Opening up Attention: A QK Transcoder Dashboard for Feature Pairs

Multi Head Attention in Transformers continues to be an elusive topic in mechanistic interpretability. One way of trying to decompose MHA into interpretable parts is applying SDL (Sparse Dictionary Learning) methods to the attention mechanism. However, despite the progress that has been made, we still don't have a firm grasp on attention patterns and their formation.

Additionally, there is an absence of comprehensive dashboard views for the attention mechanism and features present therein. Feature vis offers a tool to visualize and investigate features present in the residual stream and in the MLP, but neglect the attention mechanism.

To try to shed light on why/how attention patterns form and to visualize features that answer these questions, I introduce a query-key transcoder dashboard to be able to easily find and investigate monosemantic features in MHA.

The query-key transcoder architecture resembles Wynroe and Sharkey’s bilinear sparse dictionary learning approach. The primary contribution of this post is an interactive feature-pair dashboard for every layer/head in GPT-2 small (similar in spirit to Anthropic's SAE dashboard, but dedicated to query-key feature pairs), and an exploratory investigation using it. This is meant to be an interim, independent research report for investigating and visualizing the QK circuit of GPT-2.

Replacing dense queries and keys with sparse features

A sparse autoencoder learns to express a sparse set of features for models' internal activations. Transcoders work similarly, however instead of their reconstruction target being the same input activations, the output is the activations after some LLM transformation (usually MLP blocks, but in this case it's the query and key blocks). Specifically, the query-key transcoder acts as a replacement for the QK computation. They are two separate transcoders; the query transcoder learns a sparse set of features for every token and the key transcoder does the same. In practice L0 (the number of features active per token) ranges from about 10-20 features. Generally, earlier layer transcoders have L0 on the small side of this range and later layers require more features to reconstruct the attention pattern.

The 2 encoders read the LayerNormed residual stream vector, the query side encoder learns the query features an the key side the key features:

The transcoders each have 3-dimensional decoders of shape F × 12 × 64. This way, every feature has a decoder vector for each head:

One key characteristic of this architecture is feature activations are shared across all heads in a given attention layer. The three-dimensional decoder does not assign each feature to exactly one head. A feature can write to every head, and each head's decoder vector determines how much a particular feature contributes to the reconstructed attention pattern for that head.

T is sequence length; FQ and FK are the dictionary sizes.

The transcoder learns by replicating GPT 2's attention patterns and the original queries and keys are likewise used in the auxiliary training losses, as described below

The Loss Function

The main training signal from the loss function is Ljoint; the KL divergence between the original attention pattern and the reconstructed attention pattern. The reconstructed attention pattern is obtained by multiplying the query decoder activations with the key decoder activations, softmaxing, and dividing by sqrt(d_head).

The mixed loss sum of LQ and LK represents the auxiliary training losses. One problem with only having Ljoint is the query and key decoders can compensate for one another. To circumvent this issue, you implement a loss function with the original key and reconstructed query as well as the reconstructed key and original query to against the original attention pattern to ensure that both learned queries and keys remain faithful to the underlying model.

The last two terms in the loss function represent the sparsity and ghost gradients respectively. Sparsity penalizes too many features activating and maintains a low L0. Ghost gradients temporarily revive features that have been dead for too long and provide a learning signal in the direction of reconstruction error. This is the key architectural difference between my transcoder implementation and the bilinear sparse dictionary learning approach by Wynroe and Sharkey’s bilinear sparse dictionary learning approach. They cite that 80%-95% of features die without using any neuron resampling, whereas using ghost gradients to resample dead neurons, only 4%-12% of features die out.

ws(t) and wg(t) are the warmup functions. Since a warmup learning rate is implemented it also makes sense to scale the sparsity and ghost gradient functions by this same warmup to have cohesion between the scaled learning rates.

Interpreting the QK Feature Pair Dashboard

A layer-11 dashboard view, for an induction feature pair. Prompt: “I always loved visiting Aunt Sally. Whenever I was feeling sad, Aunt”. The selected edge is Aunt→ Sally in head 9.

1) Prompt and Token Selection

