![[Training a 200B LLM from Scratch] 6/25: How Is Attention Trained?](/_next/image?url=%2Fblog%2Ftraining-200b-part-6-cover.png&w=3840&q=75&dpl=dpl_6PvZJYM25xw2xGK8pSA7Dt43f6ZF)
[Training a 200B LLM from Scratch] 6/25: How Is Attention Trained?
There is a sentence in our training material: “It was raining outside, so I picked up my umbrella.”
Temporarily cover “umbrella.” The model sees “It was raining outside, so I picked up my…” and assigns different probabilities to what comes next. We then compare its prediction with the answer in the original text. That gives us one training exercise.
But the answer is just “umbrella.” How does it teach Attention to look back at “raining”? We will follow this one exercise to see which numbers inside the model change, and how the program decides which direction to change them.
First, Follow One Training Cycle
The text is first split into small pieces called tokens, then converted into numbers. Those numbers pass through layers of computation and eventually become a prediction of the next token. For this explanation, “umbrella” stands for the answer we want to predict. With an actual tokenizer, one word may correspond to several tokens.
The training program has the original text, so it knows the answer. It checks how much probability the model assigned to the correct answer, then calculates a number that measures the prediction error. This number is called the loss. A smaller loss means that this prediction is closer to the training target.
Next, the program works backward through the calculation to find which numbers can be adjusted, and in which direction, to help reduce the loss. After the adjustment, the model takes on the next batch of exercises. Training repeats these four steps: predict, compare, propagate backward, and update.
Blue A in the figure is Attention. It retrieves information from the visible context. Gold B processes that information further. Part Four introduced them; if you are starting here, these two jobs are enough to remember. The backward arrow passes through both A and B because both parts learn.
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.
What Is Attention Learning?
“Picked up” on its own tells us about an action. Put it back into “It was raining outside, so I picked up my…” and the action has a context. Attention learns how to bring useful context to the current position.
Inside the model, each position has a sequence of numbers that represents its current state. Attention retrieves information, then B processes it, updating those numbers. A later layer's representation of “picked up” can therefore include the rainy background. Moving right through the sentence makes more preceding text available. Moving down through the network keeps changing how each position represents that information.
Think of a stack of information cards. Each card has clues that help you find it, and content that you can retrieve. The current position comes with something it wants to look for. These three things are called 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 each K to obtain different match strengths. These become proportions used to retrieve and mix the corresponding V values. The position containing “raining,” for example, may provide context for the action. Training learns these matching methods. Actual Q, K, and V are numbers; there is no manually written “weather” field inside them.
The calculations that produce Q, K, and V each have adjustable numbers called parameters. They determine how the model expresses a search need, supplies clues, and prepares content. These are the parameters that Attention learns during training. When the next sentence arrives, the same parameters calculate new Q, K, V, and mixing proportions. The Transformer paper describes this structure.
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?
We do not separately tell the model to “give rain 60% of your attention.” The training program has the next token from the original text and the calculation that just produced the prediction.
Imagine one Attention parameter as an adjustable dial. Ask: if we turn this number up slightly, will the final loss increase or decrease? If it decreases, we have found a direction worth trying. We can ask the same question about every other parameter.
Actual training uses the relationships between operations to calculate this influence directly. A parameter change might first change Q. Changing Q changes the mixing proportions. Those proportions change the retrieved information, which finally affects the answer's probability. Following this chain backward from the final loss tells us how to adjust the parameter.
This process is called backpropagation. The calculated effect of a small parameter change on the loss is called a gradient. It lets the program handle many parameters together, without turning every dial separately and running the entire model again. PyTorch's explanation shows how this automatic differentiation works.
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.
Watch One Number Actually Change
To see this dial clearly, we will shrink the calculation to two input positions, one Attention unit, and two candidate answers. We call the correct candidate “umbrella.” This small model lets us calculate one complete parameter update. Its numbers illustrate the learning process.
On the first calculation, the correct candidate has a probability of 56.09%. We want that probability to increase. For one Attention parameter that produces Q, we choose an initial value of 0.5. Backpropagation finds that increasing this parameter slightly would help reduce the current loss.
Its gradient is approximately −0.20637. For this example, we choose a learning rate of 0.1 to control the size of each adjustment. Taking 0.1 times the gradient in the direction of lower loss increases the parameter by about 0.02064, from 0.5 to 0.52064. The parameters that produce K and V, and Attention's output parameter, are adjusted together. Recalculate the same input with all the updated parameters, and the correct candidate's probability becomes 57.10%.
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.
This is the change that one learning step leaves behind: the saved parameters have changed, so the next calculation retrieves and processes information differently. Q, K, V, and attention proportions are recalculated with the new parameters. The complete starting values, update calculation, and runnable program are in “Expand the Small Calculation” below.
The Whole Model Learns Together This Way
The full model has many layers. All of them influence the final prediction, so backpropagation follows their calculations backward to find the gradients of their parameters. Every Attention layer, every B layer, and the parts that turn text into numbers and final states into predictions participate in training.
We have focused on just one answer, “umbrella.” Actual training can turn a whole sentence into many exercises: predict “raining” after “It was,” or “umbrella” after “It was raining outside, so I picked up my.” Each position can read itself and preceding positions. The answer is used to score the prediction afterward.
Positions within the same layer use states already produced by the previous layer, so many positions can be calculated in parallel. The layers themselves proceed in order. Prediction losses across a batch are combined to determine the parameter update. This lets the training program use GPU parallelism while preserving each exercise's reading boundary.
After many sentences, the parameters gradually learn which clues help prediction, where to retrieve information, and how to process it. We supply the answers in the text; the model learns the whole calculation needed to predict them. Each adjustment needs an appropriate size, and we use another set of material to check whether a particular relationship has been learned well.
For our 200B model, the main training process for Attention is now clear: use parameters to retrieve information and make predictions, then use prediction loss to update those parameters. The next choice is how much material to give it at once, and how far back it can connect. The expandable sections below retain the explanations of parallel training, memory, and long-context design to support that choice.
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
Start with four illustrative tokens. The input positions are t1, t2, t3, and t4. The training text also contains the following token, 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 training program knows the complete sentence, but sets a reading boundary for every position. Inside blue A, Q at t2 can compare only with K at t1 and t2 and retrieve their V values. Although t3 and t4 are in the same batch of inputs, they remain invisible to t2. This boundary is called a causal mask. The answer sits beyond the boundary and is used to calculate the error after the prediction.
Within one layer, every position receives states already produced by the previous layer. These positions can therefore perform their Q/K/V calculations and attention together. Position t4 uses t1 through t4 at this layer's input. The states then enter the next layer and continue updating. Layers proceed in order; within a layer, many positions practice in parallel. 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
During this forward calculation, each layer produces its own Q, K, and V. They participate in attention and, together with other intermediate results, form later states. Backpropagation needs some of this information to calculate how to adjust the parameters.
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.
This is an inference-storage calculation. In a teaching configuration with 40 layers, 8 K/V groups per layer, 128 numbers in each group's K and V, and 2 bytes per number, one position needs:
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 200B 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 200B Model to Use Longer Context: Compute Cost, 2026 Architectures, and an Experiment Plan
How Much More Work Does Doubling the Training Length Require?
First, consider only full causal attention inside blue A. The first position can read one position, the second can read two, and the third can read three. Continue to position 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.
We also need to distinguish one sequence from the whole project. If the training total stays at 4T tokens, doubling sequence length roughly halves the number of sequences. In the small example, two eight-position sequences contain sixteen tokens and 36 plus 36, or 72, pairs. One sixteen-position sequence contains the same sixteen tokens but 136 pairs. The first arrangement has no attention links between the two sequences. The second adds links across that former boundary.
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?
Full attention gives each position the most direct connections, with a cost that grows with length. We can adjust blue A to preserve important relationships while reducing what each step must process.
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 200B Model?
Start with three experimental lengths: 8K, 32K, and 128K tokens. Use the same tokenizer and data with comparable sources and quality. Give each length the same cumulative training-token budget, reducing the number of sequences as length increases. Then adjust sequences per batch and gradient accumulation so that every parameter update processes the same total number of tokens. This lets us compare what a longer visible range brings for the same training-material budget.
These three lengths are experimental options. The eventual choice also depends on what the model should handle: everyday conversation, complete meeting transcripts, or technical documents spanning many pages. Training sequences need the corresponding long-distance relationships, and evaluation material needs sufficiently long dependencies, so the model can learn and demonstrate that capability.
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.
Returning to this article's main thread: Attention uses parameters to produce Q, K, and V, retrieve information, and contribute to predictions. Prediction loss travels backward through the calculation, helping the program update those parameters. Repeated practice gradually teaches the model how to find and use preceding text. We can then put candidate structures and material lengths into training experiments, letting measured quality and cost determine the 200B model's specific configuration.
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.