Destillation von LLM-Feedback für das Theorembeweisen mit Lean

Zusammenfassung

Das Nachtraining für Modelle für logisches Schließen kombiniert typischerweise überwachtes Feintuning mit bestärkendem Lernen mit verifizierbaren Belohnungen, am häufigsten mit GRPO. Allerdings leidet dieser Algorithmus unter spärlichen Belohnungen, eingeschränkter Exploration und Modenkollaps. Aufbauend auf aktuellen Arbeiten zur Selbst-Destillation schlagen wir Feedback-Destillation vor, eine Trainingsmethode, bei der das Modell auf Tokenebene trainiert wird, seine eigene, durch privilegiertes Feedback eines Sprachmodells konditionierte Verteilung nachzubilden. Feedback-Destillation bietet Überwachung auf Tokenebene und kann externes Wissen einbringen. Bei der Evaluierung unserer Methode für Lean4-Theorembeweise stellen wir fest, dass Feedback-Destillation eine größere Vielfalt in generierten Trajektorien beibehält als GRPO, was zu einer höheren Policy-Entropie und besserem pass@k-Scaling führt. Die beiden Methoden sind komplementär: das Initialisieren von GRPO mit einem Feedback-Destillation-Checkpoint übertrifft jede der beiden Methoden einzeln. Alles in allem deuten unsere Ergebnisse auf einen vielversprechenden Weg hin, das Nachtraining für komplexes logisches Schließen zu verbessern.

English

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.