Week 6 — From tokens to a small language model: understanding prediction through one example

Throughout this week, the tokenizer splits on spaces. The vocabulary, corpus, and dimensions stay fixed across sections. During training, the model learns from six current-token → actual-next-token tasks together. During generation, it appends just one token at a time.

A token is a basic unit that a tokenizer passes to a model. It need not be a whole word: it might be a character, a subword, or a byte. For this week's hand calculations, our tokenizer splits on spaces. The vocabulary is the finite list of allowed tokens. A token ID is a token's integer address in that list.

Different tokenizers can turn the same text into different numbers of tokens. For example, learning might be one token or might split into learn and ing. Models read and generate tokens, so the splitting method affects sequence length, how much text fits in the context window, and the computation needed for training and inference. GPT-style systems commonly use subword or byte-based methods; Week 9 examines these in more detail.

The language-model training task is not to repeat the input unchanged. It is to predict the next token from the text available so far. The original text therefore supplies both the questions and their correct answers: each position predicts the token immediately after it. This is called self-supervised learning because the labels come from the text itself, rather than someone labeling every sentence by hand.

Now put the three independent sequences in a batch: row 0 is you like AI, row 1 is we like you, and row 2 is you study AI. B = 3 means that we process three samples together. They are not concatenated and do not supply context to each other. Each contributes T = 2 prediction positions, giving 3×2 = 6 tasks in total.

Do not treat [B,T,C] as a formula to memorize. Think of a table with three indexing steps: choose a training sample, choose a prediction position within that sample, then read the continuous-valued features at that position. Start with just one sample.

Here T = 2 means that this sample supplies two next-token prediction positions. The original sentence has N = 3 tokens; offsetting the targets gives T = N−1 = 2 tasks. T is not a universal constant. It is the number of input positions per sequence in the current batch.

A batch puts several independent samples in one tensor so the CPU or GPU can compute them together. In our training setup, we then average the losses from their prediction positions. This does not concatenate three sentences into one long sentence or pass information from sample 0 to sample 1. In this course's Transformer, attention never mixes different samples in the batch.

Now add embedding. inputs[b,t] is still just an integer address. The embedding lookup retrieves a floating-point vector of length C for every address. B and T remain because we still need to identify the sample and position. A new final axis of length C holds the internal representation at that position.

The colon in x[b,t,:] means all C features at this position. For example, x[1,0,:] contains the four numbers for we in sample 1. C is the representation width allocated to each position. We deliberately choose 4 for hand calculation; real models usually use much larger widths. Individual coordinates do not have predefined semantic names; training adjusts them.

An axis tells you which question an index answers; its size tells you how many choices that index has. Shape [3,2,4] has rank 3 because it has three indexing axes: b, t, and c. It contains 3×2×4 = 24 scalar values. Rank 3 does not mean three values.

An embedding looks up a learnable continuous vector using a discrete ID. Our input table is nn.Embedding(V,C) = nn.Embedding(5,4): five vocabulary rows with four learnable features per row. The values below are specified teaching initializations, not human-written definitions of words.

Using IDs directly as numerical features would impose nonexistent relationships such as we = 4 being twice AI = 2. The ID only says where to look. The embedding provides floating-point features that later layers can combine, compare, and adjust through gradients.

For example, you selects E[0] = [0.20,-0.10,0.70,0.30], while we selects E[4] = [-0.70,0.40,0.30,0.60]. Training adjusts these coordinates through loss calculation, backpropagation, and an optimizer update.

A one-hot vector acts as a selector: here it selects row 4. An embedding lookup does not need to allocate a large vector of length V; it reads E[id] directly. The two methods produce the same forward result.

We now predict the token after the prefix you. Suppose the input embedding and context processor produce h = [0.20,-0.10,0.70,0.30] at this position. To keep the hand calculation simple, this step uses h equal to the input embedding of you. In a full Transformer, h has usually incorporated context and passed through several layers of computation.

The output head stores a row of output weights and a bias for each vocabulary candidate. Each row is a learnable scoring rule: weight the same four features of h according to that candidate's rule, sum them, and add its bias.

For like, the full calculation is 0×0.20 + 1×(−0.10) + 2×0.70 + 2×0.30 + 0.10 = 2. None of these operations looks up the correct answer. The current h and the scoring weights for like simply combine to produce a relatively high score.

By hand, we calculate five dot products separately. W_out stacks the five scoring rules into rows, so W_out @ h + bias produces all five scores at once. The bias vector [5] supplies a learnable baseline offset for each candidate, independent of the current h.

