K5: inference-time levers at matched generated tokens

Status: harness built and CPU-tested (2026-10-19); no GPU run yet. Plan: docs/SOTA_PLAN.md v3.1 section 3 (K5, "run FIRST"), docs/PREREG_SOTA.md section 4 (K5 MDE 1.3 pp estimator test) and section 5 (K5 is its own Holm family).

Question: at the SAME total generated tokens per problem, does Lean-feedback correction (L1) or the agentic loop (L6: proof bank + APOLLO-style subgoal repair + compiler feedback) solve more dev problems than restart sampling?

1. What is built

File	Role
`scripts/p2/k5_levers.py`	Orchestrator: stage plan (exact command lines), resumable runner, per-problem token ledger, budget equaliser, proof bank + decontamination + retrieval, route/assemble rounds of the agentic arm, matched-budget report with paired tests
`align/rl/envs/subgoal_repair.py`	Pure-Python subgoal extraction / reassembly; Lean only through `align.verify.lean` (read-only: `Worker.header_env` / `Worker.command`, the same calls as `verify_generations.collect_feedback`)
`extras/vllm_rollout/lean_prompts.py`, `lean_gen.py`	New prompt format `goedel_bank` (Goedel prompt preceded by retrieved verified proofs; with no examples it is byte-identical to `goedel`). The `goedel` path is untouched (its byte-pin test passes).
`tests/test_k5_levers.py`, `tests/test_subgoal_repair.py`, `tests/fixtures/k5_baseline_proofs.jsonl`	30 + 22 tests; fixture = 14 failed and 6 accepted real proofs from the P2 baseline file, verified with `--feedback`

2. Arms and the matching rule

All arms: Goedel-Prover-V2-8B (rev dfd02e62), T = 1.0, top_p = 0.95, seed 0, cap 40,960, same dev problems, same frozen verifier.

restart: n independent round-0 attempts per problem (also the roots of correct, so the arms are paired and the roots are generated once).
correct (class B): each failed root gets up to 2 Goedel-format repairs (lean_gen --round 1/2 --corrections-from, --trajectory-cap-tokens 40000, --repair-solved so every trajectory has a complete cost record, default --no-feedback-policy skip).
agentic (class C, exploratory), R = 4 rounds: round 0 = bank-conditioned attempt (goedel_bank, k = 3 retrieved proofs); rounds 1..3 = subgoal, feedback, subgoal (configurable --routes; a design choice, not tuned). A subgoal round falls back to feedback for parents with nothing extractable.

Unit and cost. The unit of every arm is a trajectory. Its cost is the completion tokens of every node up to the end of the round of its first success (all nodes of a round are paid, branches included), or all its tokens if it never succeeded. Subgoal samples (accepted or not) are charged to the trajectory that spawned them, through the assembled node's trajectory.completion_tokens. Prompt/prefill tokens are not in the budget (same convention as the class-B ledger); they are reported separately in box-hours.

