Mechanistic Permutability: Match Features Across Layers

Every layer speaks its own dictionary

A sparse autoencoder turns a hidden state into a short list of interpretable ingredients: it encodes $f(x)=\sigma(W_{\mathrm{enc}}x+b_{\mathrm{enc}})$, keeps $f$ sparse, and reconstructs $\hat{x}=W_{\mathrm{dec}}f+b_{\mathrm{dec}}$ with a feature count $F$ far larger than the hidden size $d$. Trained this way, an SAE untangles superposed, polysemantic activations into features a person can read.

But an SAE is trained per layer, so a 26-layer model gives us 26 unrelated dictionaries. Whether feature 4,071 in layer 19 has a counterpart in layer 20 is exactly the question interpretability needs answered to follow a computation through depth, and until now answering it meant running data through the model and correlating activations. We wanted the alignment to fall out of the weights themselves.

Fold the thresholds into the weights

The Gemma Scope SAEs we study, released by Lieberum and colleagues, use a JumpReLU activation $\operatorname{JumpReLU}(z)=z\odot H(z-\theta)$, where $H$ is the Heaviside step and $\theta$ holds a learned threshold per feature. Those thresholds are not decoration. Plotting them across depth, we find $\theta$ grows in step with the norms of the hidden states, so it quietly encodes how big each feature's activations really are.

That motivates parameter folding. We rescale $W'_{\mathrm{enc}}=W_{\mathrm{enc}}\operatorname{diag}(1/\theta)$ and $b'_{\mathrm{enc}}=b_{\mathrm{enc}}\odot 1/\theta$, push the scale into the decoder as $W'_{\mathrm{dec}}=W_{\mathrm{dec}}\operatorname{diag}(\theta)$, and set $\theta'=1$. The autoencoder computes exactly the same function afterwards, but each decoder column now carries its feature's true magnitude, so comparing columns compares reconstructions rather than directions of arbitrary length.

Matching is a linear assignment problem

Our hypothesis is that adjacent layers learn largely the same features, shuffled into different slots. Aligning two layers then means finding the permutation $P^{(A\to B)}$ that minimizes the mean squared error between folded decoder columns, which is equivalent to maximizing $\langle P,(W^{(A)}_{\mathrm{dec}})^{\top}W^{(B)}_{\mathrm{dec}}\rangle_F$. That is a classic linear assignment problem, the same machinery Ainsworth and colleagues used to align independently trained networks, and it needs no forward passes at all.

To check that low error really means same meaning, we had an external model judge the matches blind. GPT-4o mini compared Neuronpedia descriptions of paired features, about 1,600 comparisons over 100 feature paths spanning 16 layers, labeling each pair same, maybe, or different. Folded matching beat the unfolded variant on both the error metric and this semantic judgment, and pairs with low decoder MSE were consistently the ones judged identical.

Alignment works, but not everywhere

Past layer 10 of Gemma 2, matched features are overwhelmingly judged the same or closely related, and similarity decays gracefully with distance: starting from layers 10, 12, or 20, a feature stays recognizable for roughly five layers. Composed permutations $P^{(A\to C)}\approx P^{(B\to C)}P^{(A\to B)}$ track the exact assignment closely for nearby layers before drifting, which is what our composition hypothesis predicted.

Before layer 10 the picture inverts: almost no pairs are judged identical. The explanation turned out to be mundane and useful. The canonical Gemma Scope set picks, per layer, whatever sparsity the release happened to standardize, and in the early layers the average number of active features lurches between neighbors, 141 at layer 2 against 59 at layer 3. Two dictionaries describing the same state at wildly different sparsity carve it into incompatible pieces.

The sparsity story is causal, not incidental: when we re-selected SAEs so neighboring layers have similar $\ell_0$, explained variance improved significantly in the early layers too. Matching compares dictionaries, and dictionaries can only agree if they are written at a comparable granularity.

A matched dictionary can stand in for a block

If matched features really carry the same meaning, they should be exchangeable in the computation itself. We test this by skipping a block: encode $f^{(t)}=\sigma(W^{(t)}_{\mathrm{enc}}x^{(t)}+b^{(t)}_{\mathrm{enc}})$, permute $\hat{f}^{(t+1)}=P\,f^{(t)}$, and decode with the next layer's decoder to get $\hat{x}^{(t+1)}$, never running the transformer block in between.

From layer 10 onward the cross-entropy change is minimal, smallest on OpenWebText, with the clearest degradation at the final layer on code. Sparsity matters here too: sweeping SAEs of different density at layers 20 and 21, explained variance peaks near $\ell_0\approx 70$. Too sparse and the reconstruction starves; much past 100 active features and the assignment starts pairing noise with noise.

What a permutation cannot say

We match bijectively: every feature gets exactly one partner, even when the honest answer is none or several. High-MSE pairs flag some of that, and encoder-only matching is markedly worse than decoder-based variants, but features that split, merge, or genuinely appear mid-network need a richer formalism than a permutation.

Within those limits the method holds up beyond its home model: the same fold-then-match recipe transfers to Llama 3.1 8B. What we take away is simple. The thresholds everyone treats as a training artifact are a measurement of feature scale, and once they are folded into the weights, feature identity across depth becomes something you can compute in a spreadsheet-sized optimization, no data required.