ceselder/kimi-k3-olens-reconstructor

Activation reconstructors for Kimi-K3 -- text span -> layer-59 activation direction (LoRA r512 + affine maps)

The text-to-activation half of the Kimi-K3 oracle lens: given a short span of tokens, predict the (mean-centred) layer-59 activation h_sum of the frozen moonshotai/Kimi-K3 at the position the span continues. Two LoRA reconstructors ("AR", the olens h2hpfx recipe ported from Qwen3.6-27B via GLM-5.3; the second is the first continued onto a second harvest) and six closed-form affine maps on the frozen base. lora_ar_h59_bos_2M/ is the reward model of the lens RL in ceselder/kimi-k3-olens-lens (reward = clamp(cos(AR([BOS] + bullet), h_sum - mu), 0)^2); lora_ar_h59_bos_1M/ rewarded the early RL run. Held-out centred cosine 0.356 at the bullet's last token (4-16 tokens), 0.424 at 16-token reads (GLM-5.3 block-51 AR: 0.340 / 0.391; Qwen3.6-27B layer-42 AR: 0.400 at 16 tokens); the best frozen-base affine map 0.293. Part of the collection Kimi K3 olens with the lens repo. Written 2026-10-25 from NOTES.md and the eval files of ceselder/kimi-olens (local git); numbers are exact copies.

Figures

Held-out centred cosine of lora_ar_h59_bos_1M / _2M every 50 steps of their single pass (a), and per-prefix cosine of the final checkpoints next to the GLM-5.3 and Qwen3.6-27B reconstructors (b).

As the RL reward: held-out reconstruction of the lens's greedy bullets under lora_ar_h59_bos_2M rises through step 300 while WorkspaceBench is flat after step 25-100 (both main arms).

Vector PDFs sit next to the PNGs under figures/.

1. What is in this repo

folder	kind	input convention	rows seen	held-out centred cos, bullet spans (harvest-1 set, 1,390 rows)	same, harvest-2-matched set (7,417 rows)	FVE (h1 / h2 sets)	1-token read	8-token read	16-token read
lora_ar_h59_bos_1M/	LoRA r512 reconstructor (ar1, step 987)	[BOS] + the bullet's exact ids (4-16 tokens)	1,010,688 (one pass over harvest 1)	0.3552	0.3541	0.143 / 0.142	0.156	0.344	0.424
lora_ar_h59_bos_2M/	LoRA r512 reconstructor (ar1x, step 2117) = _1M continued onto harvest 2	same	2,167,808 (+1,157,120 harvest-2 rows, one pass)	0.3564	0.3629	0.145 / 0.149	0.153	0.343	0.424
affine_var_bos/	affine map, variable-length span	[BOS] + the bullet (4-16 tokens), read at its last token	1,009,224	0.2647	-	0.080	-	-	-
affine_var_bare/	affine map, variable-length span	the bullet alone	1,009,224	0.2612	-	0.078	-	-	-
affine_k12_bos/	affine map, exactly 12 tokens	[BOS] + the first 12 sampled tokens	1,008,073	0.2822 (1,388 rows)	-	0.089	-	-	-
affine_k12_bare/	affine map, exactly 12 tokens	the 12 tokens alone	1,008,073	0.2793	-	0.088	-	-	-
affine_k16_bos/	affine map, exactly 16 tokens	[BOS] + the first 16 sampled tokens	1,007,402	0.2934 (1,388 rows)	-	0.095	-	-	-
affine_k16_bare/	affine map, exactly 16 tokens	the 16 tokens alone	1,007,402	0.2912	-	0.094	-	-	-

