0views0forks
Quite far, as long as you keep the optimizer state out of it. That buys a 1,048,576-feature sparse autoencoder on a single GPU.
Sparse autoencoders, or SAEs, are how interpretability researchers pull a language model’s internal activity apart into individual, namable features: this one fires on legal boilerplate, that one on Python tracebacks. The more features in the dictionary, the more of the model’s repertoire it can resolve, and dictionary size is capped by GPU memory. So the practical question is how cheap you can make the arithmetic before the autoencoder gets worse.
We answered it on one layer each of two small models, Gemma-2-2b and Qwen2.5-3B, across eleven precision recipes and three sparsity designs. The answer has a clean shape. Storing the weights and running the matrix multiplies in low precision costs essentially nothing. Storing the optimizer’s running averages in low precision is where things break, and it breaks on rare, unusually large inputs rather than on the average case. One switch in how the optimizer rounds recovers most of the loss.
Eleven recipes, one autoencoder design, one layer of Gemma-2-2b, everything else held fixed: 32,768 features, 500 million training tokens. The recipes differ in which quantities get fewer bits — the weights, the matrix multiplies, or the optimizer state that Adam keeps for every parameter.
The measurement that separates them is reconstruction error on a token class deliberately held out of training: sequence-start positions, where the model’s activations are far larger than usual. Each recipe is scored against the full-precision autoencoder on exactly those inputs.
Two things follow. Precision in the weights and the matmuls is not the constraint: bf16 weights are as good as fp32 everywhere we looked, and so are fp8 matmuls in both the tensorwise and rowwise variants. The constraint is the optimizer state, and it is not a matter of bit width alone — 4-bit state (24.41×) is barely worse than 8-bit (23.46×), while fp8 state is worse than both (65.68×), and 8-bit state with stochastic rounding is an order of magnitude better than plain 8-bit state at the same bit width.
Where exactly the boundary sits matters for what you do. Halving the optimizer state from 32 to 16 bits costs almost nothing (1.220× against 1.000×, from a single seed), so the rule is not “keep the optimizer in full precision” but “do not go below 16 bits without stochastic rounding.”
Measuring the ladder at 32,768 features is not a special case. Across a 16× range of dictionary width and an 8× range of training data at this site, probing quality is flat: every cell lands between 0.7747 and 0.7918 AUROC. Since quality does not vary with width in this range, the ladder should carry over to the widths a practitioner would actually train, with the caveat that this flatness is measured at one site of one model.
Quality is half the decision; the other half is what you get back. All figures below are measured on one H100 with Gemma-2-2b layer 12 and 4096-token steps, as peak allocated memory and median seconds per step.
The simplest comparison is at 262,144 latents with the ordinary dense step, where full precision still fits: full precision takes 49.57 GiB and 0.1787 s per step, bf16 weights with 8-bit optimizer state take 22.46 GiB and 0.1269 s. That is 45% of the memory at 1.41× the speed. Quantization here is not a speed-for-quality trade, and not a memory-for-speed trade either: the step is smaller and faster at the same time.
At the headline width of 1,048,576 latents every scheme needs the recomputed step, which makes the schemes directly comparable. Full precision does not fit at this width at all.
| Scheme, 1,048,576 latents, recomputed step | Peak memory | Seconds per step |
|---|---|---|
| bf16 + fp8 optimizer state | 45.51 GiB | 0.7607 |
| bf16 + int8 state | 45.54 GiB | 0.8021 |
| bf16 + int8 state, stochastic rounding | 45.54 GiB | 0.8146 |
| fp8 rowwise matmuls + int8 state | 45.54 GiB | 0.8429 |
| bf16 + bf16 state | 54.38 GiB | 0.8014 |
| bf16 + int4 state | 77.25 GiB | 0.9868 |
| full precision | does not fit | — |
Two things to read off the table. Stochastic rounding, which recovers most of the quality lost to 8-bit state, costs 0.0125 s per step and no memory at all, so there is no reason to train 8-bit state without it. And 4-bit state is the odd one out: it needs more memory than 8-bit, 77.25 GiB against 45.54 GiB, despite storing less, because of how the library handles its compressed buffers during the update rather than anything about 4-bit arithmetic (the accounting is in the appendix).
Recomputation is the separate saving. Instead of storing the step’s widest intermediates, it recomputes them chunk by chunk. At matched width and scheme (bf16 + int8 state):
| Latents | Dense step | Recomputed step | Memory | Time |
|---|---|---|---|---|
| 524,288 | 44.63 GiB / 0.2552 s | 22.93 GiB / 0.3537 s | 51% | 1.39× |
| 786,432 | 66.80 GiB / 0.3842 s | 34.23 GiB / 0.5051 s | 51% | 1.31× |
The two savings are independent and compose. Cheap arithmetic makes the step both smaller and faster; recomputation trades time back for memory.
Two pieces of evidence bracket the conclusion, and neither is the conclusion on its own. The first is SAEBench’s sae-probes k-sparse probing benchmark, the standard evaluation the field uses, across 113 classification tasks, plus RAVEL as a causal check; it finds the savings nearly free. The second is reconstruction error on rare, high-magnitude tokens; it finds a cost the first does not see. Both are reported below.
The comparison was 72 matched SAEs — eleven precision recipes, two models, three seeds each, all at 32,768 features and 500 million tokens — scored on SAEBench’s sae-probes k-sparse probing benchmark across 113 classification tasks. The yardstick for a real effect is the spread between full-precision runs that differ only in random seed: 0.0056 AUROC on Gemma, 0.0040 on Qwen. Anything inside that band is indistinguishable from rerunning the baseline.
8-bit optimizer state sits inside it, −0.0027 AUROC on Gemma and −0.0015 on Qwen, neither surviving multiple-comparison correction (q = 0.37 and 0.68). Half-precision weights are equally clean. A causal test on Gemma — can you steer the model’s behavior through the SAE’s features? — agrees: the 8-bit SAE’s cause score differs by −0.0126 (p = 0.43) and its disentanglement score by +0.0116 (p = 0.31), noise in both directions, with a third score, isolation, differing by +0.0359 in the 8-bit autoencoder’s favour but not holding up as an improvement (see the appendix).
This instrument does find the floor of the ladder. 4-bit optimizer state degrades probing on both models, consistently across the 113 sae-probes tasks rather than through a few of them. Eight bits is where the saving stops being free even by this forgiving yardstick.
| Change to the training recipe | Probing effect, Gemma | Qwen | Verdict at q < 0.05 |
|---|---|---|---|
| bf16 weights, full-precision optimizer state | −0.0010 | +0.0009 | no effect |
| 8-bit optimizer state | −0.0027 | −0.0015 | no effect |
| 4-bit optimizer state | −0.0055 | −0.0082 | degrades, both models |
How fine a difference can this benchmark resolve? Two runs of one identical configuration — same width, same data, same recipe — differ by 0.0084 AUROC, while the two autoencoders themselves agree almost exactly: 0.000123 in explained variance and 0.086 in active features per token. Most of that variation therefore comes from fitting the probes, not from training. The consequence for the central claim is worth stating plainly: the measured 8-bit cost of 0.0027 AUROC is smaller than this benchmark can resolve, so the honest reading is “below the benchmark’s resolution” rather than “exactly zero” — which is why the rare-token measurement earns its place as a second, more sensitive check.
Sparse probing averages over hundreds of thousands of ordinary tokens. The tokens where 8-bit state fails are a fraction of a percent of the stream, so the average absorbs them. Scoring the same autoencoders on the held-out high-magnitude classes is what produces the 23–66× gaps in the figure above.
What the benchmark sees also depends on the autoencoder design. Repeating the comparison on Matryoshka BatchTopK, which learns the same top-k rule over nested dictionary widths, the benchmark does register 8-bit optimizer state: 0.0061 AUROC, the one cell of twelve that survives correction (q = 1.88×10−4). Not because the effect is larger there, but because Matryoshka is more reproducible — its full-precision seed band is 0.000545 AUROC against 0.0030 for the reference design, 5.6× tighter, so a small effect clears it. On those same runs, at matched sparsity (1.0015× spread in realized L0), the excluded token class is reconstructed 42.81× worse. Even where the benchmark registers something, the two measurements disagree by orders of magnitude about how much.
Both of those comparisons are anchored: the 23.46× on the reference architecture and the 42.81× on Matryoshka come from runs at identical width (32,768 features), identical token budget and identical seeds, three seeds per setting, so neither can be explained by one setting having effectively trained further. The int8 autoencoders also fit their training distribution very slightly worse than full precision, not better, 0.806 against 0.815 explained variance. Stated carefully: at matched training, int8 optimizer state is 23 to 43 times more brittle on rare high-magnitude tokens.
How fine can the benchmark be? Twelve autoencoders, all trained at one precision setting (bf16 weights, int8 optimizer state), varying only dictionary width and token budget, make that concrete. Read on the benchmark they are one flat sheet, spanning 0.0171 AUROC in total against a smallest detectable difference of 0.0284. Read on rare tokens the same twelve span 690-fold. The grid is a statement about the benchmark’s resolution, not a comparison between precision settings.
Read as reconstruction error on the held-out high-magnitude tokens, those twelve cells span 7.98× to 5,508.4× their own error on ordinary tokens, a 690-fold range, rising with both width and data. That number is each SAE compared against itself, not against a full-precision twin — no full-precision arm was trained at these budgets — so it measures how far off-distribution robustness falls behind on-distribution accuracy as you scale, and the 72-run comparison above is what pins the low-precision share of it. Rare-token error grows with width and token budget on its own, independent of precision: full precision alone goes from 2.77× to 17.34× as width goes from 32,768 to 131,072 features. The measurement is sensitive to training scale as well as to precision, which is why the precision claim rests on the matched comparisons rather than on this grid.
The failure is specific to large inputs, not to unfamiliar ones. One held-out class has activations smaller than an ordinary token’s, and there all eleven precision recipes land between 0.94× and 1.24× of full precision. Novelty alone breaks nothing. Above ordinary magnitude it changes: 8-bit optimizer state reconstructs Gemma’s sequence-start tokens 23.5× worse than full precision, its late special tokens 27.2× worse, and Qwen’s sequence-start tokens 5.59× worse. Magnitude decides whether the failure appears, not how large it gets: Qwen’s sequence-start class has the largest activations of the four and the smallest penalty, so the sizes are not comparable across models.
These magnitudes are not exotic. Over two-million-token samples, Qwen’s sequence-start activations run at 54.84× the ordinary-token norm and Gemma’s at 13.72×. Every sequence the model processes begins with one.
Stochastic rounding in the optimizer pulls every affected class back to within 2.20× of full precision, and on the Matryoshka runs it removes 99.0% of the excess error. It costs no memory and 0.0125 s per step at the million-feature width, about 1.6%.
A third design learns a separate threshold per feature instead of taking a fixed number of the largest activations. With everything else identical — same data, same recipe, same sparsity coefficient — the three precisions converge to different sparsity: 17.32 active features per token in full precision, 69.76 with 8-bit optimizer state, 86.29 with 8-bit state plus stochastic rounding, a 4.98× spread against a target of 100. Nothing in the training loss announces this. You ask for one sparsity and get another, and every downstream comparison is then confounded, which is why this design could not answer the reconstruction question at all.
The cause is the quantized optimizer state on the threshold parameter itself. Giving that one parameter full-precision Adam state, while everything else stays 8-bit, restores the trajectory: L0 27.13 at a budget where the fully full-precision run sits at 19.59 and the fully 8-bit run at 331.28. Instrumenting the step shows why. Measured as realized update divided by learning rate, the threshold parameter moves 1.041 per step with full-precision moments and 0.0783 with quantized ones, about 13.3× smaller: the updates are being rounded away before they can raise the threshold. This mechanism rests on one seed at reduced budget, so treat it as a well-instrumented direction rather than a settled number.
The concrete payoff: 1,048,576 features training on a single H100 with Gemma-2-2b resident on the same GPU and data streaming through it, confirmed over 200 real training steps at 56.50 GiB peak memory and 0.7831 seconds per step.
Two ideas share the credit, one doubling each. 8-bit optimizer state is the first. Recomputation is the second: with the recomputed step, full precision reaches 524,288 features at 56.32 GiB, so at equal engineering effort low precision buys exactly one further doubling.
Quality was measured up to 524,288 features, so the million-feature run is a demonstration that it trains, not evidence that it is good.
The first attempt at cashing the saving found nothing. Eleven precision recipes, from full 32-bit everything down to 4-bit optimizer state, all hit the same width ceiling on the same GPU.
It was a bug. A handful of small quantities were deliberately kept in full precision for numerical safety, a bias here, a normalization term there. Each was tiny, but wherever one of them touched the single widest tensor in the computation it dragged that whole tensor up to full precision with it. Every recipe was paying the full-precision memory bill on the one tensor that mattered. Casting those safety quantities back down before they meet the wide tensor halved peak memory at the old ceiling, 22.46 GiB against 42.72 GiB, and doubled the confirmed trainable width to 524,288 features.
For anyone training SAEs on a memory budget: bf16 weights, 8-bit optimizer state, stochastic rounding switched on, and any learned threshold’s optimizer state kept in full precision. That buys the width, keeps the probing score and the causal behavior indistinguishable from full precision, removes most of the rare-token cost, and avoids the sparsity failure. 4-bit optimizer state is where it stops paying.
One methodological point generalizes past SAEs. When you change how a model is trained, check a metric that stresses rare inputs as well as the standard average. The rare-token failure here was invisible to sae-probes’ 113 tasks and to RAVEL on the reference design, and the sparsity failure does not show up in the training loss either.
| Setting | Value |
|---|---|
| Sites | Gemma-2-2b layer 12 (2304 dimensions) and Qwen2.5-3B layer 18 (2048 dimensions), residual stream; both language models run in bf16 |
| Data | uncopyrighted Pile, streamed and tokenized on the fly, context 1024, first sequence position dropped on both models so each excludes equivalent content, special tokens excluded; document separator is the beginning-of-sequence token for Gemma and the end-of-text token for Qwen, which has none |
| Token budget | 500M tokens per run for the precision comparison; 454,549,504 for the width and data experiments |
| Autoencoder | BatchTopK, k = 100, 32,768 features for the main comparison; decoder initialized to norm 0.1; activations rescaled by decoder norm; dead-feature auxiliary loss at coefficient 1.0; threshold EMA rate 0.01 |
| Optimizer | Adam, learning rate 3e-4, betas 0.9 and 0.999, decayed over the final 20% of each run; batches of 4096 tokens; seeds 0, 1 and 2 for every precision cell |
| Other architectures | Matryoshka BatchTopK at nested widths 2048/8192/32768. JumpReLU in the Anthropic form: tanh sparsity penalty weighted by decoder norm, pre-activation loss coefficient 3e-6, tanh scale 4.0, bandwidth 2.0, initial threshold 0.1, sparsity coefficient calibrated to 0.625, plus the input normalization that recipe requires and the others do not, which is why JumpReLU numbers are compared only within that architecture |
| Precision settings | eleven combinations of weight precision (fp32, bf16), optimizer-state precision (fp32 with master weights, bf16, int8, int4, fp8) and fp8 matrix multiplication (tensorwise, rowwise, rowwise with high-precision gradient), plus one stacking fp8 matmuls with int8 state and one adding stochastic rounding to int8 state. In every low-precision setting a short list of sensitive operations stays in fp32 — decoder-norm reduction, threshold EMA, biases, dead-feature loss, error accumulation — with their per-feature results cast down before touching the large activation tensor |
| Wide runs | the step is recomputed over feature blocks of 65,536 instead of storing intermediates; batch and global top-k pool are unchanged, so the objective is identical |
| Hardware and software | one H100 80GB per run, 79.18 GiB usable; torch 2.11.0 with CUDA 13.0, torchao 0.17.0, SAELens 6.46.1 |
| Measurement | SAEBench’s sae-probes k-sparse probing benchmark: 113 binary tasks over 334,092 prompts per model at k of 1, 2 and 5, compared per task against the same autoencoder’s full-precision counterpart; a probe on the raw residual stream as the ceiling (0.9261 Gemma, 0.9155 Qwen); SAEBench core with KL and cross-entropy enabled for the reconstruction metrics; RAVEL on Gemma as the causal check; and, added by us rather than part of the suite, the rare-token measurement, defined as reconstruction error on a held-out token class over that same autoencoder’s error on ordinary tokens, across four classes spanning 0.776× to 51.6× the ordinary-token norm, the below-ordinary one being the control |
RAVEL result. Run on Gemma only, comparing full precision against bf16 weights with int8 optimizer state, three seeds per arm across six attribute cells. Cause score differs by −0.0126 (p = 0.43), disentangle score by +0.0116 (p = 0.31), isolation score by +0.0359. The isolation difference reaches significance on an unpaired test (p = 0.018) but not on the paired test across cells (p = 0.31), and with three seeds and three metrics untested for multiple comparisons we do not claim it as an improvement. The verdict RAVEL supports is no detectable cost to causal attribute separation, and nothing more.
Autointerp result. SAEBench autointerp ran on eight autoencoders at 200 features each. Scores fall between 0.860 and 0.889 across every precision setting on both models, against a per-feature spread of roughly 0.12. At that scale it cannot separate the settings, so it neither supports nor contradicts the other evaluations; it is reported here only so the panel is complete.
The stored state genuinely is smaller: 4.78 GiB against 8-bit’s 9.14 GiB at 1,048,576 features, so the advertised saving is real. Peak memory is what limits width, and there the 4-bit setting sits 36.07 GiB above what its stored state predicts. Two full-size 32-bit copies of Adam’s two moment buffers over 4.83 billion parameters come to 36.02 GiB, a match to within 0.14%.
Every other setting’s peak decomposes as parameters plus stored state plus an identical 27.36 GiB of gradients and activations, so the excess belongs to this one path: during each update it expands both compressed buffers to full 32-bit precision simultaneously instead of a block at a time. That transient, not the stored state, sets the ceiling. It is a fixable detail of the quantized-optimizer library rather than a property of 4-bit arithmetic, and fixing it would not change the recommendation, since 4-bit state is ruled out on quality.