Self-Attention Mechanism

Self-attention is a mechanism for computing a context-dependent representation of each element in a sequence by letting it attend to every other element, including itself. Given an input sequence of vectors $x_1, \dots, x_n$ (typically token embeddings, possibly with positional information added), self-attention produces an output sequence $z_1, \dots, z_n$ of the same length, where each $z_i$ is a weighted combination of transformed versions of all the $x_j$. The mechanism is organized around three learned linear projections of each input vector, borrowed from the vocabulary of information retrieval:

• Query ($q_i$): a representation of what token $i$ is "looking for" in the rest of the sequence.
• Key ($k_j$): a representation of what token $j$ "advertises" about itself, used to be matched against queries.
• Value ($v_j$): the actual content of token $j$ that gets aggregated once a match is found.

The similarity between a query and a key — computed via a Dot Product Similarity — determines how much of value $v_j$ flows into the output for token $i$. Formally, the attention weight $\alpha_{ij}$ measures "how relevant is token $j$ to token $i$," and the output is

The word "self" distinguishes this from cross-attention (where queries come from one sequence and keys/values from another, as in encoder-decoder attention): in self-attention, queries, keys, and values are all derived from the same input sequence. This is what lets a single sentence contextualize its own words using only its own content, without any external sequence to attend to. The next natural refinement — dividing the dot product by $\sqrt{d_k}$ for numerical stability — is called Scaled Dot-Product Attention, and running several attention computations in parallel with different learned projections is Multi-Head Attention; both build directly on the mechanism described here.

Self-attention proceeds in five conceptual stages, applied identically and independently at every position, then combined: 1. Project each input embedding $x_i \in \mathbb{R}^{d}$ into a query, key, and value vector using learned weight matrices $W^Q, W^K \in \mathbb{R}^{d \times d_k}$ and $W^V \in \mathbb{R}^{d \times d_v}$:

2. Score every pair $(i, j)$ by taking the dot product of query $i$ with key $j$:

3. Normalize the scores for each fixed $i$ across all $j$ using softmax, turning raw similarity scores into a probability distribution that sums to 1:

4. Aggregate the value vectors weighted by these attention probabilities:

5. Output the sequence $z_1, \dots, z_n$, which now encodes each token's meaning in context — ready to be passed to a feed-forward layer or stacked into another attention layer. In practice, steps 1–4 are implemented as dense matrix operations over the whole sequence at once rather than token-by-token loops, which is what makes the mechanism efficient on parallel hardware. ```

# Self-Attention over a sequence of n token embeddings, each of dimension d

# X: n x d matrix of input embeddings (rows are tokens)

# WQ, WK: d x dk learned projection matrices

# WV: d x dv learned projection matrix

# Returns Z: n x dv matrix of contextualized outputs function SELF_ATTENTION(X, WQ, WK, WV): n, d <- shape(X) Q <- X @ WQ # n x dk (each row is a query vector) K <- X @ WK # n x dk (each row is a key vector) V <- X @ WV # n x dv (each row is a value vector) scores <- Q @ transpose(K) # n x n, scores[i][j] = qi . kj

# (Optional but standard: scale by sqrt(dk) — see Scaled Dot-Product Attention) scores <- scores / sqrt(dk)

# Row-wise softmax: each row becomes a probability distribution over the n tokens for i in 1..n: row <- scores[i] m <- max(row) # for numerical stability exps <- [exp(s - m) for s in row] total <- sum(exps) A[i] <- [e / total for e in exps] # A[i][j] = alpha_ij Z <- A @ V # n x dv, Z[i] = sum_j A[i][j] * V[j] return Z ``` Two edge cases deserve explicit attention (pun intended):

• Padding tokens. When sequences in a batch have different lengths, padded positions must be masked out before the softmax, typically by adding $-\infty$ (or a very large negative number) to their score, so their post-softmax weight becomes 0 and they contribute nothing to any $z_i$.
• Causal masking. In autoregressive (decoder-only) models like GPT, token $i$ must not attend to tokens $j > i$ (future tokens), since that would leak information the model is supposed to predict. This is enforced by setting $e_{ij} = -\infty$ for all $j > i$ before the softmax — a lower-triangular mask.

``` function CAUSAL_MASK(scores): # scores: n x n for i in 1..n: for j in (i+1)..n: scores[i][j] <- -infinity return scores ```

Time complexity. For a sequence of length $n$ with model dimension $d$ (and query/key/value dimension roughly $d$ as well), the dominant costs are:

• Computing $Q$, $K$, $V$: three matrix multiplications of size $(n \times d)(d \times d)$, costing $O(n d^2)$.
• Computing $QK^\top$: an $(n \times d)(d \times n)$ multiplication, costing $O(n^2 d)$.
• Softmax over the $n \times n$ score matrix: $O(n^2)$.
• Computing $A V$: an $(n \times n)(n \times d)$ multiplication, costing $O(n^2 d)$.

The total is $O(n^2 d + n d^2)$. For typical settings where $n$ is on the order of hundreds to thousands and $d$ is fixed (e.g., 512, 768, 4096), the $O(n^2 d)$ term dominates as sequences grow, giving the well-known quadratic-in-sequence-length scaling of self-attention. This is the central efficiency bottleneck of Transformers at long context lengths, motivating a large body of follow-on work on sparse, linear, and approximate attention variants.

Space complexity. The full attention matrix $A$ (and the intermediate scores) requires $O(n^2)$ memory in the naive implementation, in addition to $O(nd)$ for the embeddings and projections. This $O(n^2)$ memory is often the practical limiting factor before the compute cost is, since it must be materialized (or at least streamed) for the softmax normalization.

