Inference and learning in sparse autoencoders as natural gradient flow

Hadi Vafaii; Tejas Rao; David Chanin; Thomas Fel; Jacob L. Yates; Bruno Olshausen; David Klindt; Dileep George

cs.LGOct 5, 2026

Inference and learning in sparse autoencoders as natural gradient flow

Authors: Hadi Vafaii, Tejas Rao, David Chanin, Thomas Fel, Jacob L. Yates, Bruno Olshausen, David Klindt, Dileep George, +1 more

Organizations: Astera Institute · UC Berkeley · UCL · Goodfire AI · CSHL

Abstract

Sparse autoencoders are widely used to uncover interpretable features in neural networks, yet reliable recovery remains difficult when features overlap or activate infrequently. These challenges involve both inferring which features explain an input and learning the dictionary that represents them. Here, we unify inference and dictionary learning as natural-gradient flows on a shared variational free energy. We instantiate this framework as BeFOND, an encoder-free sparse coding model with closed-form inference and learning dynamics. We show how recurrent explaining away reduces interference between overlapping features, while Fisher preconditioning can compensate for the slow learning of rare features. On synthetic data, BeFOND improves dictionary recovery and rare-feature detection, with a growing advantage over amortized baselines as superposition increases. On language-model activations, it improves single-feature concept detection and selective intervention, outperforming pretrained reference SAEs with substantially less training data. Its feature quality continues to improve with dictionary width, whereas the evaluated baselines largely plateau. Together, these results show how improving inference and learning within a unified probabilistic framework can make better use of data and dictionary capacity to interpret and intervene on neural representations.

Figures & tables

Figure 1: Feature quality and sample efficiency on Gemma-2-2B activations. a,c : Training dynamics at

524

k width. b,d : Final performance across dictionary widths. Details in section 5 .

Figure 2: Heatmaps show sprinkler probabilities under exact inference, converged

\mathrm{Be}

FOND with a fixed prior, and a numerically optimized single-projection encoder. The marked observations have identical lawn wetness but different roof evidence: neutral (white, halfway dry / wet) or wet (blue). Fixed-prior inference and encoder fitting are detailed in section C.2.3 .

Figure 3: $\mathrm{Be}$ FOND maintains the highest recovery, and without dead latents ( $K=4{,}096$ ). Means and 95% confidence intervals over 3 seeds. Exact values are reported in section C.3.1 and Table 4 .

Figure 4:

\mathrm{Be}

FOND outperforms all evaluated baselines on dictionary recover (MCC) and feature detection (Macro F1, Rare F1) in SynthSAEBench (

d=768

K=K_{\star}=16{,}384

). Rare F1 covers the least frequent quarter of features. Bars show means over 3 initialization seeds with 95% confidence intervals, evaluated on

10^{6}

held-out samples. Table 5 gives individual-seed results.

Appendix figures & tables18 assets

Supplementary material from the paper’s appendix.

Appendix

Algorithm	Metric $\mathbf{G}$	Loss $V$
Gradient Descent ( Cauchy, 1847 , GD,)	$\mathbf{I}$ (Euclidean)	Generic task loss
Natural Gradient Descent ( Amari, 1998 , NGD,)	Fisher–Rao	Generic task loss
Bayesian Learning Rule ( Khan and Rue, 2023 , BLR,)	Fisher–Rao	$\mathcal{F}$ (free energy)

Appendix

Table 1: Special cases of eq. 57 under different choices of metric

\mathbf{G}

and loss

V

. The BLR row uses exponential-family variational beliefs in natural parameters.

Panel of Figure 5 — Figure 5: Online inference with a rolling prior. The learned initial prior initializes the first belief. Under identity state dynamics, prediction carries each posterior into the next reference (green arrows). The observation model is shown by black arrows and inference by red dashed arrows. The fixed-input case presents the same observation at each step.

