Source-linked AI summary
Distilling LLM Feedback for Lean Theorem Proving
Gaetan Narozniak, Gérard Biau, Rémi Munos, Ahmad Rammal, Pierre Marion
TL;DR
Post-training with SFT and GRPO faces sparse-reward, exploration, and mode-collapse limitations. The paper introduces Feedback Distillation, which matches a student to its own feedback-conditioned distribution at the token level, and finds greater diversity and stronger pass@k scaling in Lean theorem proving; combining it with GRPO achieves 75% versus 59% pass@1 for GRPO alone.
Problem
SFT+GRPO is limited by costly data generation, off-policy training, sparse rewards, and insufficient exploration or answer diversity.
Method
Feedback Distillation trains a student to match token-level distributions from a teacher conditioned on feedback about the student’s attempt and verifier output.
Results
Feedback Distillation preserves higher policy entropy and improves pass@k scaling; on LeanWorkbook, Qwen3.5-9B reaches 75% pass@1 with Feedback Distillation+GRPO versus 59% with GRPO alone.
Takeaways & Limitations
Feedback Distillation and GRPO provide complementary benefits, with Feedback Distillation offering a promising mechanism for injecting external knowledge while maintaining diversity.
Takeaways & Limitations
The current results do not yet match state-of-the-art SFT+GRPO pipelines trained at scale, and training stability remains a challenge.
Abstract
from arXiv · showhide
Post-training for reasoning models typically combines supervised fine-tuning with reinforcement learning from verifiable rewards, most commonly with GRPO. However, this algorithm suffers from sparse rewards, limited exploration, and mode collapse. Building upon recent works on self-distillation, 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. Feedback Distillation offers token-level supervision and can inject external knowledge. Evaluating our method for Lean4 theorem-proving, we find that Feedback Distillation maintains greater diversity in generated trajectories than GRPO, yielding higher policy entropy and better pass@k scaling. The two methods are complementary: initializing GRPO from a Feedback Distillation checkpoint outperforms either method alone. All in all, our results suggest a promising avenue to improve post-training for complex reasoning.
1 Introduction
The paper motivates Feedback Distillation as an alternative to SFT+GRPO limitations and evaluates it for Lean 4 theorem proving. It reports higher diversity and stronger pass@k scaling, with further gains when GRPO is initialized from Feedback Distillation.
- Motivation: GRPO can collapse onto already high-reward modes rather than increasing answer diversity or exploration.This is especially problematic for formal mathematics, where exploring large lemma libraries is important.
- Motivation: Feedback Distillation uses privileged feedback to provide token-level credit assignment and can operate beyond directly verifiable reward settings.The method distills a prompting strategy into the model’s weights while optionally leveraging verifier outputs.
- Contribution: Feedback Distillation conditions feedback on the model’s attempt and verifier output, using either a frozen model copy or a third-party LLM.This distinguishes it from self-distillation approaches using raw environment signals or ground-truth solutions.
- Results: In Lean 4, Feedback Distillation maintains higher policy entropy and achieves better pass@k scaling than GRPO, indicating greater output diversity.The evaluation is conducted in the Lean 4 formal mathematics environment.
- Results: 59% pass@1 test accuracy with GRPO alone increased to 75% with Feedback Distillation+GRPO on LeanWorkbook.The result uses Qwen3.5-9B and a LeanWorkbook test split.
- Analysis: The paper also analyzes feedback sources, training instability, and qualitative consequences of Feedback Distillation as on-policy knowledge transfer.These analyses are presented as part of the method’s broader evaluation.
2 Related work
Related work places the paper at the intersection of self-distillation for LLMs and reinforcement learning for formal mathematics. Existing self-distillation methods use privileged environment outputs or ground-truth solutions, while formal provers commonly combine synthetic-proof SFT with GRPO.
- Self-distillation for LLMs: Self-distillation trains a student to match a same-architecture teacher conditioned on privileged information.Prior work uses environment outputs or ground-truth solutions as the privileged signal.
- Self-distillation for LLMs: Feedback Distillation extends self-distillation by using feedback from a frozen model copy or a third-party LLM.The paper also analyzes properties of self-distillation in this setting.
- RL for formal mathematics: State-of-the-art formal mathematics provers commonly use supervised fine-tuning on synthetic proof data followed by GRPO-based reinforcement learning with verifiable rewards.The related-work discussion references progress on MiniF2F and PutnamBench.
3 Feedback Distillation
Feedback Distillation trains a student on attempted proofs while a teacher receives textual feedback about those attempts. The method uses a token-level divergence objective, an EMA teacher, and optional top-K vocabulary truncation.
- Setup: The method post-trains on problems without solutions, such as Lean statements without proofs, using student-sampled attempted proofs.Problems x are sampled from dataset D and attempts y from the student policy.
- Teacher and feedback: The teacher is an augmented student conditioned on feedback F(y), where F maps an attempted response to textual feedback.The feedback model may be a stronger third-party LLM or a frozen copy of the model.
- Teacher updates: The teacher weights are updated by exponential moving average every five gradient steps.The update is µ ← αµ + (1 − α)θ, and α affects stability and performance.
- Training objective: Feedback Distillation uses a per-token KL divergence between teacher and student distributions along trajectories sampled from the student.The objective compares distributions at each position in the sampled sequence.
- Training objective: Because the teacher distribution is independent of θ, the KL divergence and cross-entropy have the same gradient with respect to θ.The paper therefore uses the equivalent cross-entropy loss in practice.
- Implementation: Top-K truncation reduces computational cost by restricting the vocabulary sum to the K teacher-high-probability tokens.The implementation uses K = 25.
4 Capabilities of Feedback Distillation for Lean Theorem Proving
In Lean 4 theorem proving, Feedback Distillation learns efficiently while preserving trajectory diversity, and its combination with GRPO achieves the strongest performance. It also supports faster early sample-efficient learning and broader tactic and lemma discovery, though GRPO-based training remains unstable.
- Experimental setting: Lean 4 provides a challenging testbed with vast tactic and lemma search, verifiable proofs, and low initial MiniF2F accuracy.The setup uses Qwen3.5-9B and Qwen3-8B, with initial MiniF2F scores of 11% and 2%, respectively, and Mathlib contains over 300,000 declarations.
- Experimental setting: Feedback Distillation equips the model with Lean tools for writing files, compiling proofs, and searching Mathlib lemmas.A problem is solved when the generated Lean file compiles and proves the statement.
- Performance and diversity: GRPO initialized from Feedback Distillation checkpoints outperforms GRPO alone and either method alone, indicating complementary benefits.The comparison uses Feedback Distillation checkpoints after 200, 300, and 400 training steps.
- Performance and diversity: Feedback Distillation maintains higher policy entropy than GRPO and scales better with pass@k when both methods have similar pass@1 accuracy.The higher entropy is reported to reflect answer diversity rather than uniform noise.
- Learning efficiency: Feedback Distillation discovers distinct Lean tactics and Mathlib lemmas faster and is more sample-efficient than GRPO during early training.For Qwen3.5-9B, GRPO eventually catches up before crashing; the authors defer a more thorough sample-efficiency assessment to future work.
- Training stability: GRPO and GRPO over Feedback Distillation exhibit training instabilities, although the combined method remains superior to GRPO alone across tested model sizes.The reported advantage is not explained solely by GRPO instability, and GRPO shows fewer instabilities with Qwen3-8B.
5 Analysis of Feedback Distillation
The analysis compares feedback sources and examines how Feedback Distillation behaves under self-feedback, EMA updates, fixed feedback, and critique-based feedback. Results show complementary gains over GRPO, internalization of feedback, and trade-offs between learning speed, stability, and output behavior.
- Other sources of feedback: Claude and ground-truth feedback outperform Lean compiler output, while Claude feedback induces higher entropy than fully specified ground-truth solutions.The authors attribute the entropy difference to hints and corrections rather than complete answers.
- Other sources of feedback: Combining GRPO with Self-Feedback Distillation achieves higher peak performance than GRPO alone, but remains below performance with Lean compilation feedback.Self-Feedback Distillation uses the same model with frozen weights, isolating gains from the model’s own analysis.
- EMA and training instability: EMA teacher updates let the teacher track the student, preventing earlier learned feedback from being penalized as training incorporates new signals.More aggressive EMA updates learn faster initially but create greater instability and limit peak performance.
- Qualitative comments: Fixed feedback is internalized into the student’s weights: “Think” increases response length, “Be concise” decreases it, and student-teacher KL divergence vanishes.After training, the student behaves as if the feedback were appended to every prompt.
- Qualitative comments: More structured feedback prompts improve stability and final performance but slow early learning.Critique-form feedback also suppresses epistemic verbalization less often than richer privileged information such as ground-truth solutions.
- Qualitative comments: Feedback Distillation provides on-policy knowledge transfer through natural-language feedback, avoiding the need to share the feedback model’s tokenizer or host its logits.When the feedback model is stronger, it can transfer information without supplying a complete solution.
6 Conclusion
The conclusion presents Feedback Distillation as a promising mechanism for improving post-training despite not yet matching scaled SFT+GRPO pipelines. It identifies training stability and broader task generalization as important directions for future work.
- Conclusion: Feedback Distillation does not yet match state-of-the-art SFT+GRPO pipelines trained at scale.The authors nevertheless view it as addressing limitations through external knowledge, maintained diversity, and fine-grained credit assignment.
- Future work: Training stability remains a challenge, motivating more robust stabilization mechanisms for longer training runs.The conclusion also identifies extension beyond Lean theorem proving to other formal and informal reasoning tasks as future work.
A.2 Tool use ablation
The tool-use ablation compares Feedback Distillation and GRPO with and without Lean-tool access. Tool use improves performance for both methods, with higher MiniF2F and training accuracy reported in each case.
- Tool use ablation: Tool use improves MiniF2F accuracy and training accuracy for both Feedback Distillation and GRPO with Qwen3-8B.The authors interpret iterative interaction with the Lean compiler as a driver of performance.
- Tool use ablation: The feedback model receives the formal statement, the model’s attempt including tool calls and outputs, and the final Lean compiler output.The feedback prompt is designed to generate advice from the complete tool-interaction trajectory.
B Loss derivation
The loss derivation uses stop-gradient sampling so optimization differentiates only through the student distribution, making the KL gradient equivalent to a cross-entropy gradient. Additional learning-rate checks show similar GRPO peak performance.
- Loss derivation: Stop-gradient sampling makes the gradient depend only on the student distribution inside the KL divergence.The sampled trajectory y is treated as fixed during differentiation.
- Loss derivation: The KL divergence and cross-entropy have the same gradient when the teacher distribution does not depend on θ.This equivalence motivates the gradient-equivalent loss used in practice.
- Learning-rate robustness: All three tested GRPO learning rates lead to similar peak performance on LeanWorkbook and MiniF2F.The evaluated rates are around 10^-7, including 3 · 10^-7 and 3 · 10^-8.
C.2 Results with Qwen3-8B
For Qwen3-8B, the reported conclusions match those for Qwen3.5-9B: GRPO initialized from Feedback Distillation outperforms GRPO alone, while training behavior depends on EMA and feedback-prompt design.
- Results with Qwen3-8B: GRPO initialized from Feedback Distillation outperforms GRPO alone for Qwen3-8B.The comparison uses Claude Opus 4.6 as the Feedback Distillation feedback model.
- EMA and training instability: Lower EMA interpolation values accelerate early learning but increase instability and limit peak performance.α = 1 corresponds to freezing the teacher weights.
- Feedback formulation: The study evaluates four Claude Opus 4.6 prompt styles with increasing structure and guidance while holding other hyperparameters fixed.The styles progress from a minimal request through formatting, self-contained feedback, and tool-use or truncation guidance.
- Feedback formulation: More structured feedback prompts produce slower initial gains but more stable training and higher peak performance.Prompts 3 and 4, which encourage generic self-contained feedback, appear especially beneficial.
C.4 Epistemic verbalization and response length
Critique-based feedback preserves epistemic verbalization and is associated with longer responses than ground-truth feedback, consistent with greater diversity in generated trajectories.
- Epistemic verbalization: Ground-truth feedback sharply reduces counts of representative epistemic words early in training.The measured categories include “check,” “wait,” and a combined epistemic total.
- Epistemic verbalization: Language-model critique feedback preserves the student’s epistemic verbalization, whether generated by the student itself or a stronger model.The preservation is especially evident when Claude provides the feedback.
- Response length: Critique-based feedback is associated with longer responses than ground-truth feedback.The paper connects this pattern with maintaining diversity rather than converging to short, confident proofs.
D Feedback examples
The feedback examples provide concrete Lean 4 proof strategies, emphasizing case splits, auxiliary lemmas, solver guidance, divisibility handling, and iterative compilation checks.
- Feedback examples: Absolute-value inequalities are approached by splitting sign cases and simplifying with abs_of_nonneg and abs_of_nonpos before linarith.The examples use le_or_lt to reduce absolute values to linear expressions.
- Feedback examples: The examples recommend writing and checking Lean files iteratively to catch syntax and tactic failures early.Suggested tools include lean_write_file and lean_check_file.
- Feedback examples: Natural-number division examples avoid field_simp, establish divisibility or use Nat.div_mul_cancel, and then search bounded values with omega or interval_cases.The suggested workflow first derives a bound such as n ≤175.
- Feedback examples: Nonlinear arithmetic examples supply auxiliary square-nonnegativity and positivity lemmas to nlinarith and use linarith or ring for algebraic transformations.The feedback also recommends avoiding division over reals and deriving intermediate polynomial identities.