Xiao et al., MIT, Meta AI, CMU and NVIDIA (2023) — ARXIV-2309.17453
Key takeaways
- The phenomenon. Beyond the bottom two layers, every layer and head of
Llama-2-7B attends heavily to the first few tokens regardless of what they
say. Substituting the first four tokens with linebreaks changes almost
nothing: 5.40 perplexity with the real tokens, 5.60 with
\n. So it is the absolute position that matters, not the content. - Why it happens, and this is the part the record needs. Softmax cannot assign zero to everything: the scores are normalised to sum to one, so a head with nothing it needs to attend to must still put its mass somewhere. Models learn to dump that surplus on whatever every query can see, which under causal masking is the initial tokens. The authors name that role the attention sink. The same argument had been made for quantization outliers, with SoftMax-off-by-One proposed as the fix.
- What it costs when the sink is evicted. Window attention drops the initial tokens’ KV once the window slides past them, which removes a large part of the softmax denominator and shifts the whole attention distribution. Llama-2-13B: 5.40 perplexity at 4+1020 cache, 5158.07 at 0+1024. The failure is not gradual.
- The remedy, for a model already trained. Keep the KV of four initial tokens permanently and roll the rest. One or two is not enough; four is, and more adds little. Positions are assigned within the cache rather than from the original text, which is what makes it work — for RoPE the keys are cached before the rotation and rotated at use. Up to 22.2× faster per token than sliding-window-with-recomputation, at comparable memory.
- The remedy, for a model you are training. Prepend one learnable sink token. Trained that way, streaming perplexity is stable with the sink alone (18.01 at 1+1023) where the vanilla model still needs four initial tokens. A zero-valued sink is not enough — it has to be learnable.
- A caveat the paper volunteers. Larger caches do not monotonically lower perplexity: Llama-2-7B is best at 4+2044 and worse at 4+4092. The models do not fully use the context they are given.
Standing in the anthology
The description of a phenomenon this record was already recommending people remove. SOTA-134 is filed on the finding that a head-specific output gate improves quality and eliminates the attention-sink pattern, and until this note the phrase appeared nowhere else in the corpus — the record advised removing something it had never defined.
Read against SOTA-134, the two are answers to the same pressure from opposite ends. This paper accepts that the softmax must shed mass and gives it a designated place to go; the gate lets a head scale its whole output down instead, so there is nothing to shed. That is why the gate removes sinks rather than relocating them, and it is also why sinks are not simply a defect: in an ungated model the sink is load-bearing, and evicting it costs three orders of magnitude of perplexity.
Filed from LIT-139’s citation of it at DeepSeek-V4 §2.3.3, via the reference pass in #40.
SepLLM (LIT-748) is the most direct later test of this design: it keeps StreamingLLM’s sink tokens and rolling window, adds the KV of every separator that scrolls out, and runs StreamingLLM as its baseline at a matched cache budget throughout, coming out ahead by seven points on GSM8K-CoT at the same 47% of KV and passing a needle-in-a-haystack test StreamingLLM fails. Its ablation also reproduces this paper’s finding that the initial tokens cannot be dropped.