Tokenization And Embeddings

LLMs do not see words. They see integer token ids produced by a fixed tokenizer, then map those ids through an embedding table. Subword tokenization is the practical compromise between character-level flexibility and word-level vocabulary size: common strings can become single tokens, while rare names, typos, code, or multilingual text can still be represented as shorter pieces.

1. Freeze the contract: model weights, tokenizer files, special token ids, and chat templates must travel together.
2. Budget with tokens: context length, cost, prefill latency, and KV-cache growth depend on token count, not character count.
3. Test round trips: decode(encode(text)) should preserve the intended text for the supported input domain.

Byte pair encoding starts from small symbols, repeatedly finds a frequent adjacent pair, and adds a merge rule that treats that pair as a new symbol. Modern production tokenizers add important details around bytes, Unicode, regex pre-tokenization, and special tokens, but the core merge loop is small enough to implement from memory.

from collections import Counter

def learn_bpe_merges(words, num_merges):
    splits = [tuple(word) + ("</w>",) for word in words]
    merges = []
    for _ in range(num_merges):
        pair_counts = Counter()
        for pieces in splits:
            pair_counts.update(zip(pieces, pieces[1:]))
        if not pair_counts:
            break
        best = max(pair_counts, key=pair_counts.get)
        merges.append(best)
        merged = "".join(best)
        next_splits = []
        for pieces in splits:
            out = []
            i = 0
            while i < len(pieces):
                if i + 1 < len(pieces) and (pieces[i], pieces[i + 1]) == best:
                    out.append(merged)
                    i += 2
                else:
                    out.append(pieces[i])
                    i += 1
            next_splits.append(tuple(out))
        splits = next_splits
    return merges, splits

1. Check determinism: define a tie-break rule for equally frequent pairs so training is reproducible.
2. Check reversibility: keep enough boundary information to detokenize without silently deleting spaces or punctuation.
3. Check unknowns: byte-level tokenizers can represent arbitrary text; wordpiece-style systems need an explicit unknown-token policy.

Chat models are still next-token models. The role list [system, user, assistant] must be rendered into the exact control tokens, separators, and generation prompt expected by the checkpoint. Hugging Face treats the chat template as part of the tokenizer via apply_chat_template, and OpenAI's token counting guidance makes the same production point: message wrappers, tool schemas, and model-specific formatting count toward the prompt budget.

messages = [
    {"role": "system", "content": "Answer from the provided sources."},
    {"role": "user", "content": "Summarize the KV cache tradeoff."},
]

input_ids = tokenizer.apply_chat_template(
    messages,
    add_generation_prompt=True,
    tokenize=True,
    return_tensors="pt",
)

1. Pin the template: log tokenizer revision, chat template, special token ids, and whether a generation prompt was added.
2. Count the rendered prompt: measure tokens after chat formatting and tool/schema insertion, not from raw message text.
3. Test stop behavior: verify assistant turns stop at the model's expected end token instead of leaking into the next role marker.

Adding a special token is more than editing a tokenizer file. Hugging Face's tokenizer docs expose special tokens as tokenizer-level metadata, while the model docs expose resize_token_embeddings for changing the embedding table after vocabulary growth. In a decoder-only LM with tied embeddings, the input lookup and output projection often share weights, so a new row must be initialized, trained, saved with the tokenizer, and tested before the model can reliably read or generate that token.

special = {"additional_special_tokens": ["<tool_call>", "</tool_call>"]}
added = tokenizer.add_special_tokens(special)
if added:
    model.resize_token_embeddings(len(tokenizer))

tool_id = tokenizer.convert_tokens_to_ids("<tool_call>")
assert tool_id != tokenizer.unk_token_id
assert model.get_input_embeddings().weight.size(0) == len(tokenizer)

if model.get_output_embeddings() is not None:
    assert model.get_output_embeddings().weight.size(0) == len(tokenizer)

1. Resize immediately: call the model resize method after tokenizer growth so every new token id has an embedding row and, for LM heads, a logit row.
2. Initialize deliberately: newly added rows start untrained, so they need fine-tuning data or a documented initialization policy before they carry meaning.
3. Save together: publish or checkpoint tokenizer files, special-token config, model weights, and chat template as one versioned bundle.
4. Test generation: verify the model can both consume the new marker in prompts and assign finite logits to it when it is allowed as an output token.

An embedding layer is a lookup table from token id to vector. During language-model training, gradients update only the rows used in the batch. Many decoder-only models tie the input embedding matrix to the output projection so the same token-vector space is used to read tokens in and score next-token logits.

class TinyTokenModel(torch.nn.Module):
    def __init__(self, vocab_size, d_model):
        super().__init__()
        self.tok_emb = torch.nn.Embedding(vocab_size, d_model)
        self.block = TransformerBlock(d_model)
        self.norm = torch.nn.LayerNorm(d_model)
        self.lm_head = torch.nn.Linear(d_model, vocab_size, bias=False)
        self.lm_head.weight = self.tok_emb.weight  # tied embeddings

    def forward(self, input_ids):
        x = self.tok_emb(input_ids)
        x = self.block(x)
        return self.lm_head(self.norm(x))

1. Shape check: input ids are [batch, time]; embeddings become [batch, time, d_model]; logits become [batch, time, vocab].
2. Resize carefully: adding special tokens requires resizing embeddings and initializing new rows before fine-tuning.
3. Keep ids stable: changing token ids after training changes which vectors the model reads and writes.

1. Count: measure token counts for short chat, long RAG context, code, emoji, and non-English text with the target tokenizer.
2. Round trip: encode and decode prompts with spaces, newlines, tabs, quotes, and special tokens.
3. Budget: reserve context for system text, tools, retrieved chunks, user input, and expected output before serving.
4. Version: log tokenizer name, model revision, special tokens, and chat template with training and eval runs.

Sources: Sennrich et al., subword BPE for rare words, SentencePiece, OpenAI tiktoken, OpenAI Cookbook token counting, Hugging Face chat templates, Hugging Face tokenizer docs, Hugging Face model resize docs, PyTorch Embedding docs, OpenAI token counting cookbook, Press and Wolf on tied embeddings.