Experiment	Settings
$\mathrm{Be}$ FOND best / continuous	Seed 0; batch 2,048; 88,779 updates; 181.82M decoder and 18.18M prior examples; linear LR $0.02\to 10^{-5}$ , no warmup; $\beta=1.6$
Inference / decoder	Bernoulli; ETD1 $h=\infty$ ; one inner step; uniform training depth 1–10, test depth 200. Prior Fisher, current moments, no damping; $\Delta t=1$ ; bias and variance fixed
Initialization / constraints	Gaussian scale 0.1443375673; prior logits $-4.7538364662$ ; soft decoder norm 3.5 with coefficient 100; prior mean cap $3\times 10^{-4}$ and floor $10^{-5}$
Prior / matrix functions	Prior update every 10 batches on 2,048 separate examples; depth $\leq 200$ , KL threshold 0.01, rejection threshold 1, window 10. Taylor limit 32 applications; Chebyshev tolerance $10^{-5}$ , maximum degree 2,048
Original baselines	BatchTopK/JumpReLU/Matryoshka: target $L_{0}=35/45/20$ , batch 1,024, LR $3\times 10^{-4}$ , 200M decoder examples. MP/ReLU-L1: target $L_{0}=35$ , 200,000,512 decoder plus 19M monitoring examples. MF: 20 inference steps in training and testing, batch 1,024, 195,313 updates (200,000,512 examples), cosine schedule with dictionary / other-parameter LR $0.001/0.003$ , damping 1 on the first step then 0.2, decoder norm cap 4
Modified baselines	Target $L_{0}=45/40/35$ for BatchTopK/JumpReLU/Matryoshka; MinFire or resampling and learning-rate decay

Appendix

Table 2: SynthSAEBench training settings. Decoder and prior examples are counted separately. Full-width runs use

d=768

K=K_{\star}=16{,}384

; superposition runs use

K=4{,}096

	Neutral roof		Wet roof
Method	Rain	Sprinkler	Rain	Sprinkler
Exact posterior	0.292138	0.292138	0.960707	0.033017
$\mathrm{Be}$ FOND	0.228996	0.228996	0.974111	0.019506
Single-projection encoder	0.217514	0.139225	0.958493	0.064003

Appendix

Table 3: Posterior activation probabilities at the marked observations in Figure 2 , rounded to six decimal places. Neutral and wet correspond to

\bm{x}=(2,1.5)^{\top}

and

\bm{x}=(2,3)^{\top}

, respectively.

		MCC			Dead latents
Method	$d$	Seed 0	Seed 1	Seed 2	0	1	2
BatchTopK	1536	0.6447949409484863	0.6418958902359009	0.6367465853691101	178	217	206
	1024	0.7172077894210815	0.7101552486419678	0.715070366859436	78	80	66
	768	0.7237797379493713	0.7247610688209534	0.7268125414848328	49	47	39
	512	0.6935576796531677	0.6925675868988037	0.6897946000099182	37	38	55
	384	0.6238665580749512	0.6256771087646484	0.6227874159812927	56	82	58

Appendix

Table 4: Individual-seed results for Figure 3 (

K=4{,}096

K_{\star}=16{,}384

). MCC values retain the recorded precision. Dead latents are counted on

10^{6}

held-out examples using each model’s native activity rule (

r_{i}>0.5

for

\mathrm{Be}

FOND and Mean-field + Adam).

