Aprendizaje por Refuerzo sin Conexión para Razonamiento de Varios Pasos en LLM

Offline Reinforcement Learning for LLM Multi-Step Reasoning

December 20, 2024

Autores: Huaijie Wang, Shibo Hao, Hanze Dong, Shenao Zhang, Yilin Bao, Ziran Yang, Yi Wu

cs.AI

Resumen

Mejorar la capacidad de razonamiento de múltiples pasos de los modelos de lenguaje grandes (LLMs) con aprendizaje por refuerzo (RL) fuera de línea es esencial para adaptarlos rápidamente a tareas complejas. Si bien la Optimización de Preferencias Directas (DPO) ha demostrado ser prometedora en alinear los LLMs con las preferencias humanas, es menos adecuada para tareas de razonamiento de múltiples pasos porque (1) DPO se basa en datos de preferencias emparejadas, que no están fácilmente disponibles para tareas de razonamiento de múltiples pasos, y (2) trata todos los tokens de manera uniforme, lo que lo hace ineficaz para la asignación de crédito en tareas de razonamiento de múltiples pasos, que a menudo vienen con recompensas escasas. En este trabajo, proponemos OREO (Optimización de Razonamiento en Offline), un método de RL fuera de línea para mejorar el razonamiento de múltiples pasos de LLMs. Basándose en ideas de trabajos anteriores de aprendizaje por refuerzo de máxima entropía, aprende conjuntamente un modelo de política y una función de valor optimizando la Ecuación de Bellman suave. Mostramos en principio que reduce la necesidad de recopilar datos emparejados y permite una mejor asignación de crédito. Empíricamente, OREO supera a los métodos de aprendizaje fuera de línea existentes en bancos de pruebas de razonamiento de múltiples pasos, incluidas tareas de razonamiento matemático (GSM8K, MATH) y control de agentes encarnados (ALFWorld). El enfoque se puede extender a un marco de múltiples iteraciones cuando hay recursos adicionales disponibles. Además, la función de valor aprendida se puede aprovechar para guiar la búsqueda en árbol de forma gratuita, lo que puede mejorar aún más el rendimiento durante el tiempo de prueba.

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.