Distilling LLM Feedback for Lean Theorem Proving

summary

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

In short

Feedback Distillation is a training method that teaches a model to match its own output distribution based on privileged feedback from another language model. This technique improves reasoning in Lean theorem proving by maintaining greater trajectory diversity and higher policy entropy compared to standard methods like GRPO. It suggests that models can internalize self-reflective abilities for better complex reasoning.

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 used across episodes

Reinforcement Learning (RL) — explained in 25 episodes
Retrieval-Augmented Generation (RAG) — explained in 20 episodes
Sample Complexity — explained in 20 episodes
Flow Matching — explained in 18 episodes
Federated Learning — explained in 17 episodes
Quantum Fisher Information (QFI) — explained in 17 episodes

This episode discusses

The paper

Distilling LLM Feedback for Lean Theorem Proving · Read on arXiv

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

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.

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.

More episodes

← Home