Figure 6: Lateral suppression supports feature recovery under superposition. Hollow markers denote early-stopped excitatory runs (

K=4{,}096

Figure 7: The natural-inference/natural-learning recipe reaches higher recovery with fewer samples. Panels show MCC, Macro F1, and Rare F1 on SynthSAEBench (

d=768

K=K_{\star}=16{,}384

). N/E denotes inference/learning geometry: natural (N) or Euclidean (E). Recipes are tuned separately ( section C.1.2 ).

Figure 8: Sparsity and recovery versus

\beta

on SynthSAEBench (

d=768

K=4{,}096

K_{\star}=16{,}384

). Colors indicate the mean initial-prior cap, lines show the mean over two seeds, and shading shows their range. Each fit uses 12,000 updates and is evaluated on 8,192 examples after 5,000 inference steps.

Figure 9:

\mathrm{Be}

FOND retains the highest MCC, Macro F1, and Rare F1 after baseline modifications aimed at preventing inactive features. Hatched bars denote modified recipes (

d=768

K=K_{\star}=16{,}384

, seed 0); numerical results appear in Table 5 .

Method	Seed	MCC $\uparrow$	Macro F1 $\uparrow$	Rare F1 $\uparrow$	$L_{0}$	Dead $\downarrow$
BatchTopK	0	0.534523	0.445173	0.299436	34.833798	1,577
	1	0.535311	0.444452	0.299666	35.031285	1,733
	2	0.537442	0.446642	0.296203	34.948331	1,707
	Mean	$0.5358\pm 0.0038$	$0.4454\pm 0.0028$	$0.2984\pm 0.0048$	34.94	1672.33
JumpReLU	0	0.483953	0.326567	0.190439	45.148172	8,069
	1	0.481728	0.326672	0.184864	44.826161	8,090

Appendix

Table 5: Full-width SynthSAEBench results (

d=768

K=K_{\star}=16{,}384

). Individual rows report initialization seeds 0–2; mean rows give pointwise 95% Student-

t

confidence intervals for the three recovery metrics, as in Figure 4 . Each evaluation uses

10^{6}

held-out examples.

L_{0}

and dead counts in mean rows are arithmetic means. The lower block retains the modified recipes (

+

) and continuous-run checkpoint, each at seed 0. The

\mathrm{Be}

FOND runs use 181.82M decoder-update plus 18.18M prior examples.

Figure 10: Recovery across feature-frequency quintiles (

d=768

K=K_{\star}=16{,}384

, seed 0). Dashed lines show original baselines, and “+” denotes their modified recipes. Coverage is the fraction of features matched above cosine 0.8. Quintiles use each run’s 500K-example frequency reference. F1 uses 1M held-out examples. Random draws differ between some runs.

Figure 11: Feature recovery throughout training on unwhitened SynthSAEBench (

d=768

; 16,384 learned and ground-truth features), using plain baseline recipes and seed 0. (a) Dictionary mean correlation coefficient (MCC). (b) Macro F1 over ground-truth features after positive-cosine one-to-one matching. (c) Rare F1 over the least frequent quartile, defined by a shared independent 500K-example frequency reference (seed 315031). Every checkpoint is evaluated on 1M held-out examples (seed 49100). BeFOND uses 200 inference steps with activity

p>0.5

; baselines use positive native activations. The x-axis counts decoder-update examples and is piecewise linear: 1M examples at 0–10M, 5M at 10–50M, and 25M beyond 50M occupy equal horizontal distances. Cross-region slopes therefore do not have a common per-sample interpretation. Baseline points through 10M, after 10M through 50M, and beyond 50M come from separate matching-recipe runs with their original full-budget learning-rate schedules; they do not form a single continuous trajectory. Overlapping checkpoints are selected by these fixed windows, not by score. BeFOND uses one continuous run. Samples are rounded up to full batches. Baselines finish at approximately 200M decoder-update examples and BeFOND at 181.82M plus 18.18M separate prior-update examples, which are excluded from the x-axis. Baseline auxiliary exposure is not counted on this axis. Final measured

L_{0}

values are BatchTopK: 34.83, JumpReLU: 45.15, Matryoshka: 19.98, Matching Pursuit: 29.65, ReLU-L1: 35.09, BeFOND: 35.70. One seed is shown; no uncertainty intervals are estimated.

	Sparse probing						RAVEL
Model	$k=1$			$k=5$
	Acc. $\uparrow$	AUC $\uparrow$	F1 $\uparrow$	Acc. $\uparrow$	AUC $\uparrow$	F1 $\uparrow$	Disentangle- ment $\uparrow$	Cause $\uparrow$	Isolation $\uparrow$
Dictionary width: 16k
BatchTopK (target $L_{0}=400$ )	${.767}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.002}}$	${\mathbf{.798}}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.003}}$	${.761}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.004}}$	${\mathbf{.814}}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.002}}$	${\mathbf{.862}}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.002}}$	${\mathbf{.812}}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.002}}$	${.729}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.005}}$	${.692}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.019}}$	${.767}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.017}}$
Matryoshka BatchTopK ( $L_{0}=200$ probing; $400$ RAVEL)	${.768}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.003}}$	${.794}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.003}}$	${.762}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.002}}$	${.810}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.002}}$	${.858}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.002}}$	${.808}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.002}}$	${.731}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.003}}$	${.695}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.013}}$	${.767}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.015}}$
Mean-field + Adam ( $T=20$ inference steps)	${.761}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.003}}$	${.761}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.003}}$	${.749}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.004}}$	${.801}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.002}}$	${.842}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.001}}$	${.798}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.002}}$	${.696}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.003}}$	${.639}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.006}}$	${.753}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.006}}$

