PaperlayerView exact revision on arXiv

Source-linked AI summary

Do New Attention Mechanisms Actually Fix Attention Sinks at Million-Token Context?

Sara Rizwan, Samaanah Abdus Salam

arXiv:2609.08574v1cs.CLcs.AI

TL;DR

The paper asks whether attention-sink fixes and long-context retrieval remain effective at larger context lengths. Using SinkProbe and controlled small-model variants, it finds that the training objective produces the sink, gating does not reproduce its published effect at this scale, and sink mass, activations, and position bias can separate.

  • Problem

    Long-context models can retain tokens in their context window yet retrieve early information vaguely or rely disproportionately on recent pages, motivating tests of attention sinks and position-dependent recall.

  • Method

    The paper applies SinkProbe to controlled model variants, measuring sink mass, massive activations, position-resolved recall, and related behavior across token-mixing and depth mechanisms.

  • Results

    The training objective produces the sink rather than the architecture, gating did not reduce it at this scale, and sink mass, massive activations, and position bias moved independently.

  • Takeaways & Limitations

    Removing attention sinks does not by itself establish that a model can read information throughout a long context, so sink mass and position-dependent retrieval should be evaluated separately.

  • Takeaways & Limitations

    The pilot omits the delta-rule correction from Kimi Delta Attention and reaches 768 tokens, far below the million-token window, limiting claims about the exact Kimi kernel and frontier-scale behavior.

Abstract

from arXiv · show

Long context language models now advertise windows of one million tokens, but two habits limit how much of that window is used. Attention heads with nothing useful to read still spend their budget on the first token, which is called the attention sink, and where a fact sits in the context changes whether the model finds it. Gated attention cut first token attention from 46.7 percent to 4.8 percent at NeurIPS 2025, and Kimi K3 pairs that idea with Kimi Delta Attention and Attention Residuals behind a one million token window, eight times past the range where these diagnostics have been reported. This paper asks whether the fix survives that jump. We build SinkProbe, a suite that measures sink mass, massive activation, position resolved recall and the recency gap, and apply it to four small models that differ only in how they mix tokens and depth. Three results follow. The training objective produces the sink, not the architecture. Gating did not reproduce its published effect at our scale. Sink mass, activations and position bias moved independently. Code, data and the measurement protocol are released at https://github.com/sararizwan7/Attention-Mechanisms-in-1M-Context-Window

I. INTRODUCTION

Long-context models can have million-token windows yet still underuse their context because attention sinks and position-dependent recall separate nominal capacity from usable memory. This paper introduces SinkProbe and a controlled architecture ladder to test whether recent mechanisms fix both habits together.

  • The two habits: Attention sinks force heads with nothing useful to read to place their attention budget somewhere, usually the first token.The first token can receive 46.7 percent of attention in a reported 15-billion-parameter model.
  • The two habits: Position-dependent recall creates a second problem: evidence near the middle is harder to retrieve than evidence near either context end.The recency gap measures the difference between recall near the question and recall at the opening.
  • Research gap: The paper addresses the gap between sink diagnostics reported to 128 thousand tokens and million-token architectures reported without those diagnostics.Kimi K3 combines gated attention, Kimi Delta Attention and Attention Residuals, but its million-token report uses benchmark scores rather than mechanism-level measurements.
  • Research question: The paper separates the mechanisms rather than treating sink removal and even context reading as one outcome.Its stated question is whether gated and linear mechanisms make the sink disappear, make context reading even, and produce those outcomes together.
  • Contributions: SinkProbe collects sink mass, massive activation, position-resolved recall and the recency gap into one runnable diagnostic suite.The suite is paired with a controlled ladder of four architectures that differ in token and depth mixing while other training factors are held fixed.

II. RELATED WORK AND THE GAP IT LEAVES

