One GPT, Two Complete Paths: Training Updates and Token-by-Token Generation

An input passes through embeddings, attention and an MLP, eventually producing logits. Training and generation use the same kind of forward network up to this point. What happens next depends on whether the true next-token labels are available, and what the output will change.

Training compares logits with labels, computes a loss and gradients, then updates parameters through an optimizer. Generation holds parameters fixed, selects an ID from the last position’s logits and appends it to the context. The training loop changes parameters; the generation loop changes the input sequence.

Fix the model and the meaning of its numbers

The running example is a GPT-2-style teaching model with the full vocabulary and a smaller network for inspecting individual values.

It uses learned absolute position embeddings, Pre-LN, GELU and tied input/output weights. There are 810,960 independent trainable parameters: 50,257×16=804,112 in the token table and 16×16=256 in the position table, plus the remaining layers. The tied token table counts once.

B denotes batch size and T the current sequence length. Most derivations omit the batch dimension: T×16 is actually 1×T×16 here. V_vocab denotes vocabulary size; attention’s Value matrix is V.

The article connects the principles. The animations display a previously saved teaching computation; the browser replays and expands its values. Training shows the first forward pass and first update. Inference uses parameters after 120 updates. The architecture matches, but the parameter snapshots differ, so identically named matrices need not contain identical values. Small arrays labeled “arithmetic example” explain formulas separately.

Connect the whole neural network

Follow the input down through two Blocks. Each contains attention and an MLP, with a residual path around each branch. Attention grids represent reading relationships between positions. MLP circles represent feature dimensions; circles on the vocabulary output side represent candidate IDs. Connecting lines represent linear transformations. Only some dimensions or candidates are drawn; computation uses the full labeled sizes. Line widths do not encode actual weight magnitudes.

The walkthroughs include the same full-network diagram: click a module to open its calculation stage. Block 2 shows the complete architecture while exposing only its saved input and output values.

Training starts by shifting the answer one position

The fixed GPT-2 tokenizer encodes Hi 你! into five IDs. The data preparation procedure explicitly appends EOS at the document boundary:

Text IDs: [17250,   220, 19526,   254,     0]
With EOS: [17250,   220, 19526,   254,     0, 50256]

x:        [17250,   220, 19526,   254,     0]
y:        [  220, 19526,   254,     0, 50256]

Each position predicts the item after its prefix. The real sequence supplies all five targets at once. A causal mask is still necessary: position 2 must not read position 3 and copy its answer.

This trains the language model. BPE merges, vocabulary entries and ID meanings have already been fixed; backpropagation does not retrain the tokenizer. EOS is an explicit data convention here, not something the original GPT-2 encode() automatically appends.

Jump between main stages, then step through the current operation. Within the network, choose a position, attention head, output feature and summation term. Tables retain complete dimensions and scroll horizontally on narrow screens.

One forward pass: parameters versus computed activations

Look up two tables and add their rows

The token ID chooses a row of E; the position chooses a row of P:

E: 50,257×16       P: 16×16
id_t → E[id_t]    t → P[t]

X⁽⁰⁾[t] = Dropout(E[id_t] + P[t])
Output: [1,T,16]

Positions are the deterministic integers 0,1,…,T−1 in the token sequence. The values inside P[t] are trainable.

Arithmetic example: [0.2,−0.1,0.5,0.3] + [0.01,0.02,−0.03,0.04] = [0.21,−0.08,0.47,0.34]. The result still has four features; the animation adds two 16-feature vectors. The same ID retrieves the same token vector at different positions, but receives a different position vector.

Training from scratch initializes embeddings and linear weights according to the initialization rules, then updates them during learning. Linear biases start at zero and LayerNorm starts with γ=1, β=0; initialization does not make every parameter random. Reinitializing each batch would discard what earlier updates learned.

LN1 controls feature scale; the following projection produces QKV

Let X be the current Block input. LayerNorm treats each position’s 16 features independently:

μ_t  = (1/D) Σ_d X[t,d]
σ²_t = (1/D) Σ_d (X[t,d] − μ_t)²
U[t,d] = γ[d] × (X[t,d] − μ_t) / √(σ²_t + ε) + β[d]

Here ε=10⁻⁵. Statistics run along features, not across the five positions. Learned γ and β mean the final output need not have exactly zero mean and unit variance. LayerNorm definition.

LN1 produces U, which then enters three projections:

Q = U W_Q + b_Q
K = U W_K + b_K
V = U W_V + b_V

With fixed parameters, a different input leaves W_Q unchanged but generally changes Q. The factors in QKᵀ are computed from the current input.

An implementation can combine these projections into one 16→48 operation, then split Q, K and V. Each 16-feature result splits into two eight-feature heads. The split is along features; each head still processes every allowed position. Equations here use row vectors, U W+b. PyTorch stores Linear.weight as output-by-input, so the combined weight is stored as 48×16.

Attention weights positions, then combines V

For one head, Q, K and V each have shape T×8:

S[i,j] = Σ_d Q[i,d] K[j,d] / √8
S_masked[i,j] = S[i,j] if j ≤ i; otherwise −∞
A[i,:] = softmax(S_masked[i,:])
A_used = Dropout(A)
H_head[i,d] = Σ_j A_used[i,j] V[j,d]

i is the query position and j the source position. Dividing by √8 controls the scale of a dot product with eight features: if the terms are roughly independent and similarly scaled, their sum’s variance grows with dimension and its standard deviation grows with the square root.

The diagonal is allowed because the current input is known and predicts the next item. Masking changes match scores; it neither deletes input tokens nor permanently zeros future rows of V.

An arithmetic example without Dropout makes the sequence explicit:

Scaled score row 2: [2, 1, 3, 4, 5]
Mask the future:    [2, 1, 3, −∞, −∞]
Softmax:            [0.244728, 0.090031, 0.665241, 0, 0]

One feature of V:   [10, 20, 30, 40, 50]
Output feature:     0.244728×10 + 0.090031×20 + 0.665241×30
                  ≈ 24.2051

Q/K determine the reading proportions; V provides the values. A T×T weight matrix times T×8 values produces T×8 outputs.

A softmax row sums to one. Training then applies Dropout, so the row actually multiplying V need not still sum to one. In inference Dropout is off and A_used=A.

Which vector does each residual addition restore?

The two heads each produce eight features for the same position. Concatenation keeps both sets, producing 16 features. The subsequent W_O learns how to mix them:

C = Concat(H_head1, H_head2)       T×16
Δ_attn = Dropout(C W_O + b_O)     T×16
R = X + Δ_attn                   T×16

The first addition restores the Block input X from before LN1. Now retain R while the next branch reads it:

U₂ = LN2(R)                      T×16
Z₁ = U₂ W_up + b_up              T×64
G  = GELU(Z₁)                    T×64
Z₂ = G W_down + b_down           T×16
Δ_mlp = Dropout(Z₂)              T×16
B₁ = R + Δ_mlp                   T×16

The second addition restores R from before LN2. Δ_mlp is the branch update; B₁ is the complete Block output.

Expanding gives B₁=X+Δ_attn+Δ_mlp, but Δ_mlp depends on R. These are not two independent branches both reading the original X. A residual connection preserves the main path and learns an adjustment to its representation; it is distinct from the prediction error in a loss.

Why 16→64→16 cannot simply become one layer

The MLP uses the same parameters independently at every position. The 64 dimensions are hidden features, not tokens. Removing the intermediate nonlinearity would give:

(U₂ W_up + b_up) W_down + b_down
= U₂ (W_up W_down) + (b_up W_down + b_down)

Two affine maps collapse into one. GELU prevents this collapse. Its exact form is GELU(x)=xΦ(x), where Φ is the standard normal cumulative distribution function: GELU(1)≈0.8413, while GELU(−1)≈−0.1587. It acts elementwise, has no trainable parameters and performs no sampling. Negative inputs do not all become zero. GELU and its approximation options.

