SAE convergence comparison on CUB ConvNeXt Stage-4 activations

Public research artifact: trained sparse-autoencoder checkpoints and the cached Stage-4 activation dataset from the full-data convergence comparison (topk_selector_comparison). Everything needed to reproduce the reported Pareto results is included — no backbone forward pass is required to load the data.

checkpoints/
  topk_l0_16/  topk_l0_32/                  # fixed TopK, per-epoch best + last
  variable_topk_l0_16/  variable_topk_l0_32/
  batch_topk_l0_16/  batch_topk_l0_32/
  jumprelu_l0_16/  jumprelu_l0_32/          # selected coefficient trial (+ selection.json)
  standard_l0_16/  standard_l0_32/          # selected coefficient trial (+ selection.json)
  coefficient_trials/                        # every trial of the coefficient search
    jumprelu/coefficient_*/training/
    standard/coefficient_*/training/
data/
  stage4_maps.pt                             # 847 MB, memory-mappable activation cache
  normalization.pt                           # training-only per-channel mean/std
  data_metadata.json                         # cache provenance and split sizes
metadata/                                    # metrics CSVs, calibration, plots, execution manifests
SHA256SUMS                                   # checksums for every checkpoint and data file

Each run directory contains checkpoint.pt (best validation objective), last_checkpoint.pt (final epoch), history.csv and termination.json. Convergence runs and baseline trials both stop on a validation plateau (min 50 epochs, patience 30, relative delta 1e-4, LR halving to 1e-6, 500-epoch cap). They are validation-based early-stopping artifacts, not proofs of optimizer convergence.

Checkpoints

Full-data convergence runs

Model	Nominal L0	Directory	Best epoch / total	val MSE	val L0	sha256 (best)
TopK	16	`checkpoints/topk_l0_16/`	57 / 87	0.459257	15.9999	`b91c404ee644e815…`
TopK	32	`checkpoints/topk_l0_32/`	72 / 102	0.380467	32.0000	`af2afcd525eaa319…`
Variable TopK (learned linear selector)	16 (nominal)	`checkpoints/variable_topk_l0_16/`	96 / 114	0.404410	16.4075	`9a83e9b1dfded742…`
Variable TopK (learned linear selector)	32 (nominal)	`checkpoints/variable_topk_l0_32/`	104 / 134	0.318392	32.7348	`ce03557c4e544433…`
BatchTopK	16	`checkpoints/batch_topk_l0_16/`	103 / 133	0.381870	16.8384	`c7ba7cc6f4a0b402…`
BatchTopK	32	`checkpoints/batch_topk_l0_32/`	107 / 137	0.298611	33.3092	`c6718238c759d489…`

Coefficient-matched baselines (best validation-objective epoch)

Model	Nominal L0	Directory	Best epoch / total	val MSE	val L0	sha256 (best)
JumpReLU	16	`checkpoints/jumprelu_l0_16/`	122 / 152	0.395566	15.8839 coef=0.00423084	`9d96346535e9f535…`
JumpReLU	32	`checkpoints/jumprelu_l0_32/`	117 / 147	0.309862	32.3392 coef=0.00235609	`f08ddf9cb9068c72…`
Standard ReLU + L1	16	`checkpoints/standard_l0_16/`	55 / 85	0.458149	15.6162 coef=0.00469905	`a7cca5f4985a3a9c…`
Standard ReLU + L1	32	`checkpoints/standard_l0_32/`	127 / 157	0.373885	32.7543 coef=0.00337395	`addbdc4c3b777421…`

jumprelu_l0_* and standard_l0_* point at the trial selected by validation L0 within 5% of the target (test data never used for selection). All coefficient trials are preserved under checkpoints/coefficient_trials/. JumpReLU and Standard ReLU monitor validation MSE plus their fixed sparsity penalty, so their best epoch is not necessarily the lowest-MSE epoch.

Full SHA-256 checksums (best and last):