LoRA folders: adapter_config.json + adapter_model.safetensors (13.7 GiB, fp32 r512 adapter on layers 0..59; base_model_name_or_path set to moonshotai/Kimi-K3), value_head.pt (fp32 Linear(7168, 7168) state dict, identity-initialised), centre_mean.pt ({"mu": [7168], "read_layer": 59, ...}, the centre = mean h_sum over the first 20,000 training rows; bitwise equal between the two folders), ar_config.json (prefix ids, span source, read layer). Affine folders: affine.pt (W [7168, 7168], mu_x, mu_y, MU, read, prefix_ids, read_layer, target) and meta.json (ridge sweep on VAL, full held-out eval). affine_results.json at the root holds every map's fit + eval. "Bullet spans" = Kimi's own 4-16-token continuation read at its last token; "k-token read" = centred cosine of the LoRA AR at content position k over rows with >= k tokens (all prefixes 1..16 are regressed, so every k is a trained read). Harvest-1 set = the kept rows of finance_001578 (never trained on); harvest-2-matched set = the kept rows of FineFineWeb file 4149 (never trained on, same distribution as the extension data). Both ship with activations in ceselder/kimi-k3-olens-heldout-activations; the per-checkpoint eval rows are its reconstructor_heldout_evals.parquet.

2. How they are wired