Prior work explains attention sinks, proposes gating and linear mechanisms, and develops long-context evaluations, but these lines stop at different combinations of scale, architecture and diagnostics. The unresolved gap is direct mechanism-level measurement at the lengths advertised by newer models.

  • Attention sinks: Streaming studies found that retaining a few initial tokens prevents perplexity collapse even though those tokens carry little semantic content.This motivated the structural interpretation of attention sinks.
  • Attention sinks: Attention sinks have been linked to massive activations and to optimization and data, while competing interpretations predict the same measurable signature.The paper therefore emphasizes measurement over choosing between the null-position and staging-area interpretations.
  • Gated attention: Gated attention reduced first-token attention from 46.7 percent to 4.8 percent, while a shared gate reduced activations but left sink mass at 30.1 percent.The comparison supports treating sink mass and activation magnitude as separable diagnostics.
  • The remaining gap: Prior evidence covers dense softmax models up to 128 thousand tokens, whereas newer million-token architectures report benchmark scores without the corresponding diagnostics.This mismatch is the gap the paper joins with SinkProbe.
  • Linear and residual mechanisms: Linear attention replaces a growing key-value cache with a fixed-size recurrent state, while delta rules and forget gates address state editing and decay.Attention Residuals separately replace uniform depth accumulation with learned weighting over earlier representations.

E. Measuring Long Context Behaviour

The paper turns long-context behavior into separate measurements for sink concentration, hidden-state magnitude, attention concentration and position-dependent retrieval. These diagnostics distinguish where attention goes from how evenly information is recovered.

  • Position-resolved recall: Position-resolved recall measures retrieval accuracy as evidence depth varies, complementing scalar sink and activation metrics.Together, the diagnostics separate attention concentration from the location-dependent ability to recover evidence.
  • Sink mass: Sink mass measures how much of a softmax head’s mandatory attention budget lands on the first position.An even-reading stack scores near 1/T; at T = 160, a value near 0.5 represents roughly eighty times the even share.
  • Sink mass: The analysis also reports the worst layer because stack averages can hide concentration in a single layer.The first query position is skipped because it can attend only to itself.
  • Massive activation: Massive activation tracks the largest absolute hidden-state value entering each layer and is averaged over the stack.It is reported separately because activation magnitude can change without sink mass changing.
  • Attention concentration: Entropy measures whether attention is concentrated without assuming that concentration occurs at the first position.This catches sinks that move away from the initial token.

E. Position Resolved Recall and the Recency Gap

The paper defines the recency gap as a signed comparison between recall near the question and recall near the context opening, then evaluates it in a controlled architecture ladder. This separates position bias from sink and activation behavior.

  • Recency gap: The recency gap compares recall near the question with recall when evidence appears at the context opening.Evidence depth d ranges from 0 at the opening to 1 immediately before the question.
  • Recency gap: A positive ∆r indicates preference for recent evidence, while a value near zero indicates more even context reading.The estimate uses a Wilson score interval because it is a difference of two binomial proportions.
  • Why separate measurements: Long-context behavior reflects multiple pressures: quadratic softmax cost, sequence dilution, and depth dilution in uniform residual streams.These pressures motivate measuring computational cost, token-level attention behavior and depth mixing separately.
  • Research question: The study tests whether sink mass and recency gap move toward zero together, treating them as two claims requiring two measurements.SinkProbe produces σ(M,T), µ(M,T), p(d;M,T) and ∆r(M,T) for context lengths reaching the tested window.
  • Controlled architecture ladder: Four models differ only in token and depth mixing, with width, depth, heads, feed-forward size, vocabulary, data, optimizer, schedule and seed policy held fixed.The ladder includes softmax, softmax with gate, a hybrid linear-global stack and the same hybrid with attention residuals.
  • Controlled architecture ladder: Each neighboring step adds exactly one mechanism, making the ladder’s differences architectural ablations by construction.This design avoids attributing differences to unrelated changes in training or model scale.

C. The Linear Layer We Use

