DEV Community

jidonglab
jidonglab

Posted on

Why resize_token_embeddings Breaks Your LLM Fine-Tune

You add <|plan|> and <|critique|> to the tokenizer, call model.resize_token_embeddings(len(tok)), run 2,000 SFT steps, and the loss curve looks perfect. Then at inference the model emits <|critique|> in the middle of a sentence, or — worse — never emits it at all and just writes the word "critique" instead. Nothing in the training logs told you anything was wrong.

The bug is that resize_token_embeddings writes random numbers into the new rows, and for a token that appears in a small fraction of your examples, those random numbers are still mostly there when training ends. Initialization is not a warm start here. It is the answer.

TL;DR

  • resize_token_embeddings appends rows drawn from N(0, config.initializer_range) (0.02 by default) unless your transformers version supports mean_resizing, which fits a multivariate normal to the existing rows instead.
  • A new token's logit is w_new · h. With a random row, that is noise of roughly 0.02 · ||h|| — mid-pack in the vocabulary, not at the floor, and it varies per position.
  • At a typical SFT learning rate of 1e-5 and a few hundred occurrences, the new row moves very little. Whatever you initialize is roughly what you ship.
  • With LoRA, new rows get zero gradient unless you set modules_to_save=["embed_tokens", "lm_head"] — which costs over 10 GB of optimizer state on a 128k vocab.
  • Initialize from semantically related existing tokens (input row and output row separately), add small noise to break symmetry between multiple new tokens, and verify the pre-training rank of each new token before launching the run.

What does resize_token_embeddings actually write into the new rows?

It allocates a bigger nn.Embedding (and a bigger lm_head if the model is untied), copies the old weights into the top rows, and initializes the remainder with the model's _init_weights. For Llama-family configs that is normal_(mean=0.0, std=config.initializer_range) with initializer_range = 0.02.

Recent transformers releases added a mean_resizing argument (default True) that instead samples new rows from a multivariate normal fitted to the mean and covariance of the existing embedding matrix. That is a real improvement over N(0, 0.02) — but it is a distributional fix, not a semantic one, and it silently no-ops on older versions pinned in a lot of training images. Check what your version actually does before assuming you are covered.

If you pass pad_to_multiple_of=64, you also get up to 63 rows that no tokenizer id ever maps to. config.vocab_size grows to include them, so your serving stack samples over them too.

Why does a new token's logit start mid-pack instead of at the floor?

Because the LM head has no bias term. The logit for token t at a position with final hidden state h is exactly the dot product w_t · h, and a zero-mean random w_new gives a zero-mean random logit with standard deviation std · ||h||.

Run the numbers for an 8B model. The head is applied after the final RMSNorm, so ||h|| ≈ √d · γ̄, and with d = 4096 that is about 64 · γ̄. Multiply by std = 0.02 and you get a logit that wanders roughly ±1.3·γ̄ around zero, position by position. Top logits in a converged LM typically sit 10–20 above the bulk, so the new token will not be sampled at step 0 — but it is not suppressed either. It floats somewhere in the middle of a 128k vocabulary, at a different rank at every position.

Mean initialization is cleaner than it looks, and the reason is worth stating precisely. If you set w_new = (1/V) Σ w_i, then by linearity logit_new = (1/V) Σ logit_i — exactly the mean logit at that position, every time. That is a stable, well-below-top starting point, which is why it is the default heuristic. The problem is that it is semantically empty: the new token starts equidistant from everything.

Why is initialization ~90% of the final embedding?

Two compounding reasons.

First, occurrence count. Embedding rows only receive gradient from positions where that token actually appears (input side) or is the target (output side). If <|critique|> shows up in 3% of a 20k-example dataset, that row sees on the order of a few hundred gradient events across the run, while every attention projection sees every token of every batch.

Second, bf16 rounding. bfloat16 has 8 bits of mantissa precision, so consecutive representable values are about 0.4% apart. An Adam step at LR 1e-5 changes a weight by roughly 1e-5; on an embedding entry of magnitude ~1e-2, that is a 0.1% relative change — smaller than half an ulp. It rounds straight back to the original value. Mixed-precision training normally saves you here because the fp32 master copy accumulates those steps, but memory-saving setups that keep embeddings in pure bf16, or that offload the master copy inconsistently, will silently drop every update to a rare row.

The combination means a rarely-occurring new token is, in practice, frozen near its init.

Why does LoRA leave new tokens permanently random?

Because LoRA adapters attach to linear layers you list in target_modules, and embed_tokens / lm_head are usually not in that list. target_modules="all-linear" does not include the embedding table. So the new rows receive no update at all, for the entire run, and you ship a token whose meaning is N(0, 0.02).

peft_config = LoraConfig(
    r=32,
    lora_alpha=64,
    target_modules="all-linear",
    modules_to_save=["embed_tokens", "lm_head"],  # without this, new rows never train
)
Enter fullscreen mode Exit fullscreen mode

