Speaking Thai, Part 2: Four Frozen Embeddings
/ 6 min read
Table of Contents
I hit this bug while building CreSTAL. As we write the research paper, I’m publishing some of the debugging work that sits behind it. I used Claude Fable to turn my notes and training logs into this account.
I was fine-tuning a model to generate structured scientific hypotheses about materials. Given a compound like TmMgHg2 and a proposed element substitution, it should reason about the effect on the band gap and answer in a strict format:
<hypothesis>The substitution of mercury with oxygen in TmMgHg2 to form MgO2Tmis expected to increase the band gap due to...</hypothesis><modification>['substitute', 'Hg', 'O']</modification>Downstream code parses those tags, so if they don’t come out, the pipeline breaks.
The reasonable-sounding setup
By default, a tokenizer shreds <hypothesis> into pieces: <, hypothesis, >. That means the model has to learn to emit a three-token sequence just to open a tag, and any drift mid-sequence corrupts the format. The standard fix is to register the tags as special tokens so each one becomes a single, atomic token:
special_tokens = ["<hypothesis>", "</hypothesis>", "<modification>", "</modification>"]tokenizer.add_special_tokens({"additional_special_tokens": special_tokens})model.resize_token_embeddings(len(tokenizer))Add the tokens, resize the embedding matrix so the new IDs have rows to live in, run SFT. I’d done everything the tutorials say to do, which is exactly why the bug that followed was so disorienting.
The symptom: no tags, and… Thai?
After SFT, the model’s outputs contained none of my tags. Worse, they contained tokens I never asked for:
ตัดThe substitution of mercury (Hg) with nitrogen (N) in TmMgHg2 ......could affect the outcome.เหลือ塚['substitute', 'Hg', 'N']ปริมาณ<|im_end|>The content was fine: coherent reasoning about electronegativity and band gaps. But right where my <hypothesis> and <modification> tags should have been, the model was emitting Thai tokens (ตัด, เหลือ, ปริมาณ) and the occasional CJK character. It was as if the model knew something belonged in those slots, reached for it, and grabbed a random token off the shelf instead.
If this symptom looks familiar, it’s the same saga as Why Is My Model Suddenly Speaking Thai?. That post chased the inference-and-decode half of the failure. This is the training-side half: the deeper reason the tag embeddings had nothing to say.
The investigation: four suspiciously identical rows
I dumped stats on the embedding matrix (the norm, mean, and std of each token’s embedding vector), comparing my newly added tokens against tokens from the original vocabulary:
Token 151669: norm=0.3613, mean=-0.000259, std=0.011295Token 151670: norm=0.3614, mean=-0.000259, std=0.011295Token 151671: norm=0.3613, mean=-0.000258, std=0.011295Token 151672: norm=0.3614, mean=-0.000259, std=0.011295
Token 100: norm=1.0396, mean=-0.000850, std=0.032492Token 500: norm=0.9730, mean= 0.000002, std=0.030422Token 1000: norm=1.0908, mean= 0.000883, std=0.034092Token 5000: norm=0.7705, mean=-0.001563, std=0.024040Two things jump out. My four special tokens are nearly identical to each other: same norm, same mean, same std, down to the fourth decimal place. And they’re roughly 3x smaller in norm than tokens the model actually trained on.
That first one is the tell. When you call resize_token_embeddings, the new rows are initialized from the existing embedding distribution: in recent versions of transformers, essentially the mean of the existing embeddings, plus a little noise. Four tokens initialized the same way, all sitting at the mean of the distribution. That’s fine as a starting point, if training ever touches them.
Mine were still sitting there, untouched, after a full SFT run.
The root cause: LoRA never trains what you don’t target
LoRA only updates the modules you list in target_modules. Everything else is frozen. My config targeted the usual attention and MLP projections:
lora_config = LoraConfig( r=16, lora_alpha=32, target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], # embed_tokens and lm_head: nowhere to be found)The embedding layer and the LM head weren’t in the list. So during training:
- The attention and MLP layers adapted to my SFT data, including learning that some token should appear at the start of the response.
- The four embedding rows for my special tokens stayed frozen at their initialization: four nearly-identical vectors carrying zero information.
- The LM head rows for those tokens were equally frozen, so the model could barely assign them probability even when the context screamed for them.
The model had learned the shape of my format without ever being able to learn the tokens of my format.
This also gave me a plausible explanation for the Thai. The new rows were still clustered near their shared initialization, while the model emitted unrelated, high-ID multilingual tokens at the positions where the tags belonged. The logs prove that the tag rows stayed effectively untrained; they do not prove the exact geometry that made those particular replacement tokens win. The important diagnostic was simpler: the expected IDs were absent, and the rows responsible for representing and predicting them had never moved.
The fix: modules_to_save
The fix is one line. PEFT’s modules_to_save marks layers to be fully unfrozen and trained (and saved with the adapter), which is exactly what a resized embedding matrix needs:
lora_config = LoraConfig( r=16, lora_alpha=32, target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], modules_to_save=["embed_tokens", "lm_head"],)After retraining, I re-checked both layers: the embedding matrix (tokens → vectors) and the LM head (hidden states → logits). The weights barely moved. There was no dramatic rewrite of the embedding space, but it didn’t take much, because the four rows just needed to differentiate from each other and drift far enough from the low-signal mean region to become expressible. The output went from Thai-sprinkled soup to exactly what I wanted:
<hypothesis>The substitution of mercury with oxygen in TmMgHg2 to form MgO2Tmis expected to increase the band gap due to several structural and chemicalfactors...</hypothesis><modification>['substitute', 'Hg', 'O']</modification><|im_end|>Takeaways
-
If you add tokens to a tokenizer and train with LoRA, their
embed_tokensandlm_headrows must be trainable and saved. In PEFT,modules_to_save=["embed_tokens", "lm_head"]is the straightforward way to do that when those layers would otherwise remain frozen. The rows need to train and survive adapter saving, though not every setup needs this exact option. -
Untrained embeddings don’t go missing; they get substituted. You won’t see an absent tag. You’ll see a bizarre stand-in token, often a rare multilingual one, sitting exactly where your tag should be. If your fine-tuned model suddenly speaks a language that isn’t in your data, check your embedding layer before you check your dataset.
-
Embedding stats are a 30-second diagnostic. Print norm/mean/std for your added tokens next to a few original-vocab tokens. If your new tokens are near-identical clones of each other with smaller norms, they never trained. This check would have saved me an entire debugging session, and now it runs before every training job I launch.
-
The loss curve won’t tell you. It looked fine the whole time. Only decoding samples and reading them surfaced the problem.