How attention gives a language model context

A naturalist writes: The ranger saw the bat by the bank. The sentence has two ambiguities. Bat might mean an animal or a piece of sports equipment. Bank might mean a river edge or a financial institution. Then comes a second sentence: Its wings folded against the reeds as water reached the muddy shore. A human reader now favors the animal and the riverbank. The surrounding words change how the first sentence is understood.

This made-up example is a guide to attention, not a trace of any particular model. Only the token split and the numbers in the worked example are invented for teaching; the research claims that follow refer to cited sources. In a Transformer, attention lets one position collect weighted information from other positions. That operation is central to many language models, while the rest of the network turns those mixed representations into predictions.

The model receives tokens

Text first passes through a tokenizer. Its output consists of tokens, which may be whole words, punctuation marks, word pieces, or other units. The exact split depends on the tokenizer and model. Work on subword units showed how smaller pieces can help neural translation handle rare words; it did not establish one tokenizer for every system. For the example, treat each written word in The, ranger, saw, the, bat, by, the, bank as one token. Real tokenizers may split them differently, so this list is only a teaching device. Sennrich, Haddow, and Birch's 2016 paper describes one approach to subword units.

In a real tokenizer, a familiar word may be one token while an unusual word is split into pieces. Heatmap row and column indices refer to tokens, not spaces or whole written words. Checking the tokenizer output helps prevent reading one subword position as if it represented the entire word.

Each token is mapped to a vector, a list of numbers. In the original Transformer, the token representation is combined with a positional encoding; other architectures can represent position differently. A vector is not a label such as animal, river, or noun. Its coordinates are useful because training adjusts them and later computations to predict language. A person generally cannot read one coordinate as a plain-language feature.

As the sequence passes through layers, a token's representation is transformed. It can carry information gathered from neighboring positions. Position information also matters: bat near bank and bank near bat have different order. The network has to preserve order while mixing information across the sequence.

One attention calculation

A Transformer layer projects the representations into queries, keys, and values. The terms name mathematical roles. A query is not a natural-language question, and a key is not a dictionary entry. For one token, the layer compares its query vector with the key vectors of positions it is allowed to consider. The resulting similarity scores are scaled and passed through softmax. Softmax converts the scores into nonnegative weights that sum to one across the available positions for that query and head.

The layer uses those scalar weights to combine value vectors. Query and key projections determine where information may be drawn from; the value projections determine what vector content is available to mix. Suppose one query assigns weights 0.70, 0.20, and 0.10 to three positions. For a teaching example, let their value vectors be (2, 0), (0, 1), and (1, 2). The mixture is 0.70(2, 0) + 0.20(0, 1) + 0.10(1, 2). Multiplying gives contributions of (1.4, 0), (0, 0.2), and (0.1, 0.2), which add to (1.5, 0.4).

The weights are scalars, while the values carry vector content. Attention therefore does more than point at a word: it takes weighted pieces of the value representations and adds them. These numbers are hand chosen, not measurements from a trained model.

Each query repeats this calculation for the positions it can see, so a layer produces one weight distribution for each query and head. A heatmap usually places these distributions in rows and available key positions in columns; a mask can make some columns ineligible. Each score compares one query with one key. A larger score gives that position a larger share of the normalized weights, but the score does not directly name a grammatical role. The map is specific to the input, layer, head, and mask being displayed.

If all three value vectors were (1, 0), then any weights that sum to one would produce the same result, (1, 0). The map of weights could differ while the mixture stayed identical. In a real layer, value vectors and later computations also matter, so a heatmap shows only one part of the calculation.

In the original paper, the operation is called scaled dot-product attention. It forms query-key scores, divides them by the square root of the key dimension, applies softmax, and combines the values. As the vectors grow in dimension, an unscaled dot product can grow in magnitude; softmax may then concentrate too sharply on a few scores. Dividing by the square root of the dimension helps keep the scores in a useful range. Vaswani and colleagues introduced the Transformer in 2017, in a paper first submitted on 12 June 2017 and published at NeurIPS 2017. The full architecture paper is a 2023 revision of that 2017 work.

A layer uses multiple attention heads. Each head has learned projections and can form its own pattern of weighted mixing. The model combines the head outputs through an output projection, returning them to the layer's representation space so later computation can blend their contributions. A head can contribute to a pattern, but calling it the model's grammar detector or animal detector requires evidence about a specific trained system. Patterns can combine across heads and layers.

After the layer's other operations, its updated representations become input to the next layer. That layer computes fresh queries, keys, and values from states that already carry information gathered earlier. A stack repeats this exchange across positions and transformations.

Attention sits inside a larger architecture