Logits matter to Softmax because of the differences between candidates. If one logit exceeds another by 1, its Softmax probability is e¹ ≈ 2.718 times as large. Raising only the logit for like increases its probability relative to the others. Adding 100 to every logit leaves all differences unchanged, so it leaves the probability distribution unchanged too.

Start with just two scores, [0,1]. Exponentiating gives approximately [1,2.718], whose sum is 3.718. Dividing by this sum gives approximately [0.269,0.731]. Both entries are nonnegative and they sum to 1. Our five-candidate table uses the same calculation with three more entries.

The preceding section calculated logits = [0,2,1,-1,0] from h, W_out, and bias. These scores can be compared, but they are not probabilities: 2 does not mean 200%, and −1 does not mean a negative probability. Sampling and the probability interpretation of cross-entropy require a valid distribution: nonnegative entries summing to 1.

Step 1 is exponentiation: e⁰ = 1, e² ≈ 7.389, e¹ ≈ 2.718, and e⁻¹ ≈ 0.368. Exponentiating a finite negative logit still gives a positive number. That candidate therefore retains a chance, although its relative weight is smaller.

Step 2 is normalization. The five exponential weights sum to approximately 12.475. Dividing each by that same sum gives approximately [0.080,0.592,0.218,0.029,0.080]. The exact probabilities sum to 1; these displayed rounded values sum to approximately 1.

For logits shaped [B,T,V] = [3,2,5], dim=-1 applies Softmax separately to the five vocabulary candidates at each [b,t] position. It does not mix sentences or positions into the same denominator.

In the two-candidate example, adding 100 to both scores gives [100,101] without changing the probabilities: the difference is still 1. Implementations often subtract the maximum first, turning [0,1] into [-1,0], to avoid very large exponentials. Softmax depends on score differences, not an absolute zero point.

Softmax has produced five candidate probabilities. Suppose the actual next token is like. We first select the model's probability for this correct answer, p_correct ≈ 0.592. We want loss to approach 0 as that probability approaches 1, and the penalty to grow as it approaches 0.

Cross-entropy does more than check whether argmax is correct. If the answer is like, increasing its probability from 0.51 to 0.90 still reduces loss, even if it was already the top candidate. For several positions, mean loss averages the negative log-probability of the correct target at each position.

Notation reminder: y = 1 can denote the integer ID of the correct class. In the one-hot formula, yᵢ denotes entry i of the corresponding target vector: 1 only when i = 1, otherwise 0. It does not split the integer 1 into many ones.

F.cross_entropy combines log-Softmax with negative log-likelihood in a numerically stable way. Its score input is therefore raw logits. This course uses hard labels: torch.long target IDs from 0 to V−1, with one correct class per task.

This code follows the previous hand calculation exactly. [1,5] means one task with five candidate scores; [1] supplies the one correct class ID it needs. Once the single-task calculation is clear, arrange the batch's six positions as six tasks.

API clarification: PyTorch cross-entropy also supports floating-point class-probability targets with the same shape as the logits, including soft labels. Saying it never accepts one-hot targets is therefore inaccurate. This course uses integer hard labels to keep the six tasks and six answers easy to match: [BT,V] logits and [BT] long targets.

The preceding calculation followed task [0,0]: the visible text is you and the correct next token is like. A context representation h is passed to the output head to produce logits z = [0,2,1,−1,0]. Softmax then gives p ≈ [0.0802,0.5923,0.2179,0.0295,0.0802]. We will not repeat the dot products or exponentials here. Instead, ask what each of these five numbers means for language.

Read P(xₜ₊₁ | x≤ₜ) as: after seeing prefix x≤ₜ, what is the probability of next token xₜ₊₁? In practice, the model does not return just this one number. It returns a distribution p over the whole vocabulary. For cross-entropy, we select the entry pᵢ corresponding to the correct target.

We can now read the main computation precisely: the context processor produces h; the output head calculates z = W_out h + b_out; and p = Softmax(z) gives the conditional probabilities of the candidates. h holds features computed from context, z contains unnormalized candidate scores, and p is the distribution that answers next-token probability questions.

This exercise combines the earlier sections. Original numbers and links are preserved, but the reading order has changed: first work through inputs/targets, batches, embeddings, logits, Softmax, and cross-entropy, then connect them here. Do not try to memorize a single diagram full of unfamiliar terms.

