How Attention Lets AI Models Focus on the Right Words

Attention is the step inside a modern language model that decides which earlier words matter for the word being processed. Instead of reading a sentence strictly one word at a time, the model scores how relevant every word is to every other word, then blends the information accordingly. This one idea sits at the heart of the Transformer design that today's chatbots are built on.

The problem attention was meant to solve

Before 2017, the leading systems for translating or generating text were recurrent neural networks. These read a sentence in order, squeezing everything seen so far into a running summary. Long sentences were hard, because details from the start had to survive many steps of squeezing.

According to the Wikipedia article on attention, a 2014 paper by Bahdanau and colleagues added an attention mechanism to these networks so that, when producing each output word, the model could look back at the input words that mattered most. That helped with long sentences. In 2017 eight researchers then published the paper "Attention Is All You Need" on arXiv (submitted 12 June 2017), proposing the Transformer, which is based solely on attention and drops recurrence entirely.

Queries, keys and values

The standard version is called dot-product attention. Every token (a word or word piece) is turned into three lists of numbers, called vectors, using three learned sets of weights. They are named the query, the key and the value.

• The query says what this token is looking for.
• The key says what each token offers to be found by.
• The value is the information a token passes along if it is chosen.

For one token, the model compares its query with the key of every token by multiplying the vectors together. A large result means a good match. These scores are then put through a function called softmax, which turns them into weights that are positive and add up to one. Finally, the model adds up the value vectors, each multiplied by its weight. A word with a high weight contributes a lot to the result; a word with a tiny weight contributes almost nothing.

The Wikipedia article writes the whole operation as softmax(QKᵀ / √d_k) · V. The division by the square root of the key size keeps the scores from growing so large that a single token swamps the rest of the softmax.

Self-attention and several heads

When the queries, keys and values all come from the same piece of text, it is called self-attention. Every position can look at every other position in the input, which the Wikipedia article calls all-to-all attention.

A Transformer does not run attention just once. It runs several copies side by side, called heads, each with its own learned weights. Their outputs are joined together and mixed by another learned step. Having several heads lets the model track different kinds of relationships at the same time, although what each head actually learns is something researchers study rather than something designers choose.

Order and masking

Plain attention has no sense of word order; shuffle the words and the maths treats them the same. Transformers therefore add positional information to each token, so the model knows where each word sits.

Chatbots that write text one token at a time also use a mask. When the model predicts the next word, it must not peek at words that come after it in the training text. According to the Wikipedia article, this causal masking is done by setting the scores for later positions to negative infinity before the softmax, which makes their weights zero.

Why it is fast to train but costly for long text

The 2017 paper reported that the Transformer was more parallelizable than earlier designs, because all positions can be processed at once rather than one after another. It reported a score of 28.4 BLEU, a standard translation measure, on the WMT 2014 English-to-German task, and 41.8 BLEU for a single model on English-to-French after 3.5 days of training on eight GPUs. Those were the authors' results in 2017, and later systems have moved far beyond them.

The cost is that every token is compared with every other token. As Wikipedia puts it, the size of the attention matrix is proportional to the square of the number of input tokens. Double the length of the text and the number of comparisons roughly quadruples. This is one reason long inputs are expensive and why the length a model can handle, its context window, is limited.

Engineers work around this in several ways. One named in the article is FlashAttention, which computes the same attention results in blocks small enough to fit in the fast memory on a GPU, reducing memory use without changing the answer.

What attention is and is not

Attention weights are useful for describing what a model is doing, but they are not a full explanation of why it gives a particular answer. A Transformer also contains other layers, such as feed-forward layers and skip connections, that do much of the remaining work. Attention is the part that moves information between tokens; the rest transforms it.

The Wikipedia article notes that attention has since largely replaced recurrent networks in frontier models such as BERT, T5 and GPT. That is why the word turns up so often when people describe how chatbots work, and why a basic grasp of queries, keys and values goes a long way.

Sources

• Attention (machine learning), Wikipedia)
• Attention Is All You Need, arXiv:1706.03762