qwen2.5-7b-instruct-nla-L20

Qwen2.5-7B-Instruct NLA — residual stream, block 20

A Natural Language Autoencoder (NLA) for Qwen/Qwen2.5-7B-Instruct: a pair of models that read a residual-stream activation and write a natural-language explanation of it, then reconstruct the activation back from that text.

• AV (verbalizer) — a LoRA on the base model. An activation is injected at a marker token (norm-matched, à la Karvonen et al.) and the AV writes an <explanation>…</explanation> of it.
• AR (reconstructor) — a 21-block truncation of the base model plus a linear head that maps the explanation text back to the activation vector.

Trained with EasyNLA (built on nanoNLA): SFT warm-start, then on-policy GRPO where the AV is rewarded by how well the AR reconstructs the activation from its words.

Results

Held-out, doc-disjoint (a crc32 bucket on doc_id, so every row of a held-out document is excluded from training — a row-level split leaks badly here because the corpus is row-shuffled and each document contributes ~10 rows).

FVE = fraction of activation variance explained, against a predict-the-mean baseline. The RL number rose from 56.4% → 72.9%; extraction rate stayed at 100% for essentially the whole run (no format collapse), and KL from the SFT reference rose smoothly to ~0.95 with no runaway.

Extraction contract

⚠️ The layer convention is off-by-one relative to naive hidden_states[K] indexing. layer_index=20 hooks model.model.layers[20] and captures its output, which equals hidden_states[21] (index 0 is the embedding output). Verified numerically: cosine 0.9999 against hidden_states[21] vs 0.914 against hidden_states[20].

Every checkpoint ships an nla_meta.yaml sidecar carrying this contract (marker token ids, prompt templates, scales). The trainers assert against it — the AR's depth is derived from extraction.layer_index + 1, not hardcoded.

Repository layout

av_sft/iter_*/          AV warm-start LoRA (final: iter_0003834)
ar_sft/iter_*/          AR warm-start LoRA + value head (final: iter_0003834)
merged/av_hf/           AV SFT merged to bf16 HF   (regenerable: merge_lora_to_hf.py)
merged/ar_hf/           AR SFT merged to bf16 HF   (validated: FVE 62.0%, == the LoRA ckpt)
rl_vllm/iter_*/         RL AV LoRA every 25 steps (final: iter_000400)
                          adapter_model.safetensors — the trained policy
                          reference/                — frozen SFT copy, the KL reference
rl_vllm/critic_latest/  RL co-trained AR reconstructor (full weights + value_head)

To use the trained NLA you need rl_vllm/iter_000400 + rl_vllm/critic_latest. The RL adapter is a LoRA on the raw base model (RL continued the SFT adapter via --av-adapter), so it does not require merged/av_hf. The intermediate iter_* snapshots are included for training-dynamics and ablation work.

Usage

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
from nla.models import NLACriticModel          # pip install -e git+https://github.com/asherps/EasyNLA

BASE = "Qwen/Qwen2.5-7B-Instruct"
REPO = "<local clone of this repo>"

tok = AutoTokenizer.from_pretrained(BASE)
av = PeftModel.from_pretrained(
    AutoModelForCausalLM.from_pretrained(BASE, dtype=torch.bfloat16),
    f"{REPO}/rl_vllm/iter_000400",
)                                              # verbalizer: activation -> text
ar = NLACriticModel.from_pretrained(
    f"{REPO}/rl_vllm/critic_latest", torch_dtype=torch.bfloat16,
)                                              # reconstructor: text -> activation

Inject the activation at the marker token with nla.injection.karvonen_inject_in_residual; see scripts/show_nla_generations.py in EasyNLA for a complete example.

peft must be <0.19 (e.g. 0.18.1) if you load adapters under torch.distributed: peft 0.19's set_peft_model_state_dict imports EmbeddingParallel from transformers.integrations.tensor_parallel, which does not exist in transformers 4.57.x, and the call is guarded by dist.is_initialized() — so it breaks only in distributed runs.

Data provenance

Text, prompts, and gold explanations come from ceselder/qwen3-8b-nla-L24-finefineweb-100k (corpus: m-a-p/FineFineWeb, 100k docs; explanations by Claude Sonnet 4.6).

Those fields are model-agnostic and were reused unchanged: the explanations describe the source text, and the prompts store an <INJECT> placeholder rather than a literal marker. Qwen2.5-7B-Instruct also tokenizes the corpus identically to Qwen3-8B and shares the marker token id, so n_raw_tokens and the doc splits carry over exactly.

Only the activations were recomputed, by forwarding Qwen/Qwen2.5-7B-Instruct over detokenized_text_truncated and taking the block-20 residual stream at the final token. Rows whose stored text ends mid-multibyte-character (~0.003%) were dropped rather than repositioned, since their text cannot reproduce the original token boundary.

Training rows: 245,344 (AV) / 245,327 (AR) / 499,828 (RL), one epoch of SFT each.

Training setup

8×A100-80GB. SFT is 8-GPU data-parallel (~1 h 40 m for both stages); RL is 400 steps of GRPO with per-rank vLLM rollouts (7 h 10 m).

License

Apache-2.0, matching the base model and the source dataset.