Appendix

Table 6: SAEBench results on whitened Gemma-2-2B layer-12 activations after

49{,}999{,}872

training tokens. Entries are mean

\pm

sample standard deviation over three training seeds for

\mathrm{Be}

FOND and five for each baseline. For SAE baselines, we selected target

L_{0}

values by mean top-1 accuracy, and then used it for all probing columns, and the model selected by mean disentanglement supplies all RAVEL columns. Parentheses give both values when they differ. Sparse probing measures how well

k=1

k=5

features predict concepts, averaged equally across 113 tasks: accuracy is the fraction of correct predictions; AUC measures how well positives rank above negatives across thresholds; F1 balances precision (correctness of positive predictions) and recall (coverage of positives). RAVEL measures whether interventions change a targeted attribute ( cause ) while preserving other attributes ( isolation ), and disentanglement is their mean. Higher is better throughout. Within each width and metric, bold marks the highest mean, selected before rounding.

Figure 12: Performance as a function of inference depth. Sparse probing top-1 accuracy (left), top-5 accuracy (middle), and RAVEL disentanglement (right) for

\mathrm{Be}

FOND on Gemma-2-2B layer-12 activations, evaluated over

T=1,\ldots,40

inference steps at dictionary widths of

16

65

131

k, and

524

k. All evaluations use checkpoint

12{,}207

, corresponding to

49{,}999{,}872

training tokens. Lines show means and shaded regions indicate

\pm 1

sample standard deviation across three training seeds. Sparse probing peaks within the first few steps, whereas RAVEL generally benefits from longer inference.

Method / metric	Settings
$\mathrm{Be}$ FOND	Widths 16,384–524,288; exponentially sampled training depth (mean 16, cap 100; Appendix C.1 ); FP32 inference/parameters; prior Fisher with current moments and no damping; $\Delta t=1$ ; decoder solver matrix_free_phi1_auto
Mean-field + Adam	Whitened inputs, 50M tokens, 20 inference steps, readout $p>0.5$ . After 2,000 updates, clip the dictionary gradient at three times its running-mean norm and skip batches whose negative ELBO exceeds five times its running mean, capped at approximately 20% of recent batches
BatchTopK / Matryoshka	Whitened inputs; batch 4,096; 12,207 updates (49,999,872 tokens); constant LR $3\times 10^{-4}$ ; FP32 SAEs, BF16 Gemma; SAELens 6.51.1. ZCA fitted on 500M training tokens. Target $L_{0}$ values in Table 8 were selected on a separate pilot
Baseline initialization / losses	Decoder norm initialized to 0.1; decoder bias subtracted before encoding; activations rescaled by decoder norm; auxiliary-loss coefficient 1; threshold EMA rate 0.01
Matryoshka levels	4,096, 16,384, and the full width (16,384 appears once for the smallest model); Matryoshka auxiliary loss enabled
Gemma Scope	Release google/gemma-scope-2b-pt-res , checkpoint paths layer_12/width_{w}/average_l0_{s} , with $(w,s)$ given by each row of Table 9 . Training: 4B tokens at 16K, 8B at 32K–524K, and 16B at 1M ( Lieberum et al., 2024 )

