gemma-3-12b-it-nla-L40-wildchat

Gemma-3-12B-IT NLA @ block 40 — WildChat-adapted

A natural language autoencoder (NLA) for google/gemma-3-12b-it at block 40, adapted from web text to conversational activations with 400 further steps of on-policy GRPO on WildChat-1M.

This is a continuation, not a from-scratch model. It resumes achand45/gemma-3-12b-it-nla-L40 at rl_vllm/iter_000400 and keeps training on chat activations. The AV is the same LoRA lineage (r=128, rsLoRA, continued via --av-adapter); the AR is that release's co-trained critic, further co-trained here.

It is the deep arm of a three-layer study. The siblings L24-wildchat and L32-wildchat are the same procedure at blocks 24 and 32, run against the same corpus file and therefore the same held-out documents.

Why

An NLA trained on FineFineWeb is measured on web-text activations. Chat activations are a different distribution, and the parent model loses a lot of ground on them. This run quantifies that drop at block 40 and how much of it RL recovers.

So the parent drops ~13pp when pointed at chat activations, and 400 steps of WildChat GRPO recover +7.8pp — a little over half the transfer gap, and slightly less than the other two depths recovered.

The depth picture

Two things fall out of running all three:

1. Block 32 is a spike, not a trend. It leads on chat activations by ~9pp, and it already led on FineFineWeb in the parent release (68.6% vs 60.9% and 59.0%). The advantage is inherited, not created by chat adaptation — consistent with 32 being the depth the recipe was developed against.
2. Blocks 24 and 40 are indistinguishable after adaptation — 54.7% and 54.2% on the mean-of-final-10 basis — despite sitting 16 blocks apart and starting from different FFW levels.

The chat-transfer drop itself is remarkably uniform across depth: −14.4, −13.5 and −12.9pp. Whatever makes chat activations harder for these verbalizers does not particularly care which layer you read.

⚠️ FVE is not on a common scale across depths. The predict-the-mean baseline is 0.0400 here, 0.0299 at block 32 and 0.0049 at block 24, so equal FVE percentages do not represent equal absolute reconstruction quality. This model's raw ar_mse is ~0.019 against block 24's ~0.0027 — an order of magnitude apart, brought together only by the baseline division. The gains are each measured against their own baseline and are the comparable quantity.

Results

Held-out, doc-disjoint (every row of a held-out document is excluded from training; the corpus is row-shuffled and each document contributes ~10 rows, so a row-level split leaks).

Predict-the-mean baseline on this eval set: 0.0400 (train-set baseline 0.0405).

⚠️ Read these as a band. Eval sampling runs at temperature 1.0, and the parent release measured repeated evals of identical weights spreading ~5 points. The +7.8pp gain clears that; differences under ~5 points here do not — including the gap between this model and the block-24 sibling, which should be read as a tie.

The curve converged early and then held flat: 46.1% → ~54% by step 690, and the last 100 steps moved within ±1pp. The 55.8% peak at step 690 is band-top rather than a level the run settled at.

Entropy rose 1.94 → ~2.2 and settled; KL from the SFT reference rose smoothly 0.93 → ~1.08. 400 steps in 18h39m on 4×H100 at ~166 s/step. No non-finite gradients, no format collapse, no OOMs. Training-step extraction dipped to 98–99% late in the run (a few malformed samples under temperature-1.0 rollouts, absorbed by the reward floor); every held-out eval stayed at 100%.

⚠️ Not a like-for-like continuation of the parent

The parent ran GRPO at batch_prompts: 256; this run used 128. Everything else matches — lr 1e-4, critic lr 8e-5, KL β=0.01 (k3), group 8, temp 1.0, 256 max new tokens. So this run saw 51,200 prompts against the parent's 102,400 per equivalent step count, with noisier per-step gradients. Absolute FVE here is not directly comparable to the parent's published figures. It is comparable to the block-24 and block-32 siblings, which used the same 128.

Extraction contract

Unchanged from the parent — the sidecar (rl_wildchat/nla_meta.yaml) is authoritative and the trainers assert against it.

Because magnitude is discarded by the per-row rescale, FVE measures direction only.

Training data

The same 20,000 WildChat-1M conversations used for the block-24 and block-32 siblings (English, non-toxic, ≥400 chars, selected from 45,560 scanned), rendered with the Gemma-3 chat template and forwarded to get block-40 residuals at 10 sampled token positions each → 200,000 rows, ~150,000 after the doc-disjoint holdout.

All three depths point at the same corpus file on disk, deliberately. Document ids are derived as {corpus path}:{split}:{row index} and the holdout is a hash of that id, so a byte-identical copy at a different path would have produced a different held-out set and made the depths non-comparable.

