How Small Can We Go? Calibrating Activation-Cache Compression for SAE Training
Organizations: Stanford University
Abstract
To train a sparse autoencoder (SAE), you need a massive cache of model activations. We look into the effects that compression has on an SAE in two ways: we keep the original fp32-trained SAE and pass the compressed cache into the same SAE to see how many feature activations stay the same (the frozen-SAE test), and we retrain an SAE on the compressed cache and see how many of the original features still exist (the retraining test). We found that in the retraining test, there is no clear scale for retained features. On a GPT-2 cache, rounding everything to fp16 keeps 60.6% of features, but changing one value in a million by the smallest step fp32 allows (a change eight orders of magnitude smaller) keeps only 70.5%. Meanwhile, changing the random seed keeps 23.0%. For the frozen-SAE test, the frozen SAE's correlation with the original increases in order of how precise the codecs are, but plain distance gives nearly the same ordering. We also see the same frozen reading correspond to different percentages retained at different sites. This leaves the question of the right measurement. We build one by calibrating: retraining on identical bytes sets the ceiling, and a reseed sets the floor. We then measure the effect of each compressed codec against its own site's floor. With TopK SAEs on GPT-2 small and Pythia-1.4B at research-scale budgets, naïve per-token int4 loses more features than a reseed everywhere we tested, as do PCA and random projection on GPT-2. Standardizing each channel before quantizing brings int4 back above the floor in terms of retained features, at the same raw file size. At our smallest budget, a bf16 forward pass, which some pipelines use to save memory, loses slightly more features than rotated 8-bit storage. Calibrating a site costs four training runs and two evaluation passes.
Figures & tables
Appendix figures & tables34 assets
Supplementary material from the paper’s appendix.
Appendix
| codec | B/tok | survived % | FVU | mean | stable % | Hung. cos | frozen % |
|---|---|---|---|---|---|---|---|
| fp16 | 1536 | 0.1355 | 0.866 | 95.7 | 0.817 | 100.0 | |
| bf16 storage | 1536 | 55.6 | 0.1355 | 0.853 | 93.3 | 0.799 | 100.0 |
| rot-int8 | 772 | 0.1356 | 0.845 | 94.2 | 0.786 | 100.0 | |
| int8 | 772 | 46.6 | 0.1356 | 0.821 | 91.5 | 0.751 | 100.0 |
| int8-o8 | 780 | 50.8 | 0.1355 | 0.837 | 93.7 | 0.773 | 100.0 |
| rot-int6 | 580 | 0.1355 | 0.823 | 92.1 | 0.752 | 100.0 |
| codec | survived % | floor | stable % |
|---|---|---|---|
| fp16 (bit-identical ceiling) | 100.0 | 6.35 [6.00, 6.71] | 100.0 |
| rot-int8 | 43.0 | 2.73 [2.58, 2.87] | 88.0 |
| int8 | 32.9 | 2.09 [1.98, 2.20] | 84.9 |
| rot-int4 | 25.8 | 1.64 [1.55, 1.73] | 81.6 |
| int4 | 2.8 | 0.18 [0.15, 0.20] | 13.7 |
| seed floor | 1.00 | (def.) |
| codec | survived % | floor | stable % |
|---|---|---|---|
| fp16 (bit-identical ceiling) | 100.0 | 5.44 [5.17, 5.70] | 100.0 |
| rot-int8 | 38.4 | 2.09 [1.99, 2.18] | 89.3 |
| int8 | 27.7 | 1.51 [1.44, 1.57] | 83.4 |
| rot-int4 | 26.1 | 1.42 [1.35, 1.48] | 83.0 |
| int4 | 0.8 | 0.05 [0.03, 0.06] | 5.1 |
| seed floor | 1.00 | (def.) |
| codec | GPT-2 | retr.% | L8 | retr.% | L16 | retr.% |
|---|---|---|---|---|---|---|
| fp16 | 0.99997 | 60.6 | 1.00000 | 100.0 | 1.00000 | 100.0 |
| rot-int8 | 0.99919 | 53.8 | 0.99922 | 43.0 | 0.99924 | 38.4 |
| int8 | 0.99632 | 46.6 | 0.99488 | 32.9 | 0.98916 | 27.7 |
| rot-int4 | 0.98536 | 37.7 | 0.98678 | 25.8 | 0.98696 | 26.1 |
| int4 | 0.84005 | 15.6 | 0.56333 | 2.8 | 0.15633 | 0.8 |
| seed floor | — | 23.0 | — | 15.7 | — | 18.4 |
| codec | survived % | mean |
|---|---|---|
| fp16 | 26.4 | 0.716 |
| rot-int8 | 25.3 | 0.708 |
| rot-int4 | 26.1 | 0.708 |
| int4 | 13.7 | 0.623 |
| seed floor | 23.7–24.2 | 0.699 |
| budget | stab. thresh | fp16 % | rot-int8 % | rot-int4 % | gap pp |
|---|---|---|---|---|---|
| (33M visits) | 0.900 | 95.7 | 94.2 | 88.8 | 7.0 |
| (98M) | 0.936 | 95.7 | 94.7 | 93.5 | 2.2 |
| (328M) | 0.954 | 95.7 | 96.3 | 94.2 | 1.6 |
| task | codec | acc | direction cos |
|---|---|---|---|
| sst2 | fp16 | +0.0011 | 1.000 |
| sst2 | rot-int8 | +0.0000 | 1.000 |
| sst2 | rot-int4 | +0.0011 | 0.949 |
| sst2 | int4-o16 | -0.0080 | 0.778 |
| sst2 | int4 | -0.0608 | 0.489 |
| sst2 | pca256 | -0.0046 | 0.500 |
| gpt2_l6 | |||||
|---|---|---|---|---|---|
| n alive | 4860 | 2927 | 1421 | 654 | 201 |
| floor % | 23.1 | 27.0 | 33.7 | 40.1 | 47.5 |
| rot-int4 floor | 1.63 | 1.50 | 1.38 | 1.28 | 1.13 |
| pythia_l8 | |||||
| n alive | 6573 | 4053 | 1523 | 506 | 100 |
| sparsity | seed floor % | rot-int8 % | floor | stable-only % | stable frac % | |
|---|---|---|---|---|---|---|
| 49.1 | 1.72 | 92.8 | 2,912 | 21.0 | ||
| 53.8 | 2.34 | 94.2 | 4,860 | 15.9 | ||
| 47.4 | 3.23 | 97.1 | 8,269 | 10.1 |
| treatment | B/tok | rel-L2 | frozen | survived % | floor | stable % | Hung. cos | FVU |
|---|---|---|---|---|---|---|---|---|
| fp16 storage | 1536 | 0.00020 | 0.99997 | 60.6 | 2.63 | 95.7 | 0.817 | 0.1355 |
| rot-int8 storage | 772 | 0.00551 | 0.99919 | 53.8 | 2.34 | 94.2 | 0.786 | 0.1356 |
| int8-o8 storage | 780 | 0.00928 | 0.99868 | 50.8 | 2.21 | 93.7 | 0.773 | 0.1355 |
| bf16 forward | 3072 | 0.00980 | 0.99819 | 51.2 | 2.22 | 94.2 | 0.775 | 0.1356 |
| int8 storage | 772 | 0.02660 | 0.99632 | 46.6 | 2.02 | 91.5 | 0.751 | 0.1356 |
| seed floor | — | — | — | 23.0 | 1.00 | (def.) | — | — |
| codec | raw B/tok | zstd B/tok | ratio |
|---|---|---|---|
| fp16 | 1536 | 1419.0 | 1.08 |
| int8 | 772 | 510.9 | 1.51 |
| rot-int8 | 772 | 709.1 | 1.09 |
| int6 | 580 | 431.9 | 1.34 |
| rot-int6 | 580 | 557.0 | 1.04 |
| int4 | 388 | 146.4 | 2.65 |
| codec | B/tok read | disk MB/s | decode MB/s | loader tok/s |
|---|---|---|---|---|
| fp32 | 3072 | 3094 | 10416 | 383k |
| fp16 | 1536 | 3085 | 4748 | 445k |
| int8 | 772 | 2743 | 754 | 319k |
| rot-int8 | 772 | 2917 | 242 | 205k |
| int6 | 580 | 2796 | 208 | 228k |
| rot-int6 | 580 | 2998 | 113 | 147k |
| GPT-2 L6 | |||||
|---|---|---|---|---|---|
| floor (survival) | 59.4 | 42.4 | 32.9 | 23.0 | 10.9 |
| fp16 | 1.45 | 1.83 | 2.15 | 2.63 | 3.78 |
| rot-int8 | 1.39 | 1.71 | 1.97 | 2.34 | 3.19 |
| rot-int4 | 1.22 | 1.39 | 1.50 | 1.63 | 1.90 |
| int4 | 0.83 | 0.79 | 0.75 | 0.68 | 0.61 |
| perturbation | rel-L2 | frozen | survived % | pp/decade |
|---|---|---|---|---|
| zero (bit-exact) | 0 | 1.000000 | 100.0 | — |
| noise | 1.00e-02 | 0.998551 | 51.2 | +7.82 |
| noise | 1.00e-03 | 0.999847 | 59.0 | +2.38 |
| fp16 (codec) | 2.03e-04 | 0.999969 | 60.6 | +0.07 |
| noise | 1.00e-04 | 0.999985 | 60.6 | +1.19 |
| noise | 1.00e-05 | 0.999999 | 61.8 | +1.69 |
| arm | rel-L2 | all alive % | seed-stable % |
|---|---|---|---|
| noise | 1.01e-07 | 66.0 | 96.4 |
| noise | 1.00e-06 | 63.5 | 97.0 |
| noise | 1.00e-05 | 61.8 | 96.1 |
| noise | 1.00e-04 | 60.6 | 95.3 |
| noise | 1.00e-03 | 59.0 | 95.9 |
| noise | 1.00e-02 | 51.2 | 93.4 |
| criterion | perturbed % | floor % | margin (pp) | below ceiling (pp) |
|---|---|---|---|---|
| dense | ||||
| 95.6 | 83.8 | +11.8 | 4.4 | |
| 88.3 | 59.4 | +28.9 | 11.7 | |
| 76.2 | 32.9 | +43.3 | 23.8 | |
| (ours) | 66.0 | 23.0 | +43.0 | 34.0 |
| 46.3 | 10.9 | +35.4 | 53.7 |
| codec | direction cos | causal cos | effect/floor |
|---|---|---|---|
| ag_news (floor 0.9995 / 0.9998) | |||
| fp16 | 1.0000 | 1.0000 | 0.002 |
| rot-int8 | 1.0000 | 1.0000 | 0.007 |
| rot-int4 | 1.0000 | 1.0000 | 0.053 |
| pca256 | 0.9998 | 0.9999 | 0.276 |
| int4 | 0.7901 | 0.8294 | 13.713 FAIL |
| arm | B/tok | rel- | frozen | survived % | floor |
|---|---|---|---|---|---|
| naïve int4 | 388 | 0.4553 | 0.84005 | 0.69 | |
| std-int4 (standardize only) | 388 | 0.1033 | 0.98466 | 1.65 | |
| rot-int4 (standardize rotate) | 388 | 0.0999 | 0.98536 | 1.64 |
| codec | lost/seed % | P(lost s1 lost s0) % | indep. % | lost all 3 % | indep. % |
|---|---|---|---|---|---|
| rot-int8 | 6.1 | 17.8 | 6.3 | 0.13 | 0.023 |
| rot-int4 | 11.1 | 25.3 | 12.0 | 0.90 | 0.134 |
| int4 | 41.0 | 75.5 | 42.9 | 24.16 | 6.876 |
| codec | % | % | % | Hungarian % |
|---|---|---|---|---|
| fp16 | 90.1 / 46.1 | 95.7 / 54.0 | 97.3 / 57.9 | 80.7 / 52.7 |
| rot-int8 | 86.0 / 38.0 | 94.2 / 46.2 | 98.8 / 50.5 | 76.4 / 44.9 |
| rot-int4 | 73.8 / 19.9 | 88.8 / 28.0 | 95.5 / 33.3 | 62.9 / 27.7 |
| int4 | 41.3 / 2.9 | 58.9 / 7.4 | 73.6 / 11.3 | 26.9 / 11.1 |
| site | frozen | rel- | cosine | |
|---|---|---|---|---|
| GPT-2 L6 | 15 | 0.982 | 0.979 | 0.982 |
| Pythia L8 | 5 | 1.000 | 1.000 | 1.000 |
| Pythia L16 | 5 | 1.000 | 1.000 | 1.000 |
| codec | CV AUC | mean act | peak act | rate |
|---|---|---|---|---|
| rot-int6 | 0.836 | 0.758 | 0.713 | 0.552 |
| rot-int4 | 0.844 | 0.755 | 0.711 | 0.570 |
| int4-o16 | 0.831 | 0.749 | 0.711 | 0.564 |
| int6 | 0.838 | 0.747 | 0.709 | 0.573 |
| condition | same effect | same | effect cos | logp | top-50 ctx | |
|---|---|---|---|---|---|---|
| identical bytes | 100.0 | 100.0 | 100.0 | 1.000 | 1.000 | 1.00 |
| 1 ulp on of values | 73.0 | 79.7 | 94.5 | 0.971 | 0.966 | 0.82 |
| fp16 | 63.7 | 68.0 | 89.5 | 0.942 | 0.935 | 0.78 |
| std-int4 | 38.7 | 38.7 | 73.3 | 0.842 | 0.816 | 0.62 |
| naïve int4 | 15.7 | 12.7 | 46.8 | 0.651 | 0.584 | 0.38 |
| reseed, seed 1 | 26.0 | 22.7 | 53.8 | 0.748 | 0.687 | 0.48 |
| condition | lost, | set EV | |||
|---|---|---|---|---|---|
| identical bytes | 100.0 | 100.0 | 100.0 | — | — |
| 1 ulp on of values | |||||
| fp16 | 60.1 | 65.6 | 67.3 | 14.0 | 0.91 |
| rot-int4 | 37.7 | 46.7 | 50.0 | 14.8 | 0.93 |
| std-int4 | |||||
| naïve int4 | 15.6 | 22.4 | 26.2 | 8.5 | 0.73 |
| regime | SAE | agreement % | floor | stable % | fixed-774 % |
|---|---|---|---|---|---|
| replay (20M, 5 passes) | floor (fp32 reseed) | 1.00 | — | — | |
| fp16 | 1.47 | ||||
| std-int4 | 1.16 | ||||
| fresh (98.5M, one pass) | floor (fp32 reseed) | 1.00 | — | — | |
| std-int4 | 1.16 |
| candidate | tokens (M) | stored GB | agreement % | ref |
|---|---|---|---|---|
| fp32, fresh (reference) | 98.5 | 303 | 33.4 (33.2–33.7) | 1.00 |
| fp16, replayed | 20.0 | 30.7 | 32.6 (32.0–33.0) | 0.974 |
| std-int4, more text | 79.0 | 30.6 | 32.6 (31.4–33.3) | 0.974 |
| fp32, replayed † | 20.0 | 61.4 | 32.8 (31.8–33.2) | 0.980 |
| naïve int4, replayed † | 20.0 | 7.8 | 15.9 (15.7–16.1) | 0.476 |
| format | decode | local | busy | volume | busy |
|---|---|---|---|---|---|
| fp32 | CPU (paper) | 117.0 | 0.83 | 75.1 | 0.99 |
| GPU | 118.5 | 0.51 | 119.5 | 0.53 | |
| fp16 | CPU (paper) | 114.6 | 0.89 | 115.1 | 0.84 |
| GPU | 121.0 | 0.28 | 119.7 | 0.33 | |
| std-int4 | CPU (paper) | 92.9 | 0.99 | 50.1 | 1.00 |
| CPU, no matmul | 92.6 | 0.99 | 90.9 | 0.99 |
| codec | frozen mean | model | statistic | crossing % |
|---|---|---|---|---|
| fp16 | 1.0000 | 0.736 | 0.644 | – |
| bf16 | 0.9998 | 0.888 | 0.734 | – |
| int8 | 0.9963 | 0.989 | 0.843 | – |
| int8-o8 | 0.9987 | 0.966 | 0.818 | – |
| int6 | 0.9832 | 0.991 | 0.861 | – |
| int6-o8 | 0.9948 | 0.979 | 0.838 | – |
| target | positive % | statistics | geometry | both |
|---|---|---|---|---|
| , (fit) | 15.9 | 0.846 | 0.829 | 0.868 |
| 24.9 | 0.787 | 0.741 | 0.794 | |
| 32.9 | 0.762 | 0.696 | 0.752 | |
| fresh text | 24.9 | 0.793 | 0.738 | 0.794 |
| 21.0 | 0.817 | 0.801 | 0.837 | |
| 10.1 | 0.899 | 0.842 | 0.902 |
| change | step | agreement % | distance | +2,000 steps |
|---|---|---|---|---|
| , one batch | 0 | 64.7 | 0.6561 | 0.2338 |
| 4,000 | 99.9 | 0.0830 | 0.0452 | |
| 6,000 | 100.0 | 0.0301 | 0.0301 | |
| 7,000 | 100.0 | 0.0198 | – | |
| 7,500 | 100.0 | 0.0107 | – | |
| 7,800 | 100.0 | 0.0041 | – |