gemma-4-26b-a4b-nla-l20-sonnet46

Natural Language Autoencoder — gemma-4-26B-A4B-it (layer 20), Sonnet-4.6 warmstart

A Natural Language Autoencoder (NLA) for google/gemma-4-26B-A4B-it: a pair of LoRA-adapted models that map the residual-stream activation at layer 20 (2/3 depth; d_model = 2816) to a natural-language description and back.

• av_lora/ — AV (activation verbalizer): activation -> text. Reads an injected layer-20 activation and writes a 2-3 feature description (<explanation>...</explanation>).
• ar_lora/ — AR (activation reconstructor): text -> activation. A truncated 21-layer backbone + Linear(2816, 2816) value head that reconstructs the layer-20 vector from the explanation text. ar_lora_value_head.safetensors bundles the LoRA + value head.
• rl_av_lora/ — the AV after GRPO RL (best on-policy checkpoint, step 60). Optimised so the AR can reconstruct from the AV's own sampled explanations.
• rl_ar_lora/ — the AR co-trained during RL (the live reward model).

Injection = embedding REPLACEMENT (paper-style, not nanoNLA's default add-at-layer-1): the activation is L2-normalised to injection_scale = sqrt(d_model) ≈ 53 and overwrites the marker token's input embedding. Built with nanoNLA.

Results

FVE = 1 − mse_reconstruction / mse_predict-the-mean (both vectors L2-normalised; higher = the explanation captures more of the activation's directional content).

RL lifts the on-policy FVE from −2.0% to +11.6%, closing ~70% of the gap to the 17.4% teacher ceiling — and ~9× the prior add-at-layer-1 run (which peaked at +1.3%). The teacher-vs-on-policy gap exists because SFT only trains the AV to imitate Sonnet (cross-entropy), while its sampled outputs are noisier and out-of-distribution for the teacher-trained AR; RL optimises the real objective directly.

Data

Explanations from ceselder/nla-matryoshka-warmstart-sonnet46 (Claude Sonnet-4.6). Each row's truncated input_text was tokenised with the gemma-4 tokenizer and the layer-20 activation extracted at its last token; 448,010 (activation, explanation) pairs total.

Sidecar

Each adapter ships an nla_meta.yaml with the prompt templates, the injection marker token id (㈜ / 246566), its canonical neighbours, injection_scale, and mse_scale — load those, never hardcode them.

Training code

training_code/ has the last-token extractor, the datagen-stages driver, and the AV/RL launchers (origin/main nanoNLA hyperparameters: LoRA r128/α16, lr 3e-5 SFT / 1e-5 RL, eff-batch 64; gemma-4 needs quant=none bf16 for SFT since bitsandbytes can't 4-bit the MoE experts).