First identify just one task: [0,0] has input you and target_id = 1, meaning like. To check the output head, the hand-calculation example supplies h = [0.20,−0.10,0.70,0.30] and fixed output weights that produce z = [0,2,1,−1,0]. Here h is a given input to a numerical demonstration, not a claim that a full Transformer always sets h equal to the embedding. Weeks 7–8 will open up the component that calculates h.

After vocabulary-axis Softmax, like has probability approximately 0.5923, giving this task loss ≈ 0.5237. If another task sees only you at the same position, the same deterministic model must produce the same distribution. If its target is study, select approximately 0.02949 from that distribution instead, giving loss ≈ 3.5237. The target does not enter the prediction computation.

Flattening only changes the arrangement: [0,0], [0,1], [1,0], [1,1], [2,0], and [2,1] become rows 0–5 in that order. It does not concatenate contexts from different sentences. Attention has already processed each batch row separately before flattening.

EX06-Integration: without looking at the answer, explain three things. Which task does flat_logits[4] represent? Why can h and the embedding have the same shape but different values? After backward, which persistent parameters receive gradients and which integers do not?

Bigram means a pair of adjacent tokens. This model learns one relationship: given current token x_t, which token tends to follow it in the training data? For example, it can learn that like follows we. It does not examine text preceding the current token. It is not a replacement for a complete GPT, but a minimal baseline whose calculation we can inspect fully.

The previous Bigram can predict by lookup, but a newly created table has not learned what follows you. Training adjusts that table by repeatedly using the actual next tokens as labels, encouraging suitable candidates to receive more probability. The loop is simply the program structure that repeats these operations.

The training-script body below runs after the BigramLanguageModel class in Section 13. Each step clears previous gradients and makes a fresh prediction using the same tasks and current parameters. Create the model and optimizer only once, outside the loop.

Think of generation as a small loop. The model reads the current prompt and produces a probability for every vocabulary candidate. A generator selects one token and adds it to the prompt. The next prompt has changed, so this simple implementation computes the next h, logits, and probabilities again. Feeding a selected output back as the next input is autoregressive generation.

Now compare training and generation. Training text already supplies the actual next tokens, so many input/target positions can be prepared and scored together. During generation, the actual next token is unavailable: we must choose and append the first next_id before we know the input for round 2.

For positive finite temperature τ, use softmax(last_logits / τ): τ = 1 preserves the distribution, τ < 1 sharpens it, and τ > 1 flattens it. Our teaching vocabulary has no end-of-sequence token, so max_new_tokens bounds the loop. This small generation fragment assumes a valid nonempty input and a nonnegative integer token count, and leaves the model in eval mode; later chapters introduce a stricter generation interface.

Both prompts end in ID 1, so the Bigram retrieves W_bigram[1,:] for both. This is not an accidental training result: with a current-token-only lookup into the [V_vocab,V_vocab] table, the earlier you or we has no computational path to affect the final-position output.

Our fixed-vocabulary table has only 5×5 = 25 logits. It can learn one-step transitions, making it a useful teaching and data-pipeline baseline, but cannot represent different continuations for the same current token under different longer prefixes.

The ideal Bigram's mean-loss infimum on these six tasks is therefore 4×ln(2)/6 ≈ 0.462, not zero. Finite logits can only approach probability 1. The causal Mini GPT in Week 11 can distinguish [you,like] from [we,like], leaving only the two conflicting targets for the identical prefix [you]. Its ideal lower limit is consequently 2×ln(2)/6 ≈ 0.231. The difference comes from available context, not inconsistent arithmetic, and does not promise that a particular small model will train to either limit.

PPL = exp(mean token loss), the reciprocal of the geometric mean of the correct-token probabilities. For probabilities 0.5 and 0.125, mean loss is [ln(2)+ln(8)]/2 = ln(4), so PPL = 4. It is not the arithmetic mean of their reciprocals, (2+8)/2 = 5, nor does it mean the model literally considered four equally likely tokens. Different tokenizers change the units being scored, so perplexity alone cannot directly compare their language capability.

Self-attention lets each position dynamically combine information from positions it may see. A query supplies matching features for the current position; a key supplies matching features for a candidate position; a value supplies the content that will actually be combined. These are typically different learned projections of the same hidden representations, not three extra tokens. The descriptions are useful intuitions, not predefined semantic labels on the coordinates.

Prompts A = [you,like] and B = [we,like] have the same token at position 1, but attention can read their different position-0 inputs. Their final logits can therefore differ. This supplies a path for using longer context; it does not guarantee a particular answer before training or establish understanding on its own.