The section defines a decaying linear layer, adds block-level attention residuals, and specifies the controlled retrieval probe and objective control used to test it.

  • Linear layer: The linear layer maintains a fixed-size state with channel-wise input-dependent decay and a full-rank output gate.The decay is bounded below for stable chunked computation in single precision.
  • Linear layer: Decay initialization spans approximately 0.78 to 0.999 so channels jointly preserve order and long-range memory.Fast decay tracks recent events, while slow decay carries content across long spans.
  • Linear layer: The pilot omits Kimi Delta Attention’s delta-rule correction, so it measures a decaying fixed-size state rather than the exact Kimi kernel.The paper reports that this simplification has a larger-than-expected cost.
  • Attention residuals: Block Attention Residuals split eight layers into two four-layer blocks, reduce each block to one representation, and retain the token embedding as an always-reachable source.Each layer attends over available sources using a learned pseudo-query and RMSNorm.
  • Retrieval probe: Each retrieval sequence plants one key-value pair at a controlled depth among distractors, then scores whether the model returns the matching value.The design uses distinct distractor keys and reports blind-chance and in-context-value floors.
  • Objective control: The objective control removes the auxiliary next-token term while leaving the answer-position task, testing whether position-wide output pressure creates the sink.The same control is evaluated within the SinkProbe pipeline, whose thresholds are fixed before large runs.

F. Statistical Treatment

The evaluation fixes statistical decisions and model settings in advance, uses controlled synthetic retrieval sequences, and reports uncertainty across seeds and within seeds.

  • Uncertainty: Wilson score intervals replace plain standard errors for every binomial accuracy.Trial counts are chosen for a target half width at the worst-case proportion.
  • Uncertainty: Every result uses three seeds, with cross-seed spread reported separately from the within-seed interval.These quantities represent different sources of variation.
  • Pre-registration: Pass-fail thresholds for released weights are fixed before full-scale runs, and intermediate results are reported as inconclusive.The thresholds are based on published values rather than measurements from this study.
  • Data construction: The haystack repeats a random 24-symbol motif while randomizing half its positions, making auxiliary next-token prediction solvable but nontrivial.A calibration found sink mass of 0.095 with 15 percent randomization versus 0.438 with 50 percent, but the setting choice is not treated as a reported result.
  • Controlled models: All four models share the same configuration, differing only in their token- and depth-mixing rules.The small scale is chosen to make controlled retraining practical rather than to estimate frontier-model scores.
  • Evaluation design: Evaluation extends to eight times the training length, while aggregate depths are sampled uniformly and profiles use a ten-point grid.Each length pools 384 sequences across three seeds; depth profiles use 32 sequences per point per seed.
  • Retrieval probe: Retrieval uses 12 keys, 12 values, 3 distractor pairs, and 44 filler symbols, with blind chance at 8.3 percent and the in-context-value floor at 25 percent.Only performance above the second floor counts as retrieval evidence.
  • Diagnostic collection: Sink, activation, and related diagnostics are recorded on the first evaluation batch because averages over heads, positions, and layers stabilize with short sequences.The diagnostics are treated as tightly determined by 16 sequences of a few hundred tokens.

D. Cache Growth for the Kimi K3 Layer Mix

The Kimi K3 layer mix reduces context-dependent cache growth through mostly linear layers, while the pilot tests its mechanisms under small-scale controlled conditions and exposes important scope limits.

  • Cache structure: Of 93 Kimi K3 layers, 69 are linear and 24 are global, so 74.2 percent of the stack contributes nothing to cache growth.Linear-layer state size depends on head geometry rather than context length, whereas global-layer caches grow with context.
  • Cache cost: Under stated assumptions, the hybrid stack needs about 24 GiB of cache at one million tokens versus about 2.5 TiB for a dense stack, roughly a 107-fold difference.Below about 8 thousand tokens, the fixed state is the larger term.
  • Scope: The pilot runs at 768 tokens with approximately one-million-parameter models, so it can compare diagnostic movement but cannot establish effect sizes at frontier scale.The paper explicitly declines to extrapolate these findings to 2.8-trillion-parameter models or one-million-token runs.
  • Evaluation scope: The study reports sink formation, mechanism effects, length extrapolation, position dependence, and total cost as separate evaluation questions.This ordering distinguishes architectural diagnostics from retrieval behavior and resource accounting.
  • Objective control: The objective control compares identical softmax stacks differing only in whether loss covers every sequence position or only the answer position.This isolates the training-objective contribution without changing architecture, data, or seeds.
  • Baseline diagnostics: The baseline’s 31.6 percent sink mass at training length is roughly thirty times an even split and comparable in order to the 46.7 percent average reported for a much larger model.Its worst layer reaches 63.3 percent, while baseline seeds vary from 0.089 to 0.455.
  • Hybrid limitation: The hybrid rows fail retrieval, scoring 6.0 and 5.7 percent against 8.3 percent blind chance and the 25 percent in-context-value floor.The paper attributes this limitation to omitting the delta-rule correction, while retaining their other diagnostic measurements.

