PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Dan Braun, Lucius Bushnaq, Stefan Heimersheim, Jake Mendel, Lee Sharkey

arXiv:2501.14926v4cs.LGstat.ML

TL;DR

Mechanistic interpretability lacks a clear way to decompose neural-network parameters into mechanisms, despite the goal of understanding learned internals. APD decomposes parameters into components optimized for faithfulness, minimality, and simplicity, and recovers ground-truth mechanisms across three toy settings.

  • Problem

    It remains unclear how to decompose neural-network parameters into mechanisms, which complicates efforts to understand network internals.

    Source: §1 · p. 3

  • Method

    APD decomposes network parameters into components optimized for faithfulness, minimality, and simplicity, using input-specific attributions to identify mechanisms used for each input.

    Source: §1 · pp. 3–4

  • Results

    APD recovers ground-truth mechanisms in toy models of superposition, compressed computation, and cross-layer distributed representations.

    Source: §1 · p. 5, §5 · p. 20

  • Takeaways & Limitations

    The toy experiments support parameter-space decomposition as a route to identifying minimal circuits in superposition and conceptualizing features.

    Source: §5 · p. 20

    The initial method may have worst-case memory cost scaling like O(N^2), and APD often required extensive hyperparameter tuning in these experiments.

    Source: §4.2.1 · p. 18, §4.2.2 · p. 19

Abstract

from arXiv · show

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.

1 Introduction

The paper addresses the difficulty of identifying neural-network mechanisms when they do not align with individual architectural components, and asks what makes a parameter decomposition useful for interpretability. It introduces APD to optimize parameter components for faithfulness, minimality, and simplicity, and reports recovery of ground-truth mechanisms in three toy settings.

1 Introduction

  • Sparse dictionary learning decomposes hidden activations into sparse directions, but its limitations leave unclear how those directions reveal underlying mechanisms.The cited concerns include reconstruction errors, uncertain interpretability of sparsity, and unexplained feature geometry.Source: §1 · p. 3
  • Neuron-by-neuron descriptions can be accurate yet unnecessarily long and polysemantic, failing to reflect deeper mechanistic structure.Source: §1 · p. 3
  • The proposed criteria are parameter faithfulness, minimality on the training distribution, and simplicity of each component.The paper motivates these criteria through minimum description length; together they are intended to capture distinct, basic computational units.Source: §1 · p. 3
  • APD decomposes network parameters into components optimized for the three criteria, seeking a shorter description of mechanisms.Its parameter-faithfulness criterion is stricter than behavioral faithfulness because different parameters can exhibit the same behavior.Source: §1 · pp. 3–4
  • The approach assumes that any given input need not use all network mechanisms, so unused mechanisms can be ablated without affecting computation.The paper illustrates this with stored knowledge implemented across multiple neurons and layers.Source: §1 · p. 4
  • APD recovers ground-truth mechanisms in toy models of superposition, compressed computation, and cross-layer distributed representations.The reported cases cover individual input features, individual functions, and functions spanning multiple layers.Source: §1 · p. 5

2 Method: Attribution-based Parameter Decomposition

The method treats mechanisms as parameter-space components and optimizes their decomposition for faithfulness, minimality, and simplicity. APD estimates component attributions, trains a small active set to reproduce target behavior, and penalizes unnecessary component complexity.

2.1 Defining ‘mechanism space’ as parameter space

  • Parameter space can represent subcomputations, mechanisms distributed across architectural components, and computations on features in superposition.The paper also states that parameter-space vectors can support multidimensional computations.Source: §2.1 · p. 5

2 Method: Attribution-based Parameter Decomposition

  • APD directly decomposes network parameters into components, aiming to minimize the description length of mechanisms used per training data point.The method is presented as optimizing the components for faithfulness, minimality, and simplicity.Source: §2 · p. 5, §2.2 · p. 6

2.2 Identifying networks’ mechanisms using Attribution-based Parameter Decomposition

  • Experiments use batch top-k to select a fixed number of active components per batch, avoiding per-sample selection but introducing other issues.The attribution calculation uses one target-model forward pass and one backward pass per output dimension.Source: §2.2 · p. 6
  • Biases are not decomposed; they are treated as their own component active for every input, with their decomposition left for future work.Source: §2.2 · p. 7
  • For faithfulness, APD trains the sum of parameter components to approximate the target parameters; for minimality, it estimates component attributions and trains the top-k attributed components to match target outputs.The attribution step approximates causal importance with gradient attributions; in some cases, hidden activations are also matched to constrain computations.Source: §2.2 · pp. 6–7
  • The simplicity objective penalizes components spanning more ranks or layers than necessary, using an Lp norm of singular values in practice.Source: §2.2 · pp. 6, 8

