语言生成建模为Token级 Markov 决策过程（Token-Level MDP）详解-CSDN博客

本文链接：https://blog.csdn.net/shizheng_Li/article/details/147276945

语言生成建模为Token级 Markov 决策过程（Token-Level MDP）详解

为了详细解释原文中将 语言生成建模为标记级 Markov 决策过程（Token-Level MDP） 的含义及其作用，我们需要从 Markov 决策过程（MDP）的基本概念入手，结合语言生成任务的特点，逐一剖析原文中定义的 MDP 元组 $\mathcal{M} = (\mathcal{S}, \mathcal{A}, \mathbb{P}, R, d_0, \omega)$ 的每个组件，并说明这种建模方式在强化学习（RL）框架下的意义和用途。以下是全面的解释，涵盖理论背景、建模细节、优势和应用场景。

Paper: https://arxiv.org/pdf/2504.05118

1. Markov 决策过程（MDP）简介

Markov 决策过程（MDP） 是一种数学框架，用于建模序列决策问题，特别适合描述智能体（agent）与环境（environment）交互并通过学习策略（policy）最大化累积奖励（cumulative reward）的场景。MDP 的核心假设是马尔可夫性，即未来的状态和奖励仅依赖于当前状态和动作，而与历史无关。MDP 通常由以下五元组定义：

$\mathcal{M} = (\mathcal{S}, \mathcal{A}, \mathbb{P}, R, \gamma)$

$\mathcal{S}$ ：状态空间（State Space），描述智能体可能处于的所有状态。
$\mathcal{A}$ ：动作空间（Action Space），智能体可执行的所有动作。
$\mathbb{P}$ ：状态转移概率（Transition Dynamics）， $\mathbb{P}(s_{t+1} | s_t, a_t)$ 表示在状态 $s_t$ 执行动作 $a_t$ 后转移到状态 $s_{t+1}$ 的概率。
$R$ ：奖励函数（Reward Function）， $R(s_t, a_t)$ 或 $R(s_t, a_t, s_{t+1})$ 表示执行动作后获得的即时奖励。
$\gamma$ ：折扣因子（Discount Factor）， $\gamma \in [0, 1]$ ，平衡短期和长期奖励。

在强化学习中，智能体的目标是学习一个策略 $\pi(a_t | s_t)$ ，最大化期望累积奖励：

$J(\pi) = \mathbb{E}_\pi \left[ \sum_{t=0}^\infty \gamma^t R(s_t, a_t) \right]$

原文将语言生成任务建模为一个标记级 MDP，并扩展了标准定义，添加了初始状态分布 $d_0$ 和终止动作 $\omega$ ，以适应自然语言处理（NLP）中的生成过程。

2. 语言生成作为标记级 MDP 的建模

原文将语言生成任务（例如给定提示 $x$ 生成响应 $y$ ）建模为一个标记级 MDP，定义为元组：

$\mathcal{M} = (\mathcal{S}, \mathcal{A}, \mathbb{P}, R, d_0, \omega)$

下面逐一解释每个组件的含义及其在语言生成中的具体体现。

2.1 状态空间（State Space, $\mathcal{S}$ ）

定义：
- 状态空间 $\mathcal{S}$ 包含所有可能的标记序列，这些序列由提示 $x$ 和已生成的响应标记组成。
- 在时间步 $t$ ，状态 $s_t$ 定义为：
  $s_t = \langle x_0, \ldots, x_m, y_0, \ldots, y_t \rangle$
  其中：
  - $\langle x_0, \ldots, x_m \rangle$ 是提示（prompt），由 $m + 1$ 个标记组成。
  - $y_0, \ldots, y_t$ 是截至时间步 $t$ 已生成的响应标记。
含义：
- 状态 $s_t$ 表示生成过程中的“上下文”，包括完整的提示和当前生成的响应部分。
- 例如，若提示为“Solve the equation”，已生成响应为“First, let’s”，则状态可能是：
  $s_t = \langle \text{Solve}, \text{the}, \text{equation}, \text{First}, \text{,}, \text{let’s} \rangle$
特性：
- 状态空间是离散的，因为标记来自固定词汇表 $\mathcal{A}$ 。
- 随着生成进行，状态长度逐渐增加（从 $m + 1$ 到 $m + t + 1$ ），反映了序列的动态增长。

