final_olens: an OLens for Qwen3.6-27B (lens B: raw, NLA-style)

An OLens reads any residual-stream vector of Qwen/Qwen3.6-27B (revision 6a9e13bd6fc8f0983b9b99948120bc37f49c13e9), at any of its 64 layers, and writes what that vector encodes as text.

Everything needed is in this repo. olens_lensB.py is a self-contained loader that needs only transformers, peft and safetensors. Its outputs match our training code exactly: AR predictions and lens logits differ by 0.0 on the same inputs.

folder	what	trainable params
lens_rl_step600/	the lens, highest WorkspaceBench (.683): residual vector → text	117M (LoRA)
lens_rl_step300/	the lens at RL step 300: fewer hallucinations, WorkspaceBench .673	117M
lens_rl_step150/	the lens at RL step 150 (.664)	117M
inverter/	the same lens before RL (supervised only)	117M
ar/	the AR: text + layer → predicted residual direction. Used to score readouts and as the RL reward.	143M (LoRA + head)

Each folder holds model.safetensors (trainable tensors only; the 27B base model is loaded from the Hub) and export.json (the exact training configuration). Total download is about 1.5 GB plus the base model.

Quickstart

pip install "torch==2.9.1" "transformers==5.11.0" "peft==0.19.1" safetensors "flash-linear-attention==0.5.1"
huggingface-cli download agu18dec/final_olens --local-dir final_olens && cd final_olens

import torch
from olens_lensB import load_lens, load_ar, residual, read, ar_direction

lens = load_lens("lens_rl_step600/model.safetensors")          # needs ~60 GB GPU memory

# a residual of Qwen3.6-27B: block-40 output at the last token of a chat turn
h = residual(lens, "Who wrote the opera Carmen?", layer=40)    # torch [5120]
for r in read(lens, h, layer=40, n=4):                         # 4 sampled readouts
    print("---\n" + r)

# or read your own vector: h = output of decoder block L (before the final norm), shape [5120]
# read(lens, h, layer=L)

ar = load_ar("ar/model.safetensors")                           # optional, another ~60 GB
d = ar_direction(ar, "Bizet", layer=40)                        # predicted residual direction
print(torch.cosine_similarity(d, h.cuda(), dim=0))

Conventions (these must match training):

• Residual at layer L: the output of decoder block L (0–63), before the final RMSNorm. residual() captures it with a forward hook. In HF's output_hidden_states it is hidden_states[L+1] for L < 63; at L = 63 the HF tensor is already normalised, so use the hook.
• Rendering: text is rendered with the chat template and enable_thinking=False.
• Injection: the lens's prompt contains one slot token ㈜. Its input embedding is replaced by 1000 · h/‖h‖.
• Sampling: temperature 1, top-p 1, up to 128 new tokens. Sampling several readouts per vector reads more of it: about +.08 WorkspaceBench at 4 samples per vector for this recipe family.
• AR input: the AR reads [Layer NN] followed by the text. It was trained on spans of ≤ 32 tokens, so score at most the first 32 tokens of a readout.

Exactly what it was trained on

All data comes from Qwen3.6-27B's own replies (temperature 1, thinking off) to user turns from WildChat. Conversations are split 70 / 20 / 5 / 5 by a hash of the first user turn (AR train / inverter train / RL / eval), so the four stages never share a conversation. Every set is balanced exactly over the 64 layers and over span lengths.

stage	examples seen	what one example is	schedule
AR	118,302,336 = 205,386 steps × 576	a reply span of 1–32 tokens (uniform; ~16.5 on average, ≈1.95B span tokens) + one layer + the true residual at the token just before the span	LoRA r16/α32 all-linear + LayerNorm→Linear head; AdamW lr 1e-4, constant for 171,155 steps, then linear decay to 0 over 34,231 steps (20%); one pass
inverter	1,200,128 = 1,172 steps × 1,024	a text crop (1–32 tokens) from the inverter split + the AR's predicted vector for it at its layer (crop pool 1.42M, < 1 pass)	LoRA r16/α32; lr 3e-4 constant for 977 steps + 195 decay steps; loss = cross-entropy on the crop given the injected vector
RL	4,800 activations, 307,200 readouts = 600 steps × 8 activations × 64 samples (checkpoints at 150 / 300 / 600)	a TRUE residual from the RL split at its layer (balanced over 64 layers); 64 sampled readouts, ≤ 128 tokens each	GRPO (group-normalised advantages), lr 5e-5, KL 0.005 (k3 estimator, reference = the inverter). Reward = NLA's centred FVE of the whole readout through the frozen AR: 1 − (1 − cos²)/(1 − ‖μ‖²)

Held-out evaluation sets:

• AR: 4,096 spans × 64 layers.
• Inverter: 16,384 crops.
• RL: 1,024 activations.

All come from the eval split.

Which checkpoints these are

• ar/: the end of the AR's learning-rate decay (step 205,386 = 171,155 constant + 34,231 decay).
• inverter/: the end of the inverter's decay (step 1,172 = 977 + 195).
• lens_rl_step{150,300,600}/: RL at a constant learning rate (no decay), at steps 150, 300 and 600.

How well the AR fits (held-out, 4,096 spans × 64 layers):

metric	value
raw cosine² (mean)	.770, vs .643 for predicting each layer's mean direction
NLA-style FVE above the mean, (raw − floor)/(1 − floor)	.289 mean; .234 at L20, .316 at L40, .455 at L60
centred FVE (mean removed)	.207
held-out loss 2(1 − cos)	.258

Training curves

	start	end
AR held-out loss 2(1 − cos)	.32	.258 (after decay)
inverter held-out CE (nats/token)	1.75	1.197 (after decay)
RL held-out NLA FVE	.376	.423 (step 150) → .446 (step 600)
RL KL to the inverter (rolling median)	.01	.7 (step 150) → ~1.1 (step 600); one outlier step (75) hit 368, where the k3 estimate blew up on a single token

Results (WorkspaceBench, all-Opus headline_23, 1 readout per cell; CI ≈ ±.017)

	score
inverter (before RL)	.594
lens, RL step 150	.664
lens, RL step 300	.673
lens, RL step 600	.683 (vs step 150: +.015, paired 95% CI [+.001, +.029])
same recipe with a whitened AR + whitened reward ("lens A")	.669
s3d reference lens	.605 (first bullet) / .726 (4 distilled bullets)

Tradeoff with more RL: the hallucination eval (1 − hallucination rate) goes .657 (step 150) → .683 (step 300) → .617 (step 600), so step 600 states false things more often. Step 300 beats step 150 on both the headline and hallucination.

By eval at RL step 600 (normalised pass rate; the headline is the mean over 23 families). Per-eval 95% CI ≈ ±.09 at 100 items:

eval	items	inverter	lens (RL 600)	s3d (1 bullet)
typo	100	0.99	1.00	1.00
multilingual	100	0.98	0.99	0.98
basic_readout	100	0.94	0.98	0.92
multihop	100	0.96	0.98	0.93
association	100	0.64	0.91	0.65
user_modeling	100	0.96	0.91	0.95
relational_multihop	100	0.81	0.90	0.86
multilingual_typo	100	0.63	0.87	0.71
role_bound_association	100	0.86	0.86	0.91
moral_rationale	200	0.81	0.85	0.68
basic_readout_mt	100	0.64	0.82	0.49
typo_mt	100	0.71	0.79	0.80
brew_intermediates	100	0.67	0.72	0.68
hallucination (1 − hallucination rate)	149	0.73	0.62	–
multilingual_mt	100	0.17	0.58	0.21
multi_concept_directed_modulation	27	0.52	0.52	0.64
chain_intermediates	120	0.45	0.44	0.47
buggy_code (score/10)	49	0.39	0.41	0.44
poetry	100	0.29	0.39	0.32
conjunctive_association	100	0.33	0.38	0.40
arithmetic_intermediates	596	0.18	0.38	0.20
multihop_mt	100	0.26	0.36	0.15
directed_modulation	100	0.33	0.34	0.33
multilingual_multihop	98	0.16	0.32	0.19

How it relates to NLA

• Same as NLA:
  • the reconstructor (AR) loss is unit-vector MSE = 2(1 − cos);
  • no norm is predicted;
  • the reward is a cosine reconstruction score (NLA's −log 2(1 − cos) gives the same WorkspaceBench score: .663 vs .663);
  • the activation enters by embedding replacement;
  • RL is GRPO with a KL penalty to the starting policy.
• Different from NLA:
  • the AR is trained first, on real text, and frozen during RL, so the lens cannot co-adapt a private code with its reconstructor;
  • the lens is warm-started by inverting the AR (no external summaries);
  • one lens covers all 64 layers;
  • readouts are continuation-style text, not tagged explanations.

Files

olens_lensB.py            self-contained loader: load_lens, load_ar, residual, read, ar_direction
lens_rl_step600/          the lens (highest WorkspaceBench)
lens_rl_step300/          the lens at step 300 (fewer hallucinations)
lens_rl_step150/          the lens at step 150
inverter/                 the lens before RL
ar/                       the AR (reconstructor / reward)
training_curves.png       AR / inverter held-out curves, smoothed RL curves
wsb_by_eval.png           WorkspaceBench by eval