gemma-4-e2b-nla-L23-ar-v0_1-paraphrase-invariant

Gemma-4-E2B NLA AR (Activation Reconstructor) — v0.1 (paraphrase-invariant)

LoRA adapter (+ a 1536→1536 linear head) for google/gemma-4-E2B that takes a natural-language explanation produced by the matched Activation Verbalizer and reconstructs a 1536-dimensional activation vector intended to round-trip (by cosine similarity) to the original L23 residual-stream activation.

This is the v0.1 paraphrase-invariant Activation Reconstructor. It continues training from the v0.0.1 AR with an added auxiliary objective that pushes the reconstruction to be stable under meaning-preserving paraphrases of the input explanation, directly targeting the AR's structural-projection behavior (see "What paraphrase-invariance means" and Limitations). Trained end-to-end on a single 4 GB consumer GPU (NVIDIA GTX 1650 Ti Max-Q) following a customized variation (see below) of the methodology of Fraser-Taliente, Kantamneni, Ong et al. 2026 (Transformer Circuits).

Customizations vs the source methodology

A consumer-hardware variation, not a faithful reproduction: a LoRA adapter (+ a 1536→1536 linear head) over a 4-bit NF4-quantized frozen base rather than full-fine-tune bf16; a single 4 GB GPU with micro-batch 1 + gradient accumulation and a modest SFT-step budget; an SFT-only released pair (Phase-4 GRPO explored separately, not shipped); a paraphrase-invariance auxiliary loss added in v0.1 (not part of the source recipe); and added AV-side evaluations beyond round-trip cosine (see the matched MODEL_CARD_AV.md). Note round-trip cosine on this AR is largely a structural-projection metric, not per-row faithfulness — see Limitations.

What the paired verbalizer reads from an activation (round-trip ceiling)

The reconstructor is one half of a round-trip: an activation is verbalized by the paired Activation Verbalizer, and this AR maps that explanation back to an activation. So the round-trip can only be as content-specific as the verbalizer's reading of the activation in the first place. A focused June 2026 evaluation measured exactly that, with a confound-free forced choice (fix the target text, swap only the injected activation, so text length and perplexity cancel; chance 0.5):

• Across domains (routing): 0.676, p = 0.005. The verbalizer reliably tracks which domain an activation came from.
• Within a domain (content): 0.456, at chance. It cannot distinguish two activations from the same domain.
• The split is stable across target lengths (16 / 32 / 48 tokens) and uniform across all five test domains; routing is led by legal at 0.93.

What this means for the AR: the v0.1 round-trip carries domain-level information well but is not expected to preserve fine-grained within-domain content, because the verbalizer it depends on does not resolve that content. This is consistent with the structural-projection behavior described below, and closing the within-domain gap is the focus of ongoing work. (n = 68 held-out activations; confound-free forced-choice likelihood metric.)

What paraphrase-invariance means (v0.1)

The v0.0.1 AR has a documented weakness: paraphrasing the input explanation moves the reconstructed cosine by only ~3% (Δcos_paraphrase = +0.014), so the reconstruction is nearly invariant to the explanation's content and mostly projects toward the activation manifold. The v0.1 AR attacks that directly during training. For each row it builds a meaning-preserving paraphrase of the explanation (Gemini-generated, same factual claims and topic, different vocabulary and sentence structure) and trains with an auxiliary term so the reconstruction from the paraphrase matches the same target as the reconstruction from the original:

L = MSE(AR(prompt), gold) + λ · MSE(AR(paraphrase_prompt), gold),   λ = 0.5

The point is to make the AR map meaning, not surface form: two different wordings of the same explanation should reconstruct toward the same activation. This is an explicit attempt to move the AR off the structural-projection floor described in Limitations. It is a training-objective change, not a claim that the floor is solved — round-trip cosine on this AR remains structural-projection dominated and should not be read as per-row faithfulness.

How to use

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel
from huggingface_hub import snapshot_download
import torch
import torch.nn as nn
import numpy as np