b91c404ee644e81571a45fc9c5d5fb62669586c165246b57325e9d03ec048ae8  checkpoints/topk_l0_16/checkpoint.pt
03e385317b296dc71c76e43c3107211d49a3dfc636383321e3d7474e0bf2a47a  checkpoints/topk_l0_16/last_checkpoint.pt
af2afcd525eaa319e1a13bce7890e4853cf4db620b920d837f506794d2b435d3  checkpoints/topk_l0_32/checkpoint.pt
ec4c9ac2f1f45ed420a066f775adf25e6141ed12e699f522ae227b44950d821b  checkpoints/topk_l0_32/last_checkpoint.pt
9a83e9b1dfded742fadb71f7dc36df41fc1eee169fa08c088ff50b8c517a3d13  checkpoints/variable_topk_l0_16/checkpoint.pt
097b4fc4e2be7988b5d70e88600cc74d7434504608c1b3281c9fcd843742b2c3  checkpoints/variable_topk_l0_16/last_checkpoint.pt
ce03557c4e5444334e46a230183821f7028c1aa0162fe3535deea3b670dce8a5  checkpoints/variable_topk_l0_32/checkpoint.pt
7d2f2e2b28879bdba3a55bbc261ec8d74150442397a48dec6b4461ec3cebef85  checkpoints/variable_topk_l0_32/last_checkpoint.pt
c7ba7cc6f4a0b402e36902adbb91f6db6f01fe13e72efa069c8babc6a8201ca7  checkpoints/batch_topk_l0_16/checkpoint.pt
20ba145ec9276d56d08c6208234b290145214b47e43d04f434edcbfdf31fc2b8  checkpoints/batch_topk_l0_16/last_checkpoint.pt
c6718238c759d489ab9d3c0432f8b4520162e6119fffcf1dd2c9cf65a0ab7c67  checkpoints/batch_topk_l0_32/checkpoint.pt
f6f9897104a57138f3e030a96023618b832021f750409a09a287b109664b3ed9  checkpoints/batch_topk_l0_32/last_checkpoint.pt
9d96346535e9f5356195f0752971dcb5180ddc8057c650daa19cc9f38d6ff210  checkpoints/jumprelu_l0_16/checkpoint.pt
76ce8b2e41a3046c12be51d5a1bf9b2ab48c52a3b5930c90f18ec557a4d6acb2  checkpoints/jumprelu_l0_16/last_checkpoint.pt
f08ddf9cb9068c7277850e3121ab8ca7adfaa061dbed9b58516a4e72520ec0a6  checkpoints/jumprelu_l0_32/checkpoint.pt
db7c4f27f510cd5e28e0689db766b8dbef1bfcc5617c45d0b6446187a0cda6ed  checkpoints/jumprelu_l0_32/last_checkpoint.pt
a7cca5f4985a3a9cd42672ca0108205934cc1c8b86eff9923a4204ab3b3bdfc3  checkpoints/standard_l0_16/checkpoint.pt
5dbce6532374cf43a242d54fce3b16628b905e41ca70051301314f30b5cf7f1f  checkpoints/standard_l0_16/last_checkpoint.pt
addbdc4c3b7774210d06befa200f786eae33f3b433f3b184fd370175ad415c3f  checkpoints/standard_l0_32/checkpoint.pt
28df100b4888b758f134ec6f052f5577de79e8f13b59b7b95ba1d59453e644bc  checkpoints/standard_l0_32/last_checkpoint.pt

Dataset: CUB ConvNeXt Stage-4 activation cache

Property	Value
File	`data/stage4_maps.pt` (887,309,253 bytes, float16)
Tensors	`train_maps` `[4796, 768, 7, 7]`, `val_maps` `[1198, 768, 7, 7]`, `test_maps` `[5794, 768, 7, 7]` plus int64 `*_labels` in [0, 199]
Spatial vectors	235,004 train / 58,702 val / 283,906 test (no pooling, image-level splits)
Normalization	`data/normalization.pt`: population mean/std per channel over training sites; std floor 1e-6

Load the cache with torch.load(..., mmap=True, weights_only=True); tensors are memory-mapped and individual vectors are cast to float32. Normalization statistics were fitted on the training split only. data_metadata.json records the source path, size, mtime, split sizes and the shapes of a finite smoke batch per split.

import torch

payload = torch.load("data/stage4_maps.pt", map_location="cpu", weights_only=True, mmap=True)
train_maps, train_labels = payload["train_maps"], payload["train_labels"]   # [N, 768, 7, 7], [N]
norm = torch.load("data/normalization.pt", map_location="cpu", weights_only=True)
mean, std = norm["mean"], norm["std"]                                      # [768]

Checkpoints are plain torch.save dictionaries (model/optimizer/scheduler/ inactivity-tracker state, epoch, stopping state, source script, normalization metadata). They are not state_dict-only files: use weights_only=False and read the keys your own loader expects, or reuse the model definitions from the topk_selector_comparison.models package.

Training summary

All models share a 768 -> 3072 biased linear encoder and a bias-free linear decoder, initialized from the same seed-42 tensors, and are trained on all 235,004 training vectors (batch 128, Adam 1e-3, gradient clipping 1.0) with AuxK (64 features, weight 1/32, one full epoch of inactivity) added to every objective. Fixed and learned TopK use reconstruction MSE plus their respective sparsity terms; BatchTopK uses a batch-wide top-floor(B*k) budget with a frozen EMA inference threshold; JumpReLU uses trained per-feature thresholds with the SAELens Step/rectangular-STE formulation; Standard ReLU uses the decoder-norm-weighted L1 penalty. Reference implementation: SAELens revision ee45e7406165ce267b33d33b7371bbb63ef24db8 (MIT; the adapted baselines are controlled re-implementations of its activation rules and objectives, with AuxK added, not an unmodified recipe benchmark).

Caveats

The Stage-4 cache is legacy: it carries no ConvNeXt checkpoint identity, no image IDs, and cannot independently certify its origin or split disjointness.
Validation L0 is a soft target for every model; the Pareto comparison uses measured test L0, never a forced sparsity.
Test metrics and calibration for the plots live in metadata/; they are versioned lightweight artifacts, while the raw per-site selection traces (NPZ) stay in the local working tree.

Downloads last month: -; Downloads are not tracked for this model. How to track

Inference Providers NEW

This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

antoniotorres02
/

sae-cub-stage4-convergence