Offline Reinforcement Learning voor LLM Meerstapsredenering

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

Samenvatting

Het verbeteren van de multi-step redeneervaardigheid van grote taalmodellen (LLM's) met offline reinforcement learning (RL) is essentieel om ze snel aan te passen aan complexe taken. Hoewel Direct Preference Optimization (DPO) veelbelovend is gebleken in het afstemmen van LLM's op menselijke voorkeuren, is het minder geschikt voor multi-step redeneertaken omdat (1) DPO afhankelijk is van gepaarde voorkeursgegevens, die niet direct beschikbaar zijn voor multi-step redeneertaken, en (2) het behandelt alle tokens uniform, waardoor het ineffectief is voor credit assignment in multi-step redeneertaken, die vaak gepaard gaan met schaarse beloning. In dit werk stellen we OREO (Offline Reasoning Optimization) voor, een offline RL-methode voor het verbeteren van de multi-step redeneerbaarheid van LLM's. Voortbouwend op inzichten uit eerdere werken van maximum entropie reinforcement learning, leert het gezamenlijk een beleidsmodel en waardefunctie door optimalisatie van de zachte Bellman-vergelijking. We tonen in principe aan dat het de noodzaak om paardata te verzamelen vermindert en betere credit assignment mogelijk maakt. Empirisch overtreft OREO bestaande offline leermethoden op multi-step redeneerbenchmarks, waaronder wiskundige redeneertaken (GSM8K, MATH) en embodied agent control (ALFWorld). De aanpak kan worden uitgebreid naar een multi-iteratiekader wanneer er extra middelen beschikbaar zijn. Bovendien kan de geleerde waardefunctie worden benut om de boomzoekopdracht gratis te begeleiden, wat de prestaties tijdens testtijd verder kan verbeteren.

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.