Appendix

Table 7: Gemma training and evaluation settings. Preprocessing, checkpoint, inference depth, and readout identify an evaluation configuration.

		Sparse probing ( Kantamneni et al., 2025 )						RAVEL ( Huang et al., 2024 )
Width	Target $L_{0}$	$k=1$			$k=5$
		Acc. $\uparrow$	AUC $\uparrow$	F1 $\uparrow$	Acc. $\uparrow$	AUC $\uparrow$	F1 $\uparrow$	Overall $\uparrow$	Cause $\uparrow$	Isolation $\uparrow$
BatchTopK (16k)
16k	200	${.765}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.004}}$	${.791}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.004}}$	${.758}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.005}}$	${.810}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.002}}$	${.859}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.002}}$	${.808}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.003}}$	${.717}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.009}}$	${\mathbf{.698}}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.018}}$	${.736}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.021}}$
16k	400	${\mathbf{.767}}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.002}}$	${\mathbf{.798}}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.003}}$	${\mathbf{.761}}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.004}}$	${\mathbf{.814}}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.002}}$	${\mathbf{.862}}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.002}}$	${\mathbf{.812}}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.002}}$	${\mathbf{.729}}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.005}}$	${.692}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.019}}$	${\mathbf{.767}}\scriptscriptstyle{\color[rgb]{0.3984,0.3984,0.3984}\pm{.017}}$
Matryoshka BatchTopK (16k)

Appendix

Table 8: SAEBench results for SAEs trained on 50M tokens of whitened Gemma-2-2B layer-12 activations. Entries report mean

\pm

sample standard deviation over five training seeds (1–5), evaluated at the final checkpoint (49,999,872 tokens). Sparse-probing metrics are macro-averaged over 113 tasks; RAVEL overall is the disentanglement score. All scores are fractions, and higher is better.

L_{0}

denotes target training sparsity, fixed using a separate seed-0 pilot. Bold indicates the highest mean across

L_{0}

values within each width–architecture group, selected before rounding.

		Sparse probing									RAVEL
Width	$L_{0}$	$k=1$			$k=2$			$k=5$			Disent- angle $\uparrow$	Cause $\uparrow$	Isola- tion $\uparrow$
		Acc. $\uparrow$	AUC $\uparrow$	F1 $\uparrow$	Acc. $\uparrow$	AUC $\uparrow$	F1 $\uparrow$	Acc. $\uparrow$	AUC $\uparrow$	F1 $\uparrow$	Disent- angle $\uparrow$	Cause $\uparrow$	Isola- tion $\uparrow$
16k	22	0.7337	0.7476	0.7159	0.7529	0.7861	0.7441	0.7771	0.8243	0.7717	0.5989	0.5121	0.6857
16k	41	0.7519	0.7710	0.7366	0.7686	0.8056	0.7601	0.7893	0.8370	0.7835	0.6644	0.6163	0.7126
16k	82	0.7627	0.7857	0.7545	0.7821	0.8207	0.7768	0.8047	0.8523	0.8016	0.6960	0.6489	0.7431
16k	176	0.7679	0.7938	0.7602	0.7890	0.8262	0.7851	0.8129	0.8587	0.8119	0.7206	0.6508	0.7904

Appendix

Table 9: SAEBench evaluations of released Gemma Scope layer-12 checkpoints ( Lieberum et al., 2024 ; Karvonen et al., 2025 ) , organized by dictionary width and average sparsity

L_{0}

