| # WriteSAE checkpoints — model card |
|
|
| ## Artifact |
|
|
| Sparse autoencoders trained on cached GatedDeltaNet write activations, accompanying the paper *WriteSAE: Sparse Atoms that Substitute for Recurrent State Writes*. |
|
|
| - HF repo: [JackYoung27/writesae-ckpts](https://huggingface.co/JackYoung27/writesae-ckpts) |
| - Code: [https://github.com/JackYoung27/writesae](https://github.com/JackYoung27/writesae) |
| - Paper: see `paper/` in the code repo |
|
|
| ## Variants released |
|
|
| | variant | encoder | decoder | downstream PPL Δ (L9 H4, 0.8B) | |
| | --- | --- | --- | ---: | |
| | WriteSAE (primary) | bilinear $v_i^\top S w_i$ | rank-1 $v_i w_i^\top$ | +0.58% | |
| | FlatSAE | linear on vec($S$) | flat | +3.31% | |
| | MatrixSAE | linear on vec($S$) | rank-1 | +11.18% | |
| | BilinearSAE | bilinear | bilinear | +1.33% | |
|
|
| Each variant carries TopK sparsity (per-variant *k*; BatchTopK supported via `--use_batchtopk`). |
|
|
| ## Base models |
|
|
| - Qwen/Qwen3.5-0.8B (primary, all main-text results) |
| - Qwen/Qwen3.5-4B (cross-scale replication) |
| - Qwen/Qwen3.5-27B (cross-scale replication) |
| - fla-hub/delta_net-1.3B-100B, fla-hub/gla-1.3B-100B (cross-arch) |
| - state-spaces/mamba-2.8b (cross-arch, no SAE training; spectral audit only) |
| |
| ## Layer / head coverage |
| |
| - L9 H4 — primary substitution site (0.8B) |
| - L1 H4, L17 H4 — cross-layer firing-ordering experiments |
| - Full L9, all 16 heads — bestiary distribution and selectivity sweep |
| - 47 cells × 8 features — uniqueness controls |
| |
| ## Training |
| |
| - Architecture: rank-1 decoder atoms $v_i w_i^\top$, bilinear encoder |
| - Dictionary size: 16,384 features (`--n_features`; expansion factor configurable) |
| - Sparsity: TopK (per-variant *k*) |
| - Data: OpenWebText (`Skylion007/openwebtext`, streaming) tokenized with the Qwen3.5 tokenizer; states extracted via `experiments/extraction/extract_states.py` |
| - Split: 80/20 train/val, seed 42 (deterministic) |
| - Compute: ~200 A100-hours total across all variants |
|
|
| ## Evaluation |
|
|
| - Substitution: forward KL at firing positions under three matched-Frobenius-norm conditions (atom, ablation, random rank-1). Pooled *n*=4,851 firings (L1 1,500 + L9 1,851 + L17 1,500) at Qwen3.5-0.8B L9 H4. |
| - Closed-form factorization: per-firing logit shift predicted at median *R²*=0.98 across 200 atom×ε cells. |
| - Steering: 30 prompts × 5 install positions × 8 target tokens × 3 magnitudes = 3,600 trials. |
|
|
| ## Files on the HF repo |
|
|
| ``` |
| JackYoung27/writesae-ckpts/ |
| writesae_L{1,9,17}_H{0..15}/ # WriteSAE checkpoints |
| flatsae_L{1,9,17}_H{0..15}/ # FlatSAE controls |
| matrixsae_L{1,9,17}_H{0..15}/ # MatrixSAE controls |
| bilinearsae_L{1,9,17}_H{0..15}/ # BilinearSAE controls |
| exp_c_full_seed2026/exp_c/predictions/ # 192 prediction NPZ shards |
| memory_edit_L9_H4/ # F412 ERASE records |
| behavioral_steering_L9_H4/ # INSTALL records |
| flat_sae_svd_{gdn,mamba2,rwkv7}/ # Cross-arch SVD |
| manifest.json # tag → SHA256 + metadata |
| ``` |
|
|
| ## Loading |
|
|
| ```python |
| import torch |
| ckpt = torch.load("writesae_L9_H4_nf2048_k32_s42/best.pt", weights_only=False, map_location="cpu") |
| print(ckpt["config"], ckpt["val_mse"]) |
| ``` |
|
|
| ## License |
|
|
| - Code and SAE checkpoints: MIT |
| - Base models retain the upstream Tongyi Qianwen license (Qwen3.5) and the licenses of the FLA models (DeltaNet, GLA) and Mamba-2. |
| - We do not redistribute base-model weights; all base models are loaded from their public HF release at runtime. |
|
|
| ## Citation |
|
|
| ```bibtex |
| @article{young2026writesae, |
| title = {WriteSAE: Sparse Atoms that Substitute for Recurrent State Writes}, |
| author = {Jack Young}, |
| year = {2026}, |
| journal= {arXiv preprint}, |
| url = {https://github.com/JackYoung27/writesae}, |
| } |
| ``` |
|
|