• Input = [BOS] (id 163584) + span token ids, no chat template, read at the span's LAST token. Feed exact token ids (for lens bullets: the sampled content ids after *, whitespace included, capped at 16). The [BOS] prefix is Kimi-K3's sequence start; prefix positions are never read or scored. The _bare affine maps take the span alone (+0.003 for [BOS] everywhere).
• Model = frozen layers 0..59 of moonshotai/Kimi-K3 (embedding + 60 decoder layers, no LM head; native MXFP4 routed experts dequantized per chunk inside grouped GEMMs, bf16 attention / shared experts, Attention Residuals with block size 12) + for the LoRA folders an rsLoRA r 512 / alpha 45 adapter (scale alpha/sqrt(r) = 1.99) on the attention projections, the shared experts and the dense layer-0 MLP of those 60 layers (the lens's target set restricted to layers 0..59; ~3.68B fp32 parameters); the routed experts stay frozen. A capacity test that added r512 LoRA on the routed-expert latent projections (+650M params, same rows and schedule, 300 steps) led by +0.003 centred cosine at step 50 and +0.0006 at step 300 -- the expert pathway is not what limits these reconstructors at this stage.
• Output = value_head(h_sum_AR[last]) with h_sum = the Attention-Residuals sum after layer 59 (embedding + every sublayer output of layers 0..59, the same read vector the lens is given), a vector in the centred space: compare it with h_sum_real - mu by cosine (mu from centre_mean.pt). Only the direction is trained (cosine loss); FVE uses the best per-row non-negative scale. The RL reward is clamp(cos, 0)^2 for one bullet (= NNLS centred FVE with one atom).
• Affine maps: AR(span) = W^T (x - mu_x) + mu_y with x = the frozen base's h_sum after layer 59 at the last token of [prefix] + span read in a fresh context -- a linear read of what the frozen base already carries about the activation the span continues. var = the bullet itself (how the RL reward reads a bullet), k12 / k16 = exactly the first 12 / 16 sampled tokens.

3. How they were trained and how much data they saw

3.1 LoRA reconstructors (lora_ar_h59_bos_*)

• _1M: 987 steps x 1,024 rows = 1,010,688 rows, one pass, no repeats (every kept harvest-1 row but 197; 158 parquet parts, FineFineWeb files 4001..4048). Each row = (Kimi's own sampled continuation at a top-half-entropy position, h_sum at that position), the same rows as the lens SFT (see the lens card, section 3.1).
• _2M: _1M continued for 1,130 more steps onto harvest 2 (FineFineWeb 4101..4148; 1,157,340 de-duplicated rows in one recorded permutation, 1,157,120 consumed, none twice; ext_order.json records the stream). Hop from _1M's final adapter, value head, centre mean and optimizer state (all bitwise-checked); LR re-warmed from 5e-6 to 5e-5 over 20 steps, then linear to 5e-6 at step 2117.
• Loss: 1 - cos(value_head(h_sum_AR[pos]), h_sum - mu), mean over the span's positions (every prefix 1..n regressed to the same target), then over the step's rows. Value head Linear(7168, 7168) fp32, identity init, applied at every span position. _1M: lr 1e-4, warm-up 50, linear to 5% at step 987; AdamW (0.9, 0.95), weight decay 0, no clipping. From scratch (never warm-started from another reconstructor), no whitening, no reuse of any row.
• 9.1-9.6 s/step on one 8 x B300 node (8-stage GPipe pipeline over layers 0..59, per-layer gradient checkpointing with the MoE expert-output cache; 130-159 GB/GPU): ~2.6 h for _1M, ~2.9 h more for _2M.
• Curve (_1M, held-out centred cosine at the bullet's last token): 0.104 (init: frozen base + identity head) -> 0.294 (step 50) -> 0.323 (100) -> 0.331 (150) -> 0.342 (300) -> 0.348 (500) -> 0.353 (750) -> 0.3548 (950) -> 0.3552 (987). Flat after step ~750 on this set. _2M: a re-warm dip to 0.349 (step 1050), then 0.351 (1250) -> 0.354 (1550) -> 0.3564 (2117). What the extension did buy: +0.0088 on the harvest-2-matched set (0.3541 -> 0.3629) and a steadily falling fresh-row training loss (1 - cos 0.7145 over steps 887-987 -> 0.7084 over 1731-1831, +-0.0003), i.e. the gain is real but sits in the extension's own distribution; five bug checks on the extension (rows consumed, LR, eval parity, train loss, per-prefix table) were clean.

3.2 Affine maps (affine_*)

• Closed-form ridge from fp64 shift-centred sufficient statistics accumulated on GPU, one forward pass over the kept harvest-1 rows (158 parts), no training loop; fitted on 1,009,224 (var) / 1,008,073 (k12) / 1,007,402 (k16) rows = every qualifying kept row except the VAL slice (the first 4,096 kept rows of part_00000, never in the fit or the test).
• Ridge lambda = 0.1 x mean diag(Cxx) for every key (VAL centred cosine flat from 1e-5 to 0.03, best at 0.1, falls above 0.3); target unit(h_sum - MU), MU = mean h_sum over 20,000 training rows (raw-centred target identical to 3 decimals). Features from the vLLM engine (TP8, capture mode sum) with early exit after layer 59; ~113 min on one 8 x B300 node for all six maps.

4. How good

4.1 LoRA reconstructors (held-out, bullet spans read at the last token)

checkpoint	harvest-1 set: all 1,390 / window-final 782 / other positions 608	harvest-2-matched set: all 7,417 / window-final 876 / other 6,541	FVE (h1 / h2)	other row (far) h1 / h2	shuffled target h1 / h2
lora_ar_h59_bos_1M (step 987)	0.3552 / 0.3581 / 0.3516	0.3541 / 0.3566 / 0.3537	0.143 / 0.142	0.020 / 0.020	0.022 / 0.040
lora_ar_h59_bos_2M (step 2117)	0.3564 / 0.3579 / 0.3543	0.3629 / 0.3645 / 0.3627	0.145 / 0.149	0.018 / 0.021	0.020 / 0.041
_1M at step 294 (not shipped)	0.3409 / 0.3436 / 0.3375	0.3404 / 0.3438 / 0.3399	0.132 / 0.131	0.018 / 0.020	-
GLM-5.3 block-51 AR lora_ar_h51_sink16 (anchor, its own held-out set)	0.340 (bullet spans)	-	0.172 (16-token)	0.008	-
Qwen3.6-27B layer-42 AR (anchor; 2.4M row-steps)	0.400 (16-token spans)	-	0.186	0.001	-

Centred cosine by content position k (harvest-1 set; the k-th read uses rows with >= k tokens):

checkpoint	p1	p2	p3	p4	p5	p6	p7	p8	p9	p10	p11	p12	p13	p14	p15	p16
lora_ar_h59_bos_1M	0.156	0.210	0.248	0.277	0.301	0.319	0.332	0.344	0.355	0.367	0.374	0.381	0.391	0.400	0.406	0.424
lora_ar_h59_bos_2M	0.153	0.207	0.246	0.276	0.299	0.318	0.331	0.343	0.354	0.367	0.374	0.381	0.392	0.403	0.407	0.424
lora_ar_h59_bos_2M, harvest-2-matched set	0.155	0.211	0.250	0.280	0.302	0.322	0.338	0.352	0.362	0.373	0.382	0.388	0.394	0.403	0.411	0.415
GLM-5.3 lora_ar_h51_sink16	0.138	0.190	0.229	0.257	0.281	0.300	0.316	0.330	0.343	0.352	0.362	0.366	0.379	0.375	0.382	0.408
Qwen3.6-27B AR	0.240	0.305	0.343	0.370	0.390	0.407	0.420	0.430	0.437	0.444	0.450	0.455	0.460	0.465	0.468	0.472

The Kimi reconstructors read 16-token spans at 0.424 (GLM 0.408, Qwen 0.472) and 1-token spans at 0.156 (GLM 0.138 with its sink prefix, Qwen 0.240). The remaining gap to Qwen sits in the first few tokens.

4.2 Affine maps (held-out harvest-1 set)

map	centred cos	permuted-row control	FVE	raw cos	reward clamp(cos,0)^2	LoRA AR _1M at the same read
affine_var_bare	0.2612	0.025	0.078	0.488	0.079	0.355 (bullet)
affine_var_bos	0.2647	0.025	0.080	0.488	0.080	0.355 (bullet)
affine_k12_bare	0.2793	0.027	0.088	0.488	0.088	0.381 (p12)
affine_k12_bos	0.2822	0.027	0.089	0.488	0.090	0.381 (p12)
affine_k16_bare	0.2912	0.029	0.094	0.488	0.095	0.424 (p16)
affine_k16_bos	0.2934	0.028	0.095	0.488	0.096	0.424 (p16)

The frozen-base affine maps trail the LoRA reconstructor by 0.09-0.13 centred cosine at the same read and reach ~60% of its FVE; [BOS] adds +0.003, longer reads help. The frozen Kimi base is more linearly readable than GLM-5.3's (GLM k12/k16 affine 0.231-0.241) and about as readable as Qwen3.6-27B's (0.264-0.270 span-alone).

4.3 As an RL reward

Under lora_ar_h59_bos_2M Kimi's own 16-token continuation scores 0.145 (reward; centred cosine ~0.36) on 256 held-out rows; the SFT lens's greedy bullets 0.197 / 0.42, the release RL lens 0.337 / 0.573 at step 300 (lens card, section 5). The Qwen3.6-27B reader of Kimi activations trained against affine_var_bos (held-out cosine 0.265) gained on WorkspaceBench for 25-50 RL steps and then slid back while its reward kept rising; against lora_ar_h59_bos_1M it climbed for ~150 steps and held (reader card) -- the reward model's held-out cosine decided whether RL gains survived on the bench.

5. How to use

The lens repo ships the project's trainer stack as kimi_stack/; the eval entry point of these folders is its kimi_ar.py --eval-only (8-stage pipeline over layers 0..59 on 8 GPUs, the exact code that produced section 4):

# needs: the frozen base copied to local NVMe (1.56 TB), one 8 x B300 node, torch 2.13 / CUDA 13.0, peft 0.21, fla-core==0.6.0, triton (vllm 0.31's), tiktoken, pyarrow
torchrun --nproc_per_node 8 kimi_stack/kimi_ar.py --model-dir /local/moonshotai__Kimi-K3 --prefix bos --lora-r 512 --lora-alpha 45 \
    --eval-only ./lora_ar_h59_bos_2M --eval-files harvest1_heldout_finance_001578.parquet --out /tmp/ar_eval        # -> /tmp/ar_eval/eval.jsonl (ccos, fve, ccos_p1..p16)

Scoring arbitrary spans against arbitrary activations = the RL reward service (kimi_stack/kimi_ar_reward_service.py: load the adapter + value head + centre, forward [BOS] + span ids through layers 0..59, read the value head at the last token, clamp(cos(pred, h - mu), 0)^2). The affine maps need only torch:

import torch
d = torch.load("affine_k16_bos/affine.pt")                        # W [7168, 7168], mu_x, mu_y, MU, prefix_ids = [163584], read_layer = 59
x = ...                                                             # h_sum after layer 59 of the FROZEN base at the last token of [BOS] + 16 span tokens (vLLM + the fork, capture mode "sum")
pred = (x.float() - d["mu_x"]) @ d["W"] + d["mu_y"]                 # predicted centred direction
ccos = torch.nn.functional.cosine_similarity(pred, h_real.float() - d["MU"], dim=-1)

olens_kimi_k3_usage.py in the lens repo has capture_h59 (the fork's capture at layer 59, extra_args={"output_residual_stream": [59], "capture_positions": [...]}) and affine_reconstruct.

6. Caveats

• Trained on web text. Every row is a FineFineWeb window read at a top-half-entropy position (next-token entropy >= 1.3517 nats), plain token ids; spans are Kimi's own 4-16-token continuations. Chat, code and low-entropy positions are out of distribution.
• Research checkpoints. No safety evaluation.
• Direction only. The cosine loss leaves the output norm untrained; FVE is reported with the best per-row non-negative scale.
• Flat after step ~750 on the harvest-1 held-out set. Doubling the rows moved that set by +0.001 and the matched set by +0.009 while the fresh-row training loss kept falling; adding LoRA on the routed-expert latent projections moved it by +0.001. The open levers are the objective (direction-only cosine on one read layer), the span length (16-token reads are still rising with k) and the engine/trainer numeric gap below, not more of the same rows.
• Attention Residuals. h_sum is a derived quantity (the fork's sum fold of the AttnRes state), not a tensor Kimi-K3 ever materialises; the layer-60 attention actually reads a learned mixture (h_mix, cosine 0.80 with h_sum). Both reconstructor and lens are defined on h_sum; a ridge map to Qwen3.6-27B fitted on h_mix scored 0.651 vs 0.712 on h_sum.
• Engine/trainer gap. The targets were captured by the vLLM engine (MXFP4 x MXFP8 experts); the reconstructor runs in the bf16 HF stack with W4A16 experts. h_sum agrees between the two at cosine 0.994 mean (min 0.914) over captured positions, so a few per cent of the residual error is the stack, not the model.
• Route-cache fix. Both LoRA runs used activation checkpointing in which a MoE routing flip between forward and recompute could mis-assign one token-slot's gradient row; measured exposure was ~124 of 5.7M routed tokens in the first step after a process start and 0 afterwards (_1M ran before the fix, _2M with it).
• Licence. Derivative works of moonshotai/Kimi-K3 under the Kimi K3 License (modified MIT; see the lens card and LICENSE).

7. Provenance

• Code: ceselder/kimi-olens -- src/kimi_ar.py (LoRA AR trainer + eval-only), src/kimi_affar.py (affine maps), src/kimi_ar_reward_service.py (the RL reward service), src/kimi_hf.py + src/kimi_experts.py + src/kimi_pipe.py (trainer stack); NOTES.md is the full log. Eval files: /vol/routeA/evals/ar_heldout/eval.jsonl, /vol/ckpt/kimi_affar/results.json.
• Lens checkpoints RL'd against these: ceselder/kimi-k3-olens-lens. GLM-5.3 reconstructors for comparison: ceselder/glm-5.3-olens-reconstructor.
• Base model: moonshotai/Kimi-K3.