QK Circuit

The QK circuit's primary function is to compute the raw attention logits. For a destination token d and source token s, the score is calculated as:

• Formula: (x_d * W_Q) · (x_s * W_K)^T
• This dot product measures the semantic relevance of the source token to the destination token's query.
• The resulting scores are scaled by 1/√d_k and passed through a softmax function to produce the final attention weights.
• This mechanism allows the head to implement behaviors like copying, attending to syntax, or focusing on specific entities.

The QK interaction is a bilinear form governed by the matrix W_QK = W_Q^T * W_K. Key properties include:

• Rank Analysis: The rank of W_QK is bounded by min(d_head, d_model). A low effective rank indicates the head is only sensitive to a few feature directions.
• Singular Value Decomposition (SVD): Decomposing W_QK reveals the principal query-key directions that dominate the attention pattern.
• Split-Head View: The circuit can be analyzed by examining the W_Q and W_K matrices directly, or by studying the composite W_QK matrix for the head's overall matching function.

The QK circuit integrates positional information to control which relative distances the head attends to:

• Rotary Position Embedding (RoPE): Applies a rotation to the query and key vectors based on their absolute positions. The dot product then depends only on the relative position d - s.
• Absolute Positional Encodings: Added to the input embeddings before the QK projection, allowing the head to learn to attend to specific absolute locations.
• Attention Sinks: The QK circuit often assigns very high scores to the initial tokens (e.g., the BOS token), creating an 'attention sink' that absorbs excess probability mass.

In decoder-only transformers, the QK circuit operates under a strict causal constraint:

• Causal Mask: A mask sets all attention scores for s > d to negative infinity before the softmax, ensuring the model cannot attend to future tokens.
• Information Bottleneck: The QK circuit is the gatekeeper of information flow. It can only route information from the past, making its learned patterns critical for in-context learning and sequence modeling.
• Efficient Implementation: The causal mask is often fused with the softmax kernel in optimized implementations like FlashAttention to reduce memory I/O.

Mechanistic interpretability uses causal interventions to understand QK circuits:

• Activation Patching: Replacing the QK attention scores of a specific head with scores from a corrupted or different forward pass to isolate its functional role.
• Path Patching: A refined technique that isolates the direct effect of the QK circuit from one layer to another, bypassing intermediate computations.
• Causal Scrubbing: Formally tests hypotheses about a QK circuit by resampling activations and verifying if the circuit's proposed function is faithful to the model's behavior under all resampled conditions.

The QK circuit's design directly impacts inference efficiency through the KV cache:

• KV Cache: During autoregressive generation, the key vectors for all previous tokens are cached to avoid recomputation. The QK circuit only computes the query for the new token and scores it against the cached keys.
• Grouped Query Attention (GQA): Uses fewer key-value heads than query heads. Multiple query heads share a single key head, reducing the KV cache size while preserving most of the quality.
• Multi-Query Attention (MQA): An extreme case of GQA where all query heads share a single key-value head, maximizing cache reduction at a small cost to model capability.

The computational pathway formed by the output (O) and value (V) projection matrices. While the QK circuit determines where to attend, the OV circuit determines what information is copied from the source token to the destination token's residual stream.

• Composed of the V matrix (which projects the source token's residual stream into a value vector) and the O matrix (which projects the weighted sum of values back to the residual stream dimension)
• The OV circuit can be analyzed via Singular Value Decomposition (SVD) to reveal the principal semantic directions an attention head reads and writes

A specialized computational subunit within a transformer layer that computes a weighted sum of past token representations. Each head independently performs:

• QK computation: Calculates attention scores to determine which previous tokens are relevant
• OV computation: Copies and transforms information from attended-to tokens
• The residual stream is partitioned across multiple heads, each operating in a lower-dimensional subspace, allowing the model to simultaneously attend to different syntactic and semantic patterns

A specific type of attention head discovered through mechanistic analysis that performs in-context copying. It attends to the token immediately following a previous occurrence of the current token.

• Composed of two collaborating heads in successive layers: a previous token head that attends to the prior token, and an induction head that uses that information to attend to the token after a pattern match
• Critical for in-context learning and sequence completion tasks like [A][B] ... [A] → [B]

A position encoding method that applies a rotation matrix to query and key vectors based on their absolute positions. The key property: the dot product between a query at position m and a key at position n depends only on the relative distance (m - n).

• Encodes position information directly into the attention computation without adding positional vectors to the residual stream
• Enables the QK circuit to naturally implement relative attention patterns like 'attend to the token 5 positions back'

A causal intervention technique that replaces a model's internal activation at a specific layer and token position with a cached activation from a different forward pass. This isolates which components are causally responsible for a behavior.

• Clean vs. corrupted runs: Patch a clean activation into a corrupted run (or vice versa) and measure the change in output
• Used to trace the QK circuit by patching attention scores and observing which token relationships are necessary for correct predictions

An attention mechanism that uses a smaller number of key-value head groups than query heads. For example, 8 query heads might share only 2 key-value head groups.

• Interpolates between Multi-Head Attention (MHA) and Multi-Query Attention (MQA)
• Reduces the KV cache size during autoregressive generation while preserving most of the quality of full MHA
• The QK circuit operates with shared key projections across multiple query heads within each group