Stochastic Parameter Decomposition

• SDL decomposes activations [Cunningham et al., 2024, Bricken et al., 2023], while linear parameter decomposition instead uses sparse parameter-space components [Braun et al., 2025].Source: §1 · pp. 1–2
• APD is impractical because of high memory costs, sensitivity to the top-k hyperparameter, and reliance on attribution estimates of causal importance.Source: §1 · p. 2
• SPD factorizes each network matrix into potentially overcomplete rank-one subcomponents, which require clustering into full components when ground-truth groupings are unavailable.Source: §1 · p. 2
• SPD predicts datapoint-specific causal importance from ablatability, randomly masks subcomponents accordingly, and trains their sum to reproduce the target output while favoring low importance values.Source: §1 · p. 2
• SPD recovers ground-truth mechanisms across previously studied toy models and succeeds on two harder models that APD could not correctly decompose in practice.Source: §1 · p. 2
• Scaling SPD to larger models may require more stable training, and cross-rank or cross-layer mechanisms still need general-purpose clustering.Source: §1 · p. 3

2.1 Our approach optimizes rank-one subcomponents instead of full-rank parameter components

• SPD replaces full-rank parameter components with rank-one subcomponents localized to individual weight matrices, which may later be clustered into full components.The number of subcomponents can exceed a matrix’s minimum row or column count, allowing computations in superposition; biases are left undecomposed.Source: §2.1 · pp. 3–4

2.2 Optimizing for faithfulness

• The faithfulness objective trains the sum of the subcomponents to approximate the target model’s parameters.The loss is normalized by the target model’s total number of parameters.Source: §2.2 · p. 4

2.3 Optimizing for minimality and simplicity by learning a causal importance function to stochastically sample masks

• SPD treats minimality as requiring as few subcomponents as possible to be causally important for each training datapoint.The method learns a causal-importance function, with losses encouraging important subcomponents to preserve the original function and unimportant ones to be ablated.Source: §2.3 · p. 4

2.3.1 Requirement 1: Formally defining a subcomponent’s causal importance as the extent to which it can be ablated

• A subcomponent’s causal importance is the extent to which it can be ablated on a datapoint without changing the model output.Importance is continuous: fully unimportant components can be ablated by any amount, while partial ablatability corresponds to partial unimportance.Source: §2.3.1 · pp. 4–5

2.3.2 Requirement 2: A loss function that lets us optimize causally important subcomponents to approximate the same function as the original network

• To preserve the target function under ablation, SPD samples a finite set of random masks and trains masked-model outputs to approximate the target outputs.Evaluating all possible mask values is intractable; the method samples S values, with the loss approaching the stated objective as S increases.Source: §2.3.2 · p. 5
• An auxiliary layerwise reconstruction loss addresses noise from jointly masking every subcomponent in a forward pass.It replaces parameters in one layer at a time with stochastically masked subcomponents while computing gradients at the model output.Source: §2.3.2 · p. 5

2.3.3 Requirement 3: A loss function that encourages as many subcomponents as possible to be causally unimportant

• Independent small MLPs predict each subcomponent’s causal importance from its inner activation, and a penalty encourages the predictions to be as low as possible.The importance function works well on the paper’s toy models, but arbitrary models likely require more expressive functions.Source: §2.3.3 · pp. 5–6

2.4 Summary of training setup

• Training combines four losses and uses five non-optimizer hyperparameters, with one mask sample per training step found sufficient.The five are three loss coefficients, the p-norm in the minimality loss, and the number of mask samples S.Source: §2.4 · p. 6

3 Results

• Across models previously decomposed with APD, SPD finds ground-truth mechanisms with fewer errors and less hyperparameter tuning.The paper reports that SPD identifies known mechanisms up to a small error across the toy models.Source: §3 · p. 7

SPD Results: Toy Model of Superposition

• For TMS5−2 and TMS40−10, SPD’s MMCS and ML2R are approximately 1 and closer to target weights than APD’s, which shows shrinkage.APD’s ML2R is approximately 0.9 for both models; SPD’s results are robust to training seeds and require less tuning.Source: Table 1 · p. 7, SPD Results: Toy Model of Superposition · p. 8

SPD Results: Toy Model of Superposition with hidden identity

• In TMS5−2+ID and TMS40−10+ID, SPD recovers W’s feature directions and decomposes I into m1 subcomponents, identifying m2 + m1 non-negligible components; APD fails to recover the mechanisms.The two identity-matrix subcomponents shown for the TMS5−2+ID model sum to its rank-2 identity matrix.Source: Figure 2 · p. 9, Figure 3 · p. 9, Figure 4 · p. 9

SPD Results: Toy Model of Compressed Computation

