Machine Learning · Transformers

Self-Attention (Transformers)

Reference entry · last updated September 12, 2026

Self-attention (sometimes called intra-attention) is an attention mechanism relating different positions of a single sequence in order to compute a representation of the sequence.[1] Unlike recurrence, which processes tokens sequentially step by step, self-attention computes direct pairwise interaction scores across all positions in parallel, yielding a constant \(O(1)\) sequential path length between any two tokens.[1]

1. First Principles: Sequence Alignment and Scaled Dot-Product

Given an input matrix \(X \in \mathbb{R}^{n \times d_{\text{model}}}\) representing a sequence of \(n\) token embeddings, self-attention first applies three learned linear projections to map the inputs into Query (\(Q\)), Key (\(K\)), and Value (\(V\)) representations:

where \(W_Q \in \mathbb{R}^{d_{\text{model}} \times d_k}\), \(W_K \in \mathbb{R}^{d_{\text{model}} \times d_k}\), and \(W_V \in \mathbb{R}^{d_{\text{model}} \times d_v}\) are parameter matrices. Each token vector produces three specialized roles: a query that searches for relevant context, a key that advertises attributes for matching, and a value that contains the retrievable content.

The relevance of token \(j\) to token \(i\) is quantified by the inner product of query vector \(q_i\) with key vector \(k_j\). The matrix product \(Q K^\top \in \mathbb{R}^{n \times n}\) yields unnormalized alignment scores across all pairs. When the projection dimension \(d_k\) is large, dot products grow large in magnitude, pushing the softmax function into regions with near-zero gradients. To preserve variance, the dot products are divided by the scaling factor \(\sqrt{d_k}\):[1]

The row-wise softmax normalizes the scaled scores so that weights along each row are non-negative and sum to 1. Each output vector is a weighted sum of the value vectors, directly pooling context from across the entire sequence into an updated representation.

2. Attention Heads and Multi-Head Attention

An attention head is an independent computational unit within a self-attention layer that calculates attention patterns across tokens using its own projection matrices \((W_i^Q, W_i^K, W_i^V)\). A single attention head projects tokens into one subspace, producing a single distribution of attention weights across the sequence.

Rather than computing a single attention function with \(d_{\text{model}}\)-dimensional queries, keys, and values, the Transformer employs multi-head attention.[1] Multi-head attention runs \(h\) separate attention heads in parallel within the same layer. The dimensionality of each head is typically set to \(d_k = d_v = d_{\text{model}} / h\), keeping total computational cost similar to a single full-dimensional attention pass.

Contextual Specialization

Running multiple heads simultaneously allows the model to attend to distinct types of relational patterns at the same time:[4]

• Syntactic dependencies: Specific heads learn to connect verbs to subjects or verbs to objects (such as linking the verb chased directly to the subject dog).
• Coreference resolution: In ambiguous sentences such as The dog chased the cat because it was fast, an attention head can route high weight from the pronoun it back to the noun dog, resolving reference before subsequent feed-forward processing.
• Positional and boundary framing: Heads attend to previous or subsequent adjacent tokens, sentence delimiters, or special tokens like [CLS] and [SEP].

Output Integration

Each head \(i \in \{1, \dots, h\}\) computes its own head output \(\text{head}_i = \text{Attention}(Q W_i^Q, K W_i^K, V W_i^V)\). The outputs of all individual attention heads are concatenated along the feature dimension and mapped back to the model dimension by a linear projection matrix \(W^O \in \mathbb{R}^{h d_v \times d_{\text{model}}}\):[1]

This linear projection integrates the distinct representations discovered by each head before the activations pass into residual connections and the feed-forward sublayer.

3. Computational Complexity and KV Cache

The core computational trade-offs of self-attention govern both model scaling and serving efficiency:

Computing the matrix product \(Q K^\top\) requires evaluating an \(n \times n\) matrix of scores per head. Consequently, time and activation memory scale quadratically with sequence length \(n\). While this quadratic factor bounds context length on fixed accelerator memory, self-attention eliminates the sequential bottleneck of recurrent networks, allowing full parallelization across training tokens.[1]

Inference and Key-Value Caching

During autoregressive decoding, tokens are generated one at a time. Evaluating all positions from scratch at step \(t\) would repeat past calculations, resulting in \(O(t^2)\) operations per step. By storing the key and value vectors of previously computed tokens in a Key-Value (KV) cache, each new step needs only project the single new query vector \(q_t\) and compute attention against the cached keys and values. This reduces step latency to \(O(t)\) while requiring dedicated high-bandwidth memory proportional to sequence length and layer count.

4. Directionality and Causal Masking

Self-attention supports two primary directional regimes:

• Bidirectional self-attention: Every token attends to all positions in the sequence, both left and right. This unconstrained visibility is standard in encoder architectures like BERT,[3] making it suitable for classification, extraction, and whole-sequence embedding.
• Causal (masked) self-attention: Autoregressive decoders must not look ahead into future positions that have not yet been generated. A causal mask \(M \in \mathbb{R}^{n \times n}\) is added to the attention logits prior to the softmax operation: \[ M_{ij} = \begin{cases} 0 & \text{if } j \le i \\ -\infty & \text{if } j > i \end{cases} \] Positions with \(-\infty\) evaluate to zero probability under the softmax function, strictly preventing information leakage from future tokens.

5. Comparison with Cross-Attention

Self-attention and cross-attention use the same scaled dot-product formulation but differ in the provenance of their inputs:

• Self-Attention: The query, key, and value vectors all derive from the same sequence representation \(X\). The mechanism maps intra-sequence relationships within a single context.
• Cross-Attention: Queries derive from the target sequence (such as decoder hidden states), while keys and values derive from an external source sequence (such as encoder outputs). The mechanism aligns two distinct sequences, as in sequence-to-sequence translation or multimodal conditioning.