BASE = "google/gemma-4-E2B"
AR_REPO = "Solshine/gemma-4-e2b-nla-L23-ar-v0_1-paraphrase-invariant"
AR_TRUNCATION = 18   # capture hidden state at layer (K+1-1) = 17 (0-indexed)
D_MODEL = 1536
AR_TEMPLATE = "Summary of the following text: <text>{explanation}</text> <summary>"

bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16, bnb_4bit_quant_type="nf4")
tok = AutoTokenizer.from_pretrained(BASE)
base = AutoModelForCausalLM.from_pretrained(BASE, quantization_config=bnb, device_map="auto")
ar_local = snapshot_download(repo_id=AR_REPO)
ar = PeftModel.from_pretrained(base, ar_local); ar.eval()

# Load the trained linear head (1536 -> 1536)
head = nn.Linear(D_MODEL, D_MODEL, bias=True).to(ar.device).to(torch.float32)
head.load_state_dict(torch.load(f"{ar_local}/linear_head.pt", weights_only=True))
head.eval()

# Capture hidden state at layer L17 (1-indexed L18)
extraction_layer = ar.base_model.model.model.language_model.layers[AR_TRUNCATION - 1]

def reconstruct(explanation: str) -> np.ndarray:
    """Run an AV explanation through the AR; return the reconstructed activation vector."""
    prompt = AR_TEMPLATE.format(explanation=explanation)
    ids = tok.encode(prompt, return_tensors="pt").to(ar.device)
    captured = {"h": None}
    def hook(m, i, o):
        captured["h"] = o[0] if isinstance(o, tuple) else o
    handle = extraction_layer.register_forward_hook(hook)
    try:
        with torch.no_grad():
            _ = ar(input_ids=ids)
    finally:
        handle.remove()
    h_last = captured["h"][0, -1].to(torch.float32)
    return head(h_last).detach().cpu().numpy().astype(np.float32)

# Use
explanation = "<explanation>This activation responds to references to the Pacific Ocean and geographic regions.</explanation>"
reconstructed = reconstruct(explanation)
# Compare to the original activation:
# cosine_similarity = np.dot(original, reconstructed) / (np.linalg.norm(original) * np.linalg.norm(reconstructed))

Working end-to-end round-trip example with the matched AV: examples/round_trip_example.py in the bundled public repo.

Training setup

• Base model: google/gemma-4-E2B
• Initialization: continued from the v0.0.1 AR (LoRA adapter + linear head), then trained with the paraphrase-invariance objective
• Truncation: K+1 = 18 layers used in AR forward; hidden state captured at layer index 17 (0-indexed)
• Quantization: NF4 4-bit base + fp16 LoRA
• LoRA config: r=64, α=128, same target_modules regex as the AV
• Linear head: 1536 → 1536 with bias, trained in fp32, projects from L17 hidden state to reconstructed activation
• Auxiliary loss: paraphrase-invariance, L = MSE(AR(orig), gold) + λ·MSE(AR(paraphrase), gold), λ=0.5
• Optimizer: AdamW 8-bit, lr=1e-4
• Max length: 512 tokens
• SFT steps: 50 (paraphrase-invariance continuation)
• Training corpus: 696 rows, each with a meaning-preserving paraphrase of its explanation (Gemini-generated), derived from the v0.0.x baseline pipeline
• Hardware: 4 GB GTX 1650 Ti Max-Q

Headline numbers (v0.1)

• Round-trip cosine with the matched v0.1 AV: 0.460, above the 0.30 noise floor (v0.0.1 round-tripped at 0.438). Round-trip cosine on this AR is structural-projection dominated, not a per-row faithfulness metric — see Limitations.

Evaluation across released versions

Round-trip reconstruction cosine (right panel) is the AR's headline metric, shown across the released NLA versions alongside the AV's content-fidelity doc-retrieval (left panel): v0.0.1 round-trips at 0.438, v0.1 at 0.460, both above the 0.30 noise floor and well below Anthropic's deployed ~0.99. Round-trip cosine on this LoRA + linear-head AR is largely a structural-projection metric, not per-row faithfulness — a fixed structural component of the reconstruction tracks the activation manifold independent of the input explanation (see Limitations). Treat it as a sanity floor, not an accuracy claim. Regenerate with make_nla_eval_figure.py as new versions or evaluations land.

