Distilling LLM Feedback for Lean Theorem Proving

arXiv:2605.30861 · cs.AI · Submitted 2026-05-29 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: We're looking at the paper: "Distilling LLM Feedback for Lean Theorem Proving".

Jane: The gist: Feedback Distillation, a training method where models are trained to match their own distribution conditioned on privileged feedback produced by another language model,

Tom: First, who's behind it and why it matters.

Paper summary: Tom: So, to recap, this paper by Narozniak and colleagues is proposing Feedback Distillation as an alternative or addition to standard GRPO for post-training reasoning models.

Jane: The central thesis is that instead of relying solely on sparse rewards in reinforcement learning, we can train the model to mimic its own behavior when guided by token-level feedback from a language model.

Lu: They define this teacher as matching the student's distribution conditioned on privileged feedback, essentially injecting external knowledge directly into the training loop.

Meng: Why is this specific setup interesting for theorem proving? What makes it different from just standard supervised fine-tuning?

Tom: The difference lies in the token-level supervision and the ability to inject that external knowledge during distillation. They claim this method helps maintain greater diversity in generated trajectories compared to GRPO.

Jane: That means higher policy entropy, which is a good sign for complex reasoning tasks because it suggests the model isn't just converging on one narrow answer.

Lu: They also found that this approach can be initialized from a checkpoint created using Feedback Distillation, and that starting GRPO from there beats GRPO alone in terms of performance.

Meng: So, the method itself isn't just about getting a slightly better reward signal; it's changing how we structure the training objective for reasoning.

Tom: The evaluation was done in the Lean four formal mathematics environment, and they specifically looked at metrics like pass@k scaling to show that diversity is maintained <ref:2605.30861#pg1>.

Jane: They also spent time comparing different types of feedback, like ground truth versus Lean output, finding that certain sources are more effective for guiding the model's learning process.

Lu: The paper also discussed the training instability and the qualitative consequences of this methodology, showing that it's more than just a simple loss function adjustment.

Meng: So what’s the main thing we need to remember about Distilling LLM Feedback for Lean Theorem Proving? It seems to be about creating a better, more diverse starting point for our reasoning models.

Tom: Exactly. The implication is that we can leverage the self-reflective capabilities of large language models by training them in grounded environments like formal math.

Conclusion: Jane: We’ve covered how Feedback Distillation works, its comparison to GRPO, and the specific results they got on Lean theorem proving. Now let's think about what this means in a broader sense.

Tom: The authors conclude that this method addresses limitations of existing approaches by injecting external knowledge and maintaining diversity during post-training.

Lu: They suggest that current large language models already possess enough reasoning capabilities to analyze their own successes and failures within grounded environments like formal mathematics.

Meng: So, they see this as a way to internalize that self-reflective ability into the model's weights through this distillation process.

Jane: It opens up a path for extending this approach to other types of formal or even informal reasoning tasks outside of Lean four theorem proving <ref:2605.30861#pg1>.

Tom: The key finding is that Feedback Distillation provides a way to improve post-training by giving the model better credit assignment and external guidance during its refinement phase.

Lu: It seems like a promising direction for making AI models more robust when they have to tackle intricate problems that require multi-step reasoning.

Meta AI Research Institute of Artificial Intelligence and Research Institute for Advanced Studies Sorbonne University École des Ponts ParisTech Inria

cs.AI

Submitted: 2026-05-29

Updated: 2026-10-19

License: http://creativecommons.org/licenses/by-nc-nd/4.0/

Importance score: 83/100

The gist: The gist: Feedback Distillation, a training method where models are trained to match their own distribution conditioned on privileged feedback produced by another language model, maintains greater

Key concepts

Feedback Distillation: A training method where a model learns to reproduce its own token distribution, guided by special feedback from a teacher language model. The model is trained to match the distribution of responses it generates when conditioned on this external, privileged advice, allowing it to incorporate new knowledge effectively.
Policy Entropy: A measure of the diversity or randomness in the actions (tokens) a model takes during reasoning. Higher policy entropy indicates that the model explores a wider variety of possible proof paths and tactics, which is desirable for finding more diverse and robust solutions in formal reasoning tasks.
GRPO vs. Feedback Distillation: GRPO is a standard reinforcement learning approach often used for post-training models, but it suffers from sparse rewards and mode collapse. Feedback Distillation outperforms GRPO alone by maintaining greater diversity in generated proof trajectories, leading to better scaling on pass@k metrics in Lean theorem proving.
Self-Feedback Distillation: A variation where the feedback signal comes from a frozen copy of the model itself. This allows the model to learn from its own successes and failures within a grounded environment like formal mathematics, helping it internalize self-correction mechanisms.

