SPD meets a real grokked transformer

My first SPD clustering pilot worked on a toy model, where gate co-activation recovered the true mechanisms with an adjusted Rand index of 0.53. Toy models are forgiving, though. The fifth experiment in the series asked the question that matters: does automated mechanism recovery survive contact with a real trained transformer?

A real model with a known answer

I trained a small transformer (embedding, attention and MLP; d_model = 64) on (a + b) mod 53 with weight decay. It grokked to 100% test accuracy, meaning it learned the actual algorithm rather than memorising. Grokked modular-addition models are known to represent numbers using a handful of Fourier frequencies (Nanda et al., 2023), which provides ground truth.

Crucially, I measured that ground truth independently of SPD, directly from a Fourier transform of the embedding. The model uses about four key frequencies: 17, 11, 12 and 26, and the top five frequencies hold 91% of the spectral power.

Result: the components collapsed

I applied SPD to the MLP up-projection and asked whether the decomposed components line up with those frequencies. They didn't:

• All components collapsed onto the single dominant frequency, 17.
• Clustering against the true frequencies scored an ARI of 0.
• The cause: the importance-minimality term stayed saturated (around 1.0), so the components never specialised within the training budget on this larger layer.

As a follow-up I ran a stronger, annealed SPD with more components (48) and a longer schedule. The model still scored 100%, but the components recovered only one of the four true frequencies.

The through-line

Across the series (toy superposition models, a balanced variant, a compositional MLP and this real transformer), one finding held and got stronger with scale:

The binding constraint on automated mechanism recovery is SPD component specialization, not the choice of clustering signal.

When components map cleanly to single mechanisms, simple co-activation works. As targets get realistic, components fail to specialise, and then no clustering signal recovers the verified mechanisms, whether co-activation, attribution or causal double-ablation. The practical recommendation is to invest in specialization first: sharper L0 gates, longer or stronger sparsity schedules, and per-mechanism capacity. Clustering-signal design comes after that.

References: Bushnaq, Braun & Sharkey (2025), Stochastic Parameter Decomposition, arXiv:2506.20790 · Nanda et al. (2023), Progress measures for grokking via mechanistic interpretability, ICLR · Elhage et al. (2022), Toy Models of Superposition.