The Router Already Knows
Organizations: University of California, Los Angeles · Astrid Cybernetics
Abstract
Mixture-of-experts (MoE) models already make a sparse choice for every token: which experts should process it. We ask whether a sparse dictionary should reuse that choice. NR-SAE gives each routed expert a small dictionary and opens only those the frozen router selects. At block 15 of Gemma 4 26B-A4B, after ten million C4 training tokens and across three seeds, the expert-aligned autoencoder explains 0.518 of the variance in router-weighted expert outputs on a fixed evaluation capture. A rank-addressed control with the same parameters but sixteen times as many candidates reaches 0.481, a candidate-matched control 0.458, and an architecture-matched control with randomized expert identity 0.387. A transcoder variant also leads all three controls: 0.432, against 0.347, 0.371, and 0.118. Part of this advantage is architectural: each expert writes into a subspace of at most a quarter of the model width, so naming the expert names the subspace. How much it explains, we have not measured. At this layer and budget, the router already knows something a dictionary would otherwise have to learn. Our implementation did not train faster, and semantic interpretability, downstream loss, and causal faithfulness remain untested.
Figures & tables
| Layout | Dictionary address | Banks width | Stored features | Candidates per token |
|---|---|---|---|---|
| A NR-SAE | true expert identity | 65,536 | 4,096 | |
| F Wide rank bank | router rank, 1–8 | 65,536 | 65,536 | |
| B Rank bank | router rank, 1–8 | 4,096 | 4,096 | |
| R Random identity | permuted identity | 65,536 | 4,096 |
| Reconstruction (SAE) | Prediction (transcoder) | |||
|---|---|---|---|---|
| Layout | EV | A layout | EV | A layout |
| A NR-SAE | 0.518 | — | 0.432 | — |
| F Wide rank bank | 0.481 | 0.038 [0.035, 0.040] | 0.347 | 0.085 [0.082, 0.088] |
| B Rank bank | 0.458 | 0.061 [0.059, 0.062] | 0.371 | 0.060 [0.058, 0.062] |
| R Random identity | 0.387 | 0.131 [0.127, 0.135] | 0.118 | 0.313 [0.308, 0.320] |
Sparse Autoencoders
Sparse autoencoders decompose neural network activations into sparse features to study, interpret, and control learned representations.
See all 220 papers