Width $K$	$\mathrm{Be}$ FOND	BatchTopK	Matryoshka	Mean-field + Adam
16,384	33.125	2.719	2.991	2.714
65,536	118.959	9.202	10.737	10.157
131,072	234.665	17.935	20.075	19.914
524,288	989.321	70.121	77.382	78.331

Appendix

Table 10: Elapsed training time in seconds for 999,424 training tokens on eight NVIDIA H200 GPUs (lower is better). Each model is measured once at seed 0. Timings include

\mathrm{Be}

FOND ’s prior updates and exclude evaluation, initialization, warmup, data preparation, checkpointing, and logging. Matryoshka denotes Matryoshka BatchTopK. Mean-field + Adam uses 20 inference steps, gradient clipping, and batch skipping.

Explore similar work

Expander Sparse Autoencoders: Parameter-Efficient Dictionaries for Mechanistic Interpretability
Jul 2, 2026Rodrigo Mendoza-SmithSparse Recovery Activation Sparsity
Independent Researcher.
A Dominant Diffuse Phase in the Sparse Autoencoder Phase Diagram
Sep 9, 2026Alexis D. PlascenciaFeature Superposition Representation Disentanglement
Stop Probing, Start Coding: Why Linear Probes and Sparse Autoencoders Fail at Compositional Generalisation
Mar 30, 2026Vitória Barin-Pacela, Shruti Joshi, Isabela Camacho +2Sparse Autoencoders Sparse Coding
Mila - Québec AI Institute & Université de Montréal · Santa Clara University · Cold Spring Harbor Laboratory

Jul 2, 2026, cs.LGJ/K move · Enter open · S save

Expander Sparse Autoencoders: Parameter-Efficient Dictionaries for Mechanistic Interpretability

Rodrigo Mendoza-Smith

Independent Researcher.

Sparse Recovery Activation Sparsity Sparse Autoencoders Mechanistic Interpretability Neural Network Compression

Open paper AskPDF arXiv

Sparse autoencoders (SAEs) decompose internal activations of neural networks into sparse linear combinations of learned features by fitting an overcomplete dictionary

\mathbf{W}\in\mathbb{R}^{m\times n}

with

m<n

, and inferring a sparse code

\mathbf{x}\in\mathbb{R}^n

from

\mathbf{h}\approx\mathbf{W}\mathbf{x}

. This inference problem closely resembles the canonical setup of compressed sensing, but dense decoders requires

O(mn)

learned values, which becomes costly at large feature counts. We introduce Expander SAEs: TopK SAEs whose decoder and tied encoder are supported on a left-

d

-regular expander mask with

d\ll m

, learning only

dn

decoder values while keeping the sparse-coding problem

(m,n,k)

fixed. The same structure reduces storage and turns the matching-pursuit correlation step

\mathbf{W}^\top \mathbf{r}

in OMP into an

O(dn)

gather-and-reduce operation. Our experiments show that across Pythia-70M/160M, Qwen2.5-3B, and Llama-3.2-1B residual-stream activations, varying

d

traces a consistent storage--fidelity frontier, and that at the most compressed modern-LM setting, Qwen2.5-3B with

d=7

uses

293\times

fewer learned decoder values than the full dense decoder while retaining

84

% of dense CE-loss recovered. Control experiments show that the improved storage--fidelity tradeoff is driven by sparse, diverse decoder support structure rather than by fewer learned decoder values, and that when sparse and dense decoders are compared at matched parameter count, part of the remaining gap comes from encoder amortisation. On the theoretical side, we show that expansion and column flatness are sufficient for identifiability of noiseless

k

-sparse codes, and we derive complementary sufficient conditions under which OMP recovers the support exactly.

Abstract

Explore similar work

Expander Sparse Autoencoders: Parameter-Efficient Dictionaries for Mechanistic Interpretability

A Dominant Diffuse Phase in the Sparse Autoencoder Phase Diagram

Stop Probing, Start Coding: Why Linear Probes and Sparse Autoencoders Fail at Compositional Generalisation