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

• 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).

Results

Stage 1 — SFT (the core NLA)

• 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).

Stage 2 — RL (GRPO), rl_av_lora/

• 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.

How this was built (gemma-4 port)

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

Training code