C. What Each Mechanism Contributes

Adding mechanisms changes sink mass and retrieval behavior, but the effects are not uniformly reliable. Gating is inconclusive at this scale, while sink mass, extrapolation behavior, and position bias can move independently.

  • Sink mass: 0.316 to 0.261 to 0.183: sink mass decreases as the plain softmax stack, gated stack, and hybrid add mechanisms.The gate’s 0.055 reduction is within seed noise, whereas the hybrid’s larger reduction is confounded by failed task learning.
  • Sink mass: The gate comparison is inconclusive because its reduction is smaller than baseline seed variation.Baseline seeds ranged from 0.089 to 0.455, and one gated run exceeded two baseline runs.
  • Recall past training length: At 768 tokens, softmax stacks remain above the 25 percent retrieval floor, but the baseline result is close enough to its 8.6-point seed spread to warrant caution.The reported recalls are 26.3 and 31.8 percent at 768 tokens.
  • Sink mass: At 96 tokens the baseline places about thirty times an even share on the first position; at 768 tokens this rises to about ninety seven times.The fraction called sink mass falls with length mainly because the attention budget is spread over more positions.
  • Position bias: At 768 tokens, recall favors evidence near the question, with profiles rising from about 21 to 42 percent for the baseline and 22 to 50 percent for the gated stack.The preference reverses from the opening bias observed at the 96-token training length.
  • Position bias: Recall varies by more than twenty points with answer position, and model ordering by sink mass does not match ordering by depth profile.The gated stack has lower sink mass but a steeper position profile at 768 tokens.

F. Massive Activations

Massive activations provide architectural evidence that separates from sink behavior. Gating and attention residuals reduce activation magnitude, while cache design determines whether million-token inference is practically affordable.

  • Gating: 74.9 to 36.2: gating reduces the mean largest hidden value by 52 percent, while sink mass falls only 17 percent within seed spread.The two quantities therefore do not move together reliably.
  • Attention Residuals: 97.9 to 38.9: Block Attention Residuals reduce the mean largest hidden value by 60 percent, while sink mass moves from 0.183 to 0.202.The reduction is partly structural because a softmax-weighted average cannot exceed its largest input.
  • Scope: The models are four orders of magnitude smaller than published models, so their values are not directly comparable across scale.The comparison is intended to show direction rather than magnitude.
  • Metric separation: 1053 to 125: a published gate after value projection cuts the largest activation by a factor of eight while leaving first token attention at 0.297.A single sink headline therefore misses separable activation behavior.
  • Cache growth: 2.5 TiB versus near 24 GiB: ordinary caching and the hybrid cache model differ sharply at a million-token window.The dense estimate exceeds the memory of any single machine, while the hybrid estimate stays near 24 GiB.
  • Training behavior: The hybrid stacks flatten near chance-retrieval loss, whereas the two softmax stacks descend well below it.The hybrid loss curves match the reported failure to learn lookup retrieval.

J. Seed Behaviour and Significance