2.2 动作空间（Action Space, $\mathcal{A}$ ）

定义：
- 动作空间 $\mathcal{A}$ 对应于固定离散词汇表，即语言模型可选择的标记集合。
- 在每个时间步，智能体从 $\mathcal{A}$ 中选择一个标记作为动作 $a_t$ 。
含义：
- 动作 $a_t$ 是生成过程中的下一个标记 $y_{t+1}$ 。
- 例如，若词汇表包含 {“the”, “is”, “First”, …, “”}，则 $a_t$ 可以是任意词汇表中的标记。
特性：
- $\mathcal{A}$ 通常很大（例如，现代语言模型的词汇表可能包含数万标记）。
- 动作空间是离散的，与连续动作空间（如机器人控制）不同。

2.3 状态转移概率（Dynamics, $\mathbb{P}$ ）

定义：
- 状态转移是确定性的：
  $\mathbb{P}(s_{t+1} | s_t, a) = 1$
  其中：
  - 当前状态 $s_t = \langle x_0, \ldots, x_m, y_0, \ldots, y_t \rangle$ 。
  - 动作 $a = y_{t+1}$ 。
  - 下一状态 $s_{t+1} = \langle x_0, \ldots, x_m, y_0, \ldots, y_t, y_{t+1} \rangle$ 。
含义：
- 给定状态 $s_t$ 和动作 $a$ ，下一状态 $s_{t+1}$ 是完全确定的，即新状态是将动作 $a$ （即 $y_{t+1}$ ）追加到当前序列后形成。
- 例如，若 $s_t = \langle \text{Solve}, \text{the}, \text{First} \rangle$ ，动作 $\text{,}$ ，则：
  $s_{t+1} = \langle \text{Solve}, \text{the}, \text{First}, \text{,} \rangle$
  转移概率为 1，无随机性。
特性：
- 确定性转移简化了建模，因为语言生成不像物理环境（如机器人导航）有外部随机干扰。
- 这种设计反映了语言生成的顺序性：每个动作直接决定序列的下一个标记。

2.4 终止条件（Termination Condition, $\omega$ ）

定义：
- 生成过程在执行终止动作 $\omega$ 时结束， $\omega$ 通常是句子结束标记（end-of-sentence token），如 <eos>。
含义：
- 当智能体选择 $\omega$ 作为动作（ $a_t = \omega$ ），生成停止，轨迹（trajectory）完成。
- 例如，生成序列 $\langle \text{First}, \text{,}, \text{let’s}, \text{<eos>} \rangle$ 表示响应结束。
特性：
- $\omega$ 是词汇表 $\mathcal{A}$ 中的特殊标记，触发 MDP 的终止状态。
- 终止条件允许模型控制生成长度，适应不同任务需求（如短答案或长推理）。

2.5 奖励函数（Reward Function, $R (s, a)$ ）

定义：
- 奖励函数 $R (s, a)$ 为智能体在状态 $s$ 执行动作 $a$ 提供标量反馈。
- 在强化学习从人类反馈（RLHF）中，奖励可以基于：
  - 人类偏好：例如，通过比较生成响应的质量。
  - 任务特定规则：例如，Long-CoT 任务中验证答案正确性。
含义：
- 在 Long-CoT 任务中，奖励通常是稀疏的，仅在序列末尾（终止动作 $\omega$ ）提供：
  - 如果生成序列正确（如数学问题答案正确）， $R(s_T, \omega) = 1$ 。
  - 如果错误， $R(s_T, \omega) = 0$ 。
  - 非终止动作的奖励通常为 $R(s_t, a_t) = 0$ （ $t < T$ ）。
- 例如，在 AIME 数学推理任务中，奖励由验证器（verifier）根据最终答案判定。
特性：
- 奖励函数的设计灵活，可根据任务调整（例如，中间步骤奖励、质量评分等）。
- 稀疏奖励是 Long-CoT 的挑战之一，需通过强化学习优化策略以发现高回报路径。

2.6 初始状态分布（Initial State Distribution, $d_0$ ）