Why it's correct — the aggregation is a genuine weighted average. The softmax guarantees $\alpha_{ij} \geq 0$ for all $j$ and $\sum_j \alpha_{ij} = 1$ for each fixed $i$, since softmax is defined as normalized exponentials over a fixed index. Therefore $z_i = \sum_j \alpha_{ij} v_j$ is a convex combination of the value vectors — it lies inside their convex hull, is well-defined regardless of scale, and interpolates smoothly as scores change (softmax is differentiable everywhere), which is essential for gradient-based training.

Why the mechanism can represent useful dependencies. The scoring function $e_{ij} = q_i \cdot k_j = x_i W^Q (x_j W^K)^\top$ is a bilinear form in $x_i$ and $x_j$. Because $W^Q$ and $W^K$ are learned, the model can, in principle, tune this bilinear form to detect any linearly-expressible relationship between token pairs (e.g., subject–verb agreement, coreference cues, adjacent-word bigram patterns) — the universality argument here is the same as for any sufficiently expressive learned bilinear/linear layer: with enough capacity in $W^Q, W^K$ the score function can approximate a broad class of pairwise compatibility functions, and gradient descent shapes it to minimize the downstream task loss.

Why it beats pure recurrence for long-range dependencies. In a recurrent architecture, information from token $j$ must pass through $i - j$ sequential update steps to influence token $i$, and each step is a lossy, non-linear transformation — the effective "path length" between distant tokens grows linearly with their distance, and gradients can vanish or explode along that path (a signal-propagation argument, not merely folklore). In self-attention, the path length between any pair of tokens is exactly 1: $z_i$ has a direct additive term involving $v_j$ for every $j$, regardless of $|i-j|$. This constant path length is the structural reason self-attention captures long-range dependencies more reliably than pure recurrent models.

Tradeoffs, stated honestly. Self-attention purchases parallelism and short dependency paths at the cost of quadratic time and memory in sequence length, and — unlike recurrence or convolution — it has no inherent notion of token order: the dot products $q_i \cdot k_j$ are invariant to permuting the input positions (aside from which vectors get which index), so positional information must be injected separately (e.g., positional encodings), a design choice explored alongside Scaled Dot-Product Attention and Multi-Head Attention.

Consider a toy 3-token sequence with the sentence "cats chase mice," using contrived (but arithmetically clean) 2-dimensional embeddings after projection, so we can trace every number by hand. Suppose the query, key, and value vectors (already projected, $d_k = d_v = 2$) are:

We compute the contextualized output $z_2$ for "chase" ($i = 2$), whose query is $q_2 = (0,1)$.

Step 1 — raw scores $e_{2j} = q_2 \cdot k_j$:

Step 2 — scale by $\sqrt{d_k} = \sqrt{2} \approx 1.414$:

Step 3 — softmax. Subtract the max (1.414) for stability, exponentiate, normalize:

(Check: $0.243 + 1.000 + 0.243 = 1.486$... let's redo the sum carefully: $0.243+1.000+0.243 = 1.486$, so $\alpha_{21} = 0.243/1.486 \approx 0.163$, $\alpha_{22} = 1.000/1.486 \approx 0.673$, $\alpha_{23} \approx 0.163$.)

These sum to $1.000$ (within rounding), confirming a valid probability distribution — as guaranteed by the softmax normalization argument above.

Step 4 — weighted aggregation of value vectors:

Interpretation. The token "chase" ends up mostly attending to itself ($\alpha_{22} = 0.673$), moderately to both neighbors ($0.163$ each), because its query aligned most strongly with its own key. The output $z_2 = (2.182, 0.489)$ is no longer just the value of "chase" — it is a blend dominated by "chase" but pulled toward "cats" and "mice" in equal measure, reflecting symmetric relevance to both. If "chase" instead needed to disambiguate its subject specifically, learned weights $W^Q, W^K$ during training would shape the scores so that $\alpha_{21} \gg \alpha_{23}$, pulling $z_2$ toward the value of "cats" much more strongly — this is exactly how the mechanism learns to encode roles like subject and object once trained on real data.

The query-key-value formulation of attention, and its use as the sole sequence-mixing mechanism (dispensing with recurrence and convolution entirely), was crystallized in 2017 by researchers at Google, most prominently Ashish Vaswani, Noam Shazeer, and their coauthors, in the paper that introduced the Transformer architecture. But the intellectual lineage runs back further: attention mechanisms for aligning encoder and decoder states in neural machine translation appeared several years earlier, and the general idea of learned soft-alignment between sequence elements can be traced to differentiable memory-addressing schemes explored in the early 2010s.

Self-attention is the computational backbone of essentially every modern large language model: GPT-family decoder-only models, BERT-family encoder models, T5-style encoder-decoder models, and vision Transformers applied to images all rely on the same query-key-value mechanism described here, typically wrapped in Multi-Head Attention and stabilized via Scaled Dot-Product Attention. Production systems — machine translation services, code assistants, chatbots, document summarizers, and retrieval-augmented generation pipelines — all route their internal computation through stacks of self-attention layers. The quadratic cost analyzed above is precisely why industry research invests heavily in efficient attention approximations (sparse attention patterns, low-rank approximations, sliding-window attention, and hardware-aware implementations) to extend usable context length without a proportional blow-up in compute and memory. Understanding the exact mechanics of self-attention — not just its name — is a prerequisite for reasoning about model capacity, context-length limits, and the design tradeoffs behind every major Transformer-based system in use today.

Ask the tutor: How does scaling by $\sqrt{d_k}$ actually prevent softmax saturation, and how does multi-head attention change the query-key-value computation described here?