ceselder/nanoNLA-gemma4-rl

Gemma-4-26B-A4B NLA — RL (GRPO) checkpoint

A Natural Language Autoencoder (NLA) on google/gemma-4-26B-A4B, trained with the nanoNLA pipeline. An NLA is a pair of LoRA adapters:

• AV (Activation Verbalizer) — reads a residual-stream activation h_l (injected at layer 1, Karvonen norm-matched ADD) and describes it in natural language.
• AR (Activation Reconstructor) — reads that text description and reconstructs the activation (truncated K+1=21-layer backbone + identity-init linear value head).

This checkpoint is the final RL (GRPO) model (iter_000500), produced by warmstarting both adapters with SFT and then co-training them with reinforcement learning so the AV's on-policy descriptions become maximally reconstructable by the AR.

Results

Reconstruction quality is measured by FVE (Fraction of Variance Explained), 1 − MSE / E‖v_norm − μ‖², using the paper raw-variance baseline (the conservative definition; baseline mse_nrm = 0.6269), not the inflated mean-norm baseline.

metric	value
Held-out eval FVE @ AV-SFT init (RL step 0)	33.3%
Held-out eval FVE, RL plateau (steps 200–500)	~51–55% (peak 55.5% @ step 460)
AR reconstruction of gold summaries (ceiling)	58.1%
Train on-policy FVE (768 rollouts/step, final)	~44%

RL closes most of the gap from the warmstart init (~33%) toward the AR's gold-summary ceiling (58.1%). The train-vs-eval offset is because the held-out eval is a small, fixed 20-prompt set (an easier-than-average draw) while the train metric is a fresh 96-prompt batch each step; both use identical T=1 sampling and the same FVE formula.

Training

• Base: google/gemma-4-26B-A4B (MoE, text decoder); extraction at layer 20 (~66% depth); injection marker char ㊗ (token 249568).
• Adapters: LoRA r=128, rsLoRA, α=16 (rsLoRA scaling is essential at this rank).
• Warmstart (SFT): AV and AR each trained 1 full epoch over ~247k (text, Claude-summary, activation) rows. AR held-out FVE on gold summaries = 58.1%.
• RL (GRPO): 500 steps; 96 prompts × group size 8 = 768 rollouts/step; data-parallel across 6× B200 (gloo grad all-reduce). KL penalty β=0.01 toward the frozen AV-SFT init; actor lr 1e-5, co-trained AR lr 5e-5 (paper-faithful: AR updated by regression ‖h − AR(z)‖² each step, no backprop into the AV); sampling T=1; max 150 new tokens.

Checkpoint contents (iter_000500/)

• adapter_model.safetensors + adapter_config.json — the AV LoRA (the "default" policy adapter).
• reference/ — frozen AV-SFT-init adapter (the KL anchor used during RL).
• critic/ar_lora_value_head.safetensors + critic/ar_meta.json — the co-trained AR (LoRA + value head).
• nla_meta.yaml — the sidecar contract (marker token + neighbors, prompt templates, mse_scale, extraction layer index), copied from the AV-SFT warmstart.
• chat_template.jinja — the (patched) Gemma-4 chat template the prompts assume (the base model ships without one).

Usage

Load with nanoNLA:

from nla_inference import NLAClient
client = NLAClient(checkpoint_dir="iter_000500", base_ckpt="google/gemma-4-26B-A4B")
# inject a layer-20 activation h -> get a natural-language description
text = client.verbalize(h)

Caveats

• Research artifact / WIP. Independent reproduction of the original NLA paper's numbers is not verified — see the nanoNLA README disclaimer.
• FVE is the paper raw-variance definition (conservative). A "mean-norm" baseline would read ~10pp higher; don't compare across baselines.
• The held-out eval is only 20 prompts (noisy, ±~5pp).

Method: "Natural Language Autoencoders" (the open-model NLA recipe). Code: ceselder/nanoNLA.