oracle-encoders

oracle-encoders

Three independent per-layer oracle encoders (Stage 7, corrected one-layer-per-model design). Each maps raw text -> the 54 concept-probe detection scores of EXACTLY ONE gemma-2-2b layer, per token. Full fine-tune of Qwen/Qwen3-0.6B-Base + MLP head (1024 -> 4096 -> GELU -> 54). Original runs trained 2026-07-10 overnight; L6/L8 continuation (cont1) checkpoints trained 2026-07-10 afternoon.

Checkpoints — which encoder is which

Two generations per layer (L6/L8). layerXX/ = the ORIGINAL overnight run (stopped by an 11h wall clock, still improving). layerXX/cont1/ = the CONTINUATION: warm-started from the original best, trained on the remaining unseen corpus-scores(+overflow) shards until the eval plateau rule fired (Δ median R² < 0.005 over trailing 150M tokens). Same val shards {353,354} throughout, so all R² values below are directly comparable. For deployment, prefer cont1 for L6/L8 (strictly better on the identical heldout set); L14 was deliberately not continued.

Baseline: the deprecated 162-joint oracle reached 0.6371 — all beat it. Each metrics.jsonl holds the full eval history (per-probe / per-family R², Spearman, loss vs tokens); cont1 metrics cover the continuation segment only.

Continuation details (cont1)

• Warm start from layerXX/best_stripped.pt (no optimizer state existed): 150-step LR re-warmup to the original cosine's value at the resume step, then cosine toward a 0.1 floor at the projected full-epoch step.
• Data: the shards the original runs never saw — 338..321, 355, 356..362 (L6 also re-queued 339, which was 80% consumed at the wall-clock kill; the plateau fired before reaching it). Boundary located by exact token accounting (measured gathered/stored ratio 0.99996).
• Continuation gains were modest (+0.0037 R² both layers over ~160M extra tokens) — the overnight "still improving" slope flattened quickly.
• cont1/best_full.pt includes full Muon+AdamW optimizer state (unlike the originals) — exact resume is possible if training is ever pushed further. cont1/DONE.json records stop reason/step/tokens; cont1/logs/ has the training + prefetch logs.

Load

import torch
from transformers import AutoModel
from train_oracle_perlayer import OracleMLPHead  # concept_probes/stage7_oracle/code/

ck = torch.load("layer08/cont1/best_stripped.pt", map_location="cpu", weights_only=False)
enc = AutoModel.from_pretrained("Qwen/Qwen3-0.6B-Base", dtype=torch.bfloat16)
enc.load_state_dict(ck["encoder_state"])
head = OracleMLPHead(ck["hidden_size"], ck["head_up"], ck["K"]).to(torch.bfloat16)
head.load_state_dict(ck["head_state"])
# per-token: pred = head(enc(input_ids).last_hidden_state)  -> [.., 54] standardized scores

Outputs are in corpus-standardized units (dequantized int8 scores standardized by per-probe ClimbMix mean/std). Output dim c = concept ck["concepts"][c]. Checkpoint dict keys: encoder_state (Qwen3-0.6B, bf16), head_state (MLP), step, train_tokens, layer, layer_idx, K=54, concepts (store order), families, head_up=4096, args (+ muon_state/adamw_state in best_full.pt).

Provenance

• Objective: MSE on corpus-standardized scores (stage-7 expA), one layer per model — targets physically sliced from the stacked [n,3,54] store before training.
• Data: corpus-scores + overflow (ClimbMix shards; originals: train 320,352..321 order, cont1 continues from the kill boundary; val 353/354 always); text from karpathy/climbmix-400b-shuffle.
• Probes: concept-probes-gemma2-2b gold_probes/layer06|08|14. DoM steering scores (separate): concept-probes-corpus-scores-dom-layer8.
• Optimizer: Muon (lr 5e-3, Newton-Schulz) on 2D matrices + AdamW (1e-4) on the rest; bsz 6 docs x grad-accum 8; gemma->qwen prefix alignment; trainer: concept_probes/stage7_oracle/code/train_oracle_perlayer.py (cont1 flags: --resume --hf-subdir layerXX/cont1 --cont-anchor-mult --cont-end-step).

stage7_eval/ — evaluation artifacts + probe arrays

Checkpoint of the Stage-7 eval results and probe weight arrays (moved out of git; source dir concept_probes/5_oracle/out/, formerly stage7_oracle/out/). Relative structure is preserved under stage7_eval/.

• gold_probes_per_layer/ — the 54-probe gold banks actually used for scoring: detection (mixed-method) and steering (difference-of-means) .npz, layers 6/8/14
• probe_set.json, probe_set_arrays.npz, probe_set_dom_steering_l6_l8_l14.npz — packed probe-set definition + arrays + DoM steering vectors
• retro_metrics/ — retroactive expA/expB metric extracts (.jsonl) + wandb watch script
• figures/ — oracle-vs-probe association figure and its underlying matrices
• gate/eval JSONs — g1_corpus_stats.json, g1_natural_ref.json, g1_residual_checks.json, g2_retention.json, verify_report.json, expA_frozen_metrics.jsonl, expB_final_analysis.json, coord_fidelity.json, quant.json, scoring_config.json, bench/ (throughput/parity benchmarks)