google

/

translategemma-4b-it

like 798

Follow

Google 61.3k

TranslateGemma-4b-it produces incorrect output (predicts < token 236820 instead of Hindi)
for specific input content at exactly 50 tokens. The bug reproduces identically in
HuggingFace Transformers (PyTorch) and Candle (Rust), confirming this is model behavior,
not an implementation issue.

Gemma 3 1B (base) does not exhibit this issue, suggesting it was introduced during
TranslateGemma fine-tuning.

Minimal Reproduction

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model_id = "google/translategemma-4b-it"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32)
model.eval()

text = "The rapid advancement of artificial intelligence has transformed how we interact with technology. Machine learning models can now translate between languages."
prompt = f"<start_of_turn>user\n<translate source_lang=en target_lang=hi>\n{text}\n</translate><end_of_turn>\n<start_of_turn>model\n"

inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=True)
print(f"Tokens: {inputs['input_ids'].shape[1]}")  # 50 tokens

with torch.no_grad():
    logits = model(inputs['input_ids']).logits[0, -1]
    top_id = logits.argmax().item()
    print(f"Prediction: {top_id} ({tokenizer.decode([top_id])})")
    # Expected: Hindi token (e.g., 21290 कृ)
    # Actual: 236820 (<) 

Boundary Analysis

Testing progressively longer text shows the model works correctly up to 48 tokens,
then fails at exactly 50:

Words: 19 | Tokens:  47 |  21290 (कृ        ) OK
Words: 20 | Tokens:  48 |  21290 (कृ        ) OK
Words: 21 | Tokens:  50 | 236820 (<         ) FAIL

Full boundary sweep code

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model_id = "google/translategemma-4b-it"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32)
model.eval()

FAILURE_TOKEN = 236820

def make_prompt(text: str) -> str:
    return f"<start_of_turn>user\n<translate source_lang=en target_lang=hi>\n{text}\n</translate><end_of_turn>\n<start_of_turn>model\n"

original_text = "The rapid advancement of artificial intelligence has transformed how we interact with technology. Machine learning models can now translate between languages."
words = original_text.split()

print("Boundary sweep test")
print("=" * 70)

for i in range(1, len(words) + 1):
    text = " ".join(words[:i])
    prompt = make_prompt(text)
    inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=True)
    n_tokens = inputs['input_ids'].shape[1]
    
    with torch.no_grad():
        logits = model(inputs['input_ids']).logits[0, -1]
        top_id = logits.argmax().item()
        top_token = tokenizer.decode([top_id])
        is_failure = top_id == FAILURE_TOKEN
        
    status = "FAIL" if is_failure else "OK"
    print(f"Words: {i:2d} | Tokens: {n_tokens:3d} | {top_id:>6} ({top_token:<10}) {status}")

Key Findings

Workarounds

Two unrelated modifications fix the issue, suggesting numerical instability:

Workaround 1: Prepend Extra BOS Token

bos_id = tokenizer.bos_token_id
fixed_ids = torch.cat([torch.tensor([[bos_id]]), inputs['input_ids']], dim=1)
# Shifts all positions by 1 → now predicts Hindi correctly

Workaround 2: Replace GELU with SiLU

Modifying the MLP to use F.silu() instead of F.gelu(..., approximate='tanh')
also fixes the issue. This suggests GELU's behavior near zero may amplify
numerical instability at specific positions.

Note: The config correctly specifies gelu_pytorch_tanh, and Google's reference
implementation (google-deepmind/gemma) uses GELU exclusively, so this is not a
config error.

Investigation: GELU Variants

We tested multiple GELU modifications to identify if numerical instability was the cause:

The issue is not numerical instability or the tanh approximation.
Even exact GELU fails. Only replacing GELU with SiLU (a fundamentally different
activation curve) resolves the issue.

Hypothesis

The fine-tuned weights exhibit a pathological response at specific content + position
combinations when using GELU. Possible explanations:

1. A component of the fine-tuning pipeline may have inadvertently used SiLU
2. The RL phase (MetricX-QE / AutoMQM reward models) may have reinforced weights
   that happen to be unstable under GELU at certain positions
3. The model learned a fragile decision boundary that GELU and SiLU cross differently

The double-BOS workaround shifts RoPE positions, avoiding the problematic configuration
entirely.

Since Gemma 3 base does not exhibit this, the issue likely originates from
the TranslateGemma fine-tuning process.

Questions

1. Has this behavior been observed in internal testing?
2. Are there recommended prompt patterns to avoid this?
3. Does TranslateGemma 12B or 27B exhibit similar issues?
4. Should the double-BOS workaround be documented?