On Reparameterizing the Score Function¶
会议: NeurIPS 2026
arXiv: 2610.24188
领域: 强化学习
关键词: 强化学习, 策略梯度, 得分函数, 自回归模型, 重参数化, 方差缩减, 信用分配
一句话总结¶
本文提出历史重参数化（HR）策略梯度框架，通过将自回归生成中的条件前缀显式建模为关于模型参数的可微轨迹，并设计中心化无偏控制变量估计器（OHR），在完全不依赖可微标量奖励的前提下打通了下游对上游决策的路径级敏感度反馈，实现了高达数倍的梯度方差缩减与更稳健的长程信用分配。
研究背景与动机¶
在强化学习优化自回归序列生成（如大语言模型推理、代码生成、组合优化）时，REINFORCE、PPO 及 GRPO 等核心算法均依赖核心原语——得分函数 \(\nabla_\theta \log p_\theta(x_t \mid x_{<t})\)。然而，现有策略梯度在对时间步 \(t\) 求导时，将条件前缀 \(x_{<t}\) 视为固定的常数，完全忽略了模型参数 \(\theta\) 如何在上游塑造该历史。当环境仅提供全局序列标量奖励 \(r(x_{0:N})\)（如稀疏二元正确性）时，每一个时间步的得分项均被赋予完全相同的标量权重 \((r - b)\)。在此机制下，自回归因子分解中的深层因果链对梯度完全不可见，梯度方差随着时间步长 \(N\) 呈爆炸式增长。
自回归无环境序列生成截然不同于未知黑盒具身环境：下一个 token 的输入完全是模型自身的输出序列，前缀历史本身关于 \(\theta\) 是连续且处处可微的。核心 idea：打破冻结前序历史的人为假设，利用连续嵌入级弛豫将历史序列重参数化为可微生成轨迹，并推导具有解析闭式最优权重的中心化控制变量（OHR），在总体层面严格保证不劣于原始策略梯度的同时大幅缩减梯度方差。
方法详解¶
整体框架¶
算法从带有基线 \(b\) 的策略梯度出发，通过将条件前缀替换为带有重参数化噪声的可微轨迹 \(\chi_{<t}(\theta)\)，将梯度分解为传统的局部得分项（Score Term）和跨时间步反向传播的历史路径项（History-Pathwise Term）。为克服直接累加带来的期望偏差，论文通过中心化构造控制变量，并以闭式解计算最优混合系数。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["噪声采样 ε ~ p(ε) 与前向展开"] --> B["可微轨迹 χ_t(θ) 生成序列 x_0:N"]
B --> C["局部得分项 S_t = ∂_θ log p_θ(x_t | χ_<t)"]
B --> D["历史路径项 C_t = ∑ (∇_θ χ_j)^T (∂ log p_θ / ∂ χ_j)"]
C --> E["策略梯度项 g_PG = (r - b) S"]
D --> F["历史路径控制变量 H = (r - b) C"]
F --> G["EMA 中心化修正 H_bar = H - E[H]"]
E & G --> H["闭式最优系数 α1* = - Tr(Cov) / Tr(Var)"]
H --> I["无偏方差最优梯度估计器 g_OHR"]关键设计¶
1. 历史重参数化梯度展开与因果分解
定义连续情形下的精确重参数化生成过程 \(\chi_t(\theta) = S(\phi_\theta(\chi_{<t}(\theta)), \varepsilon_t)\)（满足数值上 \(\chi_t = x_t\)）。将其置入策略梯度条件项后，基于多元复合函数链式法则，总导数完全解耦为两部分： $\(g_{\mathrm{HR}} = (r - b) \sum_{t=0}^N \nabla_\theta \log p_\theta(x_t \mid \chi_{<t}(\theta)) = \underbrace{(r - b) \sum_{t=0}^N \partial_\theta \log p_\theta(x_t \mid \chi_{<t})}_{g_{\mathrm{PG}}^b} + \underbrace{(r - b) \sum_{t=0}^N \sum_{j<t} (\nabla_\theta \chi_j)^\top \frac{\partial \log p_\theta(x_t \mid \chi_{<t})}{\partial \chi_j}}_{H}\)$ 其中 \(H = (r - b) C\) 即为历史路径项。它直接通过 BPTT 查询模型内部雅可比矩阵 \(\frac{\partial \log p_\theta}{\partial \chi_j}\)，将下游步骤的高置信度反向传播回早期的决策点，赋予了早期 token 依据“下游后继影响”而调整的敏感度，打破了无环境 RL 无法反传状态导数的屏障。
2. 离散 Token 的连续嵌入级弛豫（Embedding-Level Relaxation）
针对离散自回归模型，传统的 Straight-Through Gumbel-Softmax 虽然在前向保持 one-hot、反向传梯度，但前向仍为硬离散值，导致在得分函数内评估时无法提供平滑的输入变化响应。作者设计了直接作用于输入词嵌入空间的连续弛豫： $\(\chi_t = \operatorname{softmax}\left(\frac{\phi_\theta(\chi_{<t})}{\tau} + g_t\right) \cdot E\)$ 其中 \(E \in \mathbb{R}^{V \times d_e}\) 为词嵌入矩阵，\(g_t \sim \operatorname{Gumbel}(0, 1)\)。该设计将下一步的输入从单一向量替换为平滑的语义空间混合向量，使下游输出对历史前缀形成处处非零且富有语义导向的导数流，无需改动任何神经网络底层架构。
3. OHR 闭式最优系数推导与方差缩减界
由于历史路径项 \(H\) 的期望在通常情况下非零，直接相加会引入严重参数偏差。作者将其改写为中心化无偏控制变量 \(\bar{H} = H - \mathbb{E}[H]\)，并构造组合估计器 \(g_{\mathrm{OHR}} = g_{\mathrm{PG}}^b + \alpha_1 \bar{H}\)。通过建立方差迹的目标优化函数并关于标量 \(\alpha_1\) 求驻点，直接导出解析最优解： $\(\alpha_1^* = -\frac{\operatorname{Tr}(\operatorname{Cov}(g_{\mathrm{PG}}^b, \bar{H}))}{\operatorname{Tr}(\operatorname{Var}(\bar{H}))}\)$ 代回后得到理论方差： $\(\operatorname{Tr}(\operatorname{Var}(g_{\mathrm{OHR}})) = \operatorname{Tr}(\operatorname{Var}(g_{\mathrm{PG}}^b)) - \frac{[\operatorname{Tr}(\operatorname{Cov}(g_{\mathrm{PG}}^b, \bar{H}))]^2}{\operatorname{Tr}(\operatorname{Var}(\bar{H}))} \leq \operatorname{Tr}(\operatorname{Var}(g_{\mathrm{PG}}^b))\)$ 相对方差削减比例恰好等于两者方差迹的广义皮尔逊相关系数平方 \(\rho^2(g_{\mathrm{PG}}^b, \bar{H})\)。
4. 在线训练中的指数移动平均（EMA）中心化机制
在实际批量（Batch Size \(M\)）训练中，若直接使用当前 Batch 的均值对 \(H\) 进行中心化，批量平均梯度中控制变量将完全抵消消失，退化回普通 PG。为此，作者引入跨优化步的指数移动平均（EMA）估计总体均值 \(\hat{\mu}_{H, k} = \rho \hat{\mu}_{H, k-1} + (1-\rho) \frac{1}{M}\sum_{m=1}^M H_k^{(m)}\)（衰减率 \(\rho=0.99\)），从而在保证训练阶段控制变量有效生效的同时逼近无偏性。
损失函数 / 训练策略¶
本方法为策略梯度优化器的方差削减技术，其核心训练优化目标为无偏方差最优梯度估计器 \(g_{\mathrm{OHR}}\)： $\(g_{\mathrm{OHR}} = g_{\mathrm{PG}}^b + \alpha_1^* (H - \hat{\mu}_H)\)$ 优化器采用 AdamW，在每个时间步通过前向生成获得完整轨迹与标量奖励，在反向传播中并行累积得分梯度与历史雅可比向量积。针对离散自回归任务，设置温度参数 \(\tau \in [1.0, 3.0]\) 控制嵌入弛豫平滑度。
实验关键数据¶
主实验¶
在 T-Maze 长程走廊任务与旅行商问题（TSP）中，对比传统 REINFORCE、基线引入学得 Critic、以及 OHR 与 RLOO 结合的方差缩减表现：
| 评估任务配置 | 评估维度 / 指标 | 基线策略梯度 (PG) | 引入 Critic / RLOO | 本文 OHR | OHR 联合基线 (OHR + RLOO) |
|---|---|---|---|---|---|
| T-Maze (走廊长度 N=6) | 梯度方差缩减倍数 | \(1.00\times\) | \(1.54\times\) | \(1.13\times\) | \(1.85\times\) |
| T-Maze (走廊长度 N=15) | 梯度方差缩减倍数 | \(1.00\times\) | \(1.70\times\) | \(4.36\times\) | \(5.82\times\) |
| T-Maze (走廊长度 N=20) | 策略收敛平均奖励 | 0.75 (多随机种子停滞) | 0.81 | 0.90 (稳定收敛) | 0.93 |
| TSP 组合优化 (规模 n=8) | 梯度方差缩减倍数 | \(1.00\times\) | \(3.08\times\) (RLOO) | \(0.94\times\) | \(3.44\times\) |
| TSP 组合优化 (规模 n=10) | 梯度方差缩减倍数 | \(1.00\times\) | \(2.45\times\) (RLOO) | \(1.27\times\) | \(3.55\times\) |
| TSP 组合优化 (规模 n=15) | 梯度方差缩减倍数 | \(1.00\times\) | \(1.73\times\) (RLOO) | \(1.12\times\) | \(1.89\times\) |
消融实验¶
离散 Token 弛豫方式在不同序列长度 \(N\) 与温度 \(\tau\) 下的方差缩减倍数消融：
| 序列长度 \(N\) | 弛豫温度 \(\tau\) | 直通估计器 (ST-Gumbel-Softmax) | 连续嵌入弛豫 (Embedding Relaxation) | 机制说明 |
|---|---|---|---|---|
| \(N = 3\) | \(\tau = 1.0\) | \(1.00\times\) (无方差收益) | \(1.21\times\) | 连续嵌入开始提供有效导数 |
| \(N = 3\) | \(\tau = 3.0\) | \(1.01\times\) | \(1.50\times\) | 高温平滑扩展局部流形 |
| \(N = 8\) | \(\tau = 3.0\) | \(1.01\times\) | \(2.85\times\) | 方差缩减倍数随序列长度单调跃升 |
| \(N = 15\) | \(\tau = 3.0\) | \(1.01\times\) | \(4.90\times\) | 打破长程方差爆炸瓶颈 |
关键发现¶
- 连续嵌入弛豫与直通估计器的根本分野：Straight-Through Gumbel 虽然反向传梯度，但前向离散硬值阻断了条件概率对上游历史的平滑响应；而连续嵌入弛豫使长序列方差缩减达到 \(4.9\times\)。
- 与优势端基线的完美正交叠加：OHR 作用于得分函数端（\(\nabla_\theta \log p_\theta\)），与作用于优势值端（\(r-b\)）的 RLOO 展现了近乎乘积级别的正交协同收益（在 TSP-10 上从单独 2.45× 提升到 3.55×）。
亮点与洞察¶
- 打破环境无关假设的机械沿袭：深刻指出自回归序列生成的输入历史并非未知黑盒状态，其前缀处处可导，完全能够内生提供信用分配信号。
- 解析最优的无偏控制变量构筑：严格推导了最小化方差迹的封闭标量系数 \(\alpha_1^*\)，确保在数学上不劣于传统策略梯度。
局限与展望¶
- 反向传播计算与显存开销：计算历史路径项需沿展开轨迹进行反向传播（BPTT），在超长文本（数千 tokens）推理中必须依赖截断 BPTT。
- 未来方向：探索在大语言模型复杂长思维链推理任务中的大规模分布式适配。
相关工作与启发¶
- vs REINFORCE / PPO / GRPO: 传统方法机械冻结历史前缀，导致长序列方差爆炸；OHR 打通历史路径梯度，在保持标量奖励黑盒的前提下大幅压低方差。
- vs RLOO / Advantage Baselines: 现有方差缩减技术集中在奖励端调整；OHR 首创在得分函数端重参数化，二者可无缝叠加。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 颠覆性提出得分函数的历史重参数化框架与方差最优控制变量闭式解。
- 实验充分度: ⭐⭐⭐⭐☆ 理论推导极为严谨完备，涵盖解析模型、GRU、Transformer TSP，唯缺超大参数 LLM 实测。
- 写作质量: ⭐⭐⭐⭐⭐ 数学推导无懈可击，动机清晰，逻辑层层递进。
- 价值: ⭐⭐⭐⭐⭐ 为强化学习在自回归生成模型中的高效训练开辟了全新的算法设计空间。