Interpretability in Parameter Space: Minimizing Mechanistic Description Length with Attribution-based Parameter Decomposition

TL;DR

This work reframes mechanistic interpretability as a problem in parameter space by introducing Attribution-based Parameter Decomposition (APD), which decomposes a network’s parameter vector θ* into a sum of parameter components that are faithful to the original network, minimally used for any input, and individually simple. APD optimizes three losses—faithfulness (component sum recovers θ*), minimality (top-k attributed components sufficing per input), and simplicity (low-rank, sparse components)—using gradient attributions to guide component activation. Demonstrations on toy models show APD recovering ground-truth mechanisms in superposition, decomposing compressed computations into per-feature components, and revealing cross-layer distributed representations, thereby supporting an architecture-agnostic view of neural mechanisms. While scalability to large models remains a challenge, APD lays groundwork for improved interpretability, robustness, and potential applications in safe model editing and mechanistic analysis across architectures.

Abstract

Mechanistic interpretability aims to understand the internal mechanisms learned by neural networks. Despite recent progress toward this goal, it remains unclear how best to decompose neural network parameters into mechanistic components. We introduce Attribution-based Parameter Decomposition (APD), a method that directly decomposes a neural network's parameters into components that (i) are faithful to the parameters of the original network, (ii) require a minimal number of components to process any input, and (iii) are maximally simple. Our approach thus optimizes for a minimal length description of the network's mechanisms. We demonstrate APD's effectiveness by successfully identifying ground truth mechanisms in multiple toy experimental settings: Recovering features from superposition; separating compressed computations; and identifying cross-layer distributed representations. While challenges remain to scaling APD to non-toy models, our results suggest solutions to several open problems in mechanistic interpretability, including identifying minimal circuits in superposition, offering a conceptual foundation for 'features', and providing an architecture-agnostic framework for neural network decomposition.

Figures (25)

• Figure 1: Decomposing a target network's parameters into parameter components that are faithful, minimal, and simple.
• Figure 2: Top: Step 1: Calculating parameter component attributions $A_c (x)$. Bottom: Step 2: Optimizing minimality loss $\mathcal{L}_{\text{minimality}}$.
• Figure 3: Results of running APD on TMS. Top row: Plot of the columns of the weight matrix of the target model, the sum of the APD parameter components, and each individual parameter component. Each parameter component corresponds to one mechanism, which in this model each correspond to one 'feature' in activation space elhage2022toy. Bottom row: Depiction of the corresponding parametrized networks.
• Figure 4: Decomposing TMS with APD.
• Figure 5: The architecture of our Toy Model of Compressed Computation using a $1$-layer residual MLP. We fix $W_E$ to be a randomly generated matrix with unit norm rows, and $W_U={W_E}^\top$.