How language models learn

I made this short animation to show how language models learn and produce text. It starts with a small network that calculates a prediction and changes its weights. The last part illustrates how a model writes an answer one token at a time.

Watch the explainer ↓Why I built it

01 / 05 · Text and context

Training

What could come next?

An example from training text

The learning task

SignalsLearned weightsTraining correction

0:00 / 1:30

The small network calculates its predictions and weight update. The language-model steps are illustrated.

Why I built this

I think it helps to know a little about what these tools are doing. If you understand what the model learned during training and what it gets from your prompt, it is easier to decide what information to give it and what to check in its answer.

I wanted a short explanation I could send to someone who is curious about how this works. You can watch it in 90 seconds or pause at each step. I hope it helps people get more comfortable trying the tools and gives them a few more questions to ask.

The small network does the math shown on screen. The text-generation example is scripted. I used 3Blue1Brown’s explanations as a reference. His videos cover much more of the detail.

Read the narration

Use this script while presenting, or read it alongside the captions. Space plays or pauses; the arrow keys move five seconds; F toggles fullscreen.

What this model simplifies

The small network has two inputs, three hidden units, and three outputs. It calculates the numbers shown in the animation. It uses tanh to transform the signals, softmax to turn scores into probabilities, and cross-entropy to measure prediction error. Backpropagation calculates the changes needed for one gradient-descent step. The starting values were chosen for this example.

In a transformer language model, each token is represented by a list of numbers. Attention combines information from earlier tokens, and feedforward layers do further calculations. These steps repeat through many layers. The small network on this page illustrates how weights are used and changed; it is not a piece taken from a transformer. The animation leaves out position information, residual connections, normalization, multiple attention heads, and the full vocabulary.

In the calculation, tanh is a nonlinear function that compresses the sum into a value between −1 and 1. Modern transformers commonly use other nonlinear functions. The basic idea of combining signals with learned weights still applies.

The animation shows one next-token prediction at a time. During training, models usually predict many positions in parallel, using only the preceding text for each prediction. Training combines errors from batches of examples before updating the weights. An update that helps one example may hurt another, so the model also needs checks on text it was not trained on.

The final answer is scripted. This page does not run a language model. In a real model, the learned weights stay fixed during ordinary answering. Attention scores depend on the current context, so they and other internal values can change as tokens are added. Product memory features and later training are separate processes. Further training can use examples and feedback to improve how the model responds.

Further reading: 3Blue1Brown’s brief introduction to language models, backpropagation explanation, and attention walkthrough.