Attention Head

The QK circuit is the computational pathway formed by the query (W_Q) and key (W_K) projection matrices. It determines where an attention head looks by computing the attention score between a destination token's query vector and every source token's key vector.

• Bilinear Form: The attention score is computed as q^T * k, a bilinear function of the source and destination token representations.
• Positional Influence: In models using Rotary Position Embedding (RoPE), the QK circuit incorporates relative position information directly into the dot product, enabling heads to learn distance-based attention patterns.
• Pattern Types: QK circuits can learn to implement previous-token heads (attending to the immediately preceding token), duplicate-token heads (attending to identical tokens), or more complex semantic matching patterns.
• Composition: The effective QK matrix is the product W_Q^T * W_K, which can be analyzed via Singular Value Decomposition (SVD) to understand the principal directions of attention.

The OV circuit is the computational pathway formed by the output (W_O) and value (W_V) projection matrices. It determines what information is copied from the source token to the destination token's residual stream.

• Linear Transformation: The OV circuit applies a linear map to the source token's residual stream state, extracting and transforming the information that will be written to the destination.
• Effective Matrix: The combined operation is captured by the matrix W_V * W_O, which can be decomposed to understand what features the head reads and writes.
• Copying vs. Transformation: Some heads act as near-identity OV circuits that copy information verbatim, while others apply significant transformations to extract specific features.
• Subspace Writing: Each head writes to a low-dimensional subspace of the residual stream. The SVD of the OV matrix reveals the rank and principal directions of this write operation.

Each attention head in a multi-head attention layer operates independently on its own low-dimensional subspace, enabling parallel computation of distinct attention patterns.

• Dimensionality Split: If the model dimension is d_model and there are H heads, each head operates on a subspace of dimension d_head = d_model / H. This forces heads to specialize.
• Independent Parameters: Each head has its own distinct W_Q, W_K, W_V, and W_O matrices, allowing it to learn a completely different attention function.
• Concatenation: The outputs of all heads are concatenated and projected back to the full model dimension by a final output projection matrix.
• Specialization Evidence: Mechanistic analysis reveals heads specializing in roles like previous-token attention, induction (in-context copying), and inhibition (suppressing specific tokens).

Through circuit analysis and activation patching, researchers have identified distinct functional classes of attention heads that emerge during training.

• Previous-Token Heads: Attend almost exclusively to the immediately preceding token. These heads are crucial for propagating information locally and are often found in early layers.
• Induction Heads: A critical mechanism for in-context learning. An induction head attends to the token following a previous occurrence of the current token pattern, enabling the model to repeat sequences.
• Duplicate-Token Heads: Attend to positions where the same token appears earlier in the sequence, useful for coreference resolution and copying.
• Inhibition Heads: Actively suppress attention to certain tokens, often implementing negative attention scores to prevent the model from attending to irrelevant or distracting information.
• Attention Sinks: Many heads allocate a large fraction of their attention mass to the first few tokens, which act as a resting place for excess attention probability.

Attention heads do not operate in isolation; they read from and write to the shared residual stream, enabling complex compositional behavior across layers.

• Reading: The input to a head's Q, K, and V projections is the residual stream state at the destination token's position, which contains the accumulated outputs of all previous layers.
• Writing: The head's output is added back into the residual stream at the destination position. This additive structure allows heads to compose linearly.
• Composition Across Layers: A head in layer L can read information written by a head in layer L-1, forming cross-layer circuits. For example, a previous-token head in an early layer can pass positional information to an induction head in a later layer.
• Logit Lens: Applying the unembedding matrix directly to the residual stream after a head's write operation reveals how that head shifts the model's next-token prediction probabilities.

The functional role of an attention head is validated through causal intervention techniques that surgically remove or modify its output and measure the effect on model behavior.

• Zero Ablation: Setting a head's output to zero to measure its necessity for a specific task. This can reveal redundancy due to self-repair.
• Mean Ablation: Replacing a head's output with its average activation over a distribution of inputs, removing task-specific information while preserving statistical structure.
• Activation Patching: Replacing a head's activation at a specific token position with a cached activation from a different forward pass. This isolates the head's contribution to a specific behavior.
• Causal Scrubbing: A formal hypothesis-testing framework that systematically resamples activations to verify if a proposed circuit faithfully explains the model's behavior under all relevant input distributions.