ceselder/muse-spark-olens-reconstructor

Activation reconstructors for Muse Spark Open -- text span -> layer-40 activation direction (LoRA r512 + per-layer heads for layers 8..56)

Built with Muse Spark Open. The text-to-activation half of the Muse Spark Open oracle lens: given a short span of tokens, predict the (mean-centred) layer-40 activation h40 of the frozen Muse Spark Open 816B-A42B at the position the span continues. Two LoRA reconstructors ("AR", the olens h2hpfx recipe ported from Qwen3.6-27B via GLM-5.3 and Kimi-K3; same recipe, rows and seed, one clean and one with a known 1.35% duplicate caveat) and two sets of per-layer heads that read h8 / h16 / h24 / h32 / h40 / h48 / h56 off the same reconstructor body (closed-form ridge, and cosine-trained from the ridge solution). lora_ar_h40_bos_ar1c/ is the reward model of every RL run in ceselder/muse-spark-olens-lens and ceselder/muse-spark-olens-qwen-reader (reward = clamp(cos(AR(<|begin_of_text|> + bullet), h40 - mu), 0)^2); cosine_layers_ar1c/ is the reward of the multilayer lens. Held-out centred cosine 0.378 at the span's last token (4-16 tokens), 0.422 at 15-token reads (Kimi-K3 layer-59 AR: 0.363 / 0.424; GLM-5.3 block-51 AR: 0.340 / 0.408; Qwen3.6-27B layer-42 AR: 0.400 at 16 tokens). Part of the collection Muse Spark Open olens. Written 2026-10-26 from NOTES.md and the runs' log.jsonl / results.json of ceselder/muse-olens (commit 46b1c0b); numbers are exact copies.

1. What is in this repo

folder	kind	input convention	rows seen	held-out centred cos, sampled spans read at the last token (2,048 kept held-out rows)	FVE	far-row control	1 / 8 / 15 / 16-token read
lora_ar_h40_bos_ar1c/	LoRA r512 reconstructor (ar1c, the clean run; the RL reward)	`<	begin_of_text	>` + the bullet's exact ids (4-16 tokens)	2,899,712 (2,831 x 1,024), one pass, (window, t)-de-duplicated	0.3784	0.1622
lora_ar_h40_bos_ar1/	LoRA r512 reconstructor (ar1; 1.35% of rows seen twice)	same	2,899,712, one pass (39,204 rows read twice: 23 harvest parts rewritten in place under the stream)	0.3783	0.1622	0.0186 (0.0325)	0.185 / 0.371 / 0.423 / 0.418
ridge_layers_ar1c/	7 closed-form ridge heads h8..h56 on the frozen ar1c body	the span through the reconstructor body (layers 0..40, adapter merged), read at the last token	300,000	h40 head 0.3703 (section 4.2 for all layers)
cosine_layers_ar1c/	7 cosine-trained heads h8..h56, initialised from the ridge solution (the multilayer RL reward)	same	1,048,576 (256 steps x 4,096, one pass)	h40 head 0.3764 (section 4.2)

Not in this repo: the dense reconstructor trained on the document's TRUE next 16 tokens at every prefilled position (ar_dense1 / ar_dense2) -- ar_dense1 was stopped at step 4,000 after a diagnosis (its held-out cosine on sampled spans drifted down 0.346 -> 0.333 while its own dense held-out rose to 0.380: document concentration in the dense stream -- ~437 correlated positions per window per block -- not a true-vs-sampled effect; at fixed positions the continuation source moves ccos by +0.003 for both ARs) and ar_dense2 (a seeded 31-positions-per-window cap, 26.5M rows) is paused at its step-2000 save. It is added here if it finishes.