• In compressed computation, SPD decomposes Win into 100 feature-specific components that match the target model’s neuron contributions closely.For one-hot inputs of magnitude 0.75, Win’s masks select only the corresponding component.Source: SPD Results: Toy Model of Compressed… · p. 11, Figure 7 · p. 12
• SPD decomposes Wout into 50 subcomponents that act together as one rank-50 component, whereas APD splits it into 100 feature-associated components.The authors judge SPD’s decomposition more accurate and attribute APD’s Wout result to a poorly chosen top-k value.Source: SPD Results: Toy Model of Compressed… · p. 11
• In compressed computation, too-small β3 yields multiple Win components per feature, while too-large β3 lowers importance values and is associated with high reconstruction losses.The intermediate setting identifies a unique Win component per feature and the correct number of Wout components.Source: SPD Results: Toy Model of Compressed… · p. 12, Figure 8 · p. 13

SPD Results for two- and three-layer models of Cross-Layer Distributed Representations

• In the two-layer model, SPD recovers 100 feature-specific input components using neurons across both layers and a single rank-50 output component spanning both layers.This demonstrates that SPD can identify mechanisms spanning multiple layers.Source: SPD Results for two- and three-layer… · p. 14
• In the three-layer model, SPD finds 102 feature-specific input components and one rank-51 output component spread across all three layers, but results are seed-sensitive.The authors suggest unfavorable random initialization as a possible source of sensitivity.Source: SPD Results for two- and three-layer… · p. 14

4 Related Work

• SPD learns a basis for parameter-space causal interventions, unlike approaches that assume a particular intervention basis [Mueller et al., 2024, Vig et al., 2020, Geiger et al., 2024a,b].Its masks are stochastically sampled according to predicted causal importances [Csordás et al., 2021, Cao et al., 2021, Zhang et al., 2021].Source: §4 · pp. 14–15
• Causal importance is not identical to causal dependence: emergent self-repair may preserve an output after an important subcomponent is ablated [Lewis, 1973, Mueller, 2024; McGrath et al., 2023].Source: §4 · p. 15
• SPD predicts input-dependent component attributions, an approach the authors say they know of no close precedent for [Mozer and Smolensky, 1988, Syed et al., 2024, Marks et al., 2024, Braun et al., 2025].Source: §4 · p. 15
• Unlike gradient-based parameter decompositions [Chrisman et al., 2025] and Fisher-information factorization [Matena and Raffel, 2025], SPD directly tests combinations of component ablations against model outputs.Source: §4 · p. 15
• SPD treats weight degeneracy across data subdistributions as vectors that can be ablated to zero, rather than merely locally degenerate [Watanabe, 2009; Wang et al., 2024].Source: §4 · pp. 15, 18

5 Discussion

• SPD is considerably more scalable and robust to hyperparameters than APD, demonstrated by decomposing deeper and more complex models.The authors hope to test the method's scaling limits and necessary adjustments in future work.Source: §5 · pp. 18–19
• SPD replaces APD's fixed top-k estimate of active components with trained causal-importance functions, leaving a more forgiving penalty coefficient β3 to tune.The authors hypothesize that this change contributes to SPD's relative robustness.Source: §5 · p. 18
• Unlike APD's top-k updates, gradients flow through every SPD subcomponent on every datapoint, which the authors say can support error correction.Top-k activation functions create discontinuities that generally tend to destabilize gradient-based training.Source: §5 · p. 18
• By avoiding a simplicity objective, SPD avoids shrinkage in parameter-subcomponent singular values and removes the tension between faithfulness and simplicity losses.Shrunken components may fail to sum to the target parameters, requiring other components to compensate.Source: §5 · p. 18
• The authors argue that systematic errors in APD's first-order attribution estimates can bias gradients, whereas SPD directly optimizes causal importance.They describe SPD's causal-importance estimates as likely better than those from even sophisticated attribution methods such as Sundararajan et al. [2017].Source: §5 · pp. 18–19
• Scaling SPD further and clustering its rank-one, layer-specific components into full components remain open tasks beyond the known-ground-truth toy models.In the general case, the authors say clustering approaches are needed to combine sparse and dense coding schemes into components that permit a minimum-length description.Source: §5 · p. 19
• The authors question whether independent, layerwise subcomponents yield the shortest functional description, since later-layer information may be needed to assess causal importance.They also note that both SPD and APD privilege mechanisms spanning individual layers and suggest exploring less layer-biased losses.Source: §5 · p. 19
• The framework could use other subcomponent structures or causal-importance functions; the authors also propose studying memorization mechanisms and intrinsically decomposed models.These are prospective directions, not demonstrated outcomes.Source: §5 · p. 19

A.1 Architecture of Causal Importance Function MLPs

