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 the input that matter for its current task. Whether the input is a chat history, an image, or an audio file, the network learns where to look for information.

The problem

Attention weights are normalized so they sum to 1. Even when no part of the input is useful to predict the next token, the network is forced to place all the attention somewhere. It can not learn to ignore everything.

Worse still, the function traditionally used for attention is Softmax which selects the largest element in a list and really emphasizes it. You can see this in the figure below, where attention isn't spread around different places, but a single token lights up as dark red.

Attention learns a workaround

In practice, older models learn a workaround, when none of the inputs are important to predict the next token. They learn to dump all the attention onto the first few tokens, 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 of a neural network that have nothing useful to attend to. They're 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.

Further developments

This also enables an architecture for very large input texts called Streaming LLMs. The model is trained on a finite window, and during inference the text streams through that window. Earlier parts of the input text drop out of the input window and so the classical attention mechanism loses the first token it would have used as an attention sink. To solve this, you can either keep the first few tokens inside the context window while the rest of the text streams by, or use an explicitly learned attention sink token that's always prepended to the context window.

Sources

[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