I’m always seeking to learn something new, something that expands my understanding or provides me with deeper insights into an alternative point of view. There is nothing better than changing your opinion on something due to solid data! While reading a book it made a statement that caught me by surprise. The idea is that token distribution will appear less random as a language model is trained. This is intuitive but this was something I needed to see for myself!
Let’s establish a common set of understanding of why this is so interesting. If you’re already versed in Entropy and Language Models feel free to skip to the results!
What is Entropy?
When I think of entropy within information science I think about surprise! What is the minimum amount of bits I need to convey information? To better understand something I like to take the extreme.
Something with mass such as a rock will fall when you drop it in a gravity well (like, you know Earth). Every time you drop a rock you need to convey the result. How many bits of information do you need? Stepping back we really don’t need any bits to convey a result, which more or less says we ran the experiment and dropped a rock (so, like 0 bits are needed to encode a result). Every time we drop a rock it will fall. There is no chance of a surprise. that is just how things will work.
Now, let’s take something highly unlikely, like taking a flight from New York, NY to San Jose, CA. It would not be much of a surprise when you land in San Jose, but while possible highly unlikely you would land in London. This would be very surprising indeed! What about somewhere Oklahoma? This is a bit more likely in that mechanical, weather, or other related event may require a flight to diverge but as it’s heading in the right direction it’s more likely than London.
What is the most efficient encoding we can use to convey these results? To better understand the math involved let’s use dice. Since we are seeking to calculate bits let’s use 8 to make our math easier.
Fair die
Let’s compare two eight sided dice. Our first die is fair with each number having equal chances being rolled:
1 = \frac{1}{8} + \frac{1}{8} + \frac{1}{8} + \frac{1}{8} + \frac{1}{8} + \frac{1}{8} + \frac{1}{8} + \frac{1}{8}Now if we wanted to transmit the dice rolled over a network and we wanted to use binary to convert this number we would need 3 bits.
| Face | Probability pi | log2(1/pi) | Contribution pilog2(1/pi) |
|---|---|---|---|
| 1 | 1/8 | 3 | 0.375 |
| 2 | 1/8 | 3 | 0.375 |
| 3 | 1/8 | 3 | 0.375 |
| 4 | 1/8 | 3 | 0.375 |
| 5 | 1/8 | 3 | 0.375 |
| 6 | 1/8 | 3 | 0.375 |
| 7 | 1/8 | 3 | 0.375 |
| 8 | 1/8 | 3 | 0.375 |
| Total | 1 | — | 3.0 bits |
Unfair die
Our second die though is not fair. It is unbalanced.
1 = \frac{1}{2} + \frac{1}{4} + \frac{1}{8} + \frac{1}{16} + \frac{1}{64} + \frac{1}{64} + \frac{1}{64} + \frac{1}{64}How many bits do we need? For our unbalanced die we had the following breakdown:
| Face | Probability pi | log2(1/pi) | Contribution pilog2(1/pi) |
|---|---|---|---|
| 1 | 1/2 | 1 | 0.500 |
| 2 | 1/4 | 2 | 0.500 |
| 3 | 1/8 | 3 | 0.375 |
| 4 | 1/16 | 4 | 0.250 |
| 5 | 1/64 | 6 | 0.09375 |
| 6 | 1/64 | 6 | 0.09375 |
| 7 | 1/64 | 6 | 0.09375 |
| 8 | 1/64 | 6 | 0.09375 |
| Total | 1 | — | 2.0 bits |
So…
Entropy therefore is the lower bound on the amount of bits needed to transmit a message. We take advantage of the probability of each scenario to encode the “surprise” of each result. Whether it be a rock falling, landing in London, or rolling a 7.
H(X) = -\sum_{i=1}^{n} p_i \log_2 p_iEntropy thus provides us with a way to measure how uncertain a distribution is. The higher the number the more uncertainty.
What is a Language Model
“All models are wrong, but some are useful.”
Language
Language is a method for encoding and transmitting information — right now, I’m using English to encode the concept of a language model. Language has grammar and follows predictable patterns, one of which is Zipf’s law: it shows up not just in human language but also in birdsong and even in certain physical phenomena.
P(rank) \approx \frac{1}{rank}In other words, the most frequent word should appear about twice as often as the second most frequent, three times as often as the third, and so on. So how does this hold up in practice? Let’s look at Alice in Wonderland and check how the frequency of each word compares to its expected rank under Zipf’s law.
It’s not a perfect fit, but the data clearly tracks the expected curve.
Tokenization
While we understand words, the transformer architecture does not! Decomposed down to its individual parts, it is a series of linear algebra operations. We therefore need to convert our words, such as “The,” into numbers. A tokenizer performs this process. A tokenizer is trained on a corpus of text. While there are different approaches to tokenizing, they all seek to be as efficient as possible.
You can split out by character, but this would not be efficient for the amount of bits required! You can also split out by word, but there sure are a lot of words, with more created daily! If we went with words, we would need some way of identifying when we don’t know one, such as using
There are multiple tokenization approaches used with language models. Two of the most common — WordPiece and BPE (Byte-Pair Encoding) — get around this UNK problem by training the tokenizer on a corpus of data and splitting words into smaller, decomposable tokens that can be reassembled into the original text. They differ in exactly how they choose which pieces to merge, but both land on the same basic idea: keep common words whole, break rare words into reusable pieces.
Bit of trivia: BPE started out as a compression technique. That Zipf’s law sure comes in handy for efficient token usage!
| Word | Relative Frequency (Zipf Rank) | Word-level | Character-level | BPE (GPT-style) | WordPiece (BERT-style) |
|---|---|---|---|---|---|
| the | extremely high | the |
t·h·e |
the |
the |
| is | extremely high | is |
i·s |
is |
is |
| dog | high | dog |
d·o·g |
dog |
dog |
| running | high | running |
r·u·n·n·i·n·g |
run ning |
run ##ning |
| unhappiness | medium | unhappiness |
11 characters | un happi ness |
un ##happiness |
| tokenization | medium-low | <UNK> |
12 characters | token ization |
token ##ization |
| xylophone | low | <UNK> |
9 characters | xyl oph one |
xylo ##phone |
| ChatGPT | low (neologism) | <UNK> |
7 characters | Chat G PT |
Chat ##G ##PT |
| antidisestablishmentarianism | very low | <UNK> |
28 characters | anti dis establish ment arian ism |
anti ##dis ##establish ##ment ##arian ##ism |
| floccinaucinihilipilification | near-unique | <UNK> |
30 characters | flocc in auc in ihilip ilific ation |
flocc ##in ##auc ##in ##ihilip ##ilific ##ation |
Each token is assigned a unique integer ID, and this mapping is reversible — tokens become IDs (encoding) for the model to process, and IDs become tokens (decoding) to produce readable text.
Training a Language Model
Transformer architecture has two training approaches. The first is the masked language model (bidirectional) which attempts to predict the missing token. This is done by taking a string such as “The cat sat on the mat” and sends to the model “The cat [MASK] on the mat” (really, removes a token) and seeks to train the model to predict the word “sat.” In actual training the method replaces a percentage of text (roughly 15% of tokens) with a [MASK], though not every selected token is literally replaced with [MASK] — some are swapped for a random token and some are left unchanged. The model has the advantage of knowing the tokens before and after the missing token. This training ensures masked language models are useful for tasks such as classification.
The second approach is a causal language model (autoregressive). With this approach the last token is removed so our “The cat sat on the mat” becomes “The cat sat on the ____” and the model is trained to predict “mat” as an example. In practice, the model predicts every token in the sequence simultaneously in a single pass, each prediction based only on the tokens before it. This training prepares causal language models for tasks that are more generative in nature such as writing a fictional story about a cat. Going forward we are going to focus on the causal language model.
A Language Model in action
Taking what we learned about how a causal model is trained we can now use its amazing predictive powers to generate some text.
Using our now trained causal language model we pass in the tokens “The cat sat on the” into the model. The text is tokenized and the next token is predicted.
Using this approach, referred to as “greedy” the next token with the highest probability will always be selected. During some of my research and exploration with language models I found greedy is not great (probably a good life lesson in there too)! The text generated would tend to “collapse” and loop back in on itself repeating n-grams (sequence of text).
To get around the weakness of the greedy approach, other next-token selection methods employ exploration guided by the probabilities. Examples include beam search which explores a few moves ahead (think exploring a maze), top-k sampling which restricts the candidates to the top k then samples randomly. With k=3 we would take the top 3 tokens and randomly select one of them. If you have ever seen the term temperature, what it’s doing is changing the probability distribution of the next token where the lower it is the closer to greedy the distribution becomes.
There are a lot of different approaches to selecting the next token. The key takeaway is that this exploration relies on a probability distribution.
So what is a language model?
To recap: a language model predicts sequences of words. The training approach varies with the intended task. Part of what makes this learning possible is language itself. Statistical regularities like Zipf’s law give models a lot of repeated structure to generalize from, which is part of why scale helps so much. (There’s plenty more I’ve left out here — embeddings alone probably deserve a post of their own.)
For causal models specifically, the ability to sample from a probability distribution over candidate next tokens is what opened the door to such a wide range of capabilities. That one mechanism — predict, sample, repeat — turns out to generalize surprisingly well. Language models now serve as building blocks across a huge range of applications: commercial chatbots, spam detection, question-answering systems, self-service IVR systems. The possibilities are nearly endless.
Experiment
Training data
The motivation with the training data was to illustrate and study the distribution of the next token. I therefore wanted a smaller vocabulary. The training data was hand-crafted with the following lines:
- This quick brown fox jumped over the lazy dogs back.
- The cat ran up the tree.
- A journey of a thousand miles begins with a single step.
- To be or not to be, that is the question.
- A bird flew away from the nest.
I randomized the sentences over a loop executed 100 times for a total of 500 sentences with a small number of words overall.
The model
I opted for a single-head, an embedding size of 16 and a single transformer block (GPT-2 architecture). This will keep the model parameter size small which will ensure the model trains quickly (both in wall-time as well as the number of epochs). The goal is not emergent capabilities but rather to observe the change in entropy across epochs.
More importantly, GPT-2 is a causal model so it lends itself to our experiment needs of observation of the next token distribution.
The tokenizer
Now that we have our model selected it informs our tokenizer. For this setup I opted to train a tokenizer versus using the already trained BPE tokenizer for GPT-2. Since the training data has limited diversity the resulting vocabulary size was 142. This smaller vocabulary not only benefits model training but also reduces the number of possible outcomes in the next-token distribution.
Results
The model is initialized with random weights and saved to disk. The model is then trained over 20 epochs using our training data of 500 sentences. After each epoch the resulting weights are saved to disk for post-training study. This provides a set of weights at each epoch for evaluation.
Let’s first compare the predicted next token directly by passing in a substring of our training data.
The model is memorizing after just a few epochs the target strings. What is interesting is that not only is the model memorizing the predicted strings but the unseen string of “The dog sat by the” has a reasonable next token of “lazy” emerging which may be a result of the way the training data was prepared and how the attention head uses position. While not the focus on this experiment it does raise another dimension to explore in the future. Which leads us to why we are here, does entropy decrease over epochs?
And indeed, we can clearly see moving from a guess to a fairly low entropy occurs with our training data which reaffirms our results from the next token plot. What is also interesting is just like we saw “lazy” emerge from the next token we see the entropy of our unseen string decrease as well.
Conclusions
This was a fun experiment I encourage others to try on their own. It’s great to get back to the basics and zero in on a single aspect of an approach and pick up a new insight or technique. What I took away from this experiment was a deeper understanding of the relationship of entropy with the next token and how I can use this measure to my advantage during training.