final_olens_whitened: the whitened OLens ("lens A") for Qwen3.6-27B
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.
This repo holds lens A, the sibling of agu18dec/final_olens (lens B). The two lenses share the data, the lens architecture, the injection, the prompt and the sampler. They differ in one respect: lens A's reconstructor (AR) was trained, and its RL reward computed, in a whitened residual space rather than the raw one. Lens A is the base lens of the final_olens ablation study (agu18dec/olens-ablations).
| folder | what | WorkspaceBench (all-Opus headline_23, 1 readout per cell) | trainable params |
|---|---|---|---|
lens_rl_step600/ |
the lens after 600 RL steps (same number of steps as final_olens/lens_rl_step600) |
.657 | 117M (LoRA) |
lens_rl_step300/ |
the lens at RL step 300: highest WorkspaceBench of this lens | .678 [.663, .695] | 117M |
lens_rl_step150/ |
the lens at RL step 150 | .669 | 117M |
inverter/ |
the same lens before RL (supervised only) | .616 [.599, .633] | 117M |
ar/ |
the lens's AR (reconstructor): text + layer → predicted residual direction, trained in the whitened space below; the RL reward | – | 143M (LoRA + head) |
whitening/ |
the frozen per-layer whitening basis (μ, W) the AR and the RL reward used | – | 6.7 GB (fp32) |
Reading 4 samples per vector instead of 1 raises the score to .758 (measured at RL step 75). Each folder holds model.safetensors (trainable tensors only; the 27B base model is loaded from the Hub) and export.json (the source checkpoint and the exact training overrides).
The AR (ar/) predicts residual directions that are compared in the whitened basis, so the basis it was trained with is included too (whitening/); see AR and whitening basis.
Quickstart
The loader is olens.py, byte-identical to the one in agu18dec/final_olens. These files have the same tensor names and shapes as that repo's lenses.
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_whitened olens.py --local-dir .
from olens import from_hub, load_lens, residual, read
lens = load_lens(from_hub("agu18dec/final_olens_whitened", "lens_rl_step300")) # ~60 GB GPU memory
h = residual(lens, "Who wrote the opera Carmen?", layer=40) # block-40 output, last token, [5120]
for r in read(lens, h, layer=40, n=4, ban_template=True): # 4 readouts, template ban on
print("---\n" + r)
Conventions (these are the same as final_olens):
- Residual at layer L: the output of decoder block L (0–63), before the final RMSNorm. In HF's
output_hidden_statesthat ishidden_states[L+1]for L < 63. - Rendering: the chat template with
enable_thinking=False. - Injection: one slot token
㈜in the lens prompt. Its input embedding is replaced by1000 · h/‖h‖. - Sampling: temperature 1, top-p 1, up to 128 new tokens.
Recommended: the decode-time template ban (read(..., ban_template=True))
At chat-template positions — the <|im_end|> that closes a turn, the \n after it, <|im_start|>, assistant, <think>, </think> — the lens sometimes writes the chat template itself (a <|im_start|>assistant\n<think> loop) or stops at once and returns an empty readout. The lens learned to write the text that follows the read position, and its training data contains no template positions; at those positions what follows is chat formatting or the end of the conversation.
read(lens, h, layer, ban_template=True) fixes this at decode time: it never samples <|im_start|>, <think>, </think> or <|endoftext|>, and it does not stop before 24 new tokens. It is off by default, so default readouts match the published evaluations below.
| WorkspaceBench headline_23, no ban | with ban | degenerate readouts (empty or looping), share of all WorkspaceBench readouts | |
|---|---|---|---|
whitened, RL step 300 (agu18dec/final_olens_whitened) |
.678 | .684 | 4.7% → 0.3% |
non-whitened, RL step 600 (agu18dec/final_olens) |
.683 | .689 | 7.5% → 0.5% |
All-Opus judge, 23 families, 1 readout per cell, CI ≈ ±.017; 190,602 readouts per lens, a third of them at template positions.
End-to-end check of this exact olens.py with the Hub weights (16 samples per cell, a chat turn read at the <|im_end|> that closes the reply): empty readouts 15/16 → 0/16 (whitened, layers 44 and 52) and 9/16, 7/16 → 0/16 (non-whitened, layers 44 and 52); at layer 60 the non-whitened lens's assistant/<|im_start|> loop (610 template tokens over 16 samples) → none. With the ban no banned token is generated and every readout has ≥ 24 tokens.
What to expect: at the end of a turn the banned readouts are readable but usually generic end-of-turn text ("The user's inquiry has been addressed…", a predicted next user message). The non-whitened lens can still spell </think> out of ordinary tokens in ~1.4% of samples.
Also optional: read(..., force_prefix="- ") starts every readout with the - bullet the lens prompt asks for (the first generated token is constrained to a space-led word, so the bullet is followed by content). Off by default; composable with the ban.
AR and whitening basis
| file | what |
|---|---|
ar/model.safetensors |
the AR (run ar.whiten.lr7e-5, decay branch decay171155, step 205,386): LoRA r16/α32 all-linear + LayerNorm→Linear head, 994 tensors, 142,947,328 params. Same architecture as final_olens/ar, so load_ar loads it. Loss 2(1 − cos) in whitened space; held-out loss 1.283 on the 4,096-span AR ruler (mean whitened cosine .358). export.json: source checkpoint + exact training overrides |
whitening/basis.safetensors |
mu [64, 5120] and w [64, 5120, 5120], fp32: per-layer mean μ_L and whitening matrix W_L for all 64 layers (layer L = output of decoder block L, before the final norm). The training basis unchanged (source sha256 prefix 8d05c083386b6924); the covariance is not included (another 6.7 GB) |
whitening/basis.json |
fit metadata: rows, ridge, shards, checksum |
Formula. For a residual x at layer L:
z = (x − μ_L) · W_L, W_L = (Σ_L + λ_L I)^(−1/2), λ_L = 0.1 · mean(diag Σ_L)
- μ_L and Σ_L (population covariance) come from 1,310,720 residuals per layer: Qwen3.6-27B's own chat replies (16 shards of the AR-training split), at the positions the AR predicts (the token just before a span). Accumulated in fp64; W_L from a symmetric eigendecomposition in fp64, stored fp32.
- W_L is symmetric, so (x − μ_L) W_L = W_L (x − μ_L). Inverse: x = z · W_L⁻¹ + μ_L.
- The AR's loss and the RL reward compare both the AR's prediction and the true residual after whitening at the item's layer. The AR's output is in residual coordinates; compare it after whitening, not by raw cosine.
Loading (with olens.py):
from olens import from_hub, load_ar, ar_direction, load_whitening, whiten, unwhiten
wb = load_whitening(from_hub("agu18dec/final_olens_whitened", "whitening",
filename="basis.safetensors"), layers=[40]) # 105 MB per layer; omit layers= for all 64
ar = load_ar(from_hub("agu18dec/final_olens_whitened", "ar"))
z_h = whiten(h, 40, wb) # h: [5120] residual at layer 40 (or [n, 5120])
z_p = whiten(ar_direction(ar, "Bizet", 40), 40, wb)
print(torch.cosine_similarity(z_p, z_h, dim=0)) # the AR's training metric
x = unwhiten(z_h, 40, wb) # back to residual coordinates (== h)
Without olens.py: safetensors.safe_open(path, "pt"), then f.get_tensor("mu")[L] and f.get_slice("w")[L] read one layer.
Round trip verified (2026-11-01): these exact files (same sha256) were downloaded from the Hub and loaded with the published olens.py, then compared with our training build (the training checkpoint + the training basis) on 64 held-out AR-ruler spans covering 43 layers: AR outputs identical (max abs difference 0.0), whitened residuals equal to 1.4e-6 relative, unwhiten(whiten(h)) = h to 1.3e-6, and mean whitened cosine AR vs true residual .3651 in both.
How it was trained
The data and splits are the same as final_olens: Qwen3.6-27B's own replies (temperature 1, thinking off) to WildChat user turns. Conversations are split by a hash of the first user turn into AR train, inverter train, RL and eval sets, and every set is balanced over the 64 layers.
| stage | checkpoint | recipe |
|---|---|---|
| whitening basis | fit once, frozen | per-layer mean μ and W = (Σ + 0.1·mean-diag(Σ)·I)^−1/2 from AR-train residual moments |
| AR | ar.whiten.lr7e-5.decay171155, step 205,386 |
LoRA r16/α32 (all-linear) + LayerNorm→Linear head. Reads [Layer NN] + a reply span of 1–32 tokens and predicts the residual at the token before the span. Loss 2(1 − cos) in the whitened space W(x − μ). Batch 576, AdamW lr 7e-5: constant for 171,155 steps, then linear decay over 34,231 steps (one pass, 118M examples) |
inverter (inverter/) |
inv.A.unit-a1k-concepts.lr3e-4.decay977, step 1,172 |
the same LoRA on the base model. The vector enters the slot as 1000 · h/‖h‖. Supervised on the AR's outputs, with the text span as the target. Batch 1,024, lr 3e-4 for 977 steps, then decay over 195 |
RL (lens_rl_step*/) |
rl.A.final.slot.concepts.s3ds.lr5e-5.kl0.005 |
GRPO from the inverter: each step uses 8 true residuals from the RL split × 64 sampled readouts. Reward = whitened fraction of variance explained when the frozen AR reconstructs the residual from the readout. lr 5e-5 constant, KL 0.005 to the inverter, temperature 1, top-p 1, 128 new tokens |
final_olens (lens B) differs in the AR and the reward. Its AR is trained with raw (uncentred) cosine at lr 1e-4, its inverter is trained on that AR's outputs, and its RL reward is an NLA-style raw FVE. Lens B's step-600 lens scores .683 on the same benchmark.
Model tree for agu18dec/final_olens_whitened
Base model
Qwen/Qwen3.6-27B