LLMマルチステップ推論のためのオフライン強化学習

Offline Reinforcement Learning for LLM Multi-Step Reasoning

December 20, 2024

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

cs.AI

要旨

大規模言語モデル（LLM）の多段階推論能力を向上させるために、オフライン強化学習（RL）を活用することは、これらを複雑なタスクに迅速に適応させるために不可欠です。直接的な選好最適化（DPO）は、LLMを人間の選好と整合させる点で有望であることが示されていますが、多段階推論タスクには適していません。なぜなら、（1）DPOはペアの選好データに依存しており、多段階推論タスクでは利用可能でないため、（2）すべてのトークンを均一に扱うため、多段階推論タスクにおける信用割り当てに効果的でなく、しばしば希薄な報酬が付随します。本研究では、LLMの多段階推論を向上させるためのオフラインRL手法であるOREO（Offline Reasoning Optimization）を提案します。従来の最大エントロピー強化学習の洞察に基づいて構築され、ソフトベルマン方程式の最適化によってポリシーモデルと価値関数を共に学習します。ペアのデータを収集する必要を軽減し、より良い信用割り当てを可能にすることを原理的に示しています。経験的には、OREOは数学的推論タスク（GSM8K、MATH）や具現化エージェント制御（ALFWorld）を含む多段階推論のベンチマークで既存のオフライン学習手法を凌駕しています。この手法は、追加のリソースが利用可能な場合には、マルチイテレーションフレームワークに拡張できます。さらに、学習された価値関数は、テスト時にパフォーマンスをさらに向上させるために、ツリー探索を自由にガイドするのに活用できます。

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.