Destilando Retroalimentación de LLM para la Demostración de Teoremas en Lean

Resumen

El entrenamiento posterior para modelos de razonamiento suele combinar el ajuste fino supervisado con aprendizaje por refuerzo basado en recompensas verificables, más comúnmente con GRPO. Sin embargo, este algoritmo sufre de recompensas dispersas, exploración limitada y colapso modal. Basándonos en trabajos recientes sobre autodestilación, proponemos Destilación por Retroalimentación, un método de entrenamiento donde el modelo es entrenado para igualar, a nivel de tokens, su propia distribución condicionada a una retroalimentación privilegiada generada por un modelo de lenguaje. La Destilación por Retroalimentación ofrece supervisión a nivel de tokens y puede inyectar conocimiento externo. Al evaluar nuestro método para la demostración de teoremas en Lean4, encontramos que la Destilación por Retroalimentación mantiene una mayor diversidad en las trayectorias generadas que GRPO, lo que resulta en una mayor entropía de la política y un mejor escalado de pass@k. Ambos métodos son complementarios: inicializar GRPO desde un punto de control de Destilación por Retroalimentación supera a cualquiera de los métodos por separado. En conjunto, nuestros resultados sugieren una vía prometedora para mejorar el entrenamiento posterior en razonamiento complejo.

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.