This model follows nanoGPT’s default nn.GELU(). The historical GPT-2 tanh approximation is not substituted into these values.

Dropout is a separate operation. For a drop probability of 0.5, input [2,4,6,8] and mask [1,0,1,0] give [4,0,12,0]. Dividing retained values by the keep probability preserves each element’s expectation. It does not permanently delete features or parameters.

Finishing one Block does not reach the loss

This model has two Blocks. The second processes B₁ with its own parameters, followed by final LayerNorm:

B₁ → Block₂ → B₂ → LN_f → H_final
                            [1,T,16]

The animations expand the first Block in detail. The second has recorded inputs and outputs but lacks equally detailed internal tensors, so its display shows its structure and boundary values. LN_f has its own γ/β. Its output remains a vector, not a probability distribution.

Vocabulary projection scores the next-token candidates

With tied input/output E:

Z = H_final Eᵀ
Z[t,j] = Σ_d H_final[t,d] E[j,d]

[1,5,16] × [16,50257] → [1,5,50257]

Here j is a vocabulary ID, no longer a position in the input. Each position produces 50,257 real numbers. Those numbers are already logits. No additional conversion makes them logits. One candidate involves a 16-feature dot product; matrix multiplication evaluates the vocabulary together.

The vectors are not normalized by length, so this is not cosine similarity. It is a learned prediction score.

Cross-entropy selects the target from an entire logit row

Training knows y and needs no sampling first:

p[t,j] = exp(Z[t,j]) / Σ_k exp(Z[t,k])
ℓ_t = −ln p[t,y[t]] = −Z[t,y[t]] + log Σ_k exp(Z[t,k])
L = (ℓ₀ + ℓ₁ + ℓ₂ + ℓ₃ + ℓ₄) / 5

Implementations use a numerically stable cross-entropy operation rather than explicitly storing softmax and then taking logs. Although the target probability is one number, its denominator depends on the entire vocabulary.

The saved first training forward pass contains:

Their mean is approximately 10.828178. These probabilities belong to different positions; applying another softmax to the five numbers would not represent the model’s distribution. A uniform vocabulary distribution gives ln(50,257)≈10.825, a useful initial scale reference.

backward has not yet updated the parameters

Differentiate the loss with respect to any logit:

∂ℓ_t/∂Z[t,j]
= −1[j=y[t]] + exp(Z[t,j]) / Σ_k exp(Z[t,k])
= p[t,j] − 1[j=y[t]]

∂L/∂Z[t,j] = (p[t,j] − 1[j=y[t]]) / 5

At position 2, the target is 254 and the gradient is approximately (0.00001788837−1)/5=−0.19999642. A direct negative-gradient update to that logit would raise it. Real training updates shared network parameters, combining contributions from multiple positions and paths.

backward() computes and accumulates gradients along the computation graph. The optimizer’s step changes the parameters. ID 220 does not become another ID; floating-point values in E[220] change.

Follow E[220,0] through one AdamW update

AdamW maintains first-moment m and second-moment v statistics for each parameter. At update s:

m_s = β₁ m_(s−1) + (1−β₁) g_s
v_s = β₂ v_(s−1) + (1−β₂) g_s²
m̂_s = m_s / (1−β₁ˢ)
v̂_s = v_s / (1−β₂ˢ)

θ_s = (1−ηλ) θ_(s−1) − η m̂_s / (√v̂_s + ε)

The final line applies to a parameter group with weight decay. Here matrix weights decay, while one-dimensional biases and LayerNorm parameters use λ=0. m/v are optimizer state, not token vectors or current activations. AdamW update rule.

The saved first update follows E[220,0] in the shared table:

The settings are η=0.02, β₁=0.9, β₂=0.95, λ=0.01, ε=10⁻⁸. Starting with m₀=v₀=0, bias correction gives m̂₁=g and v̂₁=g²:

θ₁ ≈ (1−0.02×0.01)×0.0083439769
     −0.02×0.1401842982 / (0.1401842982+10⁻⁸)
   ≈ −0.0116576916

Clear the gradients and begin another batch with the updated parameters. zero_grad() clears gradients, not parameters or m/v. Deliberate accumulation across smaller batches changes when clearing occurs.

Evaluation and saving serve different purposes

Evaluation runs a forward pass and computes loss against true labels without updating parameters. eval() switches behavior such as Dropout; no_grad() or inference mode controls gradient recording. These are separate responsibilities. PyTorch’s eval documentation.

The teaching trace repeatedly updates on Hi 你! 120 times and uses only Hi 你好! for validation. Final eval losses are about 0.001070 and 7.975595, respectively. These numbers illustrate short-phrase memorization versus validation; they do not establish Chinese-language capability.

A checkpoint stores parameters and architecture configuration. Resuming training also requires optimizer state and progress; exact continuation may additionally require random and data-loader state. Inference still needs the matching tokenizer and ID meanings.

Inference fixes parameters and changes the context

The prompt Hi encodes to [17250]. Load the checkpoint, enter eval and disable gradient recording. Embeddings, all Blocks and LN_f still run. Only the last position enters lm_head here:

H_final[:, −1:, :]: [1,1,16]
Times Eᵀ:          [16,50257]
Next-item logits:  [1,1,50257]

Position t predicts the item after it. With T known tokens, position T−1 predicts the unknown continuation of the entire sequence. Earlier positions correspond to earlier prefixes, not the first, second and third future tokens.

Training can process many positions together because the true continuation already exists. Generation must choose a new ID before the next prediction can depend on it. Parallel matrix operations within a sequence coexist with sequential dependencies between future outputs.

Turn scores into one ID

The saved run uses temperature=0.8, top-k=5, a maximum of 10 new tokens and sampling seed 7:

Raw logits → divide by 0.8 → retain top 5 → softmax → sample one ID

A positive temperature preserves ranking while changing gaps. Filtered candidates receive negative infinity; retained probabilities are renormalized. Sampling uses this processed distribution, not the original full-vocabulary softmax.

In round 1, ID 220 has a raw logit of about 16.329382, scaled to 20.411728. After filtering and normalization its probability is about 0.99998498, and it is selected. This sharp distribution reflects repeated short-phrase memorization. Selecting the highest-probability item once does not prove greedy decoding was used; greedy decoding explicitly chooses argmax.

Why can a new token arrive without a new visible character?

The generation loop appends the integer ID unchanged. It need not decode to characters, join a string and tokenize again. The five recorded rounds are:

Round 2 adds no visible character, but round 3 proceeds because the model already has ID 19526. Model context and character display are separate states: one advances with IDs, the other waits for complete UTF-8 byte sequences. UTF-8 rules determine character boundaries; the model does not make a separate boundary prediction.

Decoding an incomplete fragment independently with a replacement-error policy can display � prematurely. Incremental decoding instead buffers incomplete bytes. On EOS, the driver checks stopping conditions without displaying the marker or making another model call with EOS as input.

max_new_tokens counts newly selected IDs, excluding the prompt. This driver counts EOS, so it stops after five new tokens.

KV cache avoids recomputation, not generation dependencies

The animations recompute the current context each round and use no KV cache. A typical cached implementation processes the prompt, stores each layer’s historical K/V, then computes new positions while reading the cache. This reduces repeated work; the next generation round still depends on the selected ID.

Context never exceeds 16 here. Longer sequences need explicit window and position policies. EOS stopping is also an addition in the teaching driver; nanoGPT’s simple generate() loops over the maximum new-token count.

Put four kinds of state in their places

Training turns labels into a loss and carries its influence back to parameters. Inference turns scores into an ID and carries that ID back to the input. Following these routes identifies which computation a number belongs to and what changes next.