Terminology

Distillation — 5× in this paper · explained in 2 episodes
GRPO — 3× in this paper · explained in 6 episodes
Fine-Tuning — 1× in this paper · explained in 5 episodes
Mode Collapse — 1× in this paper · explained in 2 episodes
Reinforcement Learning — 1× in this paper · explained in 4 episodes
Self-Distillation — 1× in this paper · explained in 7 episodes
Verifiable Rewards — 1× in this paper · explained in 2 episodes

Summary

The gist: Feedback Distillation, a training method where models are trained to match their own distribution conditioned on privileged feedback produced by another language model, maintains greater diversity in generated trajectories than GRPO and can be initialized from a Feedback Distillation checkpoint to outperform GRPO alone in Lean theorem proving >

How it works

Feedback Distillation is 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 > This involves defining a teacher as π′µ,y(· x) = πµ(· x, F(y)), where F is a function that takes an attempted response y and returns a textual feedback signal (the privileged information) > The training loss is defined as the per-token Kullback-Leibler divergence between the teacher and student distributions along trajectories sampled from the student > The equivalent loss used in practice is Lˆθ = −X y t=1 X a∈V πµ(a x, F(y), y

Key Findings and Comparisons

The paper evaluates Feedback Distillation in the Lean 4 formal mathematics environment and finds that it maintains greater diversity in generated trajectories than GRPO, yielding higher policy entropy and better pass@k scaling > Furthermore, the two methods are complementary, as initializing GRPO from a Feedback Distillation checkpoint outperforms either method alone > For instance, after training on LeanWorkbook, Qwen3.5-9B achieves 75% pass@1 test accuracy with Feedback Distillation+GRPO compared to 59% with GRPO alone >

Analysis of Feedback Sources and Stability

The researchers investigate different sources of feedback, including ground truth feedback, Lean output, and LLM-generated advice from a frozen copy of the model (Self-Feedback Distillation) > They observe that Claude and ground-truth as feedback sources outperform Lean output > The use of Exponential Moving Average (EMA) updates for the teacher weights is critical, as it addresses the issue where a frozen teacher would penalize knowledge acquired from earlier feedback > The EMA interpolation parameter α controls the tradeoff between adaptivity and stability, with α = 0.9 providing the best tradeoff in their experiments >

Qualitative Observations on Feedback Formulation

The formulation of the prompt sent to the feedback model significantly impacts training performance, with more structured prompts yielding slower initial performance gains but resulting in more stable training and higher peak performance > Prompts that encourage generic, self-contained feedback (Prompts 3 and 4) appear to be especially beneficial for preserving epistemic verbalization, which is measured by the occurrence counts of uncertainty words like check and wait in training runs > Critique-based feedback is also associated with longer responses than ground-truth feedback, reflecting preserved diversity in generated trajectories rather than converging to short, confident proofs >

Ablation and Tool Use

Tool use is confirmed to improve performance for both Feedback Distillation and GRPO when training on Qwen3-8B, indicating that iterative interaction with the Lean compiler is a driver of performance > The results also show that GRPO over Feedback Distillation outperforms GRPO alone across different checkpoints, suggesting the two methods provide complementary benefits > This combination leads to faster and more sustained increase in the number of distinct tactics and lemmas invoked compared to GRPO > The work suggests a promising avenue to improve post-training for complex reasoning >

Conclusion

Feedback Distillation is viewed as a promising mechanism that addresses fundamental limitations of existing approaches by injecting external knowledge, maintaining diversity, and providing fine-grained credit assignment > The authors conclude that current LLMs appear to have sufficient reasoning capabilities to analyze their own successes and failures in grounded environments like formal mathematics, and Feedback Distillation offers a way to internalize this self-reflective ability into the model’s weights > They suggest that extending this approach to other formal and informal reasoning tasks could further demonstrate the generality of the method >

References

Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos Garea, Matthieu Geist, and Olivier Bachem. On-policy distillation of language models: Learning from selfgenerated mistakes. In The Twelfth International Conference on Learning Representations, 2024 >

Anthropic. Introducing Claude opus 4.6, February 2026 >

Charles Arnal, Gaetan Narozniak, Vivien Cabannes, Yunhao Tang, Julia Kempe, and Rémi Munos. Asymmetric reinforce for off-policy reinforcement learning: Balancing positive and negative rewards. In Advances in Neural Information Processing Systems, 2025 >

Mahmoud Assran, Quentin Duval, Ishan Misra, Piotr Bojanowski, Pascal Vincent, Michael Rabbat, Yann LeCun, and Nicolas Ballas. Self-supervised learning from images with a joint-embedding predictive architecture. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 15619–15629, 2023 >

Maciej Besta, Nils Blach, Ales Kubicek, Robert Gerstenberger, Michal Podstawski, Lukas Gianinazzi, Joanna Gajda, Tomasz Lehmann, Hubert Niewiadomski, Piotr Nyczyk, et al. Graph of thoughts: Solving elaborate problems with large language models. Proceedings of the AAAI Conference on Artificial Intelligence, 38(16):17682–17690, 2024 >

Xingyu Dang, Christina Baek, J Zico Kolter, and Aditi Raghunathan. Assessing diversity collapse in reasoning. In Scaling Self-Improving Foundation Models without Human Supervision, 2025 >

Leonardo de Moura and Sebastian Ullrich. The Lean 4 theorem prover and programming language. In Automated Deduction – CADE 28: 28th International Conference on Automated Deduction, Virtual Event, July 12–15, 2021, Proceedings, page 625–635, Berlin, 2021 >

Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre H Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent: A new approach to self-supervised learning. In Advances in Neural Information Processing Systems, volume 33, pages 21271–21284. Curran Associates, Inc., 2020 >

Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 9729–9738, 2020 >

Thomas Hubert, Rishi Mehta, Laurent Sartran, Miklós Z Horváth, Goran Žužic, Eric Wieser, Aja Huang, Julian Schrittwieser, Yannick Schroecker, Hussain Masoom, et al. Olympiad-level formal mathematical reasoning with reinforcement learning. Nature, pages 1–3, 2025 >

Jonas Hübotter, Frederike Lübeck, Lejs Behric, Anton Baumann, Marco Bagatella, Daniel Marta, Ido Hakimi, Idan Shenfeld, Thomas Kleine Buening, Carlos Guestrin, and Andreas Krause. Reinforcement learning via self-distillation. 2026 > <ref:

Improvements for AI systems

The Feedback Distillation method can be used to inject external knowledge into reasoning models by matching their distribution to a privileged feedback signal, which allows for token-level supervision. This enables the improved system to learn better formal reasoning by training on a dataset of problems without solutions (D of problems without solutions).
The system can maintain greater diversity in generated trajectories compared to standard reinforcement learning algorithms like GRPO, as Feedback Distillation maintains significantly higher entropy than GRPO throughout training. This translates directly into better test-time scaling and higher policy entropy, meaning the improved system produces a wider variety of valid proofs.
The improved system can leverage the two methods complementarily by taking checkpoints from Feedback Distillation and continuing with GRPO, as this combination outperforms GRPO alone. This suggests that initial training on Feedback Distillation provides a superior foundation for subsequent reinforcement learning refinement.
Tool use integration can be explicitly incorporated into the feedback loop, as the results confirm that iterative interaction with the Lean compiler is a driver of performance. The system can be equipped with tools like lean write file, lean check file, and rg in mathlib to perform iterative proof refinement.
The feedback prompt formulation can be optimized for stability and performance by using structured prompts, specifically those that encourage generic, self-contained feedback (Prompt 3 and 4). This leads to a significant improvement in training dynamics, as more structured prompts, which provide richer and more constrained feedback, tend to yield slower initial performance gains but result in more stable training and higher peak performance.

Abstract

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.