Destilleren van LLM-feedback voor Lean-stellingbewijzen

Samenvatting

Post-training voor redeneringsmodellen combineert doorgaans begeleid finetunen met reinforcement learning op basis van verifieerbare beloningen, meestal met GRPO. Dit algoritme heeft echter last van schaarse beloningen, beperkte exploratie en mode collapse. Voortbouwend op recent werk over zelfdistillatie stellen we Feedback Distillation voor, een trainingsmethode waarbij het model op token-niveau wordt getraind om zijn eigen distributie te matchen, geconditioneerd op geprivilegieerde feedback die door een taalmodel wordt geproduceerd. Feedback Distillation biedt supervisie op token-niveau en kan externe kennis injecteren. Bij de evaluatie van onze methode voor stellingbewijzen in Lean4 zien we dat Feedback Distillation een grotere diversiteit behoudt in gegenereerde trajecten dan GRPO, wat leidt tot een hogere beleidsentropie en betere pass@k-schaling. Beide methoden zijn complementair: het initialiseren van GRPO vanuit een Feedback Distillation-checkpoint presteert beter dan elk van beide methoden afzonderlijk. Al met al suggereren onze resultaten een veelbelovende weg om post-training voor complex redeneren te verbeteren.

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.