cds-jb/gemma-4-26b-a4b-nla-l20

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 via the self-contained 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).

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.

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.