Clustering SPD components into mechanisms

Stochastic Parameter Decomposition (SPD) breaks a network's weights into many small subcomponents, each meant to capture a piece of the network's computation. That leaves an open problem the SPD authors flagged themselves: a real mechanism is usually spread across several subcomponents, and grouping them by hand doesn't scale. This pilot, part of a seven-repository series I ran in July 2026, asks whether that grouping can be automated without being told the answer.

A model where the answer is known

To check any automatic clustering, you need ground truth. So the target is a Toy Model of Superposition (Elhage et al., 2022): five features squeezed into two hidden dimensions. Its true mechanisms are simply those five features.

I decomposed it with an SPD-style method that follows Bushnaq, Braun and Sharkey (2025). Each weight matrix becomes 20 rank-one subcomponents, each with a learned causal-importance gate. Training combines faithfulness, a stochastic-masking reconstruction loss and an importance-minimality penalty. It converges cleanly: faithfulness error around 10⁻¹⁵, and stochastic reconstruction error near zero.

Two automatic signals

I then clustered the subcomponents using two signals, without access to the labels, and only afterwards compared the clusters with the five real features:

Clustering signal	ARI	NMI
Gate co-activation	0.53	0.47
Attribution signature	0.46	0.38
Random baseline	0.00	0.00

Both signals recover real structure: an adjusted Rand index around 0.5, against 0.0 for random. SPD subcomponents can be grouped into the true features automatically. The gap from a perfect 1.0 matters too. At this sparsity the components aren't cleanly one-feature-each.

What went wrong, and what it taught

• Gate sparsity sits on a knife edge. Too much importance penalty and every gate collapses to zero; too little and nothing specialises. That reproduces a known SPD training-stability issue and argues for annealed or adaptive importance weighting.
• Co-activation beat attribution, the opposite of what I expected, which suggests fusing the two signals rather than choosing one.
• Sharper gates aren't the fix on their own. As a follow-up I tried hard-concrete (L0) gates. They increased the number of specialised components (24 of 38), but ARI fell to 0.28.

That last result is the thread running through the series: how well components specialise, not which clustering signal you pick, is what limits recovery. I wrote it up as Specialization Is the Bottleneck.

Next

The natural next steps are annealing the importance coefficient, adding a causal-interaction (double-ablation) signal fused with co-activation, and moving to cross-layer models, where mechanisms span layers and manual clustering stops being feasible. The series later tested this on a real grokked transformer. Everything in the pilot runs on a CPU in a few minutes, and the code is public.

References: Bushnaq, Braun & Sharkey (2025), Stochastic Parameter Decomposition, arXiv:2506.20790 · Braun et al. (2025), Attribution-based Parameter Decomposition, arXiv:2501.14926 · Elhage et al. (2022), Toy Models of Superposition.