LoRA folders: adapter_config.json + adapter_model.safetensors (7.7 GiB, fp32 r512 adapter on layers 0..40; base_model_name_or_path set to meta-models/Muse-Spark-Open-816B-A42B-NVFP4), value_head.pt (fp32 Linear(8192, 8192) state dict, identity-initialised), centre_mean.pt ({"mu": [8192], "read_layer": 40, ...}, the centre = mean h40 over the first 20,000 training rows; identical between the two runs), ar_config.json (prefix ids, span source, read layer), metrics.json (the run's eval curve + arguments). Head folders: heads.pt ({layer: {"W": [8193, 8192] fp32, "mu": [8192], "scale", "lam", "ccos"}} -- [x, 1] @ W = the centred prediction of h_layer, last row = bias; the cosine heads use the same layout) + results.json (lambda grid / lr grid, held-out ccos, shuffled controls per layer). "Sampled spans" = Muse's own 4-16-token continuation read at its last token; "k-token read" = centred cosine at content position k over rows with >= k tokens (all prefixes 1..16 are regressed, so every k is a trained read). The held-out rows ship with activations at all seven layers in ceselder/muse-spark-olens-heldout-activations.

2. How they are wired

• Input = <|begin_of_text|> (id 200000) + 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 position is never read or scored.
• Model = frozen layers 0..40 of Muse Spark Open (normed embedding + 41 decoder layers, no LM head; NVFP4 routed experts dequantised per chunk inside grouped GEMMs, BF16 attention / projections / norms / layer-0 experts, the per-channel gated residual, fp32 stream) + an rsLoRA r 512 / alpha 45 adapter (scale 45/sqrt(512) = 1.99) on self_attn.{q,k,v,o,gate}_proj + mlp.{pre,post}_expert_proj of those 41 layers (7 modules per layer, 1,934.6M fp32 parameters); the routed experts stay frozen.
• Output = value_head(h40_AR[last]), a vector in the centred space: compare it with h40_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). The inverter arm of the lens repo injects mu + 57.48 * unit(AR(span)) (raw-space estimate: 57.48 = median ||h40 - mu||).
• Per-layer heads: for every stored layer l in {8, 16, 24, 32, 40, 48, 56}, pred_l = [x, 1] @ W_l with x = the reconstructor body's h40-stream state at the span's last token (fp32); per-layer centre mu_l = that layer's own training-row mean (per-layer normalisation, never pooled across layers); ccos_l = cos(pred_l, h_l - mu_l). Ridge heads: closed-form, lambda relative to mean diag(X^T X) picked per layer on the held-out set (grid 0.01 .. 30; best 0.3-1). Cosine heads: the ridge solution fine-tuned with the 1 - cos loss (AdamW, lr 1e-4 picked on an 8,192-row validation slice among 1e-4 / 3e-4 / 1e-3), 256 steps x 4,096 fresh rows, one pass.

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

3.1 LoRA reconstructors (lora_ar_h40_bos_*)

