Attention Mechanism (Self-Attention)

What is the attention mechanism?

The attention mechanism is a neural network operation that decides, for each element in a sequence, how much to draw on every other element when building its representation. Instead of treating a word in isolation or only looking at its immediate neighbors, attention lets a token gather information from anywhere in the input, weighted by relevance. The word 'bank' in 'river bank' and 'savings bank' ends up with different representations because attention pulls in the surrounding words that disambiguate it.

Self-attention is the special case where a sequence attends to itself: the queries, keys, and values all come from the same set of tokens. This is the core building block of the transformer architecture and, through it, of nearly every modern large language model. The 2017 paper 'Attention Is All You Need' by Vaswani and colleagues showed that a model built almost entirely from attention, with no recurrence, could outperform the recurrent networks that dominated sequence modeling at the time.

Attention did not start with transformers. Bahdanau, Cho, and Bengio introduced it in 2014 for machine translation, to fix a specific failure of recurrent sequence-to-sequence models: those models compressed an entire input sentence into a single fixed-length vector, which lost information on long sentences. Their additive attention let the decoder look back at all of the encoder's hidden states and weight them by relevance at each step. The 2017 transformer kept that idea, dropped the recurrence entirely, and replaced additive attention with the faster scaled dot-product form.

• Attention assigns a weight to every token pair, then mixes information accordingly.
• Self-attention means queries, keys, and values all derive from the same sequence.
• It replaced recurrence as the primary way sequence models share information across positions.
• It is the central mechanism inside transformers and the LLMs built on them.

Queries, keys, and values explained

Attention works through three learned projections of each token: a query, a key, and a value. The analogy is a retrieval system. The query is what a token is looking for, the key is what each token advertises about itself, and the value is the actual content a token contributes once it has been selected. For every token, the model compares its query against the keys of all tokens to decide where to look, then collects the corresponding values.

Concretely, each input embedding is multiplied by three separate weight matrices to produce its query, key, and value vectors. These weight matrices are learned during training, which is what lets the model discover useful notions of relevance rather than having them hand-coded. The query and key vectors share a dimension so their dot product is well defined; the value vector can have its own dimension and carries the information that actually flows forward.

A useful intuition: keys and queries decide the routing, values carry the payload. Two tokens can be highly relevant to each other (large query-key match) while the information passed along (the value) is something different from the key used to find it.

• Query (Q): what the current token is searching for in the rest of the sequence.
• Key (K): what each token offers as a basis for being matched.
• Value (V): the content a token contributes once it is attended to.
• Q, K, and V are produced by three separate learned weight matrices applied to each token's embedding.

How scaled dot-product attention is computed

Scaled dot-product attention turns queries, keys, and values into an output in four steps. First, take the dot product of each query with every key to get a raw relevance score; a larger dot product means the query and key point in similar directions. Second, divide each score by the square root of the key dimension to keep the numbers in a stable range. Third, apply a softmax across each row so the scores for a given token become non-negative weights that sum to one. Fourth, use those weights to take a weighted sum of the value vectors. The result is one output vector per token, each a blend of the values the token chose to attend to.

The scaling factor matters more than it looks. When the key dimension is large, raw dot products grow large in magnitude, which pushes the softmax into a region where one weight is near one and the rest are near zero. That makes gradients vanish and training stall. Dividing by the square root of the key dimension keeps the variance of the scores roughly constant regardless of dimension, so the softmax stays in a responsive range early in training.

• Score: dot product of each query with each key (similarity).
• Scale: divide by the square root of the key dimension to stabilize gradients.
• Normalize: softmax over each query's scores produces weights that sum to one.
• Aggregate: weighted sum of value vectors yields the attended output.

Multi-head attention: many views at once

Multi-head attention runs several attention operations in parallel, each with its own learned query, key, and value projections, then concatenates their outputs and projects the result back down. Each head can specialize: one might track syntactic dependencies like subject-verb agreement, another might follow coreference (which noun a pronoun refers to), another might attend to nearby positions. A single attention pattern would force all of these into one weighting; multiple heads let the model attend to different kinds of relationships simultaneously.

The original transformer split its model dimension across heads rather than adding cost on top of it. The base model used a model dimension of 512 divided into 8 heads, giving each head a query and key dimension of 64. Because the heads share the total budget, multi-head attention costs roughly the same as a single full-width head while giving the model several independent ways to relate tokens. This is one reason the design scaled so well.

