Offline Reinforcement Learning for LLM Multi-Step Reasoning

Fuente: arXiv
Gespeichert in:
Bibliographische Detailangaben
Hauptverfasser: Wang, Huaijie, Hao, Shibo, Dong, Hanze, Zhang, Shenao, Bao, Yilin, Yang, Ziran, Wu, Yi
Format: Preprint
Veröffentlicht: 2024
Schlagworte:
Online-Zugang:
Tags: Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
_version_ 1866915080491237376
author Wang, Huaijie
Hao, Shibo
Dong, Hanze
Zhang, Shenao
Bao, Yilin
Yang, Ziran
Wu, Yi
author_facet Wang, Huaijie
Hao, Shibo
Dong, Hanze
Zhang, Shenao
Bao, Yilin
Yang, Ziran
Wu, Yi
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2412_16145
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Offline Reinforcement Learning for LLM Multi-Step Reasoning
Wang, Huaijie
Hao, Shibo
Dong, Hanze
Zhang, Shenao
Bao, Yilin
Yang, Ziran
Wu, Yi
Machine Learning
Artificial Intelligence
Computation and Language
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.
title Offline Reinforcement Learning for LLM Multi-Step Reasoning
topic Machine Learning
Artificial Intelligence
Computation and Language
url https://arxiv.org/abs/2412.16145