Kronecker Factorization Improves Efficiency and Interpretability of Sparse Autoencoders

The encoder is the part nobody optimized

A sparse autoencoder rewrites a transformer activation $\mathbf{x}\in\mathbb{R}^{d}$ as a handful of active directions from a large learned dictionary: $\mathbf{f}=\operatorname{TopK}(W_{\mathrm{enc}}\mathbf{x}+\mathbf{b})$, then reconstructs $\hat{\mathbf{x}}=W_{\mathrm{dec}}\mathbf{f}+\mathbf{b}_{\mathrm{dec}}$. Those directions are the currency of mechanistic interpretability: they are what we inspect, name, and steer.

Scaling this recipe is where the trouble starts. With dictionary size $F$ far above $d$, the decoder can exploit sparsity, and fused sparse-dense kernels already make it cheap. The encoder cannot: it multiplies every token by a dense $F\times d$ matrix before TopK has anything to select from, an $\mathcal{O}(Fd)$ cost that dominates training. Gated and Switch variants reroute the computation, but each expert still carries a dense projection inside.

Classical SAEs also give the dictionary no internal structure: feature 4,913 has no designed relationship to feature 4,914, even though real features are heavily correlated. We built KronSAE to attack both problems with one mechanism.

Compose features instead of enumerating them

We split the latent space into $h$ independent heads. Head $k$ owns two thin projections: a composition base $P^{k}\in\mathbb{R}^{m\times d}$ and a composition extension $Q^{k}\in\mathbb{R}^{n\times d}$, with $m\le n\ll d$. Their rectified outputs $\mathbf{p}=\operatorname{ReLU}(P^{k}\mathbf{x})$ and $\mathbf{q}=\operatorname{ReLU}(Q^{k}\mathbf{x})$ are pre-latents: raw ingredients rather than finished features.

Every pair of ingredients is then multiplied into a feature, giving an $m\times n$ grid per head, which is exactly an element-wise square root of the Kronecker product $\sqrt{\mathbf{p}\otimes\mathbf{q}}$. The full dictionary has $F=h\,m\,n$ entries, but each token only ever touches $h(m+n)$ projection rows, so the encoder cost falls from $\mathcal{O}(Fd)$ to $\mathcal{O}(h(m+n)d)$. TopK then runs over the concatenated grids as usual, and any sparse decoder kernel plugs in unchanged.

A differentiable AND

How the pair $(u_i, v_j)$ is multiplied matters. We combine pre-latents with an operator we call mAND, a smooth stand-in for a Boolean AND gate:

A post-latent can fire only when both of its parents fire, so the set of inputs activating feature $z_{ij}$ is exactly the intersection $\mathcal{P}\cap\mathcal{Q}$ of its parents' activation sets. The square root keeps magnitudes tame when both parents are large. In our ablation the gate earns its keep: at a 125M-token budget, mAND beats both the plain product $u\cdot v$ and the rectified product on every configuration we tested.

Reconstruction holds, absorption drops

We train against TopK, Matryoshka, and Switch baselines on Qwen-2.5-1.5B, Gemma-2-2B, and Pythia-1.4B activations under an iso-FLOPs budget: every KronSAE gets exactly the compute its dense counterpart would burn. Across dictionary sizes, sparsity levels $\ell_0\in\{16,\dots,128\}$, and layers, the factorized encoder matches the baselines' explained variance despite training far fewer parameters, and the gap narrows as dictionaries grow.

The interpretability side is where the structure pays off most. Feature absorption, where a specific feature like "Lion" swallows a general one like "starts with L", is a known failure mode of flat dictionaries. Because our post-latents fire only when their broader parents fire, a specific feature cannot silently subsume a general one; the hierarchy is architectural, not incidental.

The choice of selector matters too: swapping TopK for JumpReLU inside KronSAE degrades explained variance at every sparsity we tried, while KronSAE with TopK edges out the plain TopK SAE, 0.814 against 0.809 at $\ell_0=32$. The factorization and the hard selector appear to work as a team.

Inside the heads: AND gates we can read

Because every feature has exactly two parents, we can read the composition mechanism directly. Pre-latents turn out deliberately polysemantic: one base in head 23 mixes "-like" comparatives, spiritual vocabulary, and geographic direction words. Each extension in the head then selects one meaning, and the intersections come out cleaner than either parent, scoring up to 0.93 in our automated interpretability pipeline.

The same structure shows up statistically: features inside a head correlate far more with each other than with the rest of the dictionary, and on synthetic data with a planted block covariance, KronSAE recovers the block structure where a TopK SAE mostly loses it, with an RV similarity of 0.30 against 0.12 in the best configuration. Not every interaction is a textbook intersection, though; some compositions produce genuinely new semantics, like an instrument base and a possession extension yielding a therapy feature.

The trade-offs we accept

The design has a dial. Pushing $m$ down to 1 and the head count up maximizes explained variance but makes pre-latents so overloaded that interpretability slips; our practical recommendation is to start at $m=2$ and spend the savings on more heads. And because KronSAE simply has fewer trainable parameters, a dense TopK encoder will eventually win once training budgets grow very large; our results target the 125M-to-2B-token regime where most interpretability work actually happens.

What we take away is bigger than the FLOPs ledger. Dictionary entries do not have to be enumerated; they can be composed from reusable parts, with the composition rule chosen so we can read it. Transcoders, crosscoders, and richer gates than AND are the obvious next places to apply the same idea.