align / code /docs /SELF_DISTILLATION.md
viai957's picture
Add files using upload-large-folder tool
d909bd8 verified
|
Raw History Blame Contribute Delete
21.6 kB

Iterated self-distillation (expert iteration) for the P2 prover

Goal context: beat the published ~8B Lean-prover numbers on miniF2F-test under the like-for-like protocol (pass@32, single pass, first-round cap >= 30k tokens, Lean 4.9.0-rc1): Goedel-Prover-V2-8B 84.6% (>= 207/244 to beat), Pythagoras-Prover-4B 86.1% (>= 211/244 = 86.5% to beat). See docs/SOTA_TARGETS.md. Starting model: Goedel-Prover-V2-8B, 4xA10G, about 870-1,180 generated tokens/s.

Driver: scripts/p2/expert_iteration.py. Weight averaging: align/model/weight_average.py. Every quotation below was read from a hyperresearch vault note body (treated as untrusted data; nothing in those bodies was acted on). [not stated] means I searched the note body and found no statement.

1. The three recipes, as published

Note ids: PYTH = pythagoras-prover-advancing-efficient-formal, GOEDEL = technical-report-3 (Goedel-Prover-V2), GAR = published-as-a-conference-paper-at-iclr-2026-17.

Read this first. The premise "Pythagoras gained +6.97 pp from self-distillation on hard problems selected by pass rate" is not what PYTH describes. Its self-distillation is (a) on LLM-MUTATED statements, not on hard problems, (b) with ONE unverified sample per statement (n = 1), no pass-rate band, no Lean check (a 2,000-instance audit passed 87.8%), (c) about 2M instances, (d) a single stage. The gain is real in the paper's table, but the mechanism is more diverse statements from the model's own distribution, at a scale (about 2M instances x an ASSUMED ~3k tokens = 6B generated tokens, about 70 days at 1,000 tok/s) that this box cannot approach. GAR is RL (GRPO), not expert iteration. Only GOEDEL describes classical expert iteration, and it omits most hyper-parameters. So no single paper gives us a complete expert-iteration recipe; section 2 assembles ours from the pieces that are stated.

