Distillation des retours des LLM pour la preuve de théorèmes en Lean

Résumé

Le post-entraînement des modèles de raisonnement combine généralement l’apprentissage supervisé fin avec l’apprentissage par renforcement à partir de récompenses vérifiables, le plus souvent via GRPO. Cependant, cet algorithme souffre de récompenses clairsemées, d’une exploration limitée et d’un effondrement de mode. En nous appuyant sur des travaux récents en auto-distillation, nous proposons la Distillation par Rétroaction, une méthode d’entraînement où le modèle est entraîné à correspondre, au niveau des tokens, à sa propre distribution conditionnée par une rétroaction privilégiée produite par un modèle de langage. La Distillation par Rétroaction offre une supervision au niveau des tokens et peut injecter des connaissances externes. En évaluant notre méthode pour la démonstration de théorèmes sous Lean4, nous constatons que la Distillation par Rétroaction maintient une plus grande diversité dans les trajectoires générées que GRPO, produisant une entropie de politique plus élevée et un meilleur passage à l’échelle du pass@k. Les deux méthodes sont complémentaires : initialiser GRPO à partir d’un point de contrôle de Distillation par Rétroaction surpasse l’une ou l’autre méthode employée seule. Dans l’ensemble, nos résultats suggèrent une voie prometteuse pour améliorer le post-entraînement dans le cadre de raisonnements complexes.

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.