SAE checkpoints across SFT and PPO training
TopK sparse autoencoders for residual-stream activations of
Qwen/Qwen2.5-0.5B-Instruct trained on GSM8K. The repository provides
checkpoint-aligned SAEs for one supervised fine-tuning (SFT) trajectory and
four PPO conditions.
Available SAE families
| Family | Checkpoints | Location |
|---|---|---|
| SFT | instruct_base, sft_step{29,58,116,174,232,290,348} |
sae_sft/ |
| PPO, flexible reward | instruct_base, ppo_step{10,30,60,100,140,180,200} |
sae_flexible/ |
| PPO, strict reward | instruct_base, ppo_step{10,30,50,80,116} |
sae_strict/k64/ |
| PPO, high-KL | instruct_base, ppo_step{10,30} |
sae_kl0p025/ |
| PPO, shuffled-label control | instruct_base, ppo_step{10,30} |
sae_shuffled/ |
All main runs use TopK SAEs with 8× expansion (d_sae=7168) at layers 6,
12, 18, and 23. The normal setting is k=64; the flexible and strict L23
robustness runs use k=256 where indicated in the metric table.
Core findings
- SFT reconstruction remains strong at layers 6 and 12 across training, while raw reconstruction error increases at layers 18 and 23 after the middle SFT checkpoints.
- PPO reconstruction error is generally stable or improves over checkpoints in the available runs.
- Few SAE features are dead: the observed dead-latent fraction is at most 4.5% for the evaluated SFT/PPO runs and is zero for all L23 runs.
Checkpoint metrics
results/sae_checkpoint_metrics.csv is the canonical 118-row table. It
contains raw MSE, NMSE, loss recovery, and dead-latent fraction for every
available (training regime, chain, checkpoint, layer) combination.
| SFT layer | Raw MSE: base → final | Dead latents: base → final |
|---|---|---|
| L6 | 0.0238 → 0.0169 | 0.60% → 0.22% |
| L12 | 0.0408 → 0.0362 | 1.12% → 0.80% |
| L18 | 0.1784 → 0.1980 | 0.40% → 0.08% |
| L23 | 0.5379 → 1.1599 | 0.00% → 0.00% |
Raw MSE should be compared within a layer. Use NMSE for comparisons across layers. Dead-latent fraction is the relevant sparsity/feature-availability metric: average L0 is nearly fixed by TopK selection.
Files
| File | Contents |
|---|---|
results/sae_checkpoint_metrics.csv |
Consolidated SFT and PPO metrics. |
results/sae_mse_dead_sft.csv |
Direct SFT raw-MSE and dead-latent results. |
collation/SAE_Collation.xlsx |
Shared collation workbook; use the sae_michael sheet. |
loader.py |
Convenience loader for SAE checkpoints. |