• Each subcomponent has its own causal-importance MLP, whose hidden layer uses GELU neurons and whose output uses a hard sigmoid.The hidden width is dgate; the passage identifies the inner activation as the component input.Source: §A.1 · p. 24
• Using a separate MLP per subcomponent keeps gate-training costs low compared with training the subcomponents themselves.The gate parameters include trainable input and output weights and biases.Source: §A.1 · p. 24
• The authors chose this simple, low-cost gate design, but say SPD can use other causal-importance predictors.They contrast this flexibility with concerns in sparse dictionary learning about requiring complex nonlinearities to identify a feature.Source: §A.1 · p. 24

A.2 Avoiding dead gradients with leaky hard sigmoids

• Leaky hard sigmoids add a slope of 0.01 outside [0, 1] to avoid dead gradients in the causal-importance gates.The lower-leaky variant is used for stochastic reconstruction losses; the upper-leaky variant is used for importance minimality.Source: §A.2 · p. 24
• Forward masks use the lower-leaky sigmoid so irrelevant subcomponents can be scaled below zero, but influential ones are not scaled above one.The importance loss uses the upper-leaky sigmoid because causal-importance values must not be negative, though values above one are allowed there.Source: §A.2 · p. 24

A.3 Heuristics for Hyperparameter Selection

• Choose reconstruction and importance-minimality loss weights by checking negligible performance loss, plausible superposition noise, known mechanisms, and other sanity checks.The reconstruction weights are β1 and β2; the importance-minimality weight is β3.Source: §A.3 · pp. 24–25
• Keep stochastic reconstruction losses no larger than the noise plausibly caused by interference from inactive mechanisms in superposition.The authors suggest estimating the expected size of those interference terms.Source: §A.3 · p. 25
• When mechanisms are known, restrict hyperparameters to recover them; when none are known, the authors suggest inserting a known mechanism as a check.Their example inserts an identity matrix and checks whether SPD recovers it as a single high-rank component.Source: §A.3 · p. 25
• A decomposition should preserve target-model performance when unmasked subcomponents are summed, and some causal-importance values should reach 1 for most inputs.If they do not, the authors say β3 is likely too high.Source: §A.3 · p. 25

• Target models were trained for 10k steps with AdamW, weight decay 0.01, and a constant learning rate of 5 × 10−3.Source: §A.4.1 · p. 25
• Input features were sampled uniformly from [0, 1] with probability 0.05 and were otherwise zero.Source: §A.4.1 · p. 25
• Batch size was 1024 for TMS5−2 and TMS5−2+ID, and 8192 for TMS40−10 and TMS40−10+ID.Source: §A.4.1 · p. 25

• SPD training ran for 40k steps with batch size 4096, Adam, and a maximum learning rate of 1 × 10−3 under cosine scheduling.Source: SPD training: Common hyperparameters · p. 25
• Training data matched the target model's feature probability of 0.05, and stochastic sampling used S = 1 for both reconstruction losses.Source: SPD training: Common hyperparameters · p. 25
• Faithfulness and both stochastic reconstruction loss coefficients were 1, with one causal-importance MLP per subcomponent using 16 hidden GELU neurons.Source: SPD training: Common hyperparameters · p. 25

Model architectures

• The 1- and 2-layer models had 100 input features and 1000-dimensional embeddings; the 3-layer model had 102 input features and the same embedding dimension.The models had 50, 50, and 51 MLP neurons, respectively.Source: Model architectures · p. 25

SPD training: Common hyperparameters

• Common SPD settings used Adam with constant learning rate, batch size 2048, feature probability 0.01, S = 1 for both stochastic losses, and p = 2 for importance minimality.Source: SPD training: Common hyperparameters · p. 26

SPD training: Model-specific hyperparameters

• The 1-, 2-, and 3-layer residual MLPs used respectively 100, 400, and 500 initial subcomponents, with training budgets of 30k, 50k, and 200k steps.The 3-layer model converged around 70k steps; its importance-minimality coefficient was 0.5×10−5, compared with 1 × 10−5 for the 1- and 2-layer models.Source: SPD training: Model-specific… · p. 26

A.5 Supplementary figures

• Figures 14 and 15 compare target-weight and SPD-component similarity across the first 10 input features in 2- and 3-layer residual MLPs.Their top panels show neuron contributions; their bottom panels show learned-component contributions.Source: Figure 14 · p. 27, Figure 15 · p. 28

A.6 SPD Pseudocode

• SPD pseudocode initializes layerwise subcomponents and causal-importance MLPs, then uses activations to compute raw importance outputs and masks.Source: §A.6 · p. 29
• For sampled masks, the algorithm computes whole-model and layerwise reconstruction losses, normalizes them, combines them with faithfulness and importance-minimality losses, and updates trainable parameters.Source: §A.6 · p. 29