Yooniel/gemma-3-12b-it-nla-L32-wildchat

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

A natural language autoencoder (NLA) for google/gemma-3-12b-it at block 32, 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-L32` 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.

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 and how much of it RL recovers.

So the parent drops ~13.5pp when pointed at chat activations, and 400 steps of WildChat GRPO recover +9.2pp — roughly two-thirds of the transfer gap.

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.0299 (the parent's FFW baseline is 0.0313 — 4.5% apart, so the two are of similar scale but not the same measurement).

⚠️ 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 +9.2pp gain clears that; differences under ~5 points here do not.

Most of the gain lands by step 600 (+8pp); the last ~150 steps move within noise. Entropy rose 1.66 → ~2.0 and settled; KL from the SFT reference rose smoothly to ~1.2. No non-finite gradients, no format collapse.

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

The parent ran GRPO at `batch_prompts: 256`; this run used 128 (the tuned config for the training box's dedicated-engine layout). 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.

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

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

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)
rl_wildchat/nla_meta.yaml, run_config.yaml, optim_latest.pt

*To use the NLA you need one `iter_ adapter + criticlatest`.** `iter000750 is the best eval (64.6%) and iter_000800 the final step; they differ by 0.3pp, inside the noise band, so treat them as interchangeable — 750` is the defensible single 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_000750",
)
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-L32` — 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](https://ai.google.dev/gemma/terms) — 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 33 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.