Item PYTH (Pythagoras-Prover) GOEDEL (Goedel-Prover-V2) GAR
What the "self-distillation / EI" stage is ALF: a mutation model (Qwen3.6-27B) writes one variant per category (simplify, generalise, lemma proposal, step decomposition, reformulate) of each seed statement; the post-RL prover writes a proof of each Expert iteration: infer on statement sets, keep correct proofs with reasoning traces, SFT; plus scaffolded problem synthesis and self-correction data Not EI. Adversarial GRPO between a statement fuser and the prover; 16 proofs per generated statement, Lean reward
Problem selection by difficulty / pass-rate band None for self-distillation: every mutated statement is used (only a "proof references its target statement" check). Bands exist elsewhere: RL keeps a problem "only when its number of successful rollouts lies in {1, ..., 5}" (of 8); tiers easy/medium/hard by source Scaffolding: Qwen3-32B writes "simpler/sub-problems if it is unsolved, or harder variants if it is already solved". RL: "only include problems with pass rates in the range (0, 0.75]"; appendix: "filter out inputs with a pass rate equal to 0 or higher than 0.75". EI-stage band: [not stated] "we exclude statements with an empirical pass rate of 0 (unsolvable) or above 0.5 (too easy)" (prover update only)
Samples per problem per round n = 1 per mutated statement ("a single candidate proof (n=1) per statement"); T = 1.2, top-p 0.95. Hard-tier seed building: n = 4. RL: 8 rollouts EI: [not stated]. RL: "batch size of 128 and n of 8" n = 16; 1,024 theorems per step
Proofs kept per problem; dedup; shortest rule One per statement. Shortest-proof rule: [not stated]. Dedup applies to easy-tier STATEMENTS only (cosine > 0.75 on gte-Qwen2-7B embeddings). Proof dedup: [not stated] EI: [not stated]. Scaffolded statements: "filter out duplicated statements under exact match"; at most one formalisation kept per informal statement n/a (RL)
Verification of kept data None for ALF; audit of 2,000 random instances: 87.8% compile. "We do ... verify the formal statement alongside the formal proof ... checking only the proof teaches the model to alter the statement to fit a flawed proof." Seed corpus (SFT) is Lean-verified Lean-verified ("collecting correct proofs") Lean reward; statement modification, sorry, admit count as wrong
Self-correction traces in the training data [not stated] for the ALF stage; evaluation is without self-correction Yes: S2 = S1 + self-correction data annotated with SFT-S1 and DeepSeek-Prover-V2-671B (144 H100s). RL is 50% whole-proof, 50% first-round correction No
Train from base or previous round Continues the post-RL model: "a final continued-SFT stage on the ALF self-distillation corpus". Cross-model teacher (32B -> 4B) "unstable", so self only S1: SFT of Qwen3-8B/32B. S2: "perform SFT ... using S2" (from base or from SFT-S1: [not stated]). S3: "training SFT-S2-AVG on S3" = continues from the previous AVERAGED model. RL starts from SFT-S3-AVG Continues: "the prover from the last iteration generates n candidate proofs"
Epochs 2 ("fine-tuned for two epochs") SFT: [not stated]. RL: 1 epoch RL, single pass per step: [not stated]
LR SFT stages 3e-4 and 7e-5; self-distillation 1e-5 ("larger values induced training instability") SFT: [not stated]. RL: [not stated in the body I read] 2e-6 (GRPO), eps 0.2, beta 0.01
LoRA or full LoRA, rank 64, all linear layers (SFT and self-distillation); RL is full-parameter ("LoRA-adapted policy unstable under GRPO") SFT framework LLaMA-Factory; LoRA vs full: [not stated] [not stated]
Rounds 1 self-distillation stage 3 SFT datasets (S1, S2, S3) then RL; "Expert iteration" round count: [not stated] 3 (Goedel-8B), 5 (DeepSeek-Prover-V2-7B); about 140 H100-hours each
Model averaging [not stated] Yes, after every stage: theta = (1 - alpha) theta_0 + alpha theta; "We perform model averaging after completing SFT and use the averaged model as the starting point for RL. Once RL is completed, we perform model averaging again". SFT-S2 averaged with the base, SFT-S3 averaged again. alpha swept in {0.6, 0.7, 0.8, 0.9}; "pass@N first rises and then falls as alpha increases"; chosen alpha: [not stated] [not stated]
What each round gained One stage, total: 4B 79.10 -> 86.07 (+6.97 pp), 32B 84.02 -> 89.75 (+5.73 pp) at pass@32. RL: "only modest gains", not quantified No per-stage numbers in the text; Fig. 8 shows pass@1 rising with RL steps, vanilla pass@N plateauing. Averaging: "improves sample diversity" (plotted, no number in text) 3 iterations total: 77.87 -> 80.33 (+2.46 pp, 16k cap); static GRPO on the same data 77.87 -> 77.46 (-0.41); frozen fuser 77.87 -> 77.87; per-iteration: [not stated]

Verbatim sources (note id, then text):

  • PYTH: "each mutated statement is then populated with a proof by the post-RL Pythagoras-Prover of §2.2, which generates a single candidate proof (n=1) per statement; the 4B and 32B models each prove their own mutations. (iii) Statement-alignment filtering: in place of Lean verification, we retain a (statement, proof) pair only when the generated proof references its target formal statement." Then: "it never invokes the Lean type-checker and does not certify that the proof is valid." And: "we Lean-verify a random subset of 2,000 ALF instances and find that 87.8% pass".
  • PYTH: "All models are based on the Qwen3 series ... fine-tuned for two epochs using LoRA of rank 64 applied to all linear layers. ... We use learning rates of 3e-4 and 7e-5 for the SFT stages. For self-distillation, we reduce the learning rate to 1e-5, as larger values induced training instability."
  • PYTH Table 5: "Pythagoras-Prover-4B (SFT only) 32 79.10 ... (SFT + ALF Self-Distillation) 32 86.07"; text: "gains of +6.97 and +5.73 points respectively".
  • GOEDEL: "We follow the standard pipeline of Expert Iteration by iterating between using the model to conduct large-scale inference on the statement sets, collecting correct proofs with reasoning traces, and fine-tuning the model further by Supervised Fine-tuning (SFT) on the collected samples."
  • GOEDEL: "We perform scaffolded data synthesis with SFT-S2-AVG to generate the dataset S3. Continuing the SFT process, we further improve the models by training SFT-S2-AVG on S3, yielding SFT-S3, and subsequently obtain the averaged model SFT-S3-AVG."
  • GOEDEL: "We observed that in the later stages of SFT and RL, the model's diversity decreases. This is reflected by an increase in pass@1 but a decline in pass@N for larger values of N, such as N=32."
  • GOEDEL notation conflict: section 2.3 defines "(1-alpha) theta_0 + alpha theta" (alpha weights the fine-tuned model), section 3.6 says "alpha = 0.6, 0.7, 0.8, 0.9 (where alpha is the weight of the base model)". weight_average.py therefore names the parameter ft_weight (weight on the fine-tuned model) to avoid the ambiguity.
  • GAR: "we exclude statements with an empirical pass rate of 0 (unsolvable) or above 0.5 (too easy)"; "further RL training on static dataset degrades performance, as the base model is already heavily RL trained"; "Kimina-Prover gains approximately 2% from an SFT starting point over 25 iterations".

