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 architecture closely follows the first stage of 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

The query-key transcoder functions as a learned substitute for the QK computation; it learns sparse queries and keys of its own from a small number of active features at each token. 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.

Two distinct encoders read the LayerNormed residual stream vector, one generating the query-feature activations, the other producing the key-feature activations:

Every query feature has a decoder vector for each head; when the active decoder vectors are added together the query is reconstructed, and the key side operates in the same manner:

In this case, a head is selected and the two c terms are decoder biases. The query and key dictionaries are separate, so there is no relationship between the same feature index in the query and key transcoders; in other words, query feature 42 has no special connection with key feature 42. This represents a general bilinear decomposition; it differs from the query–key interaction at the same index that was used in the rank-one "QK diagonalization" experiments mentioned in Progress on Attention.

T is sequence length; F_Q and F_K are the dictionary sizes.

The teacher is GPT-2 frozen. The transcoder learns by replicating its attention patterns and the original queries and keys are likewise used in the auxiliary training losses, as described below

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.

What training asks the features to preserve

The primary focus is on the original attention pattern, and the loss function imposes a penalty on the difference in attention pattern between GPT-2 and the transcoders by means of KL divergence.

At the tth optimizer step t, ws(t) and wg(t) are the configured multipliers for sparsity and for the ghost gradient warmup respectively. The KL term sums over the allowed sources and then averages over the valid destinations and heads. The sparsity term sums over the features and then averages over the valid tokens.

Simply matching the pattern still allows the learned queries and keys to make up for one another. That is why the loss function also makes use of "mixed" reconstructions: combinations of learned queries with the original keys, and of the original queries with learned keys. Each of these has to predict the original attention pattern.

Two more terms are sparsity and ghost gradients. The sparsity penalty makes it costly to use a large number of features. Ghost gradients provide a temporary training path for features that are persistently inactive: even though their normal ReLU activation is zero, they can receive a learning signal aimed at reducing attention-pattern 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.

Reading the 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. Choose a prompt, destination, source, and head

For each layer I have a separate dashboard view and then I choose a head. The 'Top' setting determines how many features or pairs are displayed, without removing the other features from the reconstruction.

2. See where the destination attends

The row labelled 'Attention' shows the original GPT-2 probabilities compared with the reconstructed probabilities; the full causal attention map has the destinations on one axis and the sources on the other, and its causal triangle is the set of permitted connections.

3. Inspect the active features

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.

A large feature activation only shows that the encoder of the transcoder decided that a particular feature is present in the residual stream, but its total contribution to the attention pattern depends on the feature's alignment with that transcoder head's decoder vector as well as other sources competing for attention.

4. Inspect the feature pair heatmap and score breakdown

Each cell in the heatmap shows the contribution made by a particular feature pair to the selected source–destination score, and the display can show either the raw contributions or the centered ones. The score breakdown includes the pair of interest together with all the other pairs and the bias terms; they are added up before the softmax operation, which is why they don't add up to 1.

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:

It is the key quantity in the dashboard heatmap and calculates the raw score for the attention pattern before applying the softmax function; it shows the level of activity of a destination feature, the level of activity of a source feature, and the strength of the interaction between their write (decoder) directions.

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.

The term "absolute coverage" refers to the proportion of the total magnitude of the absolute pair-score that is shown 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, excluding the biases.

Raw score versus centered score: a score only matters relative to alternatives

If a feature pair increases by 3 (that is, the raw pre-softmax score) the score of every source then the attention probabilities remain the same after applying the softmax function. It is for this reason that a large raw contribution has no effect on determining which source is most attended to when compared to the other source positions.

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. Inspect one pair and ablate it

The panel for the selected pair displays both its activations, compatibility, contribution, and compatibility across heads; the head profile is useful since the same active pair can raise raw scores in one head while lowering them in another. After a specified removal the ablation controls then recompute the selected reconstructed attention row.

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

This is where you can test a proposed interpretation. When the corpus search finds causal source–destination edges on which both features are activated, it will display examples that are highly positive, highly negative, and typical ones.

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. The panel also checks that the reconstructed attention probabilities sum to one.

Case Study

Consider these two unfinished sentences:

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

