Apprendimento per Rinforzo Offline per il Ragionamento a Più Passaggi con LLM

Offline Reinforcement Learning for LLM Multi-Step Reasoning

December 20, 2024

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

cs.AI

Abstract

Migliorare la capacità di ragionamento a più passaggi dei grandi modelli linguistici (LLM) con apprendimento per rinforzo (RL) offline è essenziale per adattarli rapidamente a compiti complessi. Mentre l'ottimizzazione diretta delle preferenze (DPO) ha mostrato promesse nell'allineare i LLM con le preferenze umane, è meno adatta per compiti di ragionamento a più passaggi perché (1) DPO si basa su dati di preferenza accoppiati, che non sono facilmente disponibili per compiti di ragionamento a più passaggi, e (2) tratta tutti i token in modo uniforme, rendendolo inefficace per l'assegnazione del credito in compiti di ragionamento a più passaggi, che spesso presentano ricompense sparse. In questo lavoro, proponiamo OREO (Offline Reasoning Optimization), un metodo RL offline per migliorare il ragionamento a più passaggi dei LLM. Basandoci su intuizioni da lavori precedenti di apprendimento per rinforzo con massima entropia, apprende congiuntamente un modello di politica e una funzione di valore ottimizzando l'Equazione di Bellman soft. Mostriamo in linea di principio che riduce la necessità di raccogliere dati a coppie e consente una migliore assegnazione del credito. Empiricamente, OREO supera i metodi di apprendimento offline esistenti su benchmark di ragionamento a più passaggi, inclusi compiti di ragionamento matematico (GSM8K, MATH) e controllo di agenti incorporati (ALFWorld). L'approccio può essere esteso a un framework multi-iterazione quando sono disponibili risorse aggiuntive. Inoltre, la funzione di valore appresa può essere sfruttata per guidare la ricerca nell'albero gratuitamente, il che può ulteriormente migliorare le prestazioni durante il test.

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.