Decomposition of Small Transformer Models

AI summary

Overview

• Research area: Mechanistic interpretability, specifically parameter-space decomposition of neural networks (an alternative to activation-space methods such as sparse dictionary learning).
• Technical level: Advanced. The paper assumes familiarity with concepts such as polysemanticity, superposition, sparse autoencoders, rank-1 matrix factorisation, induction heads (a type of attention circuit), and KL-divergence-based reconstruction losses.
• Scope: The paper extends Stochastic Parameter Decomposition (SPD) from toy models to Transformer architectures, demonstrating it on a 2-layer attention-only induction-head model and on GPT-2-small.

What This Paper Is About

Mechanistic interpretability has largely advanced by studying activations — asking what fired on a given input. The authors argue this leaves no canonical way to factor a model itself into a small set of reusable mechanisms. Building on prior parameter-space work (Attribution-based Parameter Decomposition by Braun et al., and Stochastic Parameter Decomposition by Bushnaq et al.), they adapt SPD to Transformers, where the previous "ablatability" measure of causal importance breaks down because identical token representations at different sequence positions can carry very different importance.

Key Contributions

1. A causal importance function for sequential data. A minimal attention network (1 head, 1 layer, QKV-only) is applied across tokens before computing causal importance, using learned relative positional encodings and learned absolute positional encodings on value vectors, so the downstream γ-MLP is positionally aware.
2. Decomposition of a toy induction-head model. SPD recovers the expected 2-step circuit in a 2-layer, 1-head-per-layer attention-only Transformer.
3. New losses for faithful decomposition over few examples. A layer-wise stochastic reconstruction loss and randomised partial reconstruction losses (which evaluate only a subset of decomposed layers at a time) prevent the decomposed model from "cheating" by writing new computation into unused components.
4. Decomposition of GPT-2-small. SPD locates subcomponents corresponding to interpretable concepts such as "golf" and "basketball," and suppressing them selectively reduces the targeted prediction probability.

Main Findings

• Induction-head circuit recovered: For the toy model, subcomponents are assigned largely to the relevant positions. In layer 0, K, V and Q importances concentrate on the first s-token (values of 1.000), while the m-token position gets a Q importance of 1.000. In layer 1, Q and K concentrate on the final s-token and K on m (both 1.000), with V showing 5.053 active subcomponents at m's position.
• Sparsity of the recovered circuit: Total unique subcomponents per matrix are 1 for Q0, K0, V0, Q1, and K1, and 11 for V1.
• Near-perfect faithfulness: Faithfulness loss is 3e-9, reconstruction loss is 1e-4, and layer-wise reconstruction loss is 1e-4. Attention KL divergences are 0.385 (layer 1), 0.002 (layer 2), and 0.194 (mean).
• Positional information matters: Removing positional encodings before projecting m through the Q-matrix drops attention strength from m to s1 from 1.0 to 0.0941, even when s1's key retains positional information.
• Higher rank needed at the final position: V1 has considerably more active subcomponents for m's position; the authors attribute this to the need for m's identity among 128 tokens to be unambiguously separable in the residual stream.
• GPT-2-small compresses heavily: For the two example sentences ("Kobe Bryant plays the sport of basketball" and "Tiger Woods plays the sport of golf"), 96 active subcomponents are found in total — a 99% size reduction compared with the whole model.
• Specific factual directions located: 2 subcomponents uniquely activate in the layer 0 MLP for tokens "obe" and "Bryant", and 1 uniquely for "Woods"; suppressing their directions reduces the probability of the correct sport.
• Both names needed for Kobe Bryant: Maximally reducing output probability requires suppressing a rank-1 direction on both the first and last name, because "Kobe" carries a strong basketball association on its own whereas "Tiger" generally does not refer to a name.
• Asymmetric knowledge storage: Suppressing only the athlete→sport direction does not prevent the reverse prompt "The most famous athlete in golf is.." from correctly answering "Tiger Woods," consistent with prior work finding knowledge stored asymmetrically.
• Facts may originate earlier than previously suggested: Causal tracing (Meng et al.) indicates the fact is editable from layer 0 MLPs through layer 5 MLPs, but zeroing out entire layers shows the fact does not live where editing is easiest; the authors conclude some facts originate as early as the first MLP layer and propagate through the network.
• SPD beats greedy SVD on minimality: On the simple toy model, SPD finds a more minimal decomposition that matches the full model's output distribution than a greedy SVD low-rank approximation.

Methodology in Plain English

SPD takes a trained network's weight matrices and factorises each into many rank-1 subcomponents (the number C is a hyperparameter). For a given input, each subcomponent gets an "activity" value between 0 and 1, and the weighted sum of subcomponents reassembles an approximation of the original weight matrix. The decomposition is trained under three pressures: it must sum back to the original weights (faithfulness), it must reproduce the original model's output, and as few subcomponents as possible should be active (minimality).

