Attention mechanism

How attention works: query, key, value

Each token vector, an embedding, gives three vectors: a query, a key and a value. Vaswani and colleagues (2017) give these steps for one token:

1. Compare its query with the key of each token.
2. Divide the scores by the square root of the key length. This keeps the softmax stable.
3. Change the scores to weights with a softmax. The weights add up to 1.
4. Add the values of all tokens by those weights.

The 2017 paper uses 8 attention heads in parallel in each layer. Each head can follow a different relation between tokens.

Why attention is important for security

• One sequence for all text. The system prompt, the user text and a web page go into one transformer input. Attention compares all of them. This lets prompt injection work.
• Injected text can pull the heads. Hung and colleagues (2024) found a “distraction effect” in which some heads move to the injected instruction.
• Long inputs cost more. Two times the input gives about four times the attention work. This helps a denial of AI service.
• Attention maps are not proof. Do not use a heat map of attention as the cause of a model decision.

In practice

The 2017 paper uses attention in three positions:

Form	Queries from	Keys and values from
Self-attention	the sequence	the same sequence
Masked self-attention	the sequence	only earlier positions, for text output one token at a time
Cross-attention	the decoder	the encoder output

Subsequent work changed the formula to make it cheaper and to let it use longer inputs. Each result below is from the authors of the method.

Change	What it does	Result
FlashAttention (Dao et al., May 2022)	Accurate attention in tiles, with fewer GPU memory reads and writes	3× faster on GPT-2 at 1K tokens
Multi-query attention	One key and value head for all query heads	Faster output, but quality can decrease
Grouped-query attention (Ainslie et al., May 2023)	A small group of key and value heads	Quality near multi-head attention, speed near MQA
Attention sinks (Xiao et al., Sep 2023)	Keeps the first tokens in the cache	Four models ran on “up to 4 million tokens and more”

Meta uses grouped-query attention in its Llama 2 models of 34B and 70B parameters.

A model keeps the keys and values of earlier tokens in a cache. Xiao and colleagues say that this cache “consumes extensive memory”. They also found that models put strong attention on the first tokens, “even if they are not semantically important”.

Is attention an explanation

The 2017 paper says that self-attention “could yield more interpretable models”. Researchers do not agree on how far that goes.

	Jain and Wallace (Feb 2019)	Wiegreffe and Pinter (Aug 2019)
Position	Attention is not explanation	Prior work does not disprove it
Primary data	Weights frequently do not agree with gradient measures. Different weights can give the same prediction	The result changes with the definition of explanation. They give four alternative tests

The two papers are from 2019, before the large language models of today. The study of what occurs in a model is mechanistic interpretability.

Questions and answers

What are query, key and value in attention?

They are three vectors that the model makes from each token. The query tells what the token looks for. The key tells what the token holds. The value is the data that the token gives to the result.

Do attention weights show why a model gave an answer?

Not reliably. Jain and Wallace (2019) found that very different attention weights can give the same prediction. Wiegreffe and Pinter (2019) replied that the result changes with the definition of explanation.