Adding tokens, and initialising them properly

When it is worth it

Adding vocabulary is not a routine step. The three cases where it earns its cost:

1. A control token. A <|tool_call|> or <|thinking|> marker you want the model to emit unambiguously, which no existing token can be confused with.
2. A pad token the base model lacks, as the previous lesson described.
3. A script or domain the tokenizer handles badly. If your language tokenises at four tokens per word, extending the vocabulary with common sequences from your corpus can halve that, but this only pays off alongside continued pretraining with a substantial amount of text, because the new embeddings have to be learned from scratch.

For anything else (a product name, a set of category labels, a handful of domain terms) leave the vocabulary alone. The model represents multi-token strings perfectly well, and every added token is an untrained row in two large matrices.

The mechanics

tok.add_special_tokens({"additional_special_tokens": ["<|tool_call|>"]})
model.resize_token_embeddings(len(tok), pad_to_multiple_of=64)

The pad_to_multiple_of=64 rounds the vocabulary up to a multiple of 64. This is a throughput detail, not a correctness one: matrix dimensions that are multiples of 64 map cleanly onto tensor core tiles, and a vocabulary of 128,257 can be meaningfully slower than one of 128,320. The extra rows are never selected and cost a few megabytes.

The initialisation that matters

By default, resize_token_embeddings fills new rows with values drawn from the model’s initialiser: small random numbers, unrelated to anything the model knows.

That is worse than it sounds. The existing embedding vectors occupy a particular region of the space with a particular norm; a random vector drawn from a naive distribution sits somewhere else entirely. On the input side the model receives a signal it has no representation for. On the output side the new token’s logit is computed from a random row, which early in training can make it wildly over- or under-probable and produce large gradients.

The fix is to initialise the new rows near the centre of the existing distribution:

emb = model.get_input_embeddings().weight.data
out = model.get_output_embeddings().weight.data
n = len(new_tokens)
emb[-n:] = emb[:-n].mean(dim=0, keepdim=True)
out[-n:] = out[:-n].mean(dim=0, keepdim=True)

Better still, where the new token has a meaning expressible in existing tokens, initialise it as the mean of those tokens’ embeddings. A <|invoice|> token initialised as the average of the embeddings for “invoice” starts in approximately the right place and converges much faster.

They will not train unless you say so

With a LoRA, the embedding table and output head are frozen like everything else. Your new token’s embedding stays at whatever you initialised it to, forever.

LoraConfig(..., modules_to_save=["embed_tokens", "lm_head"])

This trains those two matrices fully and saves them in the adapter. Be aware of the size: an 8B model’s embedding table is 128,256 × 4,096 ≈ 525 million parameters, so your adapter file grows from about 160 MB to over 2 GB once both matrices are included. That is usually acceptable and it is always surprising the first time.

Tied embeddings

Many models, particularly smaller ones, tie the input embedding matrix and the output head, using the same weights for both. Check model.config.tie_word_embeddings.

When they are tied, you have one matrix, not two: resizing handles both together, modules_to_save needs only one entry, and the mean-initialisation above should be applied once. When they are untied, they are independent and both need attention. Assuming the wrong one is a quiet source of a new token that works on input and misbehaves on output.

After adding, before training

Three checks:

assert tok.convert_tokens_to_ids("<|tool_call|>") != tok.unk_token_id
assert len(tok) <= model.get_input_embeddings().weight.shape[0]
print(tok.tokenize("a <|tool_call|> b"))   # must be one token, not five

The third is the one that catches the real mistake. A special token added to the vocabulary but not registered as special will be split into its constituent characters by the tokenizer, and you will have trained the model to emit a seven-token sequence that only looks like a control token.