• Each head has its own Q, K, V projections and attends independently.
• Different heads tend to capture different relationships: syntax, coreference, position.
• Outputs are concatenated and linearly projected back to the model dimension.
• The base transformer used 8 heads over a model dimension of 512 (64 per head).

Self-attention vs cross-attention vs causal attention

These three are variants of the same operation, distinguished by where Q, K, and V come from and what each token is allowed to see. In self-attention, queries, keys, and values all come from one sequence, so the sequence relates its tokens to each other. In cross-attention, the queries come from one sequence and the keys and values from another, which is how an encoder-decoder model lets the output attend to the input (for example, a translation attending to the source sentence).

Causal attention, also called masked attention, is the variant used by decoder-only language models that generate text left to right. It masks out every future position so a token can attend only to itself and earlier tokens. This is exactly what next-token prediction requires: when predicting position five, the model must not peek at positions six and beyond. The mask is applied by setting the scores for forbidden positions to negative infinity before the softmax, which drives their weights to zero.

• Self-attention: Q, K, V from the same sequence; tokens relate to each other.
• Cross-attention: Q from one sequence, K and V from another; used in encoder-decoder models.
• Causal (masked) attention: each token attends only to itself and earlier tokens, required for autoregressive generation.
• Masking sets disallowed scores to negative infinity so softmax assigns them zero weight.

Why attention replaced recurrence

Attention won out over recurrent networks for two practical reasons: parallelism and direct access to distant context. A recurrent network processes a sequence step by step, carrying a hidden state forward, which makes training inherently sequential and slow on modern hardware. Self-attention computes all token interactions at once as matrix multiplications, which GPUs and TPUs execute in parallel. That parallelism is a large part of why training models at today's scale became feasible.

Recurrent and convolutional models also struggle to connect tokens that are far apart, because information has to pass through many intermediate steps and can degrade along the way. Self-attention connects any two positions in a single step: the path length between any pair of tokens is constant, not proportional to their distance. A token at the end of a long input can attend directly to a token at the start. This direct routing is what made long-range dependency modeling tractable and is the architectural reason transformers handle context the way they do.

• Self-attention is parallel across positions; recurrence is inherently sequential.
• Any two tokens are one attention step apart, regardless of distance.
• Direct long-range connections avoid the information decay of multi-step recurrence.
• Parallelism on GPUs and TPUs is a key enabler of large-scale training.

The cost: quadratic complexity and its limits

Full self-attention is quadratic in sequence length, which is its main scaling cost. Every token computes a score against every other token, so the attention matrix grows in proportion to the square of the number of tokens. Double the input length and the attention computation and memory roughly quadruple. This is the core reason context windows cannot be extended for free and why long inputs are expensive to process.

Two further points temper the picture. First, a lot of research targets this cost with sparse, linear, or windowed attention variants and with system-level methods like FlashAttention that reduce memory traffic without changing the math. Second, a longer context window does not guarantee that a model uses all of it well. Studies of long-context behavior find that accuracy can drop for information buried in the middle of a long input even when it fits comfortably inside the window, a pattern often summarized as 'lost in the middle.' Attention makes long context possible, but it does not make every position equally easy to recall.

• Attention compute and memory scale with the square of the sequence length.
• Longer inputs are quadratically more expensive, which constrains context windows.
• Sparse, linear, and windowed attention plus kernels like FlashAttention reduce the cost.
• A larger window does not ensure uniform recall; mid-input information can be missed.

Attention, context, and memory in practice

Attention is what makes a context window usable, but it is not long-term memory. Within one request, self-attention lets the model relate every token in the prompt to every other token, which is why an LLM can answer a question using details stated earlier in the same conversation. Once the conversation exceeds the window, or a new session starts, those tokens are gone. The attention mechanism only operates over what is currently in the context; it does not persist anything on its own.

Because attention weighting is uneven across long inputs, simply stuffing more text into the prompt is an unreliable way to give a model knowledge. Durable, retrievable memory is an engineering layer built around the model, most commonly retrieval-augmented generation, where relevant stored items are fetched with semantic search and inserted into the prompt so attention can act on them. AI memory and second-brain tools such as MemX apply this to personal knowledge: store information outside the model, retrieve the right pieces on demand, and let the model's attention do the rest within a focused context.

• Self-attention operates only over the current context window, not across sessions.
• It powers in-context reasoning but provides no persistent storage by itself.
• Uneven attention over long inputs makes prompt-stuffing an unreliable memory strategy.
• Retrieval plus a focused context is the practical path to durable, usable memory.