Tokenizers

I explain why tokenizers are useful, and I introduce the BPE tokenizer.

Tokens are the smallest units of text that language models process. At a high level, a language model takes a sequence of tokens as input and generates a sequence of tokens as output.

A tokenizer is an algorithm that takes raw text and converts it into a sequence of tokens, and back again. See Figure 1.

Figure 1. A tokenizer acts as an interface between the user and the language model.

But what should we use as tokens?

Naive solutions

Word-level tokens

Perhaps the most intuitive choice is to use words, alongside whitespace and punctuation.

Tokenization ␣ is ␣ easy !

Each word is assigned a unique token ID, forming a vocabulary. However, this approach has several serious drawbacks:

Character-level tokens

A natural response to these problems is to use individual characters as tokens. Each character in the Unicode standard can be given a unique token ID.

T o k e n i z a t i o n ␣

i s ␣ e a s y !

This approach allows the model to represent previously unseen words as sequences of known characters. However, we still have problems:

More importantly, using character-level tokens has a major drawback: it produces much longer input and output token sequences. This has several consequences:

Byte-level tokens

To reduce the vocabulary size and eliminate out-of-vocabulary issues, we can take the previous idea one step further and use bytes as tokens.

In Unicode, each characted is assigned an integer, known as a code point. For instance, the character “A” has the code point 65. We can convert these code points into a sequence of bytes, and then use the bytes as tokens.

The most prevalent way to map Unicode to bytes is to use the UTF-8 encoder (Unicode Transformation Format, 8-bit). The string “Tokenization is easy!” becomes:

84 111 107 101 110 105 122

97 116 105 111 110 32 105

115 32 101 97 115 121 33

Each byte is an integer from 0 to 255, so the vocabulary contains just 256 tokens. Furthermore, we no longer need to worry about out-of-vocabulary tokens because any future characters will still be encoded as a sequence of bytes.

The trade-off is sequence length. Since there are currently 172,808 characters in the Unicode standard, some code points must map to more than one byte. This means that all the issues we had around sequence length with a character-level tokenizer are even more exacerbated here.

Note that most of these problems live at the level of compute and memory efficiency, so it is not crazy to imagine a future where we indeed use byte-level tokens. In the meantime, however, we need better solutions.

Subword tokenizers

We are dealing with a trade-off between vocabulary size and sequence length, or the ability of a tokenizer to compress text. We want tokenizers with a relatively small vocabulary size, where each token is dense with semantic meaning and information. See Figure 2.

Figure 2. Schematic representation of the trade-off between vocabulary size and the text compression properties of a tokenizer. We seek a balance between the two, in the red region.

To achieve this, it makes sense to use subwords as tokens. These are units falling between individual bytes and full words. The question becomes: which subwords should we include in the vocabulary?

Token ization ␣is ␣easy !

The guiding principle is that common sequences should be represented by single tokens, while rare words should be decomposable into smaller, meaningful tokens. For instance, if “the” and “ing” occur frequently in the training data, it may be useful to assign them their own tokens. Rare substrings, by contrast, need not become separate vocabulary entries.

In this sense, a tokenizer is a compression algorithm. It is trained on a text corpus, independently of the language model, and then used to encode and decode text.

There is no single best solution. Several algorithms have been proposed, including WordPiece in 2012, Byte-Pair Encoding (BPE) in 2015 or Unigram in 2018. In what follows, we focus on the BPE tokenizer.

BPE tokenizer

At a conceptual level, the BPE algorithm is very easy to understand.

We start with a vocabulary of 256 tokens, one for each possible byte. We then identify the most frequent pair of adjacent tokens in the corpus, merge that pair into a new token, and repeat until the vocabulary reaches the desired size.

BPE was originally introduced as a data-compression algorithm in 1994. In 2016, it was proposed as a tokenizer at the character level. The method gained wider popularity when OpenAI used a byte-level variant for GPT-2, described in their 2019 paper.

One important addition in GPT-2 was pre-tokenization. OpenAI realized that many versions of common words like “dog” appeared in the final vocabulary in many variations like “dog.”, “dog!” or “dog?”. These tokens get completely different token IDs, even though they are very similar semantically. To avoid this, they used a regular expression to split the text into units across which merges are not allowed.

The full algorithm is shown below, with implementation details abstracted away. In practice, a naive implementation can be made substantially more efficient by caching pair counts instead of recomputing them at every iteration, parallelizing pre-tokenization across a large corpus, and maintaining a heap to efficiently identify the most frequent pair.

Pair = tuple[bytes, bytes]
NUM_BYTES = 2**8

def train_bpe(
    input_path: str, 
    vocab_size: int, 
    special_tokens: list[str]
) -> tuple[dict[int, bytes], list[Pair]]:
    """Train a BPE tokenizer on the input text file."""
    assert vocab_size >= NUM_BYTES + len(special_tokens)

    # Initialize vocabulary
    vocab = initialize_vocab(special_tokens)

    # Pre-tokenize
    pretokens, pretoken_counts = pretokenize(
        input_path,
        special_tokens,
    )

    # Perform BPE
    merges: list[Pair] = []
    pair_counts, pair_to_indices = initialize_pair_stats(
        pretokens, pretoken_counts
    )

    for token_id in range(len(vocab), vocab_size):
        if not pair_counts:
            break

        # Select the most frequent pair, 
        # breaking ties lexicographically
        pair = max(
            pair_counts, 
            key=lambda pair: (pair_counts[pair], pair)
        )
        merges.append(pair)
        vocab[token_id] = pair[0] + pair[1]

        apply_merge(
            pretokens,
            pretoken_counts,
            pair,
            pair_counts,
            pair_to_indices,
        )

    return vocab, merges

The output consists of a vocabulary of the desired size and a sequence of merges that together define how the tokenizer encodes and decodes future text. For example, the following are the first merges produced when training BPE on the OpenWebText dataset.

[
    (b' ', b't'), 
    (b' ', b'a'), 
    (b'h', b'e'), 
    (b'i', b'n'), 
    (b'r', b'e'), 
    (b' t', b'he'), 
    (b'o', b'n'), 
    (b'e', b'r'), 
    (b' ', b's'), 
    (b' ', b'w')
]