定义：
- 初始状态分布 $d_0$ 是提示 $x$ 的概率分布。
- 初始状态 $s_0$ 由提示的标记组成：
  $s_0 = \langle x_0, \ldots, x_m \rangle$
含义：
- $d_0$ 定义了任务的起点，例如从 AIME 数据集中采样数学问题作为提示。
- 例如， $\text{“Solve } x^2 - 5x + 6 = 0\text{”}$ ，则 $s_0 = \langle \text{Solve}, x^2, -, 5x, +, 6, =, 0 \rangle$ 。
特性：
- $d_0$ 允许模型处理多样化的输入提示，增强泛化能力。
- 在训练中， $d_0$ 通常是数据集的经验分布（如 AIME 2024 问题集）。

3. 建模为标记级 MDP 的意义

将语言生成建模为标记级 MDP 是一种将 NLP 任务嵌入强化学习框架的强大方法，其意义和用途体现在以下几个方面：

3.1 将语言生成转化为序列决策问题

统一框架：
- 语言生成本质上是序列决策：模型在每个时间步选择一个标记，逐步构建响应。
- MDP 提供了一个标准的数学框架，将语言生成分解为状态、动作和奖励，允许使用强化学习算法（如 PPO）优化策略。
细粒度控制：
- 标记级建模（每个标记作为一个动作）比序列级建模（整个响应作为一个动作）更精细，允许模型优化生成过程中的每个步骤。
- 例如，在 Long-CoT 任务中，模型可以学习如何生成正确的推理步骤，而不仅是最终答案。

3.2 适配强化学习优化

策略优化：
- 语言模型被视为策略 $\pi(a_t | s_t)$ ，输出给定状态（上下文）下动作（标记）的概率分布。
- 通过 MDP 建模，强化学习算法可以优化 $\pi$ ，最大化累积奖励。例如，VAPO 使用 PPO 调整 Qwen2.5-32B 的生成策略，提高 AIME 得分。
优势估计：
- GAE（广义优势估计）等 RL 技术可为每个标记计算优势 $\hat{A}_t$ ，指导模型优先选择高回报的标记。
- 在 Long-CoT 中，这意味着模型能学习哪些推理步骤对最终正确答案贡献最大。

3.3 处理稀疏奖励

稀疏奖励建模：
- Long-CoT 任务的奖励通常仅在序列末尾提供（ $R(s_T, \omega)$ ），MDP 框架通过价值函数 $V(s_t)$ 和优势 $\hat{A}_t$ 将末尾奖励传播到早期标记。
- 例如，VAPO 的长度自适应 GAE 和标记级损失利用 MDP 结构，增强了稀疏奖励下的优化效率。
探索-利用权衡：
- MDP 的策略 $\pi$ 可以通过调整（如 Clip-Higher 的 $\epsilon_{\text{high}}$ ）平衡探索（尝试新标记）和利用（选择高概率标记），缓解熵崩溃。

3.4 支持复杂推理任务

长链推理（Long-CoT）：
- Long-CoT 任务要求模型生成数百标记的推理序列，MDP 的状态 $s_t$ 动态扩展（包含所有历史标记），完美适应这种需求。
- 确定性转移 $\mathbb{P}$ 确保推理序列的逻辑连贯性，动作 $a_t$ 直接影响下一状态 $s_{t+1}$ 。
奖励设计灵活性：
- 奖励函数 $R (s, a)$ 可根据任务定制。例如，AIME 任务使用验证器评分，MDP 允许扩展到中间奖励（如步骤正确性）以进一步指导推理。

3.5 理论与实践的桥梁

理论支持：
- MDP 提供了强化学习的理论基础，保证策略收敛性和优化目标的清晰定义（如 KL 正则化目标）。
- 原文中的 RLHF 目标：
  $\pi^* = \arg \max_\pi \mathbb{E}_{\pi, s_0 \sim d_0} \left[ \sum_{t=0}^H \left( R(s_t, a_t) - \beta \text{KL}(\pi(\cdot|s_t) || \pi_{\text{ref}}(\cdot|s_t)) \right) \right]$
  直接基于 MDP 框架，平衡奖励最大化与策略稳定性。
