Оффлайн обучение с подкреплением для многошагового рассуждения в модели LLM.

Offline Reinforcement Learning for LLM Multi-Step Reasoning

December 20, 2024

Авторы: Huaijie Wang, Shibo Hao, Hanze Dong, Shenao Zhang, Yilin Bao, Ziran Yang, Yi Wu

cs.AI

Аннотация

Улучшение способности к многошаговому рассуждению крупных языковых моделей (LLM) с использованием обучения с подкреплением (RL) в автономном режиме является важным для быстрой адаптации их к сложным задачам. В то время как метод прямой оптимизации предпочтений (DPO) показал потенциал в выравнивании LLM с человеческими предпочтениями, он менее подходит для задач многошагового рассуждения из-за того, что (1) DPO зависит от данных сопоставления предпочтений, которые не всегда доступны для задач многошагового рассуждения, и (2) он одинаково обрабатывает все токены, что делает его неэффективным для присвоения заслуг в задачах многошагового рассуждения, часто связанных с разреженной наградой. В данной работе мы предлагаем OREO (Offline Reasoning Optimization), метод обучения с подкреплением в автономном режиме для улучшения многошагового рассуждения LLM. Основываясь на идеях предыдущих работ по обучению с максимальной энтропией, он совместно обучает модель политики и функцию ценности, оптимизируя мягкое уравнение Беллмана. Мы показываем в принципе, что это уменьшает необходимость в сборе парных данных и обеспечивает лучшее присвоение заслуг. Эмпирически OREO превосходит существующие методы обучения в автономном режиме на бенчмарках многошагового рассуждения, включая задачи математического рассуждения (GSM8K, MATH) и управление воплощенным агентом (ALFWorld). Подход может быть расширен до многократной структуры при наличии дополнительных ресурсов. Кроме того, изученная функция ценности может быть использована для направления поиска по дереву бесплатно, что дополнительно повышает производительность во время тестирования.

English

Improving the multi-step reasoning ability of large language models (LLMs) with offline reinforcement learning (RL) is essential for quickly adapting them to complex tasks. While Direct Preference Optimization (DPO) has shown promise in aligning LLMs with human preferences, it is less suitable for multi-step reasoning tasks because (1) DPO relies on paired preference data, which is not readily available for multi-step reasoning tasks, and (2) it treats all tokens uniformly, making it ineffective for credit assignment in multi-step reasoning tasks, which often come with sparse reward. In this work, we propose OREO (Offline Reasoning Optimization), an offline RL method for enhancing LLM multi-step reasoning. Building on insights from previous works of maximum entropy reinforcement learning, it jointly learns a policy model and value function by optimizing the soft Bellman Equation. We show in principle that it reduces the need to collect pairwise data and enables better credit assignment. Empirically, OREO surpasses existing offline learning methods on multi-step reasoning benchmarks, including mathematical reasoning tasks (GSM8K, MATH) and embodied agent control (ALFWorld). The approach can be extended to a multi-iteration framework when additional resources are available. Furthermore, the learned value function can be leveraged to guide the tree search for free, which can further boost performance during test time.