Attention Sink

In many trained decoder-only transformers, a large share of the attention in many heads goes to the first few tokens of the sequence, whatever those tokens say. Xiao et al., who named the effect, call these tokens attention sinks. The behavior matters in practice because it explains why a simple sliding window over the most recent tokens breaks a model, and why keeping a handful of initial tokens repairs it.

Why softmax forces attention to land somewhere

Attention weights for one query are a softmax over the keys it can see:

\[ a_j=\frac{e^{s_j}}{\sum_{l} e^{s_l}},\qquad \sum_j a_j = 1, \]

where \(s_j\) is the scaled query-key score for key \(j\). The weights are positive and sum to one, so the head cannot say “nothing here is relevant” by producing zeros. It must hand out all of its mass, and the output is a weighted average of value vectors. One way for a head to contribute little at a position is to concentrate attention on a token whose value contributes little; cancellation between value vectors can also yield a small output.

Causal masking makes the first token an ideal parking spot. Every later query can see it, while the first token sees only itself, so across the whole sequence it is the one key present in every softmax. Training can settle on giving it a key that attracts large scores and a value that is near zero. Xiao et al. attribute the effect to this softmax constraint. The tokens that become sinks are the early ones because of their visibility, not because of their content.

A toy trace

The program below uses one query and six keys with hypothetical scores. Key 0 is the sink, with a large score and a zero value vector. The query matches nothing among the other five keys. The program then removes the sink and recomputes.

import math
def softmax(z):
    m=max(z); e=[math.exp(x-m) for x in z]; s=sum(e); return [x/s for x in e]
logits=[4.0,0.3,-0.2,0.1,0.0,-0.1]
vals=[[0.0,0.0]]+[[1.0,0.5],[-0.8,0.9],[0.4,-0.6],[0.7,0.7],[-0.3,-1.0]]
def out(w,V): return [sum(wi*v[j] for wi,v in zip(w,V)) for j in range(2)]
w=softmax(logits)
print("with sink   weights",[round(x,3) for x in w]," output",[round(x,3) for x in out(w,vals)])
w2=softmax(logits[1:])
print("sink evicted weights",[round(x,3) for x in w2]," output",[round(x,3) for x in out(w2,vals[1:])])

Output:

with sink   weights [0.913, 0.023, 0.014, 0.018, 0.017, 0.015]  output [0.026, 0.009]
sink evicted weights [0.261, 0.158, 0.213, 0.193, 0.175]  output [0.302, 0.105]

With the sink present, 91% of the weight goes to it and the head’s output is close to zero, which is what a head with nothing to say should produce. After the sink is removed, the same logits renormalize over five irrelevant tokens, and the output becomes about eleven times longer (vector norm 0.32 versus 0.028). In this toy, removing the sink increases the irrelevant contribution. In a trained model, removing heavily attended initial tokens changes the attention distribution and can disrupt generation; the effect on each head depends on its learned values and role.

Consequence for windowed attention

A decoder with a long sequence has to store keys and values for every past token (the KV cache). The cheap fix is a sliding window that keeps only the most recent \(W\) tokens. Xiao et al. report that quality collapses once the first tokens fall out of the window, and that keeping the first few tokens together with the recent window restores it. Their StreamingLLM does exactly this and handles streams much longer than the training length without fine-tuning. The sink tokens fix the distribution of attention weights; they do not carry information about the evicted middle of the text. A fact that scrolled out of the window is not recoverable, and a model run this way produces fluent text, not memory of what it dropped.

Related but different

• Lost in the middle describes accuracy depending on where evidence sits in a long prompt. Sinks describe where attention mass accumulates, and the two can coexist without one explaining the other.
• A sink is not a special token the user inserts, although a model can be trained with an explicit learnable sink token. In models trained without one, the first ordinary tokens take the role.
• Position encoding choices influence the effect, but the mechanism above needs only softmax normalization and a causal mask.

References

• Xiao, G., et al., Efficient Streaming Language Models with Attention Sinks — https://arxiv.org/abs/2309.17453