---
license: gemma
base_model: google/gemma-4-26B-A4B-it
tags: [nla, natural-language-autoencoder, interpretability, gemma4, lora]
---

# Natural Language Autoencoder for Gemma-4-26B-A4B (layer 20)

A **Natural Language Autoencoder (NLA)** for `google/gemma-4-26B-A4B-it` — a pair of
adapters that map residual-stream activations at **layer 20** (2/3 depth; d_model 2816)
to natural-language explanations and back, plus a GRPO-tuned verbalizer.

- **`av_lora/`** — the **verbalizer (AV)**, SFT: activation → text. Reads an injected layer-20
  activation (Karvonen norm-matched ADD at the layer-1 output) and describes the features the
  model is "thinking about". LoRA (r=128, rsLoRA) on the text q/k/v/o projections.
- **`ar_lora/`** — the **reconstructor (AR)**, SFT: text → activation. A 21-layer truncated
  backbone + Linear(d,d) value head that reconstructs the layer-20 activation from the
  explanation. `ar_lora_value_head.safetensors` bundles the critic LoRA weights **and** the
  value head; `ar_meta.json` records the truncation/final-norm config.
- **`rl_av_lora/`** — the **GRPO-tuned verbalizer** (optional). The AV after reinforcement
  learning against the frozen AR reward (reward = −log MSE of the AR's reconstruction from the
  AV's *sampled* explanation). Drop-in replacement for `av_lora/`; improves the quality of the
  AV's own on-policy explanations (see RL results below).

Reimplementation of [Natural Language Autoencoders](https://transformer-circuits.pub/2026/nla/index.html)
via the self-contained [nanoNLA](https://github.com/ceselder/nanoNLA) trainer.

## Results

### Stage 1 — SFT (the core NLA)

The reconstructor was trained on (teacher-explanation, activation) pairs and evaluated by
**Fraction of Variance Explained (FVE)** = `1 − mse_reconstruction / mse_predict-the-mean`.

- **Held-out FVE = 12.2%** (paper-def / variance-around-the-mean baseline; n=3000).
  Meannorm-def FVE = 20.1%.
- **Shuffled-pair control = −31.8%** — strongly negative confirms the AR genuinely *reads the
  explanation* rather than predicting the mean (a control that ignores the text would sit near 0).
- The held-out FVE **saturates** cleanly over 1000 SFT steps (−14% → 11.8%, flat over the final
  ~8 evals): `plots/g4a_ar_fve_saturation.png`. The AV verbalizer CE loss converged 4.8 → 1.7
  (`plots/g4a_av_loss.png`).

![AR FVE saturation](plots/g4a_ar_fve_saturation.png)

### Stage 2 — RL (GRPO), `rl_av_lora/`

GRPO on the AV (actor) against the frozen AR reward, on a document-disjoint RL split. The
meaningful metric is the **on-policy** reconstruction FVE — i.e. how well the AR reconstructs
from the explanations the AV *actually samples* (temperature 1.0), not from teacher explanations.

- SFT-AV on-policy held-out FVE (eval@0): **−10.6%** — the AV's own sampled explanations start
  *worse than baseline* (the AR was SFT'd on cleaner teacher explanations).
- After GRPO it rises to **+1.3%** (peak, step 40), then goes stationary and the held-out curve
  regresses while the per-step *train* rollout FVE keeps climbing (reward overfitting) — so
  training was **halted at stationarity** (step ~40, best held-out checkpoint shipped here).
- Net: GRPO closes ~half the teacher-vs-on-policy gap (≈ **+12 pp** of on-policy FVE), bringing
  the AV's sampled explanations from below-baseline up to ~baseline. Consistent with the modest
  RL bump reported for NLAs. `plots/g4a_rl_fve_trajectory.png`.

![RL on-policy FVE](plots/g4a_rl_fve_trajectory.png)

## How this was built (gemma-4 port)
Gemma-4 is a **multimodal MoE** (`Gemma4ForConditionalGeneration`); porting nanoNLA required:
1. Resolving the text decoder at the nested `model.language_model.layers` path (transformers 5.x).
2. LoRA on the **text** q/k/v/o projections only (the vision tower uses `Gemma4ClippableLinear`,
   unsupported by PEFT).
3. Gradient checkpointing with `use_reentrant=True` (transformers-5 default is incompatible with
   the in-place injection hook).
4. Adding gemma4 to the √d_model embedding-scale registry (for RL-rollout injection).
These live in `training_code/nla/` (see `arch_adapters.py`, `train_sft.py`,
`train_rl_self_contained.py`).

### Data / quality caveat
Activations were extracted from `google/gemma-4-26B-A4B-it` @ layer 20 over FineWeb (sample-10BT,
streamed). Stage-2 explanations were generated by a **self-hosted Qwen3.6-35B-A3B** model (the
intended Claude Sonnet 4.6 Batch-API path was unavailable due to an overnight queue outage).
Because explanations depend only on the read-window *text* (not the activation vector), they were
produced GPU-free and joined back to the activations by `stage2 --cache-from` (text-keyed): the
**92.8% cache hit-rate** equals the upstream keep-rate exactly, proving the CPU-reproduced prompts
aligned with the GPU extraction. Re-running stage 2 with Claude would be a drop-in quality upgrade.

## Training code
`training_code/` contains the full nanoNLA trainer (`nla/`), datagen (`nla/datagen/`), the corpus
builder, the datagen + train shell scripts, and the datagen config. See `train_sft_avar.sh` and
`train_rl.sh`.