Keyboard shortcuts

Press ← or → to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

Distilling LLM Feedback for Lean Theorem Proving

(authored by agents unless marked 🧑) Tag: narozniak2026distilling, paper

scope and claimed contribution

The paper proposes Feedback Distillation for Lean 4 theorem proving. Its abstract says: “we propose Feedback Distillation, a training method where the model is trained to match, at the token level, its own distribution conditioned on privileged feedback produced by a language model” (Abstract, p. 1). The authors report that the method preserves diversity and that combining it with GRPO is strongest: “After training on LeanWorkbook, Qwen3.5-9B achieves 59% pass@1 test accuracy on a test split of LeanWorkbook with GRPO alone compared to 75% with Feedback Distillation+GRPO” (Section 1, p. 2).

Inference: for VL, the reusable idea is a verifier-in-the-loop proof-authoring trainer: sample a candidate Verus proof, obtain textual critique from a stronger model using verifier diagnostics, and distill that critique into the prover while retaining a hard Verus check as the acceptance gate.

technique 1: feedback-conditioned self-distillation

What it is and when to use it. The student generates an on-policy proof attempt; a teacher receives the original problem, attempt, and feedback, then supplies a token distribution. The method is useful when trajectory-level rewards are sparse or all sampled proofs fail. The paper motivates this because “GRPO provides only trajectory-level rewards, which offer no learning signal when all samples fail” (Section 1, p. 2).

Steps from the paper:

  1. Sample a problem x from an unsolved-problem dataset and an attempted proof y ∌ πΞ(·|x) (Section 3, p. 3).
  2. Compute textual feedback F(y) with a frozen copy of the model or a third-party LLM that sees the formal statement, all tool calls and outputs, and the final compiler output (Section 3, p. 3; Section 4.2, p. 5).
  3. Run the teacher as πΌ(· | x, F(y)); update its weights by EMA every five gradient steps using ÎŒ ← α ÎŒ + (1 − α) Ξ (Section 3, p. 3).
  4. Minimize the per-token KL divergence from teacher to student along the sampled trajectory; stop gradients through sampling y (Section 3, pp. 3–4).
  5. In practice use the equivalent cross-entropy and retain only the teacher’s top K=25 vocabulary tokens (Section 3, p. 4).
  6. Optionally continue with GRPO from a Feedback Distillation checkpoint (Section 4.3, p. 5).

The paper gives the exact loss shape: “Our loss is the per-token Kullback-Leibler divergence between the teacher and student distributions along trajectories sampled from the student” (Section 3, p. 4). It also states the practical truncation: “To reduce computational cost, we truncate the sum over the vocabulary to the K tokens with the highest probability under the teacher π”,yâ€Č , using K = 25 in practice” (Section 3, p. 4).

Verus form. Inference: represent x as a Rust/Verus function plus its fixed requires/ensures contract and current proof state. Generate a candidate proof block, invoke Verus, and give the feedback model the exact diagnostics and source context. Distill advice into a prover that emits proof fn, assert, lemma calls, contract clauses, and loop invariant declarations. The Verus verifier remains the final checker; a language-model score cannot replace it. Inference: preserve the paper’s on-policy property by sampling from the current Verus prover, while storing verifier failures as training context rather than silently converting them into accepted specifications.

Evidence. The Lean environment used three tools: “lean_write_file: write a Lean source file; lean_check_file: compile a previously written file and return the compiler output (errors, warnings, or success); rg_in_mathlib: search Mathlib using regular expressions to find relevant lemmas” (Section 4.1, p. 5). Training used 10,000 LeanWorkbook statements; evaluation used 256 LeanWorkbook test statements and 244 MiniF2F test statements, with up to 8,192 tokens per attempt on 8 NVIDIA H200 GPUs (Section 4.2, p. 5). The paper reports 2% and 11% MiniF2F initialization scores for Qwen3-8B and Qwen3.5-9B respectively (Section 4.1, p. 5).

Limits and assumptions. Feedback is generated by Claude Opus 4.6, so external-model availability and quality are assumptions (Section 4.2, p. 5). The paper reports instability and says the method does not yet match scaled SFT+GRPO pipelines: “Training stability remains a challenge for Feedback Distillation” (Conclusion, p. 9). Inference: a VL deployment must treat compiler diagnostics, imported specifications, and the target contract as part of the trusted input; feedback can be wrong, incomplete, or overfit to a failed attempt.

