Attention sinks in LLMs

The attention mechanism is used by every modern large language model. But it has a small flaw that recent models fix with one extra token.

The original attention

Attention (introduced by [1]) lets a network focus on the parts of its input that matter for the task at hand, whether the input is a chat history, an image, or an audio file. The network learns where to look.

The problem

But it can't learn to not look at all. Attention weights are normalized to sum to 1, so even when no part of the input is useful, the network is forced to place its attention somewhere.

In practice, models use the first few tokens for this, because the first tokens of a text often carry little information. Take this article: its first word, "the", does nothing but make the sentence grammatically correct.

The figure below (from [2]) shows layers with nothing useful to attend to dumping their attention onto the first token, the dark red column.

The solution

The fix is an explicit attention sink [2]: a dedicated token that exists only to absorb attention the network has no use for.

A Transformer runs many attention heads in parallel, each specializing in a different pattern. With an attention sink a head can now stay quiet when its specialty isn't needed instead of adding noise to the output.

[1] Dzmitry Bahdanau, Kyunghyun Cho, Yoshua Bengio: Neural Machine Translation by Jointly Learning to Align and Translate

[2] Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, Mike Lewis: Efficient Streaming Language Models with Attention Sinks