Enhancing Multi-Step Reasoning Abilities of Language Models through Direct Q-Function Optimization

📄 arXiv: 2410.09302v2 📥 PDF

作者: Kaixuan Ji, Guanlin Liu, Ning Dai, Qingping Yang, Renjie Zheng, Zheng Wu, Chen Dun, Quanquan Gu, Lin Yan

分类: cs.LG, cs.AI, cs.CL

发布日期: 2024-10-11 (更新: 2025-02-11)

💡 一句话要点

提出DQO：通过直接Q函数优化提升语言模型的多步推理能力

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 强化学习 语言模型对齐 Q函数优化 多步推理 马尔可夫决策过程

📋 核心要点

  1. 现有强化学习方法在对齐语言模型时，面临计算资源需求高或难以处理多步推理任务的挑战。
  2. DQO将响应生成建模为MDP，利用SAC框架直接优化Q函数，实现更有效的过程监督。
  3. 实验表明，DQO在数学问题解决任务上优于现有方法，验证了其作为离线强化学习方法的潜力。

📝 摘要（中文）

强化学习在使大型语言模型（LLM）与人类偏好对齐并提高其执行复杂任务的能力方面起着关键作用。然而，当前的方法要么由于使用多个模型和大量的在线采样进行训练（例如，PPO）而需要大量的计算资源，要么被构建为bandit问题（例如，DPO，DRO），这通常难以处理多步推理任务，例如数学问题解决和涉及长链思维的复杂推理。为了克服这些限制，我们引入了直接Q函数优化（DQO），它将响应生成过程形式化为马尔可夫决策过程（MDP），并利用软演员-评论家（SAC）框架来直接优化由语言模型参数化的Q函数。DQO的MDP公式比基于bandit的方法具有结构优势，从而能够更有效地进行过程监督。在两个数学问题解决数据集GSM8K和MATH上的实验结果表明，DQO优于以前的方法，使其成为一种有前途的离线强化学习方法，用于对齐语言模型。

🔬 方法详解

问题定义：现有基于强化学习的语言模型对齐方法，如PPO，需要大量的计算资源进行在线采样和训练。而基于bandit的方法，如DPO和DRO，在处理需要多步推理的任务时表现不佳，因为它们缺乏对中间步骤的有效监督。因此，论文旨在解决语言模型在复杂推理任务中，由于缺乏有效过程监督而导致的性能瓶颈问题。

核心思路：论文的核心思路是将语言模型的响应生成过程建模为一个马尔可夫决策过程（MDP），并利用强化学习中的Q函数来评估每个状态-动作对的价值。通过直接优化Q函数，DQO能够更有效地学习到最优策略，从而提高语言模型的多步推理能力。这种方法允许对中间步骤进行更细粒度的监督，克服了bandit方法的局限性。

技术框架：DQO的整体框架包括以下几个主要组成部分：1) 将语言模型的文本生成过程建模为MDP，其中状态是当前生成的文本序列，动作是下一个token的选择。2) 使用一个由语言模型参数化的Q函数来评估每个状态-动作对的价值。3) 利用软演员-评论家（SAC）算法来优化Q函数，SAC是一种off-policy的强化学习算法，适用于连续动作空间。4) 使用离线数据集进行训练，避免了在线采样带来的计算开销。

关键创新：DQO的关键创新在于将语言模型的响应生成过程建模为MDP，并直接优化Q函数。与传统的bandit方法相比，MDP公式允许对中间步骤进行更有效的过程监督，从而提高了语言模型的多步推理能力。此外，DQO使用SAC算法进行训练，SAC算法具有较好的稳定性和收敛性。

关键设计：DQO的关键设计包括：1) 状态表示：使用当前生成的文本序列作为状态。2) 动作空间：将语言模型的词汇表作为动作空间，每个token对应一个动作。3) 奖励函数：根据任务的不同，可以设计不同的奖励函数，例如，对于数学问题解决任务，可以使用答案的正确性作为奖励。4) Q函数网络结构：使用Transformer网络作为Q函数的网络结构，Transformer网络具有强大的表示能力。5) 损失函数：使用SAC算法中的损失函数来优化Q函数，包括Q函数损失和策略损失。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，DQO在GSM8K和MATH两个数学问题解决数据集上均取得了显著的性能提升。在GSM8K数据集上，DQO的准确率超过了之前的最佳方法，达到了新的state-of-the-art水平。在MATH数据集上，DQO也取得了显著的性能提升，证明了其在复杂推理任务上的有效性。这些结果表明，DQO是一种有前途的离线强化学习方法，可以有效地对齐语言模型。

🎯 应用场景

DQO具有广泛的应用前景，可用于提升语言模型在需要复杂推理的任务中的表现，例如数学问题解决、代码生成、知识问答等。该方法还可以应用于对话系统，使其能够生成更连贯、更合理的回复。此外，DQO作为一种离线强化学习方法，可以利用已有的数据集进行训练，降低了训练成本，加速了语言模型对齐的进程。

📄 摘要（原文）

Reinforcement Learning (RL) plays a crucial role in aligning large language models (LLMs) with human preferences and improving their ability to perform complex tasks. However, current approaches either require significant computational resources due to the use of multiple models and extensive online sampling for training (e.g., PPO) or are framed as bandit problems (e.g., DPO, DRO), which often struggle with multi-step reasoning tasks, such as math problem solving and complex reasoning that involve long chains of thought. To overcome these limitations, we introduce Direct Q-function Optimization (DQO), which formulates the response generation process as a Markov Decision Process (MDP) and utilizes the soft actor-critic (SAC) framework to optimize a Q-function directly parameterized by the language model. The MDP formulation of DQO offers structural advantages over bandit-based methods, enabling more effective process supervision. Experimental results on two math problem-solving datasets, GSM8K and MATH, demonstrate that DQO outperforms previous methods, establishing it as a promising offline reinforcement learning approach for aligning language models.