Train Your First Language Model: From Raw Text to Generation
Train your first language model from scratch — a character tokenizer, the six-step training loop, overfitting, and temperature/top-k text generation.
This article has an interactive companion
Don't just read it — step through it interactively, one state change at a time.
Every large language model, all the way up to the ones with hundreds of billions of parameters, runs the same core loop you can fit on one screen: read some text, predict the next token, measure how wrong you were, nudge the weights, repeat. The size and the attention machinery change. The loop does not.
So the fastest way to actually understand training is to shrink everything until you can watch every number move. Drop the tokenizer down to single characters. Drop the model down to a bigram — one character of context predicting the next. Train it on a sentence short enough to read. Nothing here is a metaphor; it is the genuine makemore/nanoGPT pipeline with the batch sizes turned down so the arithmetic stays legible.
This walk-through mirrors YeetCode's interactive deep dive, which trains a real tiny model live in the browser. We'll go from raw text to batched tensors, through the six-step training loop, into what overfitting looks like, and out the other side generating characters with temperature and top-k.
How does a character-level model turn text into training data?
A character-level model's entire vocabulary is just the distinct characters in your corpus. Build it with one line — vocab = sorted(set(text)) — and every character gets an integer index. Real Shakespeare gives 65 characters; a short phrase gives whatever it actually contains. The number is computed, never assumed.
Two maps fall out of that vocabulary: stoi (string-to-integer) turns characters into ids for the model, and itos reverses it to turn the model's predicted ids back into readable text at generation time. Models never see letters — they see integers.
The clever part is how one string becomes thousands of training examples for free. Take a window of block_size (T) consecutive ids as the context x, and the same ids shifted right by one as the target y. Position by position, the target is simply "the next character." A window of length 4 isn't one example — it's four, because every position inside it carries its own next-char prediction.
x = ids[i : i + T] # context
y = ids[i + 1 : i + T + 1] # next-char target (x shifted right by 1)Stack batch_size (B) of those windows and you get an [B, T] tensor — the exact shape the training loop's forward pass consumes. B trades memory for a smoother gradient estimate; T is how far back the model can see.
Walkthrough: batching "to be or not to be"
Take the string to be or not to be with block_size = 4 and batch_size = 3. First, the vocabulary — the sorted set of unique characters:
sorted(set("to be or not to be")) = [' ', 'b', 'e', 'n', 'o', 'r', 't']
vocab_size = 7
stoi = { ' ':0, 'b':1, 'e':2, 'n':3, 'o':4, 'r':5, 't':6 }Encode the whole string by looking each character up in stoi:
"t o _ b e _ o r _ n o t _ t o _ b e"
6 4 0 1 2 0 4 5 0 3 4 6 0 6 4 0 1 2Now slide three windows of length 4 across those ids (stride 1) and shift each target right by one:
| Window | start | x = ids[start:start+4] | y = ids[start+1:start+5] |
|---|---|---|---|
| #0 | 0 | [6, 4, 0, 1] | [4, 0, 1, 2] |
| #1 | 1 | [4, 0, 1, 2] | [0, 1, 2, 0] |
| #2 | 2 | [0, 1, 2, 0] | [1, 2, 0, 4] |
Read window #0 left to right: given t(6) predict o(4); given t o(6,4) predict _(0); given t o _(6,4,0) predict b(1); and so on. Every y value is exactly the x value one step to its right. Stack the three windows and both xb and yb have shape [B, T] = [3, 4]. That tensor is what flows into the model next.
The model: a bigram next-character predictor
The tiny model is a single weight matrix W of shape [V, V] — for our example, 7×7. To predict from a character, you look up its row. That row lookup is the forward pass:
logits = W[prev_char] # a length-V row (mathematically: onehot(prev) · W)
probs = softmax(logits) # V numbers that sum to 1Softmax turns raw scores into a probability distribution over the whole vocabulary. Cross-entropy loss then reads off the probability the model assigned to the correct next character and penalizes it for being low:
loss = mean over the batch of -log(probs[target])Here's a number worth memorizing. With small random initial weights, softmax is nearly uniform, so every character gets probability ≈ 1/V. The loss is therefore ≈ -log(1/V) = log(V). For V = 7 that's log(7) ≈ 1.946. If your untrained loss starts anywhere near log(vocab_size), initialization is healthy; if it starts wildly higher, something is miswired before you've trained a single step.
The six-step training loop
Training is the same five moves, repeated until the loss stops falling — then grab the next batch:
for step in range(max_iters):
xb, yb = get_batch('train') # next [B, T] batch
logits = W[xb] # 1. forward
loss = cross_entropy(logits, yb) # 2. loss
loss.backward() # 3. backward
clip_grad_norm_(W, max_norm) # 4. clip
optimizer.step() # 5. step
# 6. record loss, loop1. Forward. Gather the logit row for every context character in the batch. One matrix indexing op.
2. Loss. Cross-entropy over the batch, as above. This scalar is the only signal training gets.
3. Backward. Softmax composed with cross-entropy has a famously clean analytic gradient — no autograd needed to see it:
dlogits = (softmax(logits) - onehot(target)) / BThe gradient is literally "predicted distribution minus the truth." Where the model was too confident, the gradient is positive and that logit gets pushed down; the correct character's logit (where onehot is 1) gets pushed up. Accumulate dlogits into dW[prev_char] for each example.
4. Clip. Big gradients can blow a step out of the sane region. Global-norm clipping rescales the entire gradient if its L2 norm exceeds a ceiling:
coef = max_norm / (||dW|| + 1e-6) # apply only when coef < 1This preserves the gradient's direction and only shortens its length — a guardrail, not a redirect.
5. Step. The optimizer moves the weights against the gradient. Plain SGD is one line, W -= lr * dW. AdamW keeps per-parameter momentum (m) and variance (v) buffers with bias correction, adapting the effective step size per weight and decoupling weight decay. On this tiny model SGD at a learning rate around 3 already drives the measured loss down step after step — but push the learning rate to 100+ and you'll watch the same loss curve diverge upward instead. The learning rate is the single most consequential knob.
6. Next batch. Record the loss for the curve and repeat. A descending curve means the weights are learning the corpus's character statistics.
Why train loss alone lies: overfitting
A falling training loss feels like progress, but it only proves the model is memorizing the data it can see. The real question is whether it generalizes to text it hasn't seen. So you hold out a validation split and measure loss on it with a forward pass only — no gradient, no update.
Early on, train and validation loss fall together. Then they part ways: train loss keeps dropping while validation loss flattens and eventually turns back up. That upturn is the moment the model stops learning the language and starts memorizing the training set. On a bigram over a tiny corpus you see it fast, because there are only so many character pairs to memorize.
| Signal | What it tells you | Trap |
|---|---|---|
| Train loss falling | The model is fitting the data it sees | Says nothing about new text |
| Val loss falling | It's generalizing | — |
| Val loss rising while train falls | Overfitting has begun | Stop, or add data/regularization |
The practical rule: validation loss is the metric that matters, and the gap between the two curves is your overfitting gauge.
Generating text with temperature and top-k
Generation runs the exact same softmax machinery forward, one character at a time. Feed the current character, get logits, turn them into probabilities, sample the next character, append it, and repeat — autoregression. Two knobs shape the distribution before you sample:
logits = logits / temperature # scale BEFORE softmax
logits = top_k_mask(logits, k) # keep the k largest, rest → -inf
probs = softmax(logits) # distribution over the vocab
next = multinomial(probs) # a real random drawTemperature stretches or flattens the distribution. As T → 0 the softmax collapses onto its single largest logit — pure greedy argmax, deterministic and often repetitive. High T (say 2.0) flattens everything toward uniform, producing wilder, less coherent output. Around 0.8–0.9 is the usual sweet spot.
Top-k truncates the long tail. Masking all but the k highest logits to −∞ before softmax means the model can never sample an implausible character, no matter how the temperature reshapes what's left. A default like k = 8 keeps a few live options while cutting off the nonsense. The order matters: mask first, then normalize, so the surviving probabilities still sum to 1.
What to remember
- Vocabulary is computed from the text —
sorted(set(text))— and every character maps to an integer id in both directions. - A shift-by-one turns one string into many examples:
x = ids[i:i+T],ythe same ids offset by 1, stacked into an[B, T]batch. - The loop is always the same five moves: forward → cross-entropy loss → backward → clip → optimizer step, then the next batch.
- Untrained loss should start near
log(vocab_size)— a fast sanity check that initialization isn't broken. - Validation loss, not training loss, is the real score. When val turns up while train keeps falling, you're overfitting.
- Temperature and top-k are generation-time knobs only — they reshape the same softmax the model was trained with; they change nothing about the weights.
Keep going
You've now trained a base model — one that predicts the next character but does nothing you ask of it. The next articles cover how that raw predictor becomes something useful:
- LLM Fine-Tuning, Part 1: When to Fine-Tune & How LoRA Works — adapting a pretrained model cheaply.
- Supervised Fine-Tuning: How a Base Model Becomes an Assistant — teaching a base model to follow instructions.
- RLHF vs DPO: How Models Learn Human Preferences — aligning outputs to what humans actually want.
- RLVR: Reinforcement Learning with Verifiable Rewards, Explained — training against checkable correctness signals.
To watch a real tiny model train, overfit, and generate live — with every logit, gradient, and loss value computed on your own text — open the interactive deep dive on YeetCode.
FAQ
What is the simplest possible language model I can train?
A bigram character model: a single [V, V] weight matrix where predicting from a character means looking up its row and running softmax over the vocabulary. There's no attention, no embeddings, no hidden layers — just one lookup and one softmax. It's genuinely trainable with cross-entropy and SGD, so you can watch loss fall and generate text, while every number stays small enough to check by hand. It's the honest floor beneath every larger model.
Why is the target just the input shifted by one?
Because next-character (or next-token) prediction is self-supervised: the text itself provides the labels. For any position in the sequence, the "correct answer" is simply the character that actually came next, which is the same id stream offset by one. That single trick is why unlabeled text is such rich training data — a document of length N yields roughly N supervised examples for free, no human annotation required.
How do I know if my model is overfitting?
Track two loss curves — one on the training split, one on a held-out validation split you never train on. While both fall together, you're generalizing. The moment validation loss flattens and starts rising while training loss keeps dropping, the model has begun memorizing the training set instead of learning the underlying patterns. The gap between the two curves is your overfitting gauge; the upturn in validation loss is the signal to stop, add data, or regularize.
What do temperature and top-k actually change?
Both reshape the probability distribution at generation time, and neither touches the trained weights. Temperature divides the logits before softmax: low values sharpen the distribution toward the single most likely character (T → 0 becomes greedy argmax), high values flatten it toward random. Top-k keeps only the k highest-scoring characters and masks the rest to −∞ before softmax, so the tail of implausible options can never be sampled. Together they trade off coherence against diversity in the output.
Should my untrained model have a specific starting loss?
Yes — roughly log(vocab_size). With small random initial weights, softmax is nearly uniform, so the model assigns about 1/V probability to every character, and the cross-entropy of that is -log(1/V) = log(V). For a 7-character vocabulary that's about 1.95; for the 65-character Shakespeare set it's about 4.17. If your first-step loss lands near that value, initialization is healthy. If it's dramatically higher, something is broken before training even begins.
Make it stick: run this one yourself
Don't just read it — step through it interactively, one state change at a time.