Because the original model is exactly recovered when every activity equals 1, the training can always fall back on the real model. To decide which subcomponents matter on a given input, the method learns a causal importance score for each one, sampled stochastically so that subcomponents whose importance has fallen to zero can still be switched on and receive gradient signal. Reconstruction is enforced both across the whole model and layer-by-layer.

The key novelty for sequences is that importance is computed with a tiny attention network over the tokens rather than per-token independently. Consider "Sat by the river bank, the bank manager" — the two instances of "bank" would otherwise be forced to share one importance score. For the GPT-2 experiments, the authors use only two sentences but add randomised partial reconstruction losses that replace only a random subset of layers with decomposed weights, leaving the rest of the original model intact. This removes the decomposed model's ability to invent new shortcut mechanisms in layers it cannot rely on. Afterwards, they test the recovered subcomponents by suppressing their directions in the original model via orthogonal projection, confirming the effect also holds on a separate 100-example synthetic fact dataset.

Why This Matters

Impact on research. The paper is an explicit attempt to close the gap between parameter-space decomposition results on "toy models" and "real models." It shows that the machinery can be carried over to Transformers and that the resulting subcomponents act as narrow, causally relevant directions rather than broad capacity. It also introduces losses that make few-example decomposition honest rather than degenerate, which opens a path to asking targeted questions about specific model behaviours without assembling large datasets.

Potential applications (the paper does not report deployed applications; these follow from its stated framing of parameter-space objects as "clean handles for mechanistic interventions"):

• Targeted model editing. Suppressing or modifying specific rank-1 directions to change a factual association, as the authors do with the athlete→sport direction.
• Knowledge attribution and auditing. Determining where in a network a given fact or concept originates, which the paper argues may be earlier (layer 0 MLP) than causal-tracing edits suggest.
• Safety and debugging. Identifying and isolating narrow subcomponents responsible for particular behaviours, so they can be inspected or intervened on without degrading performance on unrelated inputs.
• Model comparison and compression research. Using minimality as a measure of how concisely a model's computation can be expressed, and contrasting against low-rank approximation baselines such as greedy SVD.

Industry relevance. Any organisation deploying or auditing language models has an interest in tools that locate the internal mechanism behind a specific output. The paper's key commercial caveat is that its evaluations are qualitative and limited to small-scale models, so significant scaling work remains before such methods are operational at production model sizes. The authors also note a practical limitation for sparsity comparisons: if a QK-circuit ends up with zero active subcomponents, that circuit is still used but attention becomes uniform, so they explicitly do not compare sparsity or size against pruning methods.

Future Directions

1. Scale to larger models. The experiments are restricted to small-scale models, and the authors state that scalability to larger architectures remains an open question.
2. Reduce computational overhead. The causal importance parameterisation adds cost; Table 2 shows iterations per second ranging from 12 (scalar input, sequence length 16) down to 1 (non-flex attention, sequence length 10240 with batch size reduced from 64 to 4).
3. Develop proper benchmarks. Evaluation is currently qualitative and task-specific, so general benchmarks across architectures and tasks are needed to establish generality.
4. Improve the definition of causal importance for feature interactions. The authors note ablations are not truly causal because they are subject to interactions within the model, and they explicitly plan to refine causal importance for computations involving feature interactions.

Target Audience

Mechanistic interpretability researchers and interpretability-focused engineers already familiar with activation-space methods (sparse dictionary learning, causal tracing, circuit discovery) who want to understand the parameter-space alternative. It is also relevant to researchers working on model editing and knowledge localisation, since the GPT-2-small results speak directly to where facts are stored and how they propagate. Readers seeking a gentle introduction to interpretability will find it conceptually demanding, and readers seeking production-ready tooling will find the scope deliberately preliminary.

Authors’ abstract

Recent work in mechanistic interpretability has shown that decomposing models in parameter space may yield clean handles for analysis and intervention. Previous methods have demonstrated successful applications on a wide range of toy models, but the gap to "real models" has not yet been bridged. In this work, we extend Stochastic Parameter Decomposition (SPD) to Transformer models, proposing an updated causal importance function suited for sequential data and a new loss function. We demonstrate that SPD can successfully decompose a toy induction-head model and recover the expected 2-step circuit. We also show that applying SPD to GPT-2-small can successfully locate subcomponents corresponding to interpretable concepts like "golf" and "basketball". These results take the first step in the direction of extending SPD to modern models, and show that we can use the method to surface interpretable parameter-space mechanisms.