final_olens: an OLens 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.

Everything needed is in this repo. olens.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
p3_bullets_rl_step150/ bullet variant ("P3"): the same lens distilled into readouts of up to 4 - bullets it picks itself, then RL. Highest WorkspaceBench (.761), but more hallucination (see below) 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.

The whitened lens (lens A), its AR and the whitening basis it was trained in are in agu18dec/final_olens_whitened (ar/, whitening/); olens.py loads them (load_whitening, whiten, unwhiten).

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 olens.py --local-dir .   # the loader; weights download on first use
import torch
from olens import from_hub, load_lens, load_ar, residual, read, ar_direction

REPO = "agu18dec/final_olens"
lens = load_lens(from_hub(REPO, "lens_rl_step600"))           # 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, ban_template=True):      # 4 readouts, template ban on
    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(from_hub(REPO, "ar"))                             # optional, another ~60 GB
d = ar_direction(ar, "Bizet", layer=40)                        # predicted residual direction
print(torch.cosine_similarity(d, h.cuda(), dim=0))

# the bullet variant (most information per read): same loader
bullets = load_lens(from_hub(REPO, "p3_bullets_rl_step150"))

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.

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.

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

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])
the same recipe with a whitened AR and a whitened reward .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:

WorkspaceBench by eval

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

Bullet variant: p3_bullets_rl_step150/ ("P3")

Same architecture, injection, prompt and loader as the lens above (load_lens(from_hub("agu18dec/final_olens", "p3_bullets_rl_step150"))); its readouts are a list of up to 4 - bullets, each a short phrase it expects the vector to encode.

How it was made, from this repo's inverter and AR (nothing else):

  1. Labels (own-distill). For each of the RL split's activations, the inverter samples 16 readouts on the true residual. Every line prefix of 2/4/8/16/32 tokens is a candidate bullet; the frozen AR turns each into a predicted direction (raw space). A greedy non-negative least-squares pursuit picks the candidates that together best rebuild the activation, up to 4 (mean 1.96 picked, joint raw FVE .79).
  2. SFT. The inverter is fine-tuned on those bullet targets: 376 steps x 64 (~24k examples), lr 1e-4, one epoch.
  3. RL (150 steps, same settings as the main lens): the reward is the joint reconstruction of up to 4 bullets by the frozen AR (non-negative least squares over their AR directions, raw space).
WorkspaceBench (all-Opus headline_23, same pipeline) score hallucination eval (1 − rate)
the main lens, RL step 600, 1 readout .683 .617
P3 bullets, RL step 150, 1 readout .761 [.745, .776] .527
P3 bullets, 4 readouts per cell .805 [.790, .819] .525
P3 after more RL (steps 300 / 480 / 600) .746 / .736 / .746 ~.53
ML-NLA-64 (all-layer NLA, RL step 2000), read through the same pipeline .725 [.709, .741] .002

Use P3 when you want the most information per read; use the single-readout lens when false statements matter more.

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.py                  self-contained loader: from_hub, load_lens, load_ar, residual, read (ban_template, force_prefix), ar_direction
olens_lensB.py            the same loader under its old name
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
p3_bullets_rl_step150/    bullet variant (P3): up to 4 self-chosen bullets per readout
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

More checkpoints, one controlled change each (AR training text, AR target, where RL starts, reward shape, RL length): agu18dec/olens-ablations.

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

Base model

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