![[Training a 500B LLM from Scratch] 6/25: How Is Attention Trained?](/_next/image?url=%2Fblog%2Ftraining-500b-part-6-cover-refresh-v2.webp&w=3840&q=75&dpl=dpl_HKDZtjNTn1U2zk8pGrYRnYGkyGR9)
[Training a 500B LLM from Scratch] 6/25: How Is Attention Trained?
The model’s route now reached the output. Attention retrieved context, later networks processed it, and the end scored the next token. But connecting the arrows correctly did not mean the model knew how to use them. I placed a training sentence beside the diagram: “It was raining outside, so I picked up my umbrella.”
Cover “umbrella,” and the input becomes “It was raining outside, so I picked up my…” The model guesses, and the original supplies the answer. The starting point is simple: an exercise and its correct continuation.
The difficulty appears after the answer arrives. It says only “umbrella.” It does not tell Attention to look at “raining,” or identify which layer and number to change. I decided to follow this exercise all the way backward from the output error to where learning actually happens.
After the Guess, How Does the Error Enter Training?
First, let the text make a complete pass. The sentence becomes tokens, IDs retrieve numerical representations, and layers of A and B produce a next-token prediction. We keep calling the answer “umbrella,” although a tokenizer may split a word into several tokens.
The training program holds the original text. It checks the probability assigned to the correct answer and expresses the prediction gap as a number called loss. A higher probability for the correct answer means a smaller loss for that prediction.
That number establishes the objective: adjust the network so the next prediction has lower loss. The program works backward through the calculation to find adjustment directions, updates parameters, and supplies the next batch. Prediction, comparison, backward calculation, and update form a loop.
Blue A retrieves information and gold B processes it, the same two networks as in Part Four. The backward arrow passes through both, so both receive learning signals. Keeping the whole loop in view makes it easier to place Attention’s learning when we open A.
Token: A small piece of text processed by the model. It can be a character, a word, or part of a word.
Loss: A number that measures the difference between a prediction and the training answer. Training adjusts the model to try to reduce it.
To Make Retrieval More Accurate, What Changes?
I circled the final position and the earlier rain information. “Picked up” alone describes an action; in this sentence, the action has a weather context. A’s job is to bring information useful for prediction to the current position.
Every position carries numbers representing its state. A retrieves information and B processes it, changing those numbers. In later layers, “picked up” can already carry the rain context. Moving right through the sentence makes more preceding text available; moving down through layers changes its representation. Later words can therefore carry richer context. The original text is only their starting point in the network.
“Retrieve useful information” was still too broad. The program needs something to search with, a way to compare matches, and content to bring back. I split retrieval into three parts, like a stack of information cards: the current position has a need, and each card offers clues and retrievable content. These are Q, K, and V:
Q: What am I looking for now?
K: What clues do I have that can match your need?
V: What information can I provide after a match?
Q is compared with K at each visible position to obtain match strengths. Those become proportions for retrieving and mixing the corresponding V values. The rain position can then provide context for the action at the end. The card analogy separates the roles. Actual Q, K, and V are numbers, and matching develops through training; nobody fills in a “weather” field by hand.
The question was now narrower: how can the next lookup improve? Separate adjustable parameters generate Q, K, and V, determining how to express a need, supply clues, and organize retrievable content. Training changes those parameters. The next sentence uses them to calculate new Q, K, V, and mixing proportions. The Transformer paper describes these connections.
Parameter: A number that training adjusts and saves for future calculations. Parameters are also commonly called weights.
With Only “Umbrella” as the Answer, How Does the Program Know What to Change?
A finally contains things we can change. Yet the answer is still only “umbrella”; nobody separately labels “rain” with a 60% attention share. The training program has another clue: the operations that produced the prediction.
Follow one parameter used to generate Q. If that number increases slightly, how does Q change, and what happens to the final loss? If loss decreases, there is a useful adjustment direction. Asking the same question of other parameters identifies how each affects this prediction.
The program connects these effects through the operations themselves. A parameter affects Q, Q affects matching proportions, the proportions affect retrieved information, and that information affects answer probability. The result tells us the direction and rate at which loss changes when a parameter changes slightly.
This return path is backpropagation, and the calculated influence is a gradient. It lets the program handle many parameters together, differentiating through the existing computation without rerunning the whole model separately for each parameter. PyTorch’s automatic-differentiation guide shows the process. An exercise with one answer can now guide changes to many numbers in the network.
Backpropagation: Working backward from the final loss through the existing calculation to find each parameter's influence.
Gradient: The direction and rate at which the loss changes when a parameter changes slightly.
Follow the Gradient to a Number That Actually Changes
A direction was clear, but I wanted to see a number move. I narrowed the calculation to two input positions, one Attention unit, and two candidate answers, still calling the correct one “umbrella.” Starting values, gradients, and the updated prediction could then be checked individually.
On the first pass, “umbrella” received 56.09% probability. Following the backward calculation reached a parameter producing Q, initially 0.5. Its gradient was negative: near this exercise’s current state, increasing that parameter slightly would help reduce loss.
The gradient was about −0.20637. With a learning rate of 0.1, the update takes 0.1 times the gradient in the direction of lower loss. This parameter therefore increases by about 0.02064, from 0.5 to 0.52064. Parameters producing K and V and the Attention output parameter update in the same round. After all updates, the same input receives a new prediction: 57.10% for “umbrella.”
Learning rate: A number that controls the size of each parameter adjustment. In this example, the adjustment uses the gradient multiplied by the learning rate, moving in the direction of lower loss.
The change now lives in saved numbers. The next input uses the updated parameters to generate Q, K, and V and calculate attention proportions again. What the model retains is a changed calculation method; the proportions from a particular lookup are produced afresh for each input. The expanded section below keeps the complete initial values, gradients, and program for reproduction.
One Position Updated. What About the Other Layers and Positions?
Once one parameter can update along this route, expanding the drawing to the whole model reveals the same return path in every layer. All layers affect the final prediction, so backpropagation follows their computations backward. A, B, the input representations, and the final prediction output all learn.
Remove the cover from “umbrella,” and the sentence contains many exercises. After “It was raining,” predict “outside”; after “It was raining outside, so I picked up my,” predict “umbrella.” Following only the final position made one update easier to trace. Training can use answers at many positions together.
That raises a practical question. With the whole sentence available, can earlier positions peek at later answers? A gives each position a reading boundary: itself and earlier text only. Positions in one layer read states already formed by the previous layer, allowing parallel calculation within a layer. Layers then proceed in order. Losses from the predictions combine to determine the parameter update.
Repeatedly, text supplies answers and prediction updates the parameters, gradually developing methods for retrieving and processing context. Increasing one sentence’s probability completes only a small step. Other material tests whether the relationship was learned reliably. Update size matters too; we will return to it when choosing an optimizer.
Our 500B model learns along this route as well. We can now trace a prediction error to Attention parameters. The next boundary to choose is clearer: how much preceding text should each exercise provide, and how far may the model connect it? Short material cannot provide distant information; longer input costs computation and memory. The engineering details below connect training material and model structure in the same plan.
Expand the Small Calculation: Starting Values, Gradients, and Recalculation
W_Q, W_K, and W_V are the weights that produce Q, K, and V. W_O is the output transformation weight after mixing V. This example uses −ln(probability of the correct answer) as its loss. A higher correct-answer probability gives a lower loss.
The two input states are x1 = [1, 0] and x2 = [0, 1]. We show only the next-token prediction loss at the second position. The input states and final two-candidate readout rule are fixed. Attention's seven parameters participate in this update:
| Parameter | Before update | Gradient | After update |
|---|---|---|---|
| First number in W_Q | 0.20000 | 0.00000 | 0.20000 |
| Second number in W_Q | 0.50000 | −0.20637 | 0.52064 |
| First number in W_K | 1.00000 | −0.10318 | 1.01032 |
| Second number in W_K | 0.00000 | 0.10318 | −0.01032 |
| First number in W_V | 1.00000 | −0.27331 | 1.02733 |
| Second number in W_V | −1.00000 | −0.16577 | −0.98342 |
| W_O | 1.00000 | −0.10754 | 1.01075 |
Initially, Q at the current position is 0 × 0.2 + 1 × 0.5 = 0.5. The two K values are 1 and 0, and the V values are 1 and −1. Multiplying Q by K gives match scores of 0.5 and 0. Softmax converts these scores into proportions: exp(0.5) ÷ [exp(0.5) + exp(0)] = 0.622459 for the first position, and 0.377541 for the second.
The mixed V result is 0.622459 × 1 + 0.377541 × (−1) = 0.244919. Multiply by W_O, and the output remains 0.244919. Use this number and a fixed 0 as the two candidate scores, then apply softmax again. The probability of “umbrella” is exp(0.244919) ÷ [exp(0.244919) + exp(0)] = 0.560925, and the loss is 0.578167.
We can connect the local rates of change from the second number in W_Q to the loss. From that parameter to Q, the rate is 1. From Q to the first position's attention proportion, it is 0.622459 × 0.377541. From the proportion to the mixed result, it is the difference between the V values, 1 − (−1) = 2. From the mixed result to the output score, it is W_O, or 1. From the output score to the final loss, it is 0.560925 − 1. Multiply these rates along the path:
Gradient = 1 × (0.622459 × 0.377541) × 2 × 1 × (0.560925 − 1) ≈ −0.20637
That is where the gradient in the main text comes from. The first number in W_Q has a zero gradient because the first component of the current input x2 is zero, and this example uses only the loss at the second position.
After all seven parameters are updated together at a learning rate of 0.1, the new Q is 0.520637. The two attention proportions are about 62.98% and 37.02%. The new mixed result is 0.282960. Multiplying by the updated W_O gives 0.286003, a correct-candidate probability of 0.571017, and a loss of 0.560336. The attention proportions, Q, and V are recalculated. What training saves is the updated parameters.
Download and run this small Python example to reproduce the numbers. The program also compares each analytical gradient with a numerical gradient obtained through tiny perturbations. The largest difference is approximately 1.11 × 10⁻¹⁰. Real models use wider vectors, many layers, and fuller training objectives. The principle of finding gradients along the calculation path is the same.
Softmax: A function that converts scores into nonnegative proportions summing to 1. Attention uses it to assign mixing proportions; the output uses it to obtain candidate-token probabilities.
How Can a Whole Sentence Train in Parallel? How Do Training Intermediates Differ from a KV Cache?
How Positions Divide the Work
The exercises within a sentence can run together. Take input positions t1, t2, t3, and t4, followed by t5. Each position predicts the token immediately after it:
| Current input position | Inputs this position can use | Answer to predict |
|---|---|---|
| t1 | t1 | t2 |
| t2 | t1, t2 | t3 |
| t3 | t1, t2, t3 | t4 |
| t4 | t1, t2, t3, t4 | t5 |
Here, t1 through t5 represent positions after tokenization. A character or a word can correspond to one or more tokens. Returning to our sentence, the final position in “picked up my” can use the earlier rain to learn to predict “umbrella.”
The table’s reading boundaries are implemented inside A as a causal mask. Q at t2 compares only with K at t1 and t2 and retrieves their V values. Although t3 and t4 share the input batch, they remain invisible to t2. After prediction, t3 supplies the correct answer for scoring t2, separating inputs from answers.
Now look at t4. It reads t1 through t4, using four states already available at this layer’s entrance. It does not wait for t1 to finish this layer and then for t2. The layer can generate Q, K, and V for positions together, then apply each position’s reading range. All resulting states enter the next layer. The Transformer paper implements this with masking and prediction targets shifted by one position.
This method is usually called teacher forcing: training uses the real preceding text from the material. It lets us complete many predictions from one sentence in the same calculation. Hugging Face's from-scratch training example organizes known text into these inputs and labels.
Causal mask: A restriction on what Attention can read, allowing a position to use only itself and earlier positions.
Teacher forcing: Providing the real preceding text during training so that the model learns to predict what follows.
What Backpropagation Needs to Keep
After prediction and before backward computation, GPU memory already holds intermediate results. Every layer generated its own Q, K, and V during this forward pass, using them to form later states. Some of that information must remain available to compute gradients along the same route.
These intermediate training results are called activations. Keeping them in GPU memory makes them immediately available for backpropagation. Alternatively, we can keep only some results and repeat the relevant calculations when backpropagation needs them. This is activation recomputation: more computation in exchange for less memory. PyTorch's explanation shows both arrangements.
A KV cache for token-by-token generation serves a different task. The model has already processed preceding text and reuses its old K and V as generation continues. Our ordinary pretraining processes a whole segment of known material at once and usually does not use this continually growing generation cache. Hugging Face's cache documentation also places it in inference. After training updates the parameters, K and V calculated with old parameters cannot simply be reused as results of the new ones.
Training-memory planning therefore needs to account for parameters, gradients, optimizer states, activations, and temporary working space. The KV cache needed for one generation request belongs to a separate budget for serving the model. Generating synthetic training data, or having a model produce responses during reinforcement learning, may also use a KV cache. That step is performing generation.
Activation: An intermediate result produced while the model processes its current input. Some activations are needed for backpropagation.
Activation recomputation: Saving fewer intermediate results and calculating them again when backpropagation needs them, to reduce memory use.
KV cache: Stored K and V from earlier positions, reused by later steps during incremental generation.
How the KV Cache Works During Generation
When a deployed model generates an answer, it first processes the full prompt. This stage is called prefill. It then produces one new token at a time, processes that position, and predicts the following token. This stage is called decode. With the same input prefix and model parameters, causal attention leaves old positions' K and V unchanged. They can therefore be saved and read directly by later steps.
The cache avoids repeatedly producing K and V for old positions. A new position still calculates its own Q, K, and V, compares its new Q with visible K, and retrieves the corresponding V. Reading this information still has a cost that grows with the context.
For an inference-storage calculation, take 40 layers, eight K/V groups per layer, 128 numbers each for K and V, and two bytes per number. One position requires:
40 × 8 × 128 × 2 copies (K and V) × 2 bytes = 163,840 bytes = 160 KiB
Keeping 8,192 positions for one request uses 1.25 GiB; 32,768 positions use 5 GiB; 131,072 positions use 20 GiB. These figures count only that request's KV. They exclude model weights and other working data, and do not represent training memory. Once the 500B model's actual configuration is chosen, its layer count and K/V structure will determine the serving budget.
Prefill: The stage in which a generation service processes the complete supplied input to prepare states for subsequent output.
Decode: The stage that processes a newly generated token and predicts the next one.
KiB and GiB: Storage units. One KiB is 1,024 bytes; one GiB is 1,024 cubed bytes.
Teaching the 500B Model to Use Longer Context: Compute Cost, 2026 Architectures, and an Experiment Plan
How Much More Work Does Doubling the Training Length Require?
To read longer material, first count the extra work. With full causal attention in blue A, the first position can read one position, the second two, the third three, up to N:
1 + 2 + … + N = N × (N + 1) ÷ 2
Eight positions have 36 allowed pairs; sixteen positions have 136. As N grows, doubling one sequence's length brings the pair count close to four times its original size. Each pair requires a Q/K comparison and participates in retrieving information. Training also backpropagates through these operations.
But the project sheet fixes total training at 10T tokens, so one sequence is only part of the account. Doubling length roughly halves the sequence count. Two eight-position sequences contain sixteen tokens and 36 plus 36, or 72, relationships. One sixteen-position sequence still contains sixteen tokens but has 136 relationships. Text separated by a sequence boundary can now connect.
With a fixed total token count and full attention within every sequence, doubling length therefore roughly doubles the total number of these attention pairs. Other model calculations, GPU parallel efficiency, and data movement also affect elapsed time, so we need to measure the complete training workload. Part Five's distinction between doing less computation and finishing faster still applies.
Memory also needs measurement. A longer sequence produces more intermediate states. The implementation determines which results must remain in GPU memory together. FlashAttention, for example, calculates attention in blocks, avoiding storage of the entire attention-score matrix in GPU memory while preserving exact attention. The triangle in the figure represents allowed relationships and helps count the work. Actual execution need not store it as one enormous table.
Full causal attention: Each position can read itself and every earlier position.
Sequence length: The number of tokens in one training sample. With a fixed total token count, changing the length also changes the sample count.
FlashAttention: A family of implementations that reduces attention's memory overhead through blockwise calculation and carefully arranged data movement.
How Can We Design Blue A to Read Farther?
Those distant connections can be valuable, with a cost alongside them. I returned to blue A to see whether its effort could be directed where needed: retain nearby detail and important distant connections.
One approach lets some layers inspect nearby positions closely while other layers handle more distant connections. Meta Muse Glimmer, introduced in Part Four, is a 2026 example. It repeats three local layers followed by one global layer, with a local window of 2,048 positions. In our overall diagram, this changes which positions each layer's blue A can read.
Another route first organizes a stretch of preceding text into fewer numerical representations, then selects relevant parts for attention. Nearby text can retain detailed information, while distant information takes the compression-and-selection route. Training must jointly learn how to organize information, select it, and use what both routes retrieve.
DeepSeek's V4.1-Flash report, published in September 2026, combines local attention with compressed sparse attention. The October 17 Flash-1017 update retains this architecture. The report discusses attention sharing during training and cache management after deployment separately. Here, we take the architectural choice: which preceding text keeps a detailed representation, which is compressed first, and how to select what should participate.
A smaller window reduces the nearby range directly visible to one layer. More compression requires the retained representation to concentrate more information. Selecting fewer positions makes selection accuracy more important. Training these structures therefore requires material with long-distance relationships. We then check whether the model can find dates in long documents, connect scattered evidence, and distinguish earlier statements from later revisions. Length and architecture jointly determine what it has an opportunity to learn.
Sliding window: An attention arrangement that limits a position's direct reading range to a nearby stretch of input.
Sparse attention: Selecting some of the available positions or representations to participate in attention.
Compressed representation: Retaining input information in less data for later calculation.
How Should We Choose This for Our 500B Model?
To give length selection a traceable basis, I put 8K, 32K, and 128K tokens on the experiment sheet. Keep the tokenizer fixed, sources and quality comparable, and cumulative training tokens equal. Longer sequences mean fewer sequences. Adjust sequences per batch and gradient accumulation so each update processes the same total tokens. The comparison now asks what a longer visible range adds for the same amount of material.
The choice returns to what the model needs to read. Everyday conversation, complete meeting records, and technical documents spanning many pages demand different context ranges. Training material must contain those distant relationships, and evaluation must require connecting them, for length to have a useful role.
Begin with full attention, then compare candidates that alternate local and global attention or compress and selectively retrieve context. Record training loss, long-document performance, training tokens per second, and peak GPU memory for every experiment. Throughput includes forward computation, backpropagation, and parameter updates, using the same GPU configuration. Keep settings such as activation recomputation and sequences per batch in the record too. That ties quality, time, and machine requirements to the same plan.
If we choose to increase training length in stages, list each stage's length, token count, and data mix. Calculate its duration from measured throughput, then update Part Two's schedule and Part Three's budget. The final configuration should state both the input window the model supports and the lengths it actually practiced during training.
The originally hidden “umbrella” now connects the full training route. The answer creates loss, loss propagates backward, parameters update, and the next batch retrieves context with those new parameters. What material to supply, and how long to make an exercise, become the next concrete choices.
Gradient accumulation: Processing several small batches and accumulating their gradients before one parameter update, to arrange the total training volume per update.
Training throughput: Tokens processed per unit of time during training, including forward computation, backpropagation, and parameter updates.
Peak GPU memory: The highest amount of GPU memory used during a run.
New ideas, straight to your inbox.
AI, engineering, and experiments. One or two useful emails a month.
No spam. Unsubscribe anytime.