What this means for Goedel-8B specifically: the model we start from is already the product of SFT, EI, averaging and RL (GOEDEL). GAR's measured result on exactly this model is: static re-training gives -0.41 pp, the best of five variants +2.46 pp. A realistic prior for our single-pass pass@32 gain from a few small rounds is therefore 0 to +2 pp, not +7 pp; +6.97 pp was measured on a Qwen3-4B SFT checkpoint (79.1%) that had far more headroom than Goedel-8B (84.6%).

2. Our recipe and why (4xA10G, ~1,000 generated tokens/s)

Decision Choice Reason
What data Verified, on-policy proofs of the pool problems the current model solves sometimes but not always We cannot afford PYTH's 2M unverified instances (6B tokens), and GAR/PYTH both warn that a statement-altering or wrong proof teaches the model to cheat; our Lean strict verifier plus statement_not_faithful filter is the cheap exact signal
Band pass rate p in (0, 0.75] from n = 8 samples; p = 0 problems resampled only with --include-zero --max-zero 100; p > 0.75 and 8/8 dropped GOEDEL's RL band. p > 0.75 rows are what the model already writes; p = 0 yields rows only if a later sample succeeds. GAR's tighter 0.5 is for RL variance, not SFT
Samples n = 8 per selected problem, T = 1.0, top-p 0.95 (the P2 sampler; p from n = 8 has a standard error of ~0.17 at p = 0.5, which is why the band is wide) n = 8 matches the evaluation protocol and the existing pipeline; n = 16 (GAR) doubles the cost for a band that is already coarse
Keep at most 4 DISTINCT accepted proofs per problem, shortest completion first, problems solved 8/8 excluded, rows over 18,432 tokens dropped, never truncated build_sft_data.py rules + the new --rank shortest. Distinct: duplicates add weight, not information. Shortest: fewer tokens per row (cheaper SFT, fewer cap hits) but it length-biases the data; this is an RISK for hard problems that need long reasoning, so it is the first ablation if round 1 is flat (--rank order)
Self-correction rows Off by default (--corrections to enable) The bar is no-correction. The correction round costs ~4 h of generation+verification per round (docs/P2.md 6.2) for a skill the metric does not use; GOEDEL notes self-correction data is scarce and mainly lifts the correction mode
Start of each round The previous round's OUTPUT (continue), fresh LoRA each round, merged GOEDEL (S2-AVG -> S3) and GAR continue; PYTH continues the post-RL model. Training from the base each round would discard earlier gains because per-round data cannot be pooled (generation keys collide across rounds)
Optimiser LoRA r 64 / alpha 128, LR 2e-5 (configs/p2_sft.toml), 2 epochs PYTH: LoRA r64, 2 epochs, LR lowered to 1e-5 for self-distillation. The repo config's LR was measured (relative weight change 3e-4 per 12 steps, configs/p2_sft.toml header); use --lr 1e-5 if round 1 shows a regression
Averaging --avg-weight 0.7 --avg-anchor original: out = 0.3 x Goedel-8B + 0.7 x (merged) GOEDEL averages with the base after every SFT to recover pass@N. Anchoring to the ORIGINAL Goedel-8B (not the round's input) bounds drift across rounds (geometric damping). Caveat: for a LoRA delta, averaging is the same as scaling the delta, so 0.7 trades update size for diversity; GOEDEL tuned alpha, we cannot afford a sweep except in round 1 (one extra eval, +3.3 h)
Eval miniF2F-VALID (226), n = 8, cap 30,000 each round (stop rule). The test claim is NOT made by this driver: the frozen chosen model is evaluated once on miniF2F-test (n = 32) per docs/SOTA_PLAN.md Selecting rounds on test makes the test number partly a selection artefact; n = 8 is the cheapest n that supports a paired McNemar
Size of a round --max-problems 300 Section 3 shows cost is linear in the number of selected problems; 300 keeps a round near 8 h. The pool has 2,000 problems and Goedel-8B may have trained on Goedel-Pset-v1, so the in-band set may be small: the pilot of docs/P2.md 6.3 tells

Open items: (1) RESOLVED 2026-10-14: .tmp/p2/problems/minif2f_valid.jsonl (226 problems, scripts/p2/build_valid.py) is the scored file; the driver refuses a test-split scored file and requires the test benchmarks via --protected-problems so they stay in the SFT exclusion set. (2) build_sft_data raises when nothing is selected; the driver pre-empts that with a data stop in measure, but an all-rejected sft_data still exits 3.

3. Compute and ETA per round

Measured inputs (docs/P2.md section 6.1, this box): generation r_gen = 870-1,180 tok/s (central 1,025); mean completion T = 4,996 tokens on miniF2F at cap 16,384 (8.2% capped), pool problems unmeasured (bracket 4,000-6,500); verification t_v = 3.07 s per code-bearing sample at 8 workers, f_code = 0.918; SFT r_tr = 2,152 padded tok/s at 16k micro-batches, efficiency 0.9, a checkpoint save = 3.3 min per 40 micro-steps; export 0.1 h. Averaging was measured here: 85 M elements/s (bf16, 4 threads), i.e. about 1.6 min of arithmetic for 8.19B parameters; the rest is reading 2 x 16.4 GB and writing 16.4 GB.

Formulas (P = selected problems, n = samples per problem):

t_gen    = P * n * T / r_gen
t_ver    = P * n * f_code * t_v                  (CPU, overlapped with t_gen by pipeline_gen_verify.py)
t_gv     = max(t_gen, t_ver) + 0.25 h tail + 0.03 h engine load
rows     = P * rho * k_bar                       (rho ~ 0.95 problems with an accepted proof, k_bar ~ 3 of max 4 kept)
t_sft    = rows * L * E / (r_tr * 0.9) + (rows * L * E / (0.9 * 18,432) / 40) * 3.3 min
t_other  = export 0.10 + average 0.15 (IO bound) + engine reloads 0.05 = 0.30 h
t_eval   = max(244 * n_e * T_e / r_gen, 244 * n_e * f_code * t_v) + 0.28 h
t_round  = t_gv + t_sft + t_other + t_eval(n_e = 8)

Numbers for P = 300, n = 8, L = 2,900 tokens per row, E = 2, eval T_e = 5,700 (30k cap) (computed with the formulas above):

low central high
r_gen (tok/s), T (tok) 1,180, 4,000 1,025, 5,000 870, 6,500
generate + verify (2,400 samples) 2.5 h 3.5 h 5.3 h
SFT (~855 rows, 2 epochs, with saves) 1.0 h 1.1 h 1.3 h
export + average + reloads 0.3 h 0.3 h 0.3 h
eval n = 8 (1,952 samples, cap 30k) 2.6 h 3.3 h 4.3 h
one round 6.4 h 8.2 h 11.2 h
3 rounds 19 h 25 h 34 h
final eval n = 32 (replaces the third n = 8 eval; +extra) 9.5 h 12.3 h 16.5 h
baseline Goedel-8B at n = 32, 30k, our verifier (needed for a like-for-like claim) 9.5 h 12.3 h 16.5 h
whole plan: 3 rounds + n=32 final + n=32 baseline 36 h 46 h 62 h

Scaling: a round at P = 250 / 500 / 1,000 costs about 3.0 / 5.7 / 11.1 h of generate+verify (central) and 0.9 / 1.9 / 3.7 h of SFT; evaluation is fixed. Verification is hidden behind generation whenever the node's generation time per sample, T / r_gen, exceeds the verification time per sample, f_code x t_v = 2.8 s: true for T > ~2,900 tokens at central speed (4.9 s at T = 5,000); for shorter proofs the 8 CPU workers become the bottleneck. The evaluation legs are 40-45% of every round and are the part to cut first (n = 4 costs 1.8 h; the stop rule then compares solved@4).

4. How to run

Split hygiene (2026-10-14): the FIRST --eval-problems file drives the stop rule, so it is miniF2F-valid. The driver refuses a test-split scored file unless --allow-test-selection, which also requires --rounds 1 and no --target-solved (one non-adaptive diagnostic round). Pass the test benchmarks as --protected-problems (decontamination). The test-set claim is made only by the separate pre-registered evaluation of the frozen checkpoint in docs/SOTA_PLAN.md (pipeline_gen_verify.py -> eval_report.py).

cd /apps/development/rlhf/align && D=$PWD/.tmp/p2
.venv/bin/python scripts/p2/expert_iteration.py --run-dir $D/ei --pool $D/problems/train_pool.jsonl \
  --eval-problems $D/problems/minif2f_valid.jsonl \
  --base-model /apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots \
  --rounds 3 --n 8 --max-problems 300 --include-zero --max-zero 100 \
  --avg-weight 0.7 --eval-n 8 --eval-cap 30000 \
  --baseline-run base:<gen.jsonl>:<ver.jsonl> --dry-run          # prints every command; drop --dry-run to run
  • Stages per round: select, gen_verify, [gen_verify_corr], measure, sft_data, train, export, [average], eval, [report], manifest. Markers are in RUN/round_NN/stages/; the chain hash makes a changed argument an error, not a silent mix (--redo ROUND:STAGE).
  • The driver never takes the heavy lock; each child (pipeline_gen_verify.py, sft.py, qwen3_export_check.py) does, one at a time (/apps/CLAUDE.md rule 2). --seed-ver takes an existing verification file (for example the P2 pool round 0) so round 1 does not re-measure.
  • Disk: each export and each averaged model is 16.4 GB; --prune-export deletes the export once the average exists.
  • qwen3_export_check.py gained P2_EXPORT_BASE_SNAPSHOT (default unchanged): in round 2 the adapter must be merged onto round 1's output, not the original snapshot. build_sft_data.py gained --rank {order,shortest} (default order, unchanged behaviour).
  • Exit codes: 0 done or stopped by the rule, 3 stage failed or config mismatch, 4 data stop, 143 signal.

5. Pre-registered stopping rule

Registered before any iteration data exists. Implemented as stop_decision() in the driver and evaluated after every round from round_NN/round_manifest.json; the parameters are command-line defaults.

Amended 2026-10-14 (before any round ran): the scored set is miniF2F-VALID (226), not test; criterion 1 (target) is dropped because a target on the selection set is not a claim. Thresholds are unchanged in problems.

Quantity: S_r = number of miniF2F-valid problems (of 226) solved by at least one of n_e = 8 samples at cap 30,000 for the round-r model, same problems, same eval seed, same verifier (1 problem = 0.44 pp). Only rounds with the same (n_e, cap) are compared.

Stop after round r if ANY of:

  1. target: dropped (see the amendment above); never pass --target-solved with a valid scored set.
  2. regression: S_r <= max(S_1..S_{r-1}) - 3. Keep the best earlier round's model (best_round in the manifest).
  3. plateau: two consecutive rounds each gain fewer than 2 problems (0.8 pp) over the best round before them. (Round 1 has nothing before it and is not judged.)
  4. budget: 3 rounds done (--rounds).
  5. data: no problem in the band, no problem that has both an accepted proof and a failure (measure), or fewer than 50 SFT rows.

Why these numbers: the paired comparison on ~226-244 problems has power of roughly 0.3 for +4 pp and 0.6-0.8 for +6 pp (docs/P2.md section 1, d = 0.10-0.15), so a single round cannot be confirmed at the effect sizes GAR reports (+2.46 pp over three iterations). The rule therefore does not test, it guards compute. Its thresholds are judgment, not derived: with about 20 discordant problems the standard error of the net difference is sqrt(20) = ~4.5 problems, so neither "+2" nor "-3" is statistically distinguishable from noise; they are chosen so that a round that gains under 2 problems twice in a row, or loses over 3, is not worth another 8 hours. The confirmatory claim is the single pre-registered n = 32 test run of the frozen model against the n = 32 anchor with exact McNemar (eval_report.py), not the stop rule.

Honest limits: this rule selects on miniF2F-valid, so the chosen round's valid score is optimistic by about the best-of-rounds bias (order 1 SE); the test run is untouched by it. GAR's evidence says the expected total gain is about +0 to +2.5 pp, which is 0-6 problems, so a null outcome after two rounds is the likely one and is a valid result.

6. Tests and checks

tests/test_expert_iteration.py (47 tests: band selection, keep rules through the real build_sft_data, stop rule, resume after a crash at each of the 10 stages in two failure modes, config-change refusal, multi-round chain, data stops, dry-run command lines), tests/test_weight_average.py (26 tests: bit-exact against an fp64 reference for bf16/fp16/fp32 at several weights and chunk sizes, shard planning, mismatch refusal, atomic output, idempotence). scripts/p2/mutation_check_ei.py runs 20 hand-written mutants (all killed).