Equaliser. Arms are generated generously as pools of units, then replayed at a per-problem budget B_p: units are drawn in random order (256 permutations, seed fixed per problem) until one succeeds or the next unit does not fit the remaining budget (a real run cannot know a unit's cost in advance; it would be clipped and fail, so the replay stops rather than skipping to a cheaper unit). The per-problem estimate is the fraction of orderings that solve it. B_p defaults to restart_k:8 = 8 x the mean restart attempt tokens of that problem (the spend of a pass@8 run); fixed:N is available. Curves at 0.25 / 0.5 / 1.0 x B_p are reported. A problem whose pool cannot fill the budget is underfilled; k5_levers.py status prints how many more units to generate. Contrasts (correct - restart, agentic - restart, agentic - correct) use align.eval.paired.sign_flip_test, paired_bootstrap_ci and holm (the K5 family of PREREG section 5). The K0 training-seed null does not apply (no training).

Tested arithmetic (tests/test_k5_levers.py): closed-form replay (4 equal units, 1 solved, budget for 2 -> 0.5; for 3 -> 0.75), never overspends, deterministic, stop-rule, underfill planning, restart_k budget, a 2x-costlier chain loses to restarts at equal tokens, chain/branch/assembled-node cost, refusal of unverified or partial lineages.

3. Subgoal repair (APOLLO-style)

Locate (locate_targets): every Lean error (positions index compile_view(code), translated by view_line_map) maps to the innermost failing have name : T := rhs (indent/bracket based extents, masked comments), else a main-line tactic, else a "tail" (proof ended with goals). Errors on a have header are reported unrepairable. A target containing a finer target is dropped (the outer failure shows up next round). More than 6 gaps: refused.
Sorrify (build_skeleton): replace each target by sorry.
Goal states: one REPL command on the skeleton; its sorries give the exact hypotheses and goal at each gap. The tail cut is added only if the skeleton still reports the unsolved-goals error at the theorem's by (after a parse error Lean reports a cascade that vanishes once the broken step is cut; found on a real case).
Lemma (lemma_from_goal): theorem k5_sg_<hash> (hyps) : goal := by sorry; inaccessible hypotheses (n✝) get fresh names with every later mention rewritten, instances become [inst]; let-hyps unsupported. Each lemma is checked to elaborate (check_lemma_elaborates); those that do not are not sent to the prover. They are ordinary problems for lean_gen (Goedel round-0 prompt + bank examples).
Reassemble (reassemble): the prover's tactic body is spliced inline at the gap's indentation, preceded by rename_i for inaccessible hypotheses so the body sees the lemma's names. Unproven gaps get their original text back (so a half-repaired proof is a better parent next round). Inline rather than lemma-first, because the strict verifier rejects extra declarations unless allow_helper_defs, which we do not change.
The assembled proof is only a candidate: it is verified by the unchanged strict verifier against the ORIGINAL statement.

Real-data evidence (all pass): every one of 14 real failures yields targets; skeleton + restore is the identity on all 14; 36 real have cuts of 6 accepted proofs re-inlined reproduce the proof up to whitespace; under real Lean (-m slow): goal states returned for 8 failures, at least 5 lemmas elaborate, restoring gaps gives the same REPL errors as the original, and the full cut -> lemma proven by the original body (strict-accepted) -> splice -> strict-accepted loop succeeds on 5 real proofs and on a crafted case that needs rename_i. In an exploratory run on the 14 failures, the first gap's lemma failed to elaborate for 2 (a ZMod numeral ambiguity, which stays a limitation, and a shadowed n✝, fixed by the rename rewrite).

4. Proof bank and decontamination

Source: accepted round-0 proofs of the TRAIN pool (train_pool.jsonl, ver_r0.frozen.jsonl), the shortest per problem (<= 2,500 chars). Decontamination (decontaminate_bank) uses align.data.decontam (normalised-text, signature and alpha hashes, MinHash-LSH, token n-gram containment, informal-text n-grams) against EVERY evaluation and query set passed (--eval-problems plus the dev problems themselves), drops ids equal to an eval id, and propagates ancestry. Over-flagging is the policy; dropped entries and reasons are written to bank_dropped.jsonl. Retrieval is TF-IDF cosine over statement tokens (identifiers kept), excluding the query's own problem. Tests: copies, renames, hypothesis reorderings and eval-id collisions are dropped; a near-copy of a dev problem is dropped; on the real pool, a fresh detector finds zero hits in the kept set and the kept set shares no id with any eval set. NOT yet done: the elaborated-type hash (needs the Lean env; scripts/p2/supply/elab_hash.py) is supported by the detector but not fed here, so reordered-hypothesis copies rely on MinHash/n-grams. Subgoal-lemma prompts retrieve from the same bank.

5. Design sources

APOLLO (Ospanov, Farnia, Yousefzadeh), arXiv 2505.05758: Lean compiler + LLM; repair syntax, isolate failing sub-lemmas, automated solvers, small-budget LLM calls on remaining goals, recombine and recheck. We take: failing-step isolation, per-subgoal proving, reassembly, re-verification. We do not take: the automated-solver pass (hammer tactics) or syntax-refiner agent; adding them would confound the "model" lever.
Hilbert (Varambally et al.), arXiv 2509.22819 (ICLR 2026): recursive decomposition of unsolved subgoals, verifier feedback, semantic theorem retrieval. We take: retrieval-conditioned proving and the idea that unsolved subgoals recurse (our per-round re-extraction on the assembled proof). We do not take the informal-reasoner model (needs a second, stronger model; breaks "same 8B").
DeepSeek-Prover-V2, arXiv 2504.21801: subgoal decomposition with have ... := sorry skeletons, prove subgoals separately, recombine. Our skeleton-of-sorries is the same idea, built from errors instead of a decomposing model. (Cited from memory; not re-fetched in this session.)
Goedel-Prover-V2, arXiv 2508.03613: the correction prompt and the 2-repair protocol, ported byte-faithfully in lean_prompts.py. (From memory.)
Kimina-Prover / Kimina Lean Server, arXiv 2504.11354: REPL-pool verification, which align.verify.lean mirrors. (From memory.)
OProver-8B (the plan's class-C bar: feedback ablation 91.8 vs 87.7, SR): motivation for the bank + feedback combination; no code consulted.

6. Limits (read before interpreting)

The agentic arm confounds three levers (bank, subgoal repair, feedback). Attribution needs ablations (bank off via empty bank_examples; --routes feedback,feedback,feedback), not run here.
Bank-conditioned prompts are a different conditioning of the policy than the Goedel prompt; goedel_bank is not covered by the Goedel byte-fidelity gate.
Assembled nodes carry placeholder token ids (completion_ids = [0]); their real cost is in trajectory.completion_tokens. Anything that trains on these records must exclude kind == "assembled".
The replay estimator treats units as exchangeable draws from the generated pool; correlated units (same root) are not an issue because each root is one unit, but pool size caps what budget is testable (underfilled).
The class-B arm uses Goedel's skip policy for truncated or error-less parents, so truncated roots are wasted; agentic uses the same.
Prefill/verifier CPU cost is not in the token budget; box-hours below include generation only.
Not run on GPU: route/assemble have CPU tests (assemble end-to-end with fake probe, route only via the unit pieces); route against a live pool is covered by the slow plan_subgoals tests, not by a CLI test.

7. GPU runs (dev set = miniF2F-valid + ProofNet-valid), estimated box-hours

Prerequisites: proofnet_valid.jsonl / proofnet_test.jsonl built (Phase 0a item 3; not present under .tmp/p2/problems/ at the time of writing), chain8 step pool finished (r1/pool/ver_r0.frozen.jsonl), and the verifier frozen (.tmp/restart/VERIFIER_FROZEN).

cd /apps/development/rlhf/align && . /apps/development/rlhf/env/apps_only.sh
P=.tmp/p2; K5=$P/k5/dev
PYTHONPATH=$PWD .venv/bin/python scripts/p2/k5_levers.py run --run-dir $K5 \
  --problems $P/problems/minif2f_valid.jsonl $P/problems/proofnet_valid.jsonl \
  --bank-problems $P/problems/train_pool.jsonl --bank-gens $P/r1/pool/gen_r0.jsonl --bank-vers $P/r1/pool/ver_r0.frozen.jsonl \
  --eval-problems $P/problems/minif2f.jsonl $P/problems/proofnet_test.jsonl $P/problems/mobench.jsonl \
  --n-restart 16 --n-agentic 6 --budget restart_k:8

run executes the 18 stages below in order, each inside hostguard.py run --wait (GPU stages 24 cores / 100 GB, CPU stages 16 cores / 64 GB), each with a done-marker ($K5/done/<stage>), so a re-run resumes; every generation/verification file is keyed and resumable as before. plan with the same arguments prints all 18 expanded lines without starting anything; --only <stage>... runs a subset. The first stages, expanded (paths real, $K5 = .tmp/p2/k5/dev):

# restart_r0
/apps/development/rlhf/align/.venv/bin/python /apps/development/rlhf/align/scripts/hostguard.py run --wait --cores 24 --mem-gb 100 --label k5-restart_r0 -- /apps/development/rlhf/align/.venv/bin/python /apps/development/rlhf/align/scripts/p2/pipeline_gen_verify.py --run-dir /apps/development/rlhf/align/.tmp/p2/k5/dev/pipe/restart_r0 --problems /apps/development/rlhf/align/.tmp/p2/k5/dev/dev.jsonl --verify-out /apps/development/rlhf/align/.tmp/p2/k5/dev/restart/ver_r0.jsonl --workers 8 --label k5-restart_r0 --gen-args '--model /apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots/dfd02e6271a58375dfbf3ece0175277cf6b6a89a --problems /apps/development/rlhf/align/.tmp/p2/k5/dev/dev.jsonl --out /apps/development/rlhf/align/.tmp/p2/k5/dev/restart/gen_r0.jsonl --n 32 --cap-tokens 40960 --round 0 --tp 4 --seed 0 --max-num-seqs 128 --gpu-mem-util 0.82 --label k5-restart_r0' --feedback

# correct_r1
/apps/development/rlhf/align/.venv/bin/python /apps/development/rlhf/align/scripts/hostguard.py run --wait --cores 24 --mem-gb 100 --label k5-correct_r1 -- /apps/development/rlhf/align/.venv/bin/python /apps/development/rlhf/align/scripts/p2/pipeline_gen_verify.py --run-dir /apps/development/rlhf/align/.tmp/p2/k5/dev/pipe/correct_r1 --problems /apps/development/rlhf/align/.tmp/p2/k5/dev/dev.jsonl --verify-out /apps/development/rlhf/align/.tmp/p2/k5/dev/correct/ver_r1.jsonl --workers 8 --label k5-correct_r1 --gen-args '--model /apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots/dfd02e6271a58375dfbf3ece0175277cf6b6a89a --problems /apps/development/rlhf/align/.tmp/p2/k5/dev/dev.jsonl --out /apps/development/rlhf/align/.tmp/p2/k5/dev/correct/gen_r1.jsonl --n 1 --cap-tokens 40960 --round 1 --tp 4 --seed 0 --max-num-seqs 128 --gpu-mem-util 0.82 --label k5-correct_r1 --trajectory-cap-tokens 40000 --repair-solved --corrections-from /apps/development/rlhf/align/.tmp/p2/k5/dev/restart/ver_r0.jsonl /apps/development/rlhf/align/.tmp/p2/k5/dev/restart/gen_r0.jsonl' --feedback

# agentic_r0
/apps/development/rlhf/align/.venv/bin/python /apps/development/rlhf/align/scripts/hostguard.py run --wait --cores 24 --mem-gb 100 --label k5-agentic_r0 -- /apps/development/rlhf/align/.venv/bin/python /apps/development/rlhf/align/scripts/p2/pipeline_gen_verify.py --run-dir /apps/development/rlhf/align/.tmp/p2/k5/dev/pipe/agentic_r0 --problems /apps/development/rlhf/align/.tmp/p2/k5/dev/dev_bank.jsonl --verify-out /apps/development/rlhf/align/.tmp/p2/k5/dev/agentic/ver_r0.jsonl --workers 8 --label k5-agentic_r0 --gen-args '--model /apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots/dfd02e6271a58375dfbf3ece0175277cf6b6a89a --problems /apps/development/rlhf/align/.tmp/p2/k5/dev/dev_bank.jsonl --out /apps/development/rlhf/align/.tmp/p2/k5/dev/agentic/gen_r0.jsonl --n 8 --cap-tokens 40960 --round 0 --tp 4 --seed 0 --max-num-seqs 128 --gpu-mem-util 0.82 --label k5-agentic_r0 --prompt-format goedel_bank' --feedback

# agentic_route_r1
/apps/development/rlhf/align/.venv/bin/python /apps/development/rlhf/align/scripts/hostguard.py run --wait --cores 16 --mem-gb 64 --label k5-agentic_route_r1 -- /apps/development/rlhf/align/.venv/bin/python /apps/development/rlhf/align/scripts/p2/k5_levers.py route --run-dir /apps/development/rlhf/align/.tmp/p2/k5/dev --round 1 --mode subgoal --workers 2 --bank-k 3 --sub-n 4 --pairs /apps/development/rlhf/align/.tmp/p2/k5/dev/agentic/ver_r0.jsonl /apps/development/rlhf/align/.tmp/p2/k5/dev/agentic/gen_r0.jsonl

# agentic_sg_r1
/apps/development/rlhf/align/.venv/bin/python /apps/development/rlhf/align/scripts/hostguard.py run --wait --cores 24 --mem-gb 100 --label k5-agentic_sg_r1 -- /apps/development/rlhf/align/.venv/bin/python /apps/development/rlhf/align/scripts/p2/pipeline_gen_verify.py --run-dir /apps/development/rlhf/align/.tmp/p2/k5/dev/pipe/agentic_sg_r1 --problems /apps/development/rlhf/align/.tmp/p2/k5/dev/agentic/round1/sg_problems.jsonl --verify-out /apps/development/rlhf/align/.tmp/p2/k5/dev/agentic/round1/sg_ver.jsonl --workers 8 --label k5-agentic_sg_r1 --gen-args '--model /apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots/dfd02e6271a58375dfbf3ece0175277cf6b6a89a --problems /apps/development/rlhf/align/.tmp/p2/k5/dev/agentic/round1/sg_problems.jsonl --out /apps/development/rlhf/align/.tmp/p2/k5/dev/agentic/round1/sg_gen.jsonl --n 4 --cap-tokens 16384 --round 0 --tp 4 --seed 0 --max-num-seqs 128 --gpu-mem-util 0.82 --label k5-agentic_sg_r1 --prompt-format goedel_bank'

After the GPU stages: k5_levers.py report --run-dir $K5 --budget restart_k:8 (also status, which lists under-filled problems and how many units each needs; generate those with the matching stage and a larger --n-restart/--n-agentic, then re-run report).

Estimate (k5_levers.py estimate; measured rate 522.5 tok/s end-to-end = 1.88 M tokens per box-hour, mean 5,685 tokens/attempt and 36% failure on miniF2F-test from the P2 baseline manifest; ProofNet-valid mean 9,000 tokens and 75% failure are ASSUMED; repair success 12%, repair prompts 1.15x cost, subgoal rounds 0.35 x an attempt per failed parent are ASSUMED):

Profile	restart	correct	agentic	total (generation only)
full: n_restart 32, n_agentic 8 (PREREG pass@8 from n = 32)	50	63	23	~136 box-h
lean (recommended first pass): n_restart 16, n_agentic 6	25	31	17	~74 box-h
minimal: n_restart 8, `--budget restart_k:4`, n_agentic 4	13	16	11	~40 box-h
miniF2F-valid only, lean profile	11	9	6	~26 box-h

These exceed the plan's ~30 box-h for K5 except for the miniF2F-valid-only and minimal profiles; the drivers are repairing EVERY failed root (needed for an unbiased replay) and ProofNet's long attempts. Options to cut cost, each a decision for the plan owner: miniF2F-valid first (26 box-h), then ProofNet only if informative; repair only the first 8 roots; or a stratified dev subset. Verification CPU time (6 s/proof order of magnitude at 8 workers, timeouts dominate) overlaps generation in pipeline_gen_verify.py for the GPU stages; the CPU-only stages (route, asm verify) add roughly 1-2 h of wall time per agentic round at 2-8 Lean workers.

8. Tests

cd /apps/development/rlhf/align-wt/agentic   # (or the merged checkout)
. /apps/development/rlhf/env/apps_only.sh; export PYTHONPATH=$PWD
PY=/apps/development/rlhf/align/.venv/bin/python
$PY scripts/hostguard.py run --cores 4 --mem-gb 16 --label k5-tests --wait -- $PY -m pytest tests/test_k5_levers.py tests/test_subgoal_repair.py -q            # CPU, ~5 s (slow tests deselected with -m "not slow")
$PY scripts/hostguard.py run --cores 4 --mem-gb 16 --label k5-lean --wait -- $PY -m pytest tests/test_subgoal_repair.py -q -m slow     # real Lean, ~45 s