The 2017 Transformer was an encoder-decoder model for sequence tasks such as machine translation. Its encoder built representations of the input sequence. Its decoder generated the output sequence. Encoder self-attention let an input position use information from positions on either side. The decoder used a causal mask so an output position could not see future output tokens, and it also attended to the encoded input. In self-attention, queries, keys, and values come from the same sequence; in encoder-decoder cross-attention, decoder queries compare with keys and values from the encoded input. Thus, at the bank position in an encoder, attention may draw on later words such as shore. A decoder-only model cannot use those later words while processing bank, but a later prediction can use them after they enter its prefix. The title Attention Is All You Need described attention as the central sequence-mixing operation in a model that avoided recurrence and convolution. Each layer also uses position information, a feed-forward transformation at each position, and residual and normalization paths that carry and reshape the representations.

The paper compared the computational costs of sequence-mixing methods. For dense self-attention, its table gives a per-layer cost of O(n²d), where n is the number of positions and d is the representation dimension. At fixed dimensions, doubling n creates roughly four times as many query-key score pairs. This describes the dense attention calculation, not the total cost of every model or system; real systems may use different arrangements or optimizations. Attention connects distant positions without passing information through a recurrent chain, while generation remains sequential because each new token depends on the prefix already produced. The original paper's comparison table presents these complexity and path-length comparisons.

Training and generation use different states

GPT-3, described by Brown and colleagues in 2020, is an autoregressive language model. It predicts a next token from earlier tokens. Their experiments showed that instructions or examples in a prompt could guide the model without a gradient update for each task. This is called in-context learning. Brown et al.'s 2020 paper reports that setup.

During autoregressive training, the model adjusts shared parameters to predict next tokens in training examples. Those parameters include the projections that produce queries, keys, and values. In a training sequence, each position is paired with the next token as its target, while a causal mask hides that future token from the input used for the prediction. The target supplies the error signal after the prediction; it is not information the position attends to. During inference, the model reuses its learned parameters and calculates new activations from the current prompt. A different prompt can change the activations and attention weights without changing the parameters.

Causal direction affects how to describe the ranger example. At the position of bat, a decoder-only model cannot use later words such as wings or reeds. When it predicts after reading the full two-sentence prompt, the current position can use all preceding tokens. The model does not go back and rewrite the earlier state for bat. Later positions draw on earlier ones as the prediction advances. An encoder can use input positions on both sides, so the same teaching sentence can illustrate different paths in different architectures.

At generation time, the model produces a probability distribution over possible next tokens. A decoding policy selects one, either by choosing the most probable option or sampling according to its settings. The chosen token is appended to the input, and the model predicts again. This repeats one token at a time.

What a heatmap does and does not show

The 2019 debate over attention as explanation examined particular models and tests. Jain and Wallace showed cases in NLP models where an attention distribution was not sufficient to explain a prediction. Wiegreffe and Pinter challenged assumptions in that analysis and proposed more careful tests. Both papers focus especially on recurrent attention models; together they motivate testing whether changing an attention pattern changes a prediction in a fixed model and task. They do not assign universal meaning to colored cells in later Transformers. Jain and Wallace and Wiegreffe and Pinter present the original arguments.

A model's context window describes how much input it can accept under a particular setup. Capacity alone says little about how reliably it uses each position. In tests of multi-document question answering and key-value retrieval, Liu and colleagues found that the location of relevant information affected performance in the systems they studied. Results were often weaker when useful material sat in the middle. Their first preprint appeared on 6 July 2023; the result describes those tasks and systems, not every model or prompt. The study and its task details explain the scope.

To inspect an attention explanation, record the model and version, prompt, tokenization, layer, head, and mask. Keep the task and decoding settings fixed while changing one phrase, such as the sentence that disambiguates bat, then compare the output distribution and the same attention view. This tests whether that computation changes with the cue. A factual answer still needs evidence from the cited paper or data.

Sources

• Sennrich, Haddow, and Birch, “Neural Machine Translation of Rare Words with Subword Units”, ACL, August 2016.
• Vaswani et al., “Attention Is All You Need”, first submitted to arXiv on 12 June 2017; NeurIPS 2017. The full-text HTML is a 2023 revision of the 2017 paper.
• Brown et al., “Language Models are Few-Shot Learners”, first submitted to arXiv on 28 May 2020; NeurIPS 2020.
• Jain and Wallace, “Attention is not Explanation”, NAACL, June 2019.
• Wiegreffe and Pinter, “Attention is not not Explanation”, EMNLP, November 2019.
• Liu et al., “Lost in the Middle: How Language Models Use Long Contexts”, first submitted to arXiv on 6 July 2023; journal version in TACL, 2024.