Дистилляция обратной связи от LLM для доказательства теорем в Lean

Аннотация

Посттренировка моделей рассуждений обычно сочетает обучение с учителем и обучение с подкреплением на основе проверяемых вознаграждений, чаще всего с использованием GRPO. Однако этот алгоритм страдает от разреженных вознаграждений, ограниченного исследования и коллапса мод. Опираясь на недавние работы по самодистилляции, мы предлагаем дистилляцию обратной связи — метод обучения, при котором модель обучается на уровне токенов соответствовать своему собственному распределению, обусловленному привилегированной обратной связью, генерируемой языковой моделью. Дистилляция обратной связи обеспечивает контроль на уровне токенов и может внедрять внешние знания. Оценивая наш метод для доказательства теорем в Lean4, мы обнаруживаем, что дистилляция обратной связи сохраняет большее разнообразие в генерируемых траекториях, чем GRPO, что приводит к более высокой энтропии политики и лучшему масштабированию pass@k. Эти два метода являются взаимодополняющими: инициализация GRPO с контрольной точки дистилляции обратной связи превосходит каждый из методов по отдельности. В целом, наши результаты указывают на перспективное направление для улучшения посттренировки в сложных рассуждениях.

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.