Aprendizagem por Reforço Offline para Raciocínio de Múltiplas Etapas em 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

Resumo

Melhorar a capacidade de raciocínio em vários passos de grandes modelos de linguagem (LLMs) com aprendizado por reforço (RL) offline é essencial para adaptá-los rapidamente a tarefas complexas. Embora a Otimização de Preferência Direta (DPO) tenha mostrado promessa em alinhar LLMs com preferências humanas, ela é menos adequada para tarefas de raciocínio em vários passos porque (1) a DPO depende de dados de preferência em pares, que não estão prontamente disponíveis para tarefas de raciocínio em vários passos, e (2) trata todos os tokens de forma uniforme, tornando-a ineficaz para atribuição de crédito em tarefas de raciocínio em vários passos, que frequentemente têm recompensas esparsas. Neste trabalho, propomos OREO (Otimização de Raciocínio Offline), um método RL offline para aprimorar o raciocínio em vários passos de LLMs. Com base em insights de trabalhos anteriores de aprendizado por reforço de entropia máxima, ele aprende em conjunto um modelo de política e uma função de valor otimizando a Equação de Bellman suave. Demonstramos, em princípio, que isso reduz a necessidade de coletar dados em pares e possibilita uma melhor atribuição de crédito. Empiricamente, OREO supera os métodos de aprendizado offline existentes em benchmarks de raciocínio em vários passos, incluindo tarefas de raciocínio matemático (GSM8K, MATH) e controle de agentes incorporados (ALFWorld). A abordagem pode ser estendida para um framework de múltiplas iterações quando recursos adicionais estão disponíveis. Além disso, a função de valor aprendida pode ser utilizada para orientar a busca em árvore gratuitamente, o que pode aumentar ainda mais o desempenho durante o teste.

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.