Simply choose a prompt and choose which token is the destination, and which to attend to. You can choose the layer/head you want to analyze. Top X refers to how many features present in the transcoders should be part of the analysis, it doesn't exclude the rest of the features from reconstruction, it only excludes them from the dashboard display. Top X is kept based on magnitude of feature activation.

2) Where does this destination attend?

This illustrates GPT-2's attention pattern and the reconstructed attention pattern for a given destination-source token pair. 'Selected pair's contribution across sources' determines how positively or negatively a feature pair contributes to all possible sources. The 'full causal attention map' portrays a grid-like view for both GPT-2's attention pattern and the reconstructed attention pattern.

3) Active Features at Selected Tokens

This is where you select a feature pair of interest. Oftentimes large feature activations are correlated with features that contribute highly to the attention pattern, but this is not always the case. The heatmap in the next step provides the computation behind how much a particular feature pair contributes to a specific attention pattern.

The query panel shows feature activations at the selected destination; the key panel shows them at the selected source. “L0” is the number of nonzero features at that token.

4) Query x Key Interactions

For a particular cell in the heatmap, the number shown represents the contribution made by that feature pair to the selected source-destination score. Before softmax, each source token and destination token pair has a total contribution by every feature pair and bias terms. The 'Exact decomposition of the selected source's score' shows each of this contributions.

How is a Feature Pair Computed?

To compute a feature pair, first define a compatibility score as the alignment between the query decoder direction and the key decoder direction for a particular head, h and feature pair i, j:

This is a fixed compatibility score. It says whether the two decoder directions align, oppose one another, or barely interact in this head. It can be positive or negative.

For a particular destination and source, the pair contributes:

This is the crux of the dashboard heatmap. It represents how closely the two features' decoder directions align as well as the magnitude of the features' activations.

For any source-destination token pair, each query feature can interact with each key feature. Although sparse activations reduce the number of relevant terms, 20 active query features and 20 active key features can still result in 400 pair terms. The heatmap offers a convenient method to distill the total number of pair terms down to the ones that contribute most to the attention pattern.

'Absolute coverage' simply refers to what percentage of total feature pair contribution is covered in the heatmap. For instance, 80% coverage means that the pairs displayed in the heatmap represent 80% of the sum of all the absolute pair terms, not including biases.

Centered Score vs Raw Score

The centered score view subtracts the pair’s average contribution across the legal sources:

S consists of the sources that the destination t can refer to with respect to the feature pair i, j. Centering asks whether this pair gives the chosen source a score that is higher or lower than the average score of the sources.

5) Selected Pair and Ablations

This panel is useful for ablating a feature pair and identifying how important a particular feature pair is to the reconstructed attention pattern. The panel shows the magnitude of query and key feature activation, the compatibility score, and raw contribution score to the source-destination token pair before softmax.

What exactly does “ablate the pair” mean?

For the default pair ablation, subtract that pair’s raw score contribution at every legal source in the selected destination row, then apply softmax again:

The dashboard shows the original reconstructed probability subtracted by the ablated probability; if the result is positive then removing the pair has decreased the attention given to this source whereas if it is negative then the removal has increased that attention.

The other controls answer different questions:

· Pair, selected edge only: eliminate the term of this pair at the specified source.

· Query feature: at the destination, remove the selected query feature, together with its interactions involving all the reconstructed key features and the key bias.

· Key feature: remove the selected key feature at the selected source, along with its interactions with all reconstructed query features and the query bias.

We can assess how useful a particular feature pair really is in reconstructing the attention pattern by removing the pair and then recalculating the attention. In order to check whether a pair helps a destination to attend to an earlier place name, you can see if removing it leads to a reduction in attention for that place, measure how much the change is, and find out where the probability goes instead. This makes ablation a helpful heuristic for testing interpretations, rather than blindly trusting high scores on the heatmap.

6. Look for the pair in other contexts

See other contexts where the feature pair also activated. Use this section to validate hypotheses.

7. Check the reconstruction

The diagnostics include KL divergence in pattern error, total variation distance, 'Top 1 source agrees', and score decomposition error.

Total variation distance or TV is given by:

For one attention row, TV = 0.04 means four percentage points of probability mass have to be redistributed to make the two distributions match. A small value is reassuring, but can still conceal a large relative error on a particular source, so this metric is not a direct substitute for manually inspecting the reconstructed attention pattern and ensuring fidelity. 'Top 1 source agrees' simply refers to whether the reconstructed attention pattern correctly placed the majority of attention mass onto the source that should be attended to the most. And score decomposition error checks whether the feature-pair explanation adds back up to the transcoder’s actual attention scores.

Case Study

Consider these two unfinished sentences:

a) Earlier, Maya washed the rice in Rome. Later, she wanted tob) Earlier, Maya washed the rice in Rome. Later, she travelled to

The final token is identical, the previous sentence is also identical, but the subsequent part makes either an action or a place relevant

Can we examine an attention head and see if it makes that kind of distinction?

a):

The feature pair Q1031 and K1195 involves earlier action cues

b):

Q1644 together with K322 constitutes a feature pair that involves place cues.

Section 6 check. Other indexed contexts include put → to, step → to, and debut → to for the action pair; geographical cues include Rapids → in and Jersey→ in.

So what?:

If you were to only inspect the attention pattern on layer 9 head 6 of GPT-2, you would see the attention mass on the attended token change from 'washed' to 'Rome', but there is no clear picture for why this change occurred. By attributing a single feature pair to source-destination pair 'to->washed' and one to 'to->Rome', we identify candidate vector directions in the query activation space and key activation space that 'cause' this change.

Note: HeadVis also provides feature-level QK attribution; however, this dashboard adopts a different method. It learns separate dictionaries of queries and keys in order to reconstruct the attention patterns, and combines the pairwise contributions along with the reconstruction diagnostics and ablations.

Limitations & Future Work

• A label is a hypothesis. Features may combine a number of behaviors, some appear to be monosemantic but many are polysemantic. In certain cases, clear examples can cause a general lexical cue that is polysemantic to appear as if it were a precise semantic concept.
• Features are unlabelled. I have yet to apply autointerp to the features in the transcoders.
• These ablations establish dependence inside the reconstruction, not GPT-2 itself. To make stronger claims about GPT-2, interventions in the original GPT-2 are needed, including downstream predictions, DLA, and vector steering, to test whether the mechanisms suggested by the reconstruction are genuinely causal. Tracing Attention Computation Through Feature Interactions is especially relevant: this too looks at attention formation, using the residual-stream features and producing attribution graphs by means of the model's original QK computation.
• QK is only part of attention. The QK circuit is only half the picture, the other half is the OV circuit. The next focus of this research direction will be to apply transcoders to the OV circuit of GPT-2 small to couple the 'which sources should be attended to given a destination' with 'what information moves from source to destination'. A dashboard that covers both these bases should provide a more cohesive picture of MHA.
• The layers are separate models. The dashboard represents a separate query-key transcoder for each of the 12 layers. Thus, superposition across layers of attention cannot be decomposed within the transcoders.

Related work

The architecture behind the transcoder setup and training regime comes from Decomposing the QK Circuit with Bilinear Sparse Dictionary Learning. The primary way transcoders has been used in the past is based on Transcoders Enable Fine-Grained Interpretable Circuit Analysis for Language Models which uses sparse replacements (transcoders) to find interpretable decompositions in MLP blocks. With regard to attention, the work Progress on Attention looks at a number of methods for interpreting attention, one of which is QK diagonalization. The techniques described in Tracing Attention Computation Through Feature Interactions allow feature paths, which are routed according to the original attention computation, to be traced through the model. The main difference between a query-key transcoder and this research is that Anthropic's implementation obtains the residual-stream features by means of weakly causal crosscoders, while my approach involves learning separate feature dictionaries for the query and the key specifically in order to reconstruct attention patterns. This means that training is encouraged to focus on the information that is useful for forming attention patterns, whereas in the case of residual-stream reconstruction the target is more general. The dashboard is inspired by the ARENA SAE interpretability notebook.

All WandB training runs from qk transcoders trained on layers 0-11 of GPT-2.