Attention, Derived From First Principles

Part 1 of the Transformers Under the Hood series.

To understand the Transformer architecture, one must strip away the abstraction of "attention" as a cognitive metaphor and view it strictly as a linear algebra operation. At its core, self-attention is a mechanism that allows a model to weigh the importance of different parts of the input sequence relative to each other. This process is governed by three primary components: Query (Q), Key (K), and Value (V) projections, followed by a scaled dot-product interaction and softmax normalization.

The Geometric Projection: Q, K, and V

The first step in the attention mechanism is the transformation of input embeddings into three distinct vector spaces. Let $X$ be the matrix of input embeddings for a sequence of length $n$, where each row represents a token's embedding vector of dimension $d_{model}$.

The model applies three learned weight matrices, $W^Q \in \mathbb{R}^{d_{model} \times d_k}$, $W^K \in \mathbb{R}^{d_{model} \times d_k}$, and $W^V \in \mathbb{R}^{d_{model} \times d_v}$, to project the input into Query, Key, and Value spaces respectively:

$$Q = XW^Q, \quad K = XW^K, \quad V = XW^V$$

This projection is not arbitrary. The learned matrices allow the model to define a coordinate system where semantic relationships become geometric distances. The Query vector represents what the model is looking for; the Key vector represents what the token contains; and the Value vector contains the actual information to be retrieved. By projecting into these separate spaces, the model can learn to ignore irrelevant dimensions in the raw embedding and focus on features relevant to the attention task.

The Interaction: Dot Products and Scaling

Once the projections are established, the model calculates the compatibility between every Query and every Key. This is done via the dot product. For a specific query vector $q_i$ and key vector $k_j$, the score is computed as $q_i \cdot k_j^T$.

The dot product serves as a measure of similarity. If two vectors point in the same direction, their dot product is large and positive. If they are orthogonal, the result is zero. In the context of attention, a high dot product indicates that the content of the Key matches the intent of the Query.

However, raw dot products can result in very large values, especially when the dimensionality $d_k$ is large. This poses a problem for the subsequent softmax function. The softmax function, defined as $\text{softmax}(x_i) = \frac{e^{x_i}}{\sum_j e^{x_j}}$, is sensitive to the magnitude of its inputs. When inputs are large, the exponential function saturates, causing the gradients to vanish during backpropagation. This makes learning difficult because the model cannot adjust the attention weights effectively.

To counteract this, the attention scores are scaled by the square root of the key dimension, $\sqrt{d_k}$, where $d_k$ is the dimension of the key vectors. The formula becomes:

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

This scaling ensures that the variance of the dot products remains close to 1, keeping the softmax function in its sensitive, linear region where gradients are strong. As noted in the original "Attention Is All You Need" paper, this scaling is critical for training stability in deep networks.

Normalization: The Softmax Distribution

After scaling, the attention scores undergo a softmax operation across the key dimension. This transforms the raw scores into a probability distribution. Each element in the resulting matrix represents the weight assigned to a specific value vector based on the corresponding query-key pair.

Mathematically, for a query $q_i$ and a set of keys $K$, the attention weights $\alpha_{ij}$ are:

$$\alpha_{ij} = \frac{\exp(e_{ij}/\sqrt{d_k})}{\sum_{j=1}^{n} \exp(e_{ij}/\sqrt{d_k})}$$

where $e_{ij}$ is the dot product between $q_i$ and $k_j$. The softmax function ensures two critical properties: all weights are positive, and they sum to 1. This normalization allows the model to interpret the weights as probabilities, indicating the relative importance of each token in the sequence for the current context.

Aggregation: The Weighted Sum

The final step is the weighted sum of the Value vectors. The attention output for a sequence is computed as:

$$O = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

Here, the matrix multiplication of the attention weights and the Value matrix aggregates information from the entire sequence. Tokens with higher attention weights contribute more to the output, while those with lower weights are suppressed. This operation effectively filters noise and highlights relevant context, allowing the model to dynamically focus on different parts of the input for each token.

Pitfalls

Implementing attention mechanisms correctly requires awareness of several common pitfalls:

1. Vanishing Gradients with Large Dimensions: Without proper scaling (by $\sqrt{d_k}$), large dot products cause the softmax function to saturate, leading to vanishing gradients. This stalls learning in deep models.
2. Computational Complexity: Standard self-attention has a time and space complexity of $O(n^2)$ with respect to sequence length $n$. For very long sequences, this becomes computationally prohibitive, necessitating approximations or sparse attention patterns.
3. Ignoring Positional Encoding: Since attention is permutation-invariant, it does not inherently understand the order of tokens. Failing to inject positional information (e.g., via sinusoidal encodings or learned positional embeddings) results in the model treating sequences like bags of words, losing crucial syntactic and sequential context.

Practical Takeaways

To build robust transformer-based systems, keep these mental models in mind:

• Scaling Stabilizes Gradients: Always scale by $\sqrt{d_k}$. It is not just a theoretical nicety but a practical necessity for stable gradient flow during backpropagation.
• Q/K Space Defines Similarity: The Query and Key spaces are learned projections designed to make semantic similarity computable via dot products. Treat them as distinct feature spaces, not just intermediate steps.
• V Carries Information: The Value vectors hold the actual content being passed forward. Attention weights determine how much of this content is retained, making the Value space the repository of knowledge for the current context.

FAQ

Why scale by $\sqrt{d_k}$? Scaling by the square root of the key dimension prevents the dot products from becoming too large. Large dot products cause the softmax function to output probabilities close to 0 or 1, which leads to vanishing gradients during backpropagation. Scaling keeps the values in a range where the softmax function remains sensitive and gradients are strong.

What is the role of the Value (V) vector? While Queries and Keys determine the attention weights (the "where" and "how much"), the Value vectors contain the actual information to be aggregated. The final output is a weighted sum of these Value vectors, meaning V carries the semantic content that the model retrieves based on the attention pattern.

How does attention differ from RNNs? Recurrent Neural Networks (RNNs) process sequences step-by-step, maintaining a hidden state. This creates a bottleneck for long-range dependencies and prevents parallelization. Attention, particularly self-attention, computes relationships between all tokens simultaneously (in parallel) and directly connects any two tokens regardless of distance, making it far more efficient for capturing global context.