Post
123
Just dropped a new SAE juiceb0xc0de/gemma-4-e2b-it-SAE
I built this using the Event-Aware SAE Trainer I have been developing. It trained all 35 layers in 1 run and handled the transition from KV owning layers to the KV cache at layers 15 and 16 where MLP width doubles.
I developed the trainer with two goals in mind, keep it accessible and affordable. I wanted to make SAE training simple enough that you don't require a team of frontier developers or a corporate budget. Anyone can choose a model -> hit play -> train a set of SAE's. Zero hyperparameter tweaking. Train every layer.
Optimized features:
- Single block rolling capture approx. 2.9x faster vs full forwards capture zero measured relative error
- Hoist system the model lives in CPU the layers required for SAE training stay in VRAM
- Augmented-Lagrangian λ controller with two sided nudge controller for faster, bidirectional L=0 control
- Dead feature revival preventing dead pileup saving layers from requiring another pass
- Layer-aware early stop triggers prevent training from executing maximum training steps dropping training as low as recorded 1000 steps
juiceb0xc0de/gemma-4-e2b-it-SAE
https://github.com/JuiceB0xC0de/event-aware-SAE-trainer
Drop a new discussion in the community and let me know what you think!
I built this using the Event-Aware SAE Trainer I have been developing. It trained all 35 layers in 1 run and handled the transition from KV owning layers to the KV cache at layers 15 and 16 where MLP width doubles.
I developed the trainer with two goals in mind, keep it accessible and affordable. I wanted to make SAE training simple enough that you don't require a team of frontier developers or a corporate budget. Anyone can choose a model -> hit play -> train a set of SAE's. Zero hyperparameter tweaking. Train every layer.
Optimized features:
- Single block rolling capture approx. 2.9x faster vs full forwards capture zero measured relative error
- Hoist system the model lives in CPU the layers required for SAE training stay in VRAM
- Augmented-Lagrangian λ controller with two sided nudge controller for faster, bidirectional L=0 control
- Dead feature revival preventing dead pileup saving layers from requiring another pass
- Layer-aware early stop triggers prevent training from executing maximum training steps dropping training as low as recorded 1000 steps
juiceb0xc0de/gemma-4-e2b-it-SAE
https://github.com/JuiceB0xC0de/event-aware-SAE-trainer
Drop a new discussion in the community and let me know what you think!