skip to content
Jordan Madison
Table of Contents

This post details one of many frustrations I experienced this past year in building out my research project, CreSTAL. Currently, we are writing the research paper, so here is a debug that serves as a complement. For transparency: I used Claude Fable to synthesize my written notes and training logs into this concise breakdown. Enjoy!

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 MgO2Tm
is expected to increase the band gap due to...</hypothesis>
<modification>['substitute', 'Hg', 'O']</modification>

Downstream code parses those tags, so the structure isn’t cosmetic — if the tags 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. Textbook. 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.011295
Token 151670: norm=0.3614, mean=-0.000259, std=0.011295
Token 151671: norm=0.3613, mean=-0.000258, std=0.011295
Token 151672: norm=0.3614, mean=-0.000259, std=0.011295


Token 100:    norm=1.0396, mean=-0.000850, std=0.032492
Token 500:    norm=0.9730, mean= 0.000002, std=0.030422
Token 1000:   norm=1.0908, mean= 0.000883, std=0.034092
Token 5000:   norm=0.7705, mean=-0.001563, std=0.024040

Two 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 observation 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

Here’s the thing about PEFT/LoRA that bit me: 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 happily 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.

And this explains the Thai. My untrained embeddings sat right at the mean of the embedding distribution — a crowded, low-signal region of the space where rare, low-frequency tokens (in this vocab, a lot of multilingual tokens) tend to live. When the model reached for “that token that opens the response,” the nearest things it could actually express were its neighbors in that region: ตัด, เหลือ, ปริมาณ. The Thai wasn’t random noise; it was the ghost of an untrained embedding.

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). Interestingly, the weights barely moved. No dramatic rewrite of the embedding space. But it didn’t take much: 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 MgO2Tm
is expected to increase the band gap due to several structural and chemical
factors...</hypothesis><modification>['substitute', 'Hg', 'O']</modification><|im_end|>

Small nudge, total behavior change. Such is the joy of working with PEFT/LoRA in SFT and RL.

Takeaways

If you add tokens to a tokenizer and train with LoRA, you must unfreeze embed_tokens and lm_head (via modules_to_save). LoRA’s whole premise is “freeze everything, adapt a little” — which silently includes your brand-new, never-trained embedding rows unless you say otherwise.

Untrained embeddings fail loudly, but in a costume. The failure mode isn’t “tags missing” — it’s bizarre substitute tokens (often rare multilingual ones) appearing where your tokens 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.

Trust, but decode. The loss curve looked fine the whole time. Only actually decoding samples and reading them surfaced the problem.