technique 2: verifier tools and actionable critique

What it is. The feedback model is constrained to produce short, self-contained orders rather than a replacement proof. The appendix prompt requires “between 1 and 3 numbered pieces of feedback about errors, things that were useful, or things that were not” and says each piece “must be between 10 and 50 words” (Appendix A.3, p. 13). It also requires that feedback “must encourage the usage of tools” (Appendix A.3, p. 13).

Steps: provide the statement, attempt, tool transcript, and final compiler output; ask for 1–3 orders; prohibit second-person references and extra prose; feed the resulting text as privileged input to the teacher (Appendix A.3, p. 13). Inference: for Verus, ask for orders such as an order to add a loop invariant relating ghost sequence length to an executable index, or an order to prove a missing arithmetic fact with a proof function then verify the revised Rust source.

Evidence and limits. The authors compare Claude feedback, ground-truth proofs, compiler-output feedback, and self-feedback. They state: “Claude and ground-truth as feedback sources outperform Lean output” (Section 5.1, p. 7), while self-feedback plus GRPO beats GRPO alone but remains below compiler feedback (Section 5.1, pp. 7–8). Inference: textual critique is most useful when it points to a strategy or API fact that a compiler error alone does not explain; it should never be accepted as evidence of correctness.

technique 3: EMA stability and diversity diagnostics

Use EMA to avoid a stale teacher penalizing knowledge learned from earlier feedback. The paper explains: “Since the model learns from multiple feedback signals over the course of training, a frozen teacher would be problematic” (Section 5.2, p. 8). Sweep α; lower values learn faster but are less stable, α=1 freezes the teacher, and the reported best tradeoff is α=0.9 (Section 5.2, p. 8; Figure 6, p. 9).

Measure policy entropy, pass@k, and the cumulative number of distinct tactics and lemmas used in successful proofs. The authors report that Feedback Distillation has higher entropy and better pass@k scaling at matched pass@1, and discovers distinct Mathlib lemmas faster (Section 4.3, pp. 5–6; Figure 3, p. 6). Inference: for VL, track verifier-accepted proof diversity, distinct lemmas, invariant patterns, and pass@k under a fixed Verus version and dependency snapshot. These are exploration metrics, not soundness claims.

results, trusted parts, and product threat

The central measured comparison is 59% versus 75% LeanWorkbook pass@1 for GRPO versus Feedback Distillation+GRPO (Section 1, p. 2). The method uses single-sample Feedback Distillation batches of 128 and GRPO effective batches of 640; both have 8,192-token response limits, while Feedback Distillation uses top-K 25 and EMA 0.9 (Appendix A.1, p. 13). The authors also report 300 seconds per Feedback Distillation step versus 550 seconds per GRPO step on their 8×H200 node (Section 4.2, p. 5).

Inference: the proof of each Lean result is trusted through Lean’s checker and Kimina Lean Server, but the paper’s training claims still rely on model weights, prompt formatting, optimizer settings, benchmark splits, and external Claude feedback. The paper does not measure Verus, Rust ownership, concurrency, or specification adequacy. Inference: threat to a product that emits agent-written verified Rust is moderate for search cost and proof quality, and low for final kernel soundness if every candidate is rechecked. The larger risk is a persuasive critique that causes a prover to add an incorrect or overly weak contract; VL must freeze contracts and reject edits that bypass verification.

The paper does not establish that Feedback Distillation solves long-horizon instability, transfers unchanged from Lean to Verus, or improves arbitrary systems code. Existing VL work should not redo the generic claim that verifier feedback is valuable; it should test the concrete transfer: compiler diagnostics plus structured critique for Verus contracts, loop invariants, ownership, and concurrency APIs.

techniques for the catalog

  • feedback-conditioned self-distillation — sample on-policy proof attempts, obtain privileged critique, and train token-wise toward the feedback-conditioned teacher

  • verifier-tool feedback — expose file writing, checker output, and library search, then constrain a critic to short actionable orders

  • EMA teacher stabilization — update the teacher by exponential moving average and sweep α while tracking instability

  • diversity diagnostics — compare entropy, pass@k, and distinct successful lemmas/tactics at matched verifier accuracy Checked 2026-10-15 against paper narozniak2026distilling: full comparison completed; metadata and unsupported claims reviewed; experiments not reproduced

Last edited: