Destilando Feedback de LLM para Prova de Teoremas em Lean

Resumo

O pós-treinamento para modelos de raciocínio geralmente combina ajuste fino supervisionado com aprendizado por reforço baseado em recompensas verificáveis, mais comumente com GRPO. No entanto, esse algoritmo sofre de recompensas esparsas, exploração limitada e colapso de modo. Com base em trabalhos recentes sobre autodestilação, propomos a Destilação de Feedback, um método de treinamento no qual o modelo é treinado para corresponder, no nível do token, sua própria distribuição condicionada ao feedback privilegiado gerado por um modelo de linguagem. A Destilação de Feedback oferece supervisão em nível de token e pode injetar conhecimento externo. Avaliando nosso método para demonstração de teoremas no Lean4, descobrimos que a Destilação de Feedback mantém maior diversidade nas trajetórias geradas do que o GRPO, resultando em maior entropia da política e melhor escalabilidade de pass@k. Os dois métodos são complementares: inicializar o GRPO a partir de um checkpoint de Destilação de Feedback supera qualquer um dos métodos isoladamente. Em suma, nossos resultados sugerem um caminho promissor para melhorar o pós-treinamento para raciocínio complexo.

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.