实践可行性：
- 标记级 MDP 与语言模型的架构兼容（自回归生成），便于实现。例如，Qwen2.5-32B 的输出分布直接作为 $\pi(a_t | s_t)$ 。
- VAPO 的技术（价值预训练、解耦 GAE 等）依赖 MDP 的状态-动作定义，确保训练高效稳定。

4. 建模的具体用途

基于标记级 MDP 的建模，VAPO 在 Long-CoT 任务中实现了以下具体用途：

4.1 优化复杂推理能力

AIME 2024 表现：
- VAPO 在 Qwen2.5-32B 上达到 60.4 分，超越 DAPO（50 分）和 DeepSeek-R1（47 分），得益于 MDP 框架支持的细粒度优化。
- 标记级 MDP 允许模型学习每个推理步骤的贡献，例如，如何从“设变量”到“推导公式”，而非仅关注最终答案。

4.2 缓解训练挑战

价值模型偏差：
- MDP 的状态定义（包含完整上下文）支持价值预训练，减少初始化偏差（论文 Section 4.1）。
- 例如，价值模型 $V(s_t)$ 基于 $s_t = \langle x_0, \ldots, y_t \rangle$ 估计长期回报，MDP 确保状态捕捉所有相关信息。
异构序列长度：
- 长度自适应 GAE（Section 4.2）利用 MDP 的动态状态长度，调整 $\lambda_{\text{actor}}$ 以优化短序列（低方差）和长序列（低偏差）。
稀疏奖励：
- 标记级损失和 Clip-Higher（Section 4.3）通过 MDP 的动作级优化，增强探索能力，缓解熵崩溃。

4.3 提升训练稳定性

稳定性：
- VAPO 未观察到训练崩溃，分数稳定在 60-61（论文 Section 5.2），得益于 MDP 提供的结构化优化框架。
- 确定性转移 $\mathbb{P}$ 和终止条件 $\omega$ 确保生成序列的逻辑一致性，避免无效轨迹。
效率：
- 仅需 5000 步达到 SOTA 性能（比 DAPO 少 40% 步数），因为 MDP 允许高效传播奖励信号。

4.4 可扩展性

通用性：
- 标记级 MDP 不仅适用于 Long-CoT，还可扩展到其他生成任务（如代码生成、对话系统）。
- 奖励函数 $R (s, a)$ 的灵活性支持不同任务的定制，例如，代码任务可奖励语法正确性，对话任务可奖励流畅度。
未来改进：
- MDP 框架支持引入中间奖励、动态 $\omega$ （如条件终止），进一步提升复杂任务的性能。

5. 为什么这样建模？

5.1 与其他建模方式的对比

序列级建模：
- 将整个响应 $y$ 视为单一动作，奖励基于完整序列（如语言模型评分）。
- 缺点：无法优化中间步骤，难以处理长序列推理，奖励信号传播效率低。
- MDP 优势：标记级分解允许逐一优化每个步骤，适合 Long-CoT 的细粒度推理。
无模型方法：
- 直接优化语言模型输出（如基于梯度的无 RL 方法）。
- 缺点：缺乏结构化框架，难以处理稀疏奖励或长期依赖。
- MDP 优势：提供状态-动作-奖励的明确定义，支持 RL 算法的高效优化。

5.2 适应语言模型架构

自回归生成：
- 现代语言模型（如 Qwen2.5-32B）是自回归的，基于上下文 $s_t$ 预测下一个标记 $a_t$ ，与 MDP 的状态-动作流程天然契合。
- 例如， $s_t = \langle x_0, \ldots, y_t \rangle$ 作为输入， $\pi(a_t | s_t)$ 输出动作概率。
计算效率：
- 确定性转移 $\mathbb{P}$ 和离散动作 $\mathcal{A}$ 简化了计算，适合大规模模型训练。

5.3 支持 RLHF

RLHF 兼容性：
- RLHF 需要通过人类反馈或规则定义奖励，MDP 的奖励函数 $R (s, a)$ 直接支持这一需求。
- 例如，AIME 验证器作为 $R(s_T, \omega)$ ，MDP 确保奖励信号通过 GAE 传播到早期标记。
KL 正则化：
- 原文的 KL 正则化目标（ $\beta \text{KL}(\pi || \pi_{\text{ref}})$ ）在 MDP 框架下自然实现，防止策略偏离预训练模型。