3 Experiments: Decomposing neural networks into mechanisms using APD

Across three toy settings with known ground-truth mechanisms, APD identifies components that match those mechanisms up to small error. The experiments examine superposition, compressed computation, and representations distributed across layers.

3 Experiments: Decomposing neural networks into mechanisms using APD

  • Across all three toy settings, APD identifies ground-truth mechanisms up to small error.The settings are superposition, compressed computation, and cross-layer distributed representations.Source: §3 · p. 8

APD Results: Toy Model of Superposition

  • In the superposition model, APD learns components corresponding to the ground-truth matrices whose sum equals the target weight matrix.MMCS is approximately 1 in both tested settings, while ML2R is close to 0.9, indicating some shrinkage.Source: APD Results: Toy Model of Superposition · p. 9
  • The superposition components imply an overcomplete basis in activation space, with forward-pass reading directions aligned to activations and backward-pass projecting directions aligned to gradients.The parameter vectors themselves do not form an overcomplete basis for parameter space.Source: APD Results: Toy Model of Superposition · p. 11

APD Results: Toy Model of Compressed Computation

  • With 100 input features and 50 neurons, APD learns components that implement the target computation for individual input dimensions.The target computes yi = xi + ReLU(xi), including more functions than it has neurons.Source: APD Results: Toy Model of Compressed… · p. 12, §3.2 · p. 12
  • Ablating unrelated components has little effect on output or MSE, whereas ablating components assigned to active computations causes a large MSE increase.The MSE comparison is reported across scrubbed and anti-scrubbed runs on the training distribution.Source: APD Results: Toy Model of Compressed… · pp. 13–14, Figure 8 · p. 14
  • Compressed-computation components can partially represent secondary features, producing bimodal scrubbed and anti-scrubbed losses.The authors suggest this may improve with hyperparameter or training changes, alternative losses, or rank-1 constraints.Source: APD Results: Toy Model of Compressed… · p. 14

APD Results: Toy Model of Cross-Layer Distributed Representations

  • In the two-layer setting, APD learns components using neurons from both MLP layers, but more components represent two feature computations rather than one.Batch top-k can also leave some feature computations distributed across multiple components.Source: APD Results: Toy Model of Cross-Layer… · p. 15

4 Discussion

The discussion presents APD as a parameters-first approach that may address conceptual challenges in activations-first interpretability, while emphasizing practical limitations and future work.

4 Discussion

  • APD directly decomposes parameters into faithful, minimal, and simple mechanistic components, taking a parameters-first approach.The discussion contrasts this with approaches that first decompose activation directions and then build circuits from them.Source: §4 · p. 15

4.1 Addressing issues that seem challenging from an activations-first perspective

  • APD treats components as circuits and directly targets minimal circuits in superposition, while offering a route to address several limits of activations-first methods.The discussion also argues that faithfulness constrains feature splitting, components are algorithmic steps rather than input representations, and parameter-space analysis may illuminate SAE feature geometry and distributed representations across architectures.Source: §4.1 · pp. 16–18

4.2 Next steps: Where our work fits in an overall interpretability research and safety agenda

  • The paper’s next steps are to make APD-like methods more robust and scalable, while leaving intrinsically decomposable architectures for future work.The authors say APD is currently appropriate only for toy models because of computational cost and hyperparameter sensitivity.Source: §4.2 · p. 18

4.2.1 Improving computational cost

  • O(N^2) worst-case scaling is possible because each component has memory costs similar to the target network, though lower-rank, localized, or layerwise approaches may reduce costs.The authors suggest constraining components to lower rank and fewer layers, or applying the method one layer at a time.Source: §4.2.1 · pp. 18–19

4.2.2 Improving robustness to hyperparameters

  • Extensive hyperparameter tuning was often needed, although sensitivity issues were less frequent in larger toy models and may reflect interference noise in small dimensions.Improving robustness to hyperparameters is identified as a high priority for future work.Source: §4.2.2 · p. 19

4.2.3 Improving attributions

  • APD’s gradient attributions are first-order approximations to causal ablations and may become unreliable when parameter gradients are small.The authors suggest integrated gradients, learned masks, or other attribution methods as possible alternatives.Source: §4.2.3 · p. 19

