T1E3 · Attention Residuals

Transcript

MayaA standard PreNorm transformer lets every earlier layer vote with a fixed weight of one, even when the current token needs only one earlier interpretation. Attention Residuals replaces that pileup with a selector.

LeoThat's a different axis from Kimi Linear. That episode asked how attention should move across tokens efficiently and already pointed us toward scaling. Before we pay that scaling bill, we need one final architecture footnote: how information should move across layers.

MayaAcross tokens, attention chooses which earlier words matter. Across depth, this technical report argues that a later layer should choose which earlier computations matter.

LeoPut the law firm's contracts assistant on that stack. One processing layer spots a definition, another marks an exception, another connects a liability clause, and dozens more keep annotating the same long contract.

MayaThat's the Additive Pile. In a standard residual stream, each new layer output is added to the running state. Unroll the computation and every earlier output sits there with the same fixed coefficient: one.

LeoOne fixed vote.

MayaNo election, no retrieval choice. Residual connections are wonderful gradient highways, but they also quietly define how depth mixes information. That mixing rule is still just addition.

LeoSo the definition found early and the half-useful formatting note from the middle both enter the same pile. Later layers inherit the blend, not the individual annotations.

MayaAnd PreNorm sharpens the problem. Each layer receives a normalized view, produces an update at a controlled scale, then adds it to a hidden state whose magnitude grows with depth. The new update becomes a smaller fraction of the whole.

LeoIt hasn't vanished. It's diluted.

MayaRight. The paper reports that deeper baseline layers must produce larger outputs to remain influential. The oldest useful signal may still be present, but a later layer can't reach in and retrieve it cleanly.

LeoFor contract review, that's nasty. The final risk assessment may need the early interpretation of “affiliate,” not an equal blend of every later annotation that happened to touch the clause.

MayaNow, the Selector. Full Attention Residuals, or Full AttnRes, keeps the embedding and each preceding layer output available as separate sources. Before a layer computes, it forms a softmax-weighted mixture of those sources.

LeoWhat asks the question? A normal attention query comes from the current token state.

MayaHere, each layer owns one learned pseudo-query vector. It's the layer's standing request: what kinds of earlier representations are useful to me? The possible answers are the earlier outputs themselves.

LeoBut the paper calls the choice input-dependent even though that query is fixed after training.

MayaBecause the keys come from the current input's actual layer outputs. The same learned query meets different keys for different tokens and contracts, so the weights change with content. The query is layer-specific; the selection is content-dependent.

LeoAnd large outputs could bully the selector just by being large.

MayaWhich is why the keys are RMS-normalized before scoring. Magnitude doesn't get to impersonate relevance. Then softmax makes the sources compete for a fixed pool of probability.

LeoCompetition matters. If one earlier interpretation gains weight, something else must give it up.

MayaThat's the conceptual break from the Additive Pile. The later layer isn't asking, “What has accumulated?” It's asking, “Which earlier representation should shape this token now?”

LeoDo the ablations support that story, or is “selection” just a pleasing diagram?

MayaUnder the authors' controlled runs, learned but input-independent mixing was worse than Full AttnRes. Swapping softmax for independent sigmoid gates was also worse, and removing RMS normalization hurt. The pieces that enable content-sensitive competition each mattered.

LeoWhat did the learned weights actually look like?

MayaMostly local. Layers still leaned hardest on their immediate predecessors. But some reached back to early sources, and the original embedding retained weight deep into the network. It learned selective skip paths, not random teleportation.

LeoBack at the firm, most reviewers use the latest annotation. Yet the conflicts reviewer can jump back to the first definition analysis when the current clause reuses that exact term.

MayaAnd keep the axes straight: that's attention across depth for one token's representation. It doesn't replace the across-token attention that connects page three to page eighty. It changes what each layer receives after those token relationships have been processed.

LeoFull AttnRes sounds affordable while the model is small. At scale, keeping every layer output alive becomes—

MayaThe systems invoice, and the training setup determines its size.

LeoOrdinary training already retains those outputs for backpropagation, so the added storage can be negligible.

MayaLarge runs use activation recomputation and pipeline parallelism, though. Outputs that were once discarded now have to be preserved and moved between devices.

LeoEnter the Block Shelf.

MayaInstead of shelving every layer output separately, Block AttnRes groups nearby layers. Inside a block, ordinary addition builds one summary. Across blocks, the selector attends to the completed summaries, the embedding, and the current block's partial sum.

LeoSo it gives up some fine-grained retrieval inside old blocks.