Reproducibility detail worth knowing: apply_chat_template emits a leading <bos>, and the extraction path tokenizes with add_special_tokens=True, which adds another — giving [2, 2, 105, ...]. The template's BOS is stripped before extraction so each document carries exactly one. Position sampling skips special tokens, so double-BOS would not corrupt the vectors directly, but every activation would be conditioned on a prefix the model never sees at inference.

The activation parquet is not published here. It carries detokenized_text_truncated — verbatim WildChat conversations, which are real user–chatbot logs. Redistributing that is a separate decision under WildChat's own terms, not something this model card covers.

Contents

rl_wildchat/iter_000425 … iter_000800/   AV LoRA every 25 steps (16 checkpoints)
                             adapter_model.safetensors  the trained policy
                             reference/                 frozen SFT copy, the KL reference
rl_wildchat/critic_latest/  AR reconstructor (full weights + value_head, 5 shards)
rl_wildchat/nla_meta.yaml, run_config.yaml, optim_latest.pt

To use the NLA you need one iter_* adapter + critic_latest. Use iter_000800, the final step. Evals run every 10 steps but checkpoints save every 25, so the best eval (55.8% @690) has no checkpoint of its own; the run was flat from ~step 690 and the final 10 evals span 53.5–54.8%, well inside the noise band, so 800 is the defensible pick.

The full every-25-step series is included for training-dynamics and ablation work. The RL adapter is a LoRA on the raw base model (RL continued the SFT adapter via --av-adapter), so no merged AV is required.

Usage

Identical to the parent — same contract, same gotchas.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
from nla.models import NLACriticModel      # pip install -e git+https://github.com/chand-ab/easy_nla
from nla.utils.arch_adapters import resolve_text_model

BASE = "google/gemma-3-12b-it"
REPO = "<local clone of this repo>"

tok  = AutoTokenizer.from_pretrained(BASE)
base = AutoModelForCausalLM.from_pretrained(BASE, dtype=torch.bfloat16)
av = PeftModel.from_pretrained(
    resolve_text_model(base),              # REQUIRED on Gemma-3, see below
    f"{REPO}/rl_wildchat/iter_000800",
)
ar = NLACriticModel.from_pretrained(
    f"{REPO}/rl_wildchat/critic_latest", torch_dtype=torch.bfloat16,
)

Inject the activation at the marker with register_karvonen_hook from nla.utils; it hooks decoder layer 1's residual output (layer_idx=1).

⚠️ Gemma-3 gotcha. gemma-3-12b-it loads as a multimodal wrapper nesting the language model under model.language_model.*, while these adapters are keyed on model.layers.*. Loading onto the unresolved wrapper matches zero keys and PEFT silently random-initializes instead of erroring — you get a fluent model that is not this NLA. Always pass the resolved text model. Sanity check: with one adapter attached, total params should be 12,289,797,888.

Loading the critic prints ... newly initialized: ['model.norm.weight']. Expected — the AR is trained with its final RMSNorm stripped, and NLACriticModel replaces that module with nn.Identity() immediately after loading.

peft must be <0.19 (e.g. 0.18.1) if you load adapters under torch.distributed.

Attribution

• Parent NLA: achand45/gemma-3-12b-it-nla-L40 — the SFT warm-start and the first 400 RL steps are theirs; this repo continues them.
• Training code: EasyNLA (MIT).
• Conversations: allenai/WildChat-1M, subject to its own terms.
• Method: Natural Language Autoencoders (Anthropic, 2026); Karvonen et al. norm-matched activation injection.

License

These weights are a Model Derivative of google/gemma-3-12b-it and are distributed under, and subject to, the Gemma Terms of Use — not the MIT licence of the training code. A copy of the Agreement is included as LICENSE, and the required notice as NOTICE:

Gemma is provided under and subject to the Gemma Terms of Use found at ai.google.dev/gemma/terms

Use restrictions. Your use of these weights is subject to the Gemma Prohibited Use Policy, incorporated by reference into the Terms (§3.2). If you redistribute these weights or anything derived from them, you must pass these restrictions on to your recipients as an enforceable provision, supply them a copy of the Agreement, and give notice that the weights are subject to those restrictions (§3.1).

Modification notice (§3.1). rl_wildchat/critic_latest/ is modified Gemma weights: google/gemma-3-12b-it truncated to its first 41 blocks, with the final RMSNorm stripped and a value_head added, further trained by us. The adapters under rl_wildchat/iter_*/ are new weights trained by us, not modified Gemma files, but they only function when applied to Gemma and are Model Derivatives on the same terms.