Apprentissage par Renforcement Hors Ligne pour le Raisonnement à Plusieurs Étapes en LLM

Offline Reinforcement Learning for LLM Multi-Step Reasoning

December 20, 2024

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

cs.AI

Résumé

Améliorer la capacité de raisonnement à plusieurs étapes des grands modèles de langage (LLM) avec l'apprentissage par renforcement hors ligne (RL) est essentiel pour les adapter rapidement à des tâches complexes. Bien que l'Optimisation des Préférences Directes (DPO) ait montré des promesses en alignant les LLM avec les préférences humaines, elle est moins adaptée aux tâches de raisonnement à plusieurs étapes car (1) le DPO repose sur des données de préférence appariées, qui ne sont pas facilement disponibles pour les tâches de raisonnement à plusieurs étapes, et (2) il traite tous les tokens de manière uniforme, ce qui le rend inefficace pour l'attribution de crédit dans les tâches de raisonnement à plusieurs étapes, qui sont souvent associées à une récompense clairsemée. Dans ce travail, nous proposons OREO (Optimisation du Raisonnement Hors Ligne), une méthode RL hors ligne pour améliorer le raisonnement à plusieurs étapes des LLM. S'appuyant sur les enseignements des travaux précédents en matière d'apprentissage par renforcement à entropie maximale, il apprend conjointement un modèle de politique et une fonction de valeur en optimisant l'Équation de Bellman douce. Nous montrons en principe qu'il réduit le besoin de collecter des données par paires et permet une meilleure attribution de crédit. Empiriquement, OREO dépasse les méthodes d'apprentissage hors ligne existantes sur des référentiels de raisonnement à plusieurs étapes, y compris des tâches de raisonnement mathématique (GSM8K, MATH) et le contrôle d'agents incarnés (ALFWorld). L'approche peut être étendue à un cadre multi-itérations lorsque des ressources supplémentaires sont disponibles. De plus, la fonction de valeur apprise peut être exploitée pour guider la recherche arborescente gratuitement, ce qui peut encore améliorer les performances lors des tests.

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.