4.2.4 Improving layer non-privileging

  • Penalizing matrix rank can privilege layers, potentially assigning higher rank to an equivalent computation spread across layers than to one implemented in a single layer.The authors suggest seeking layer-invariant simplicity measures instead of effective rank.Source: §4.2.4 · p. 19

4.2.5 Promising applications of interpretability in parameter space

  • If practical hurdles are overcome, parameter-space interpretability may support mechanistic anomaly detection, model editing, and unlearning, as well as study of memorization, generalization, noise robustness, and cross-architecture structures.The authors connect editing and unlearning to APD’s parameter-vector descriptions, while presenting the scientific applications as possibilities they suspect APD could help investigate.Source: §4.2.5 · p. 20

5 Conclusion

APD directly decomposes parameter space into components intended to be faithful, minimal, and simple, and identifies ground-truth mechanisms in several toy models. The authors note that scaling the method to real-world models still presents challenges.

5 Conclusion

  • APD directly decomposes parameter space into components intended to be faithful, minimal, and simple.The authors frame this as a shift from analyzing networks through their activations.Source: §5 · p. 20
  • APD identifies ground-truth mechanisms in toy models of superposition, compressed computation, and cross-layer distributed representations.The paper presents parameter decomposition as a route toward interpreting mechanisms in neural networks.Source: §5 · p. 20
  • Scaling APD to real-world models requires better computational efficiency, greater hyperparameter robustness, and more robust attribution methods.Source: §5 · p. 20
  • The authors suggest APD may help open scientific and safety-oriented applications by bringing reverse engineering of capable AI systems closer.Source: §5 · p. 20

6 Related Work

Related methods often decompose activations or rely on predefined components, whereas APD decomposes parameters into components that can represent mechanisms across layers. The paper distinguishes APD from prior approaches by its representation, attribution-based use, and minimum-description-length framing.

6.1 Sparse Autoencoders

  • APD finds sparsely used parameter components, which may span layers or act on multidimensional or nonsparse activation subspaces, rather than forming an overcomplete activation-space basis.The authors contrast compressed descriptions of causal processes with SAEs' compressed descriptions of activation datasets.Source: §6.1 · pp. 20–21

6.2 Transcoders

  • Unlike crosscoders, which decompose activations across layers, APD decomposes the parameters implementing mechanisms that can span layers.Transcoders reconstruct layer outputs from inputs, while APD decomposes the transformation in parameter space.Source: §6.1 · p. 21, §6.2 · p. 21
  • APD's component activation is based on downstream causal attribution, which filters interference terms; transcoder latents use nonlinear thresholds to filter them.Source: §6.2 · p. 21

6.3 Weight masking and pruning

  • APD resembles learning a separate parameter mask per component, but its implicit masks can take values outside [0, 1] and thus need not privilege neuron or parameter bases.Source: §6.3 · p. 22

6.4 Circuit discovery and causal mediation analysis

  • Circuit discovery typically intervenes on predefined graph components, while APD learns parameter-space components and measures complexity using an approximation to their matrix ranks.The approaches share an aim of isolating computational mechanisms but represent components and complexity differently.Source: §6.4 · p. 22

6.5 Interpreting parameters

  • Prior parameter decompositions using SVD are limited to single layers and cannot identify mechanisms in superposition; the cited factorization methods do not target minimal, simple components.Source: §6.5 · p. 23

6.6 Quanta identification

  • Quanta work clusters parameter-loss gradients, whereas APD trains parameter components; sparse attribution also implicitly encourages component parts to align with gradient components.The paper describes this as a shared intuition that different gradients reflect different active mechanisms.Source: §6.6 · p. 23

6.7 Mixture of experts

  • Like mixture-of-experts systems, APD associates mechanisms with subsets of data, but identifies active components through attribution after an initial forward pass rather than a gating function.APD mechanisms may span layers, unlike the typically single-layer experts described in the paper.Source: §6.7 · p. 23

A More detailed description of the APD method

APD represents a network as parameter components that sum to its parameters and aims to minimize the description length of the components used per training datapoint. It estimates component use with attributions, combines a faithfulness objective with a practical sparsity-and-simplicity proxy, and acknowledges limits in attribution accuracy and the parameter-linearity assumption.

A More detailed description of the APD method

  • APD defines mechanisms as parameter components that minimize per-datapoint description length while remaining faithful, simple, and sparse in use.The appendix describes the method as optimizing these properties over the training dataset.Source: A More detailed description of the APD… · p. 31

