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_states that is hidden_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 by 1000 · 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.

Downloads last month

-

Downloads are not tracked for this model. How to track
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Model tree for agu18dec/final_olens_whitened

Base model

Qwen/Qwen3.6-27B
Finetuned
(428)
this model