Chapter 17

Tokenization & Embeddings

Byte-level BPE from scratch, embedding lookups and their gradients, and weight tying.

LLMs do not read characters or words directly. They read integer token ids, turn those ids into vectors, and train the same vector table that the rest of the network sees. Tokenization chooses the units; embeddings make those units differentiable enough for gradient descent. This chapter builds that input pipeline before the transformer chapters use it.

17.1 From characters to bytes to subwords

A character tokenizer is simple and robust, but sequences are long. A word tokenizer makes sequences shorter, but every typo, name, emoji, or new language needs an unknown-token escape hatch. A subword tokenizer sits between them: frequent strings such as the or ing become one token, while rare words can still be decomposed.

The safest starting alphabet is not Unicode characters but UTF-8 bytes. Every string becomes bytes first: a uses one byte, é uses two, 世 uses three, and 🙂 uses four. A byte-level tokenizer therefore starts with 256 atomic symbols and can represent any valid text without an unknown character. The cost is that a very small vocabulary behaves like a byte model, while a very large vocabulary spends parameters on rare pieces.

Whitespace is just data in this view. The byte for a leading space may merge with the following letters, so ` cat` and cat can become different pieces. Case, accents, and Unicode normalization are also modeling choices rather than mathematical necessities. For from-scratch work, byte-level BPE is attractive because the only required preprocessor is UTF-8 encoding.

17.2 Byte-level BPE

Byte-pair encoding (BPE) learns a vocabulary greedily [sennrich2015neural]. Start with each training string as a sequence of bytes. Count adjacent token pairs, merge the most frequent pair into a new token, rewrite the corpus, and repeat until the requested vocabulary size is reached. The merge list is the model: encoding new text applies the learned merges in order, and decoding concatenates each token’s byte string before UTF-8 decoding.

Listing 17.1 Byte-level BPE training

def train_byte_bpe(texts, vocab_size):
    """Train byte-level BPE by repeatedly merging the most common pair."""
    if vocab_size < 256:
        raise ValueError("byte-level BPE needs room for all 256 bytes")
    sequences = [tuple(utf8_bytes(text)) for text in texts]
    vocab = {i: bytes([i]) for i in range(256)}
    merges = []
    while len(vocab) < vocab_size:
        counts = {}
        for tokens in sequences:
            for pair in zip(tokens, tokens[1:]):
                counts[pair] = counts.get(pair, 0) + 1
        if not counts:
            break
        pair, count = sorted(counts.items(), key=lambda item: (-item[1], item[0]))[0]
        new_id = len(vocab)
        vocab[new_id] = vocab[pair[0]] + vocab[pair[1]]
        merges.append((pair[0], pair[1], new_id, count))
        sequences = [tuple(_merge_pair(tokens, pair, new_id)) for tokens in sequences]
    return vocab, merges

Because each merge replaces two adjacent ids by one id, BPE is lossless only when the decoder keeps the byte string for every learned token. The training objective is compression by frequency, not linguistic truth. On the tiny corpus embedded in the tests, ten merges shrink the corpus from 27 bytes to 10 BPE tokens; the tests compute and assert those counts.

The greedy choice has a simple effect. If the chosen pair appears $c$ times in the rewritten corpus, that merge shortens the corpus by $c$ token positions before later merges run. Ties do not change the definition of BPE, but code should break them deterministically so tests and saved tokenizers are reproducible.

Listing 17.2 Encoding and decoding with the learned merges

def encode(text, merges):
    """Encode text by applying learned merges in order."""
    tokens = utf8_bytes(text)
    for left, right, new_id, _count in merges:
        tokens = _merge_pair(tokens, (left, right), new_id)
    return tokens


def decode(tokens, vocab):
    """Decode token ids by concatenating their byte strings, then UTF-8 decoding."""
    return b"".join(vocab[int(token)] for token in tokens).decode("utf-8")

17.3 Vocabulary-size trade-offs

A larger vocabulary shortens sequences. That lowers attention cost later, because attention compares every query position with every key position. It also enlarges the embedding table and output classifier, both roughly $Vd$ parameters for vocabulary size $V$ and model width $d$ . A smaller vocabulary shares statistics better across rare words and languages, but it gives the transformer more positions to process.

This is why tokenizer choice is a systems decision as much as a modeling decision. The same sentence can be cheap in one tokenizer and expensive in another, so perplexities are only comparable when the tokenizer and evaluation text are fixed. The unigram language-model tokenizer used by SentencePiece starts from many candidate pieces, assigns each a probability, and prunes pieces that least hurt corpus likelihood; subword regularization samples alternative segmentations during training [kudo2018sentencepiece], [kudo2018subword].

The vocabulary is usually frozen before model training. Changing it changes every input id, the embedding table shape, and the output classifier shape, so it is not a harmless data-cleaning tweak. For a tiny educational model, choose a vocabulary large enough that common substrings merge, then stop before the tables dominate the parameters.

17.4 Embeddings are gathered rows

Let $\text{id}_i$ be a token id and $\mE \in \R^{V \times d}$ be the embedding table. A lookup returns row $\mE_{\text{id}_i,:}$ . Equivalently, if $\boldsymbol{e}_{\text{id}_i}$ is a one-hot row vector,