NLAttack capability-floor evaluation

The AR reconstructs an activation from the AV's text, so what that target activation actually carries matters. It was characterized independently by the NLAttack capability-floor harness — a battery of concept-survival and emergence tests over the NLA's information bottleneck. On a held-out deception-domain set the bottleneck scores EmergenceIndex 0.601 — "established: stable, selective, generalizing representation," driven by decodability = 1.00 (a linear probe reads the injected concept off the residual activation at ceiling) and stability = 0.88 across seeds.

For the AR this fixes the target: the representation it maps back toward is well-formed and near-perfectly decodable, so the open problem in the pair is the verbalizer's surfacing of content into text, not the bottleneck the AR reconstructs. One of NLAttack's deferred axes, faithful_rank (reconstruction faithfulness with the AR in the loop), is the natural AR-side eval to add next — it is unscored here pending hard-negative minimal pairs and AR-in-the-loop setup. Regenerate with make_nlattack_v01_figure.py after a fresh NLAttack run.

What makes this release distinctive

• First non-Anthropic-team open-source NLA AR at any model scale.
• First LoRA-based NLA AR. Anthropic's published NLA ARs are full fine-tunes at bf16. This release demonstrates a LoRA adapter (r=64, α=128) + 1536→1536 linear head + AR truncation at K=18 layers over NF4-quantized Gemma-4-E2B. Shipping as LoRA + small head means the AR loads in ~0.6 GB VRAM on top of the frozen NF4 base — the entire matched (AV, AR) pair fits in 4 GB. The structural-projection properties documented below are characteristic of this LoRA-AR class at 4 GB scale; they may differ at higher AR capacity / full-FT.
• Paraphrase-invariance objective. v0.1 adds a training term that pushes the reconstruction to be stable under meaning-preserving paraphrases of the input explanation, an explicit attempt to move the AR off the structural-projection floor that round-trip cosine sits on (see Limitations).
• Consumer-GPU trainable. Fits on 4 GB laptop GPU end-to-end alongside the matched AV.

Release rationale: why this SFT pair and not a GRPO checkpoint