A.1 Linear Parameter Decomposition

  • Linear Parameter Decomposition assumes mechanisms are parameter vectors that sum to the target parameters and can be ablated to assess whether they affect outputs.The authors call this parameter linearity and report it approximately holds in their studied networks, while its validity for non-toy models remains empirical.Source: §A.1 · pp. 31–32
  • APD estimates a component's effect by gradient attribution, but first-order attributions may fail in cases such as a saturated attention softmax.The authors anticipate that more sophisticated methods, such as integrated gradients, may be needed.Source: §A.1 · p. 32

A.2.1 Idealised loss: Minimum description length loss, LMDL

  • The ideal MDL objective favors sparse component attributions and low-rank components, but its L0 and rank terms are nondifferentiable and require a tractable proxy.The description-length account can encode active components by their indices and attributions; the rank penalty controls component simplicity.Source: §A.2.1 · pp. 33–34

A.2.2 Practical loss: Top-k formulation of APD

  • The main top-k formulation selects components with the largest attributions, then checks whether a sparse forward pass reproduces the target outputs; Schatten norms approximate rank.Hidden-activation reconstruction can aid deeper-model training, but matching hidden activations everywhere may conflict with predictions about unused components' noise, leaving the balance open.Source: §A.2.2 · pp. 34–35

A.2.3 Alternative practical loss: APD formulation that uses an Lp penalty on attributions

  • An alternative Lp attribution penalty is differentiable and encourages sparse attributions, but early experiments did not work as well as the top-k implementation for unknown reasons.The authors did not thoroughly explore this alternative.Source: §A.2.3 · p. 36

B.1.1 Setup

The hand-coded gated-function task tests whether APD can separate mechanisms in a network that sums controlled trigonometric functions. Each sparse binary control bit activates its corresponding function, which the target network approximates with a two-layer residual MLP.

B.1.1 Setup

  • The target network approximates a sum of trigonometric functions, with each function activated by its corresponding sparse binary control bit.The input contains a scalar x in [0, 5] and control bits that are 1 with probability 0.05.Source: §B.1.1 · p. 37
  • The hand-crafted two-layer residual MLP assigns each of its separable mechanisms to a unique set of neurons that approximates one function.The experiment uses four functions and ten neurons per function; the coefficient ranges are summarized in Table 2.Source: §B.1.1 · p. 37, Table 2 · p. 37
  • APD is expected to partition the target weights into one parameter component for each approximated function.Figure 13 colors target weights by the function they implement and shows the APD components by their corresponding target function.Source: §B.1.1 · p. 38, Figure 13 · p. 38

B.1.2 Results

APD approximately recovers parameter components for the target network, with most components capturing one function each. The decomposition has errors, and training is sensitive to hyperparameter changes.

B.1.2 Results

  • APD is particularly difficult to train in this setting, with minor hyperparameter changes causing large divergences.The authors hypothesize that this may be because the ground-truth network was hand-coded rather than trained.Source: §B.1.2 · p. 38
  • Most APD-generated components capture one function each, yielding an approximately correct decomposition of the target network’s parameters.Most weights within a component correspond to the same target-model feature.Source: §B.1.2 · p. 38
  • Some absent target weights remain nonzero in components, while some output weights appear absorbed into other components.The authors suggest either training difficulty on a hand-coded model or a layer-dependent simplicity penalty as possible explanations.Source: §B.1.2 · p. 38

C.1 Analysis of the compressed computation target model

The residual MLP target approximates the identity-plus-ReLU computation, with outputs that are similar across input features. For a single active feature, its ReLU output undershoots the target.

C.1 Analysis of the compressed computation target model

  • For input x42 = 1, the model outputs approximately 1.6 at index 42 versus the target 2.0, while other outputs are close to 1.0.The figure compares outputs across indices for this single one-hot input.Source: §C.1 · p. 39, Figure 14 · p. 39
  • The model’s output at the active index approximates ReLU over inputs from −1 to 1 but systematically undershoots the labels.The authors expect scaling the output to match the label would increase overall MSE loss.Source: §C.1 · p. 39
  • Across all 100 input features, the model treats features approximately the same.Figure 15 overlays the outputs for all features, with colors indicating the active input feature.Source: §C.1 · p. 39

C.2 Analysis of the compressed computation APD model

The batch top-k choice trades off how many components are active against whether components represent one or multiple input-feature computations. With top-k = 1.28, extra components can split a feature’s computation across components.