MayaDeliberately. Full AttnRes offers the sharpest layer-by-layer choice. Block AttnRes compresses that history so memory and communication scale with the number of blocks, not the number of layers. Around eight blocks recovered most of the gain in the scaling experiments.

LeoThe law-firm version is a shelf of review packets: one summary from definitions, one from obligations, one from conflicts. The final reviewer chooses among packets instead of reopening every pen stroke.

MayaIn the final Kimi Linear configuration, each block summary covers six computational layers. The model carries nine completed block summaries plus the token embedding as depth-wise sources.

LeoSmaller shelf, yes. But every pipeline stage still needs the relevant shelf, which means the elegant architecture can turn into—

MayaA distributed-systems headache at every handoff.

LeoBecause the summaries still have to reach the stage that needs them.

MayaNow, the Systems Handoff. The paper's cache-based response is pipeline communication. A device keeps block summaries it already received and later transfers only the newly completed summaries.

LeoInstead of resending the whole shelf every time the work crosses a stage boundary.

MayaExactly. The companion response is a two-phase schedule. Because every pseudo-query is learned in advance and doesn't wait on the current layer's output, all the queries in one block can look at earlier block summaries together.

LeoThat's the parallel phase: read the shared shelf once for the whole group.

MayaThen comes the sequential phase. Layers walk through the evolving partial sum inside their own block, and an online-softmax merge combines that local result with the earlier-block result without changing the answer.

LeoSo batching removes repeated shelf reads, while the sequential pass preserves dependencies inside the block.

MayaUnder the authors' setup, pipeline training added less than four percent end-to-end overhead. Typical inference workloads added less than two percent latency. Those are reported engineering results, not universal constants.

LeoAnd long-context prefilling still has to store block representations for every token.

MayaThe implementation shards them across devices along the sequence dimension and works in chunks. In the paper's one-hundred-twenty-eight-thousand-token example, that cuts the per-device burden dramatically. Again, it depends on their parallel layout and kernels.

LeoHere's my objection. The architectural idea is clean; the practical claim depends on caching rules, fused merges, sharding, chunking, and a pipeline schedule. “Drop-in” is doing heroic work there.

MayaAnd my objection to that objection is that residual streams aren't a tiny detail. If fixed addition wastes effective depth across every layer, accepting some systems work can be rational—especially when the paper shows a consistent loss advantage across model sizes.

LeoConsistent in the authors' experiments. I won't call that settled architecture law.

MayaNor should we. But the scaling-law result is substantial enough to notice: at the comparison point they report, Block AttnRes reaches the baseline's loss with about one unit of compute where the baseline needs roughly one point two five.

LeoSo “one point two five times compute” means the baseline needs about twenty-five percent more training compute to match that fitted loss. It doesn't mean every AttnRes model is twenty-five percent cheaper.

MayaCorrect. It's a relationship between fitted curves over five tested model sizes, not a frontier-wide guarantee. Full AttnRes was slightly better, while the block version stayed close enough to buy back practicality.

LeoAnd the large run people will quote?

MayaThe authors integrated Block AttnRes into Kimi Linear at forty-eight billion total parameters, three billion active per token, and trained on one point four trillion tokens. Against their matched baseline, it matched or improved every reported downstream task.

LeoThose results came from the same team, same architecture family, and their evaluation suite. We don't yet have broad independent replication, a frontier-scale comparison, or evidence that every model stack tolerates the systems complexity equally well.

MayaI concede all three. The honest claim is paper-reported: selective depth mixing improved this training program, at these tested scales, with low overhead under this implementation. That's promising evidence, not a universal verdict.

LeoThen what makes it more than a benchmark bump?

MayaThe mechanism and the diagnostics line up. The baseline's hidden magnitude climbed with depth and its earliest layers took disproportionately large gradients. Block AttnRes produced a bounded, repeating magnitude pattern and spread gradients more uniformly.

LeoMeaning the selector appears to address the dilution it was designed to address, not merely stumble into a better score.

MayaYes—with a useful caveat. A fully input-dependent query performed even better in one ablation, but required an expensive projection and sequential memory access. The chosen pseudo-query is itself a systems compromise.

LeoWhich is why I land here: the architectural gain looks real enough to test, while the evidence and integration burden still make adoption an engineering decision, not a slogan.

MayaAnd that closes Topic One with the final architecture footnote we paused for. We began by learning attention across tokens, then made it efficient, and now aimed it across layers. The scaling bill is waiting: compute budgets, data, and the cost of training larger models well.

LeoIn your contract-review model, if a late reasoning layer could retrieve one earlier interpretation instead of inheriting the whole additive pile, which interpretation should it select, and what evidence would convince you the extra systems machinery had earned its place?