Budget for it. On a 128,256 × 4,096 vocab each matrix is ~525M parameters; keeping both trainable in fp32 with Adam's two moment buffers is well over 10 GB on top of your base weights. If that does not fit, the alternative is to keep the matrices trainable but register a gradient hook that zeroes every row except the new ids — you pay the memory for the parameter copy but not for a full-vocab drift, and you avoid perturbing 128k trained rows with a handful of examples.

What about tied embeddings and zeroed reserved tokens?

Two traps that bite in opposite directions.

Tied embeddings (config.tie_word_embeddings = True, common in smaller Qwen and Gemma variants): get_input_embeddings() and get_output_embeddings() return the same tensor. Any careful separate init of the input row and the output row means the second write silently overwrites the first. Assert on the flag rather than assuming.

Zeroed reserved tokens: Llama-3 checkpoints ship a few hundred <|reserved_special_token_N|> rows that are exactly zero. Reusing one instead of adding a new token sounds elegant — the vocab size never changes, no serving mismatch — but a zero input row means the first transformer layer sees a literal zero vector at that position (RMSNorm of zero is zero, and RoPE rotating zero gives zero). The token carries no information at all. And a zero lm_head row gives a logit of exactly 0.0 at every position, which is only "suppressed" if your logit bulk sits above zero. If it sits below, you have created a constant-height token that spikes to the top whenever the real distribution dips.

Also: never compute your mean over the full matrix without filtering those zero rows out.

How should you initialize a new special token?

Anchor it on existing tokens that already mean something close to what you want, separately for the input and output sides.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "meta-llama/Llama-3.1-8B-Instruct"
tok   = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16)

NEW = {
    "<|plan|>":     " plan planning strategy steps",
    "<|critique|>": " critique review feedback flaws",
}

n_old = len(tok)
tok.add_special_tokens({"additional_special_tokens": list(NEW)})
model.resize_token_embeddings(len(tok))          # no pad_to_multiple_of
assert model.config.tie_word_embeddings is False # Llama-3.1-8B is untied

W_in  = model.get_input_embeddings().weight
W_out = model.get_output_embeddings().weight

with torch.no_grad():
    # per-dimension scale of rows that were actually trained
    alive = W_in[:n_old].norm(dim=-1) > 1e-6     # drop all-zero reserved rows
    sigma = W_in[:n_old][alive].float().std(0).mean().item()

    for tokstr, anchor_text in NEW.items():
        new_id  = tok.convert_tokens_to_ids(tokstr)
        anchors = torch.tensor(tok(anchor_text, add_special_tokens=False)["input_ids"])

        v_in  = W_in[anchors].float().mean(0)
        v_out = W_out[anchors].float().mean(0)
        # symmetry breaking, scaled to the real row distribution
        noise = torch.randn(v_in.shape) * sigma * 0.1

        W_in[new_id]  = (v_in  + noise).to(W_in.dtype)
        W_out[new_id] = (v_out + noise).to(W_out.dtype)
Enter fullscreen mode Exit fullscreen mode

The noise matters more than it looks. Two new tokens initialized identically are indistinguishable to the model, and any token that is never a target only ever receives downward gradient — identical rows stay identical for the whole run.

How do you verify it before burning a training run?

Measure the rank of each new token's logit on a probe set, before you train anything. You want it comfortably outside the top few hundred, and you want it stable across positions — a rank that swings from 400 to 90,000 means the row is still noise-dominated.

@torch.no_grad()
def new_token_ranks(model, tok, new_ids, prompts):
    out = []
    for p in prompts:
        batch  = tok(p, return_tensors="pt").to(model.device)
        logits = model(**batch).logits[0, -1]           # [V]
        out.append({tok.convert_ids_to_tokens(t): int((logits > logits[t]).sum())
                    for t in new_ids})
    return out
Enter fullscreen mode Exit fullscreen mode

Two more checks before you launch:

  • len(tok) == model.config.vocab_size, and the same for whatever you serve with. A pad_to_multiple_of mismatch between training and vLLM is a silent shape error at best.
  • If you must pad, do suppression in the sampler, not the weights. There is no reliable way to write a weight row that yields a uniformly low w · h; scaling a row scales the logit toward zero, which raises it whenever the mean logit is negative. Use logit_bias, bad_words_ids, or vLLM's allowed_token_ids instead.

The short answer

resize_token_embeddings breaks fine-tunes because it treats a new token's embedding as something training will figure out, when in practice training barely touches it: the row is seen by only the small fraction of positions where the token appears, bf16 rounding eats sub-0.4% updates, and under LoRA the row gets no gradient at all unless embed_tokens and lm_head are in modules_to_save. A random row is not a neutral row — with no bias in the LM head, w_new · h lands mid-vocabulary and drifts position to position, which is exactly the behavior that produces a token the model emits at random or never learns to emit. Initialize the input and output rows from semantically related existing tokens, add noise at the scale of the real row distribution, filter zeroed reserved rows out of any statistics you compute, and check each new token's logit rank on a probe set before you spend a single GPU-hour.

Top comments (0)