sae-cub-stage4-convergence

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.

Contents

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

Coefficient-matched baselines (best validation-objective epoch)

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

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.