The Anthropic NLA recipe (Fraser-Taliente et al. 2026) has four phases: Stages 0–3 (data + labeling) → SFT (supervised fine-tune of the AV+AR pair) → Phase 4 GRPO (joint REINFORCE-style RL fine-tune of the AV with the AR's reconstruction-MSE as reward signal, plus an AR "keep-up" SFT update and a KL anchor). The published v0.0.1 and v0.1 pairs are the SFT-only output of Phases 1–3; Phase 4 GRPO was deferred at first release because it had not yet been adapted to the 4 GB hardware regime.

Between 2026-05-25 and 2026-05-29 the deferred Phase 4 was implemented and run end-to-end on the same 4 GB GTX 1650 Ti Max-Q, with alternating AV/AR loads and R=4 rollout batching to fit in VRAM. The trial swept 5 reward formulations × 4 entropy regimes across 120 rollouts. At every intermediate L2 readout (rollouts 40, 60, 80, 100, 120) the GRPO-updated AV+AR pair scored L2 cross-row-argmax = 0.100 (chance) on the n=10 held-out RL eval — the same as the SFT v0.1 baseline pair. Higher-entropy configurations additionally produced degenerate AV outputs (random Unicode tokens, whitespace, or "evasion evasion evasion…" mode collapse).

Verdict for this AR. No GRPO AR checkpoint is shipped. Within the GRPO trial the AR was updated under MSE-keep-up on the AV's rollout outputs (and, at rollout 108–120, briefly under a contrastive AR-loss variant). The post-GRPO AR's reconstruction quality on held-out activations did not improve over the SFT AR — the round-trip cosine and L2 cross-row-argmax both stayed in the same noise band as the released pair. This v0.1 paraphrase-invariant AR therefore remains the recommended Activation Reconstructor for this hardware/model class.

Research contribution. Combining the 8-attempt SFT lever sweep with the 5-readout GRPO sweep yields 14 distinct training attempts spanning the full Anthropic recipe, all converging to L2 = chance at 4 GB. The L2 ceiling at this hardware scale is robust to (a) optimizer-/loss-/scheduler-side levers within SFT, (b) reward shape (MSE vs contrastive vs contrastive-max), (c) entropy regularization (β ∈ {0, 0.1, 0.3, 1.0}), and (d) training paradigm (SFT-only vs SFT+GRPO). The structural-projection signature of the AR (Δcos ≈ 0 per per-claim ablation) is now characterized as a 4 GB-LoRA-AR property robust to GRPO updates, not an artifact of incomplete training; the paraphrase-invariance objective in this v0.1 is one direct attempt to address it. The open question — whether the bottleneck is base-model scale (2B vs 27B/70B) or the 4 GB hardware constraint (NF4 + LoRA + small contrast pool) — would be answered by a cross-model recipe-controlled retrain on Gemma-3-27B; that experiment is flagged for follow-on grant-funded work.

The v0.0.1 + v0.1 SFT AR pair therefore represents the best-coherent-pair checkpoint from a comprehensive characterization of the Anthropic NLA recipe at 4 GB, not a checkpoint that ran out of training budget before further phases could be attempted.

Limitations

NLAs can produce unexpected or incorrect explanations, and AR reconstructions can be structurally projected. Specifically for this AR:

• Round-trip cosine is largely structural-projection on this trained AR. Replicating the published §"Measuring steganography" and §"Characterizing confabulations" tests on the v0.0.1 AR: paraphrasing the input AV explanation moves the AR's reconstructed cosine by ~3% (Δcos paraphrase = +0.014); removing entire claims from the AV explanation moves cosine by ~0% per claim (Δcos = +0.001 per claim ablated). The reconstructed vector is approximately invariant to the explanation's content — the AR is largely projecting toward "somewhere in the L23 activation distribution" rather than reading the explanation. The v0.1 paraphrase-invariance objective targets exactly this gap by training the reconstruction to be stable across paraphrases, but round-trip cosine should still not be read as per-row faithfulness. For reference, Anthropic's deployed Gemma-3-27B AR reports round-trip cosine ~0.99 via the Neuronpedia API on the same shape of input; the gap to our 0.44–0.46 quantifies the hardware / methodology distance. This is a property of the AR's reconstruction at 4 GB, not a statement that the content is missing: an independent AV-side ceiling test recovers the source document from the raw L23 activation at 60% linear-probe accuracy, so the system-level per-row gap is about how the verbalizer and reconstructor read the activation — an open research question — not an absence of content in the activation (see the matched AV card and the repo README).
• This is a methodologically interesting finding about reconstruction behavior on under-trained AR architectures, not a unique pathology of this release. The same paraphrase/claim-ablation disaggregation should be measured on any NLA AR before relying on round-trip cosine as a content-fidelity proxy.
• Use this AR for: matched round-trip eval with the v0.1 AV (the cosine number is a valid characterization of the AV+AR pair as a system); replication of Anthropic's NLA validation pipeline at small scale; benchmarking AR-side improvements, including paraphrase-invariance.
• Do not use this AR for: inferring that the AV's explanation faithfully describes the activation. Use AV-side direct content-fidelity judging instead, or in addition.

Full development history and methodology retraction notes: HISTORY.md.

Citation

@article{frasertaliente2026nla,
  title={Natural Language Autoencoders},
  author={Fraser-Taliente, Kit and Kantamneni, Kshitij and Ong, Antonia and others},
  journal={Transformer Circuits},
  year={2026},
  url={https://transformer-circuits.pub/2026/nla/}
}

@misc{deleeuw2026nlagemma4e2bar,
  title={Gemma-4-E2B NLA AR (v0.1, paraphrase-invariant): a 4 GB consumer-GPU Activation Reconstructor},
  author={DeLeeuw, Caleb (SolshineCode)},
  year={2026},
  url={https://huggingface.co/Solshine/gemma-4-e2b-nla-L23-ar-v0_1-paraphrase-invariant}
}

License

CC-BY 4.0. See LICENSE in the bundled repo.