b) 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.

For the term 'wanted to', Q1031 shows an activation level of 13.69, but for 'travelled to' it drops to 1.87. Q1644, on the other hand, goes from zero to 7.82. Since the initial prefix is the same and the model is causal, the activations associated with the earlier part remain constant. The differences for these particular raw pair terms are therefore due to the query side.

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.

Two shorter examples: punctuation and suppression

Quotes at a later sentence boundary. For Alice wrote, "The parcel arrived." Bob smiled., layer 9, head 11, Q1573 × K562 supports the opening quote at the final period. GPT-2 assigns a probability of 36.84% to that quote while the reconstruction gives it 40.88%; when the pair is removed the figure drops to 25.69%. The query feature is zero whenever the prompt ends at "smiled" without the final period. However, removing the closing quote after "arrived" does not switch off the feature. Section 6 also identifies cases of quote-to-period. This serves as evidence of a punctuation-conditioned interaction between quotes and may indicate a mechanism for determining which quotation is still open.

A pair can act as a brake. For The report (written by Nora in Paris), in layer 9, head 3, for Q1902 × K2032 the closing parenthesis has a negative score contribution of −1.174; its reconstructed self-attention is 10.26%, and when the pair is removed this rises to 27.01%. GPT-2's value in this case is 8.09%. This shows the effect of ablation: when a suppressive interaction is removed, attention can increase.

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.

Appendix

Counts, densities, and display coverage

The L0 count for a single-token query includes only the positive activations; similarly, query density (which can be found under the raw pair-score distribution in section 6 of the dashboard) asks how frequently a given feature is active among the tokens being searched, and the key side is analogous.

The number of pair density edges is calculated by dividing the count of edges where both the destination query feature and the source key feature are positive by the total number of eligible causal edges. The implementation in the corpus includes self-edges and source entries from the beginning token, but excludes the trivial destination associated with the beginning token. When the window stored has a length of L and includes the beginning token, there are L(L+1)/2 − 1 eligible edges.

For the score histogram the number of bins is set at 32 for all coactive edges before the decision is made as to which of the positive, negative or typical display tabs to use. A distinct-document filter alters the examples that are displayed, not the histogram underneath. The sampling details are important when determining how common a feature's behaviour is.

If V is the set of displayed pairs and c is whichever raw or centered score view is selected, absolute coverage (given in section 4 of the dashboard) is:

The values of the decoder norms shown in the head profile indicate the extent to which a feature's write vector lies within that head. Compatibility also depends on the angle between the two vectors: if the vectors are parallel their dot product is positive, if they are in opposition the dot product is negative, and if they are at right angles the dot product is zero.

The complete decomposition (Dashboard section 4) score includes biases

The terms that involve the key bias remain constant for a given destination across different sources and therefore disappear in the softmax calculation. The interaction between the query bias and the source key features, however, does vary from source to source and has to be kept. In this implementation the bias of the query decoder is learned while the bias of the key decoder is a fixed buffer derived from the teacher. The compatibility matrix has dimensions FQ by FK for a head, but it is made up of decoder vectors that are 64 in width; its rank is therefore no more than 64.

The pattern losses

Let P(Q,K) mean scaled dot products, causal masking, and row-wise softmax. Let Q and K without hats denote the frozen teacher’s vectors. Averaging over valid destinations and heads:

KL(p‖q) is the sum taken over the sources of p log(p/q). In the implementation numerical safeguards are used in relation to zero probabilities.

For the query sparsity term, the implementation uses a smoothed square-root penalty:

and similarly for S_K. This is an L0.5-style activation penalty, not an L1 penalty and not the dashboard’s count of active features. The decoder feature directions are normalized with respect to their combined head and width dimensions in order to control the activation/decoder rescaling ambiguity. With scheduled coefficients absorbed into the lambdas, the objective is:

Why ablation effects do not sum

When dealing with one destination row, let ps represent the original reconstructed probabilities and Cs the contributions of the removed pair. The new probabilities are then:

The denominator is different for each source. That is the reason why removing the score from one source can increase the probability of another source, and why an attention-probability change is not a separate portion of the original probability.

WandB Runs:

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