• Rows = the lens SFT's rows: (Muse's own sampled continuation at a top-half-entropy position of an sft-split document, h40 at that position); 2,831 steps x 1,024 rows = 2,899,712 rows, one pass, no repeats (streamed from the harvest in round-robin blocks over the 8 shards as it was written; stream_order.json records every block). ar1c de-duplicates on (window, t); ar1 (launched before the duplicate audit) consumed the pre-rewrite version of 6 blocks and therefore saw 39,204 rows twice (the second time with a freshly sampled span) -- below our 2% restart bar, recorded as a caveat; ar1c is its clean twin with the same recipe, rows and seed. The two end within 0.0001 of each other (0.3784 vs 0.3783; identical FVE 0.1622): the duplicates made no measurable difference.
• Loss: 1 - cos(value_head(h40_AR[pos]), h40 - 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(8192, 8192) fp32, identity init. lr 1e-4, warm-up 50, linear to 5% at step 2,831; 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. mu = mean h40 of the first 20,000 streamed rows.
• 2.8-4.0 s/step on one 8 x B200 / 8 x H200 node (8-stage 1F1B pipeline over layers 0..40 with the MoE recompute cache, 52-85 GB/GPU); both streamed behind the harvest for most of their life (the round-robin advances at the slowest shard), so wall time was dominated by data waits.
• Curves (held-out centred cosine at the span's last token, 2,048 kept held-out rows, every 50 steps): ar1c 0.055 (init: frozen base + identity head) -> 0.350 (350) -> 0.3545 (500) -> 0.359 (700) -> 0.3634 (1,000) -> 0.368 (1,400) -> 0.3727 (2,000) -> 0.3767 (2,500) -> 0.3784 (2,800); ar1 0.3519 (500) / 0.3618 (1,000) / 0.3724 (2,000) / 0.3759 (2,500) / 0.3783 (2,800). Still rising at the end of the pass (+0.002 over the last 300 steps); the numbers are those of the final adapters (the last logged eval is step 2,800 of 2,831).

3.2 Per-layer heads (ridge_layers_ar1c, cosine_layers_ar1c)

• Ridge: closed-form per layer on the FROZEN ar1c body's h40-stream state at the span's last token (fp32, [x, 1] with an unpenalised bias), 300,000 sft kept rows of 48 parts (seeded file permutation), per-layer centre mu_l = that layer's own training-row mean; held-out = 4,096 kept rows of the held-out file. lambda relative to mean diag(X^T X), grid 0.01 / 0.1 / 0.3 / 1 / 3 / 10 / 30 -> best 1.0 for h8 / h16 / h24, 0.3 for h32 .. h56 (peaks inside the grid). Controls: shuffled (neighbour) = the next row (often the same document), shuffled (far) = the row n/2 away (another document). The same fit on ar1 step 566 (the validation run, not shipped) gave 0.340 / 0.330 / 0.338 / 0.352 / 0.357 / 0.377 / 0.397.
• Cosine: each ridge head fine-tuned with the 1 - cos loss on 1,048,576 fresh rows (256 steps x 4,096, one pass; lr grid 1e-4 / 3e-4 / 1e-3 on an 8,192-row validation slice, 1e-4 won at every layer): +0.006 to +0.007 centred cosine over ridge at every layer.

4. How good

4.1 LoRA reconstructors (held-out, sampled spans read at the last token; 2,048 kept rows of the held-out file)

checkpoint	centred cos	FVE	neighbour-row control	far-row control	GLM-5.3 / Kimi-K3 / Qwen anchors (own held-out sets)
lora_ar_h40_bos_ar1c (step 2,831; last eval 2,800)	0.3784	0.1622	0.0317	0.0184	GLM block-51 AR 0.340 (FVE 0.172 at 16 tokens) / Kimi layer-59 AR 0.356-0.363 (0.145-0.149) / Qwen layer-42 AR 0.400 at 16 tokens (0.186)
lora_ar_h40_bos_ar1 (step 2,831)	0.3783	0.1622	0.0325	0.0186
reward-service self-test (ar1c, Muse's own continuations <= 16 tokens, 2,048 held-out rows)	0.3776 (reward 0.163; shuffled 0.018)				the service reproduces the trainer's eval within 0.001

Centred cosine by content position k (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
ar1c final	0.185	0.239	0.276	0.303	0.324	0.343	0.358	0.371	0.385	0.395	0.402	0.405	0.407	0.417	0.422	0.417
ar1 final	0.185	0.239	0.276	0.303	0.324	0.342	0.357	0.371	0.384	0.395	0.402	0.405	0.408	0.417	0.423	0.418
Kimi-K3 lora_ar_h59_bos_2M (anchor, its own set)	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
GLM-5.3 lora_ar_h51_sink16 (anchor)	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 (anchor)	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 Muse reconstructor reads 1-token spans at 0.185 (Kimi 0.153, GLM 0.138, Qwen 0.240) and 15-token spans at 0.422 (Kimi 0.407, GLM 0.382, Qwen 0.468); at the span's own last token 0.378 vs Kimi 0.363 / GLM 0.340 (different models, tokenizers and held-out sets: indicative only). The 16-token column dips below the 15-token one for both Muse runs (fewer rows reach 16 tokens; the same dip appears for GLM at p14).

4.2 Per-layer heads on the ar1c final body (300,000 ridge rows; 4,096 held-out rows; the cosine heads on 1,048,576 more rows)

layer	ridge lambda	ridge held-out ccos	shuffled (neighbour)	shuffled (far)	ridge FVE	cosine head ccos	cosine FVE	RMS of centred h_l
h8	1.0	0.3541	0.0668	0.0084	0.1378	0.3606	0.1420	0.760
h16	1.0	0.3425	0.0833	0.0131	0.1291	0.3494	0.1330	0.808
h24	1.0	0.3487	0.1257	0.0279	0.1383	0.3546	0.1427	0.743
h32	0.3	0.3633	0.1195	0.0221	0.1521	0.3692	0.1561	0.691
h40	0.3	0.3703	0.1071	0.0193	0.1561	0.3764	0.1608	0.646
h48	0.3	0.3895	0.1082	0.0186	0.1684	0.3958	0.1736	0.614
h56	0.3	0.4088	0.1130	0.0196	0.1802	0.4145	0.1850	0.544

ar1c's own trained cosine head on h40 scores 0.3788 (far-row control 0.0189, FVE 0.1633) on the same 4,096 rows -> the ridge h40 head is 0.008 below it and the cosine-tuned h40 head 0.002 below; the later layers are the most predictable from the span (h56 0.41), h16 the least (0.35). Every layer's far-row control stays at 0.01-0.03 (the neighbour-row control is higher because neighbouring rows are often the same document). Figure: figures/layer_heads.png.

4.3 As an RL reward

Under ar1c on 256 held-out rows (greedy one bullet, reward = clamp(cos, 0)^2): Muse's own <= 16-token continuation scores **0.1706**; the fresh-init SFT lens 0.1854 (ccos 0.404); the arm-F RL lens 0.2659 (25) -> 0.2828 (50) -> 0.2973 (100) -> 0.3085 (150) -> 0.3222 (175) -> 0.3240 (225) -> 0.3249 (250) -> 0.3307 (300; ccos 0.565) = the released checkpoint (lens card). The Qwen3.6-27B reader after SFT scores 0.1985 (ccos 0.425) and 0.251 / 0.270 / 0.285 / 0.292 / 0.301 at RL steps 25 / 50 / 75 / 100 / 125 (lr 3e-6 arm; reader card). The multilayer reward service (this repo's cosine heads + the AR body) reproduces the per-layer held-out cosines in its boot self-test (results/rl/reward_selftest_mar1c_ml.json: h8 0.370 on its 2,048-row slice).

5. How to use

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

# needs: the frozen NVFP4 base copied to local NVMe (533 GB), one 8 x B200 (or H200) node, torch 2.14 / CUDA 13.0, triton 3.8, transformers 5.20, peft 0.21, pyarrow, tokenizers
torchrun --nproc_per_node 8 muse_stack/muse_ar.py --load-dir /local/unsloth__Muse-Spark-Open-816B-A42B-NVFP4 --prefix bos --lora-r 512 --lora-alpha 45 \
    --eval-only ./lora_ar_h40_bos_ar1c --eval-file heldout_rows.parquet --out /tmp/ar_eval        # -> eval ccos, fve, ccos_p1..p16 in the log

Scoring arbitrary spans against arbitrary activations = the RL reward service (muse_stack/muse_ar_reward_service.py: load the adapter + value head + centre, forward <|begin_of_text|> + span ids through layers 0..40, read the value head at the last token, clamp(cos(pred, h - mu), 0)^2; muse_ar.ar_predict(st, vh, spans) is the single call; --layer-heads cosine_layers_ar1c/heads.pt scores against any of the seven layers). The heads need only torch once you have the body's state:

import torch
H = torch.load("cosine_layers_ar1c/heads.pt")                      # {layer: {"W": [8193, 8192], "mu": [8192], ...}}
x = ...                                                             # the reconstructor body's h40-stream state at the span's last token (fp32; muse_ar.ar_predict with the raw state)
pred56 = torch.cat([x.float(), torch.ones(x.shape[0], 1)], -1) @ H[56]["W"]
ccos = torch.nn.functional.cosine_similarity(pred56, h56_real.float() - H[56]["mu"], dim=-1)

olens_muse_spark_usage.py in the lens repo has capture_h40 (the fork's capture at any layer: extra_args={"output_residual_stream": [L], "capture_positions": [...]}) and layer_head_reconstruct.

6. Caveats

• Trained on web text. Every row is a FineFineWeb window read at a top-half-entropy position (next-token entropy >= 1.664 nats), plain token ids, no BOS in the context; spans are Muse'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.
• ar1 duplicates. 1.35% of ar1's rows were trained twice (section 3.1); use ar1c unless you are reproducing the inverter arm (which inverted ar1 step 849).
• Engine/trainer gap. The targets were captured by the vLLM engine (W4A16 Marlin experts, fp32 stream); the reconstructor runs in the PyTorch stack with the same W4A16 math (lens-prompt h40 cosine 0.993, next-token KL 0.002). Muse's own bf16 sensitivity (batch-composition re-captures differ at cosine ~0.95 at layer 40) is part of the residual error.
• Dense (true next-16) reconstructor not shipped (see section 1): document concentration in the dense stream, not the target choice, held it below the sampled-span runs; the capped re-run is paused.
• Licence. Derivatives of Muse Spark Open under the Muse Model License (LICENSE.md + USE_POLICY.md, Meta's files verbatim; "Built with Muse Spark Open"). Copyright (c) 2026 Meta Platforms, Inc. for the base model.

7. Provenance

• Code: ceselder/muse-olens (commit 46b1c0b) -- src/muse_ar.py (LoRA AR trainer + eval-only + --precompute for the inverter arm), src/muse_ar_layers.py (per-layer ridge + cosine heads), src/muse_ar_reward_service.py (the RL reward service, single- and multi-layer), src/muse_hf.py + src/muse_experts.py + src/muse_pipe.py (trainer stack); NOTES.md is the full log. Logs: /vol/ckpt/{ar1,ar1c}/log.jsonl, /vol/ckpt/ar1c_layers_{final,cos}/results.json, /vol/results/rl/reward_selftest_mar1c*.json on the muse-olens Modal volume; report http://5.78.192.0/reports/view/muse-spark-olens/report.html.
• Lens checkpoints RL'd against these: ceselder/muse-spark-olens-lens. Kimi-K3 / GLM-5.3 reconstructors for comparison: ceselder/kimi-k3-olens-reconstructor, ceselder/glm-5.3-olens-reconstructor.
• Base model: meta-models/Muse-Spark-Open-816B-A42B-NVFP4 (ungated mirror unsloth/Muse-Spark-Open-816B-A42B-NVFP4). Built with Muse Spark Open.