\vx_i = \boldsymbol{e}_{\text{id}_i}\mE .\tag{17.1}

The one-hot view is useful for derivations, but real code uses integer indexing (Section B.4) so it never materializes a $V$ -dimensional one-hot vector. During backpropagation, the same row may have been used many times. Its gradient is the sum of all upstream gradients that selected it:

\bar{\mE}_j = \sum_{i:\,\text{id}_i=j} \bar{\vx}_i .\tag{17.2}

Listing 17.3 Embedding lookup and its scatter-add gradient

def embedding_lookup(indices, weight):
    """Gather embedding rows: weight[indices]."""
    return weight[np.asarray(indices)]


def one_hot_matmul(indices, weight):
    """The same lookup, written as a one-hot matrix multiply."""
    one_hot = np.eye(weight.shape[0], dtype=weight.dtype)[np.asarray(indices)]
    return one_hot @ weight


def embedding_backward(indices, grad_output, vocab_size):
    """Scatter-add output gradients into the rows that were gathered."""
    grad_weight = np.zeros((vocab_size, grad_output.shape[-1]), grad_output.dtype)
    np.add.at(grad_weight, np.asarray(indices), grad_output)
    return grad_weight

For a batch of ids with shape $B \times T$ , the lookup returns $B \times T \times d$ . The backward pass is sparse in spirit: only rows that appeared in the batch receive nonzero contributions. We still store the gradient as a dense table here because it keeps the NumPy code direct and matches the parameter shape.

17.5 Weight tying

A language model also needs an output matrix that maps hidden states to vocabulary logits. With weight tying, the output matrix is the transpose of the input embedding table: $\vz_t = \vh_t\mE^\T$ . The same vector for a token is used both to read that token and to score it as the next token. This saves parameters and often improves language models [press2016using]; the gradient into $\mE$ is then the sum of the input-lookup scatter gradient and the output-softmax matrix gradient.

In practice

Current LLMs almost always use subword tokenizers before the transformer stack. Byte-level systems avoid unknown characters and make multilingual and noisy web text easier to ingest; byte-level subwords were studied directly for neural translation [wang2019neural]. SentencePiece is popular because it trains from raw text without pre-tokenized words [kudo2018sentencepiece]. Weight tying remains a common default when the input and output vocabulary are the same [press2016using].

Key equations

\text{UTF-8 text} \rightarrow (b_1,\ldots,b_n), \qquad b_i \in \{0,\ldots,255\}

(a,b) = \argmax_{(u,v)} \operatorname{count}(u,v)

\vx_i = \boldsymbol{e}_{\text{id}_i}\mE = \mE_{\text{id}_i,:}

\bar{\mE}_j = \sum_{i:\,\text{id}_i=j}\bar{\vx}_i

\vz_t = \vh_t\mE^\T \quad \text{(tied output weights)}

17.6 Teach it

The one-sentence version. Tokenization turns text into reusable integer pieces; embeddings turn those ids into trainable rows.

An analogy. A tokenizer is a packing list: common bundles get their own label, but every bundle can still be unpacked into bytes.

At the board.

Write lowest as UTF-8 bytes.
Count adjacent pairs in low lower lowest; merge the most common pair.
Show that a lookup is a one-hot row times $\mE$ , then erase the one-hot and write E[ids].
Backpropagate two uses of the same id and add their row gradients.

Misconceptions to address.

"Tokens are words." Many tokens are word pieces, spaces, bytes, or punctuation.
"Embeddings are fixed features." They are ordinary parameters trained by gradients.
"A larger vocabulary is always better." It trades shorter sequences for larger tables and rarer pieces.

Check for understanding. If token id 7 appears three times in a batch, what happens to row 7 of the embedding gradient?

17.7 Exercises

Exercise 17.1 ★ Choosing units

For unhappiness 🙂, describe one advantage and one drawback of character, word, and byte-level subword tokenization.

Solution in Appendix D

Exercise 17.2 ★★ Lookup gradient

Derive (17.2) from the one-hot matrix view in (17.1).

Solution in Appendix D

Exercise 17.3 ★★ Tiny BPE

Train ten byte-level BPE merges on low lower lowest and newer wider. Explain why decoding is exact even when an emoji appears at test time.

Solution in Appendix D

Exercise 17.4 ★★★ Implement scatter-add

Write a NumPy function that receives token ids and upstream gradients and returns the embedding-table gradient. It must handle repeated ids.

Solution in Appendix D

References

[kudo2018sentencepiece] T. Kudo and J. Richardson. SentencePiece: A simple and language independent subword tokenizer and detokenizer for Neural Text Processing. 2018. arXiv:1808.06226
[kudo2018subword] T. Kudo. Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates. 2018. arXiv:1804.10959
[press2016using] O. Press and L. Wolf. Using the Output Embedding to Improve Language Models. 2016. arXiv:1608.05859
[sennrich2015neural] R. Sennrich, B. Haddow, and A. Birch. Neural Machine Translation of Rare Words with Subword Units. 2015. arXiv:1508.07909
[wang2019neural] C. Wang, K. Cho, and J. Gu. Neural Machine Translation with Byte-Level Subwords. 2019. arXiv:1909.03341