C.2 Analysis of the compressed computation APD model

  • The chosen batch top-k also lets APD handle rarer samples with many active input features and discourages components from representing multiple features.These are stated benefits of choosing a value large enough that active features rarely outnumber active components.Source: §C.2 · pp. 39–40
  • With batch top-k = 1.28, APD can use multiple components for part of a feature’s computation, producing low loss across features but large single-component losses for some.The setting was chosen for the one- and two-layer experiments so almost no batches had more active input features than active components.Source: §C.2 · pp. 39–40
  • Reducing batch top-k to 1 reduces the number of features with large single-component MSE, but increases components that represent two feature computations.The smaller value more often leaves fewer active components than active input features, favoring mixed computations.Source: §C.2 · p. 40

C.3 Analysis of the cross-layer distributed representations target model

The trained two-layer residual MLP target is qualitatively similar to the one-layer residual MLP target analyzed for compressed computation.

C.3 Analysis of the cross-layer distributed representations target model

  • The two-layer residual MLP target is qualitatively similar to the one-layer residual MLP target used for compressed computation.The comparison concerns the trained target models in the cross-layer distributed representations and compressed computation settings.Source: §C.3 · p. 40

C.4 Analysis of the cross-layer distributed representations APD model

The cross-layer APD model is qualitatively similar to the compressed-computation model, while batch top-k and component-rank choices affect how computations are represented. Ablation comparisons indicate that components have minimal influence on forward passes when their corresponding input feature is inactive.

C.4 Analysis of the cross-layer distributed representations APD model

  • The cross-layer APD model is qualitatively similar to the compressed-computation APD model, and components minimally influence forward passes when their corresponding input feature is inactive.The comparison uses 2-layer residual MLP settings.Source: §C.4 · pp. 40–41
  • In the 2-layer residual MLP, APD at batch top-k = 1.28 leaves some input-feature computations uncaptured by individual components; reducing batch top-k to 1 helps, but more components fully represent multiple computations.The MSE comparison is for APD trained at batch top-k = 1.28, using one-hot inputs averaged over 100k samples; it compares the target with APD using component top-k = 1 and batch top-k = 1.28.Source: §C.4 · p. 41, Figure 21 · p. 44
  • Rank-1 constraints prevent components from fully representing multiple input-feature computations; with this constraint, batch top-k = 1.28 almost completely resolves partial representations.Fully representing multiple computations would require component matrices with rank greater than 1.Source: §C.4 · p. 41

D Training details and hyperparameters

The appendix specifies target-model and APD training settings for toy models, compressed computation, cross-layer representations, and hand-coded gated functions. It notes that some hyperparameters were not rigorously ablated, while loss-term coefficients and, to a lesser extent, batch top-k can substantially affect APD outcomes.

D.1.1 TMS with 5 input features and hidden size of 2

  • For the 5-feature, hidden-size-2 TMS, the target uses AdamW for 5k steps, while APD uses Adam for 20k steps with batch top-k = 0.211.Target training uses batch size 1024 and learning rate 0.005; APD uses batch size 2048 and loss coefficients 1, 1, and 0.7 for faithfulness, minimality, and simplicity.Source: §D.1.1 · p. 44

D.1.2 TMS with 40 input features and hidden size of 10

  • On the 5-feature dataset, APD uses 50 components and trains for 20k steps with batch size 2048, batch top-k 1, and simplicity norm p=0.9.The target model trains for 2k steps with batch size 2048. APD loss coefficients are 1 for faithfulness, 10 for minimality, and 20 for simplicity.Source: §D.1.2 · pp. 44, 47

D.2 Compressed computation and cross-layer distributed representation

  • For compressed computation and cross-layer representations, both APD models use batch top-k = 1.28; the 1-layer model trains 40k steps and the 2-layer model 10k.The models use 100 input features; APD uses batch size 2048, with simplicity norm p = 0.9 for both.Source: §D.2 · p. 47
  • Several hyperparameter differences were not rigorously ablated; loss-term coefficients and, to a lesser extent, batch top-k significantly influence APD outcomes.Learning-rate schedules, training steps, simplicity norm, and batch size tended not to have a large influence.Source: §D.2 · p. 47

D.3 Hand-coded gated function model

  • The hand-coded gated-function experiment uses four functions, each approximated with 10 neurons, and trains APD for 200k steps with batch top-k = 0.2217.The data uses one input variable drawn uniformly from [0, 5] and four control bits active with probability 0.05.Source: §D.3 · p. 47
Loading 2501.14926v4…