Seed variation can exceed architectural differences, making single-run sink comparisons unreliable. The controlled protocol therefore emphasizes retraining, matched controls, and predefined interpretation rules.

  • Controls: The models are matched on optimisation, and their training-loss curves remain close across three seeds.This reduces the likelihood that diagnostic differences simply reflect unequal training quality.
  • Seed variation: Sink mass in the baseline varied fivefold across seeds, from 0.089 to 0.455, exceeding every measured architectural difference.This makes seed spread more consequential than the architecture comparison at this scale.
  • Seed variation: The auxiliary-objective control produced sink masses of 0.051, 0.046 and 0.050, a spread of 0.003.The same architecture was stable under this control but unstable under the full language-modelling objective.
  • Interpretation: The pilot reports architectural comparisons as inconclusive because retraining the same architecture can change the sink more than changing the architecture.The protocol consequently specifies controls and sample sizes rather than relying on one run.
  • Objective effects: Training the same stack for every-position prediction produced 31.6 percent first-token attention, versus 4.9 percent when answering only the final question.Recall remained unchanged, and the gap was a factor of six.

C. Cost

The cost analysis supports serving a million-token window with mostly fixed-size recurrent state, but it does not establish whether compressed contents remain usable. The pilot also finds that sink mass, activations, and position bias can diverge.

  • Diagnostic separation: Sink mass and evenness of context reading moved independently: the gated stack had lower sink mass but a steeper depth profile at 768 tokens.The lower sink therefore did not imply more even reading across the context.
  • Diagnostic separation: The gate halved largest hidden values and attention residuals reduced them by sixty percent, while sink mass remained within its seed spread or did not fall.These observations separate massive activations from sink mass.
  • Cache cost: Holding 74.2 percent of layers at fixed state size makes the million-token cache roughly one hundredth the cost of dense caching.Doubling or halving either unpublished assumption changes the total by less than a factor of two.
  • Scope of cost analysis: Fixed-size state makes the window affordable, but the cost model cannot determine whether the compressed tokens remain usable for the question asked.Recall measurements are needed to evaluate usability rather than memory capacity alone.

D. Reading the Pilot Honestly

The pilot’s conclusions are bounded by scale, task, architecture, and measurement scope. In particular, the simplified linear layer prevented hybrid stacks from learning the retrieval probe, so those rows cannot support position-bias conclusions.

  • Architecture boundary: The simplified linear layer omits Kimi Delta Attention’s delta-rule correction, and hybrid stacks consequently failed to learn the retrieval probe.This failure belongs to the stand-in rather than to Kimi K3, whose design includes the omitted term.
  • Scope boundaries: The pilot reaches 768 tokens and uses roughly one-million-parameter models, so it cannot establish diagnostic magnitudes at frontier scale.Its results show controlled mechanism changes or failures without extrapolating their size.
  • Task boundary: The synthetic key-value task isolates retrieval but omits summarisation, multi-hop reasoning, and code navigation.Performance on this probe should not be treated as coverage of all long-context uses.
  • Measurement boundary: Sink mass is defined only on softmax layers, so hybrid-model aggregates average over fewer layers than dense-model aggregates.The reported layer counts reduce ambiguity, but the comparison remains imperfectly like for like.
  • Language boundary: The suite measures symbolic sequences, while prior work reports that long-context behaviour degrades outside English.The authors therefore flag language coverage as an untested boundary.

B. Data Availability

The paper releases a reproducible diagnostic workflow alongside its findings. Its evidence supports objective-driven sink formation and dissociation among diagnostics, while retaining explicit limits on scale and mechanism fidelity.

  • Findings: The same stack produced 31.6 percent first-token attention under every-position prediction and 4.9 percent under final-question training, with unchanged recall.This supports objective-driven sink formation rather than architecture alone.
  • Findings: Sink mass, massive activations, and context-reading evenness behaved as distinct measurements rather than one phenomenon.Two mechanisms reduced activations without reducing sink mass, while lower sink mass coincided with steeper position preference.
  • Findings: The million-token cache cost is roughly one hundredth of same-depth ordinary caching, but the paper does not establish that the retained tokens remain usable.The cost result supports affordability, not memory quality.
Loading 2609.08574v1…