juiceb0xc0de/gemma-4-e2b-saes

# Gemma-4-E2B SAE Atlas

JumpReLU Sparse Autoencoders trained on every residual stream layer of `google/gemma-4-E2B-it` using an adaptive Lagrangian controller that eliminates manual per-layer hyperparameter tuning.

## What this is

A complete layer-by-layer SAE atlas for Gemma-4-E2B, trained and published live as each layer completes. Each SAE decomposes the residual stream activations at that layer into a sparse dictionary of 49,152 learned features.

## Prior work

Before training these SAEs, the model was mapped behaviorally using a neural census pipeline across 35 layers × 8 components × 16 behavior categories on 184,320 probe prompts. The results are interactive and fully queryable.

👉 [Gemma-4-E2B Brain Atlas](https://huggingface.co/spaces/juiceb0xc0de/gemma-4-e2b-brain-atlas)

The atlas identified several structural findings that informed SAE training priorities, including a three-phase behavioral leadership transition in the first four layers, a deep-layer gate sparsification event at L23-26, and a selectivity plateau at L4-L6 where neurons are 6.7× more likely to be category-selective than topic-entangled.

## Training methodology

Standard JumpReLU SAE architecture with an adaptive Augmented Lagrangian controller for automatic sparsity targeting. The controller treats L0 sparsity as a hard constraint rather than a soft penalty, using projected dual ascent to find each layer's natural KKT-point without manual tuning.

The key finding: every layer converges to a different λ equilibrium automatically. No grid search, no failed runs, no human in the loop.

## Results so far

| Layer | EV (best) | L0 (best) | dead% (final) | λ_eq | steps | status |
|---|---|---|---|---|---|---|
| 0 | 0.913 | 493 | 0.0% | 1.455e-3 | 3,501 | ✅ |
| 1 | 0.880 | 518 | 0.0% | 1.254e-3 | 3,001 | ✅ |
| 2 | 0.981 | 847 | 3.9% | 3.260e-3 | 9,501 | ✅ use best ckpt |
| 3 | 0.994 | 1112 | — | 4.650e-3 | 15,000 | ⚠️ unstable — use best ckpt |
| 4 | 0.988 | 1111 | 1.1% | 4.750e-3 | 7,001 | ✅ use best ckpt |
| 5 | 0.989 | 1217 | 0.01% | 5.000e-3 | 7,501 | ✅ use best ckpt |
| 6 | — | — | — | — | — | 🔄 training |
| 7-9 | | | | | | ⏳ queued |
| 10-34 | | | | | | ⏳ queued |

Layer 2 is notable — it required 2.2× more steps and a 2.2× higher λ equilibrium than layer 0, consistent with the entanglement cliff measured independently in the brain atlas. The controller handled it automatically.

## Architecture

- Dictionary size: 49,152 features (24× overcomplete)
- Activation: JumpReLU with learned per-feature thresholds
- L0 target: 500 active features per forward pass
- Training data: FineWeb-Edu (pre-tokenized)
- Base model: `google/gemma-4-E2B-it`

## Paper

Methodology writeup coming. The short version: this approach makes full-model SAE atlas training accessible on a single A100 for under $20 total, with zero manual tuning per layer.