大语言模型研究18——线性回归视角看强化学习策略梯度

作者: 引线小白-本文永久链接:https://www.limoncc.com/post/ef608130a2545b76/
知识共享许可协议: 本博客采用署名-非商业-禁止演绎4.0国际许可证

一、简单回顾

为了说明问题,代理目标函数暂时忽略clip、min这两个操作,,同时不采用原始的GRPO长度归一, $\frac{1}{|\bm{o}_i|}$ 会带来的长度偏差,正优势(答对)时短响应梯度更大 → 正确答案被压短;负优势(答错)时长响应被”摊薄”惩罚 → 错误响应被鼓励越写越长,这里采用DAPO的形式按总响应的总token数归一化。后面会发现如何做归一化影响甚大。

$$\begin{align}
\mathcal{L}(\bm{\theta})
=\mathbb{E}_{q\sim \mathcal{D},\bm{o}\sim \pi_{\bm{\theta}_{old} }}
\Bigg[
\frac{1}{\sum_{i=1}^G|\bm{o}_i|}
\sum_{i=1}^G\sum_{t=1}^{|\bm{o}_i|}
{\color{blue}\bigg[}
w_{i,t}\big(A_i
-\beta\mathbb{\hat{KL} }\big[\pi_{\bm{\theta} }\mid\mid\pi_{\bm{\theta}_{ref} }\big)\big]
{\color{blue}\bigg]}
\Bigg]
\end{align}$$

其中重要性比例 $\displaystyle w_{i,t}=\frac{\pi_{\bm{\theta} }(o_{i,t}\mid q,o_{i,:t})}{\pi_{\bm{\theta}_\text{old} }(o_{i,t}\mid q,o_{i,:t})}$ , $:t$ 含义是 $<t$ 。这里选 $\displaystyle k2=\frac{1}{2}\log^2\frac{\pi_{\bm{\theta}_{ref} }}{\pi_{\bm{\theta} }}$,同时梯度截断的KL约束:

$$\begin{align}
\mathcal{L}(\bm{\theta})
&=\underbrace{\mathbb{E}_{q\sim \mathcal{D} }
\frac{1}{\sum_{i=1}^G|\bm{o}_i|}\sum_{i=1}^G\sum_{t=1}^{|\bm{o}_i|}
\bigg[
\pi_{\bm{\theta}_{old} }\frac{\pi_{\bm{\theta} }}{\pi_{\bm{\theta}_\text{old} }}A_{i,t}\bigg]}_{\small\text{ 标准的强化学习目标} }
-\beta \cdot \underbrace{ \mathbb{E}_{q\sim \mathcal{D} }
\frac{1}{\sum_{i=1}^G|\bm{o}_i|}\sum_{i=1}^G\sum_{t=1}^{|\bm{o}_i|}
\bigg[\pi_{\bm{\theta}_{old} }\frac{ {\color{gray} \pi_{\bm{\theta} }} }{\pi_{\bm{\theta}_\text{old} }}\cdot\hat{\mathbb{KL} }\big[\pi_{\bm{\theta} } ||\pi_{\bm{\theta}_\textit{ref} }\big]\bigg]}_{\small\text{KL惩罚项} }\\
&=\underbrace{\mathbb{E}_{q\sim \mathcal{D} }
\frac{1}{\sum_{i=1}^G|\bm{o}_i|}\sum_{i=1}^G\sum_{t=1}^{|\bm{o}_i|}
\bigg[
\pi_{\bm{\theta} }\cdot A_{i,t}\bigg]}_{\small\text{ 标准的强化学习目标} }
-\beta \cdot \underbrace{ \mathbb{E}_{q\sim \mathcal{D} }
\frac{1}{\sum_{i=1}^G|\bm{o}_i|}\sum_{i=1}^G\sum_{t=1}^{|\bm{o}_i|}
\bigg[\mathrm{SG}(\pi_{\bm{\theta} })\cdot\hat{\mathbb{KL} }\big[\pi_{\bm{\theta} } ||\pi_{\bm{\theta}_\textit{ref} }\big]\bigg]}_{\small\text{KL惩罚项} }
\end{align}$$

这样有策略梯度

$$\begin{align}
\nabla_{\bm{\theta} }\mathcal{L}(\bm{\theta})
=\mathbb{E}_{q\sim \mathcal{D},\bm{o}\sim \pi_{\bm{\theta} }}
\Bigg[\frac{1}{\sum_{i=1}^G|\bm{o}_i|}\sum_{i=1}^G\sum_{t=1}^{|\bm{o}_i|}
\bigg[\nabla_{\bm{\theta} }\log\pi_{\bm{\theta} }\Big(
A_{i,t} -\beta\log\frac{\pi_{\bm{\theta} }}{\pi_{\bm{\theta}_{ref} }}
\Big)
\bigg]\Bigg]
\end{align}$$

为了节约符号令 $a_{i,t}=A_{i,t} -\beta\log\frac{\pi_{\bm{\theta} }}{\pi_{\bm{\theta}_{ref} }}$, 令得分函数 $\bm{s}_{i,t}=\nabla_{\bm{\theta} }\log\pi_{\bm{\theta} }(o_{i,t}\mid q,o_{i,:t})$,同时令 $c=\frac{1}{\sum_{i=1}^G|\bm{o}_i|}$。于是有:

$$\begin{align}
\nabla_{\bm{\theta} }\mathcal{L}(\bm{\theta})
=\mathbb{E}_{q\sim \mathcal{D},\bm{o}\sim \pi_{\bm{\theta} }}
\Bigg[
\sum_{i=1}^G\sum_{t=1}^{|\bm{o}_i|}
c\cdot a_{i,t}\cdot\bm{s}_{i,t}
\Bigg]
\end{align}$$

由于每个轨迹长度 $|\bm{o}_i|$是随机的。为了把两个求和符号提取出来可以令 $\displaystyle T_{max}=\max_{i \in [1,G]}|\bm{o}_i|$ 和同时使用示性函数 $\mathbb{I}\big[t\leqslant |\bm{o}_i|\big]$ 来补零,这样有:

$$\begin{align}
\nabla_{\bm{\theta} }\mathcal{L}(\bm{\theta})
=\sum_{i=1}^G\sum_{t=1}^{T_{max} }\mathbb{E}_{q\sim \mathcal{D},\bm{o}\sim \pi_{\bm{\theta} }}
\Bigg[
\mathbb{I}\big[t\leqslant |\bm{o}_i|\big]\cdot c\cdot a_{i,t}\cdot\bm{s}_{i,t}
\Bigg]
\end{align}$$

二、线性回归视角看强化学习

2.1、几个基本结论

1、回顾一下协方差公式 $\mathrm{Cov}(X,Y)=\mathbb{E}[XY]-\mathbb{E}[X]\mathbb{E}[Y]$,实际有

$$\begin{align}
\mathbb{E}[XY]=\underbrace{\mathbb{E}[X]\mathbb{E}[Y]}_{\text{背景噪声} }+\underbrace{\mathrm{Cov}(X,Y)}_{\text{线性关系} }
\end{align}$$

2、其中得分函数在理想情况下( 排除训推不一致问题[^1]):

$$\begin{align}
\mathbb{E}_{o_{i,t}\sim\pi_{\bm{\theta} }}\big[\bm{s}_{i,t}\,\big|\,o_{i,:t}\big]
&=\sum_{o}\pi_{\bm{\theta} }(o\mid q,o_{i,:t})\,\nabla_{\bm{\theta} }\log\pi_{\bm{\theta} }(o\mid q,o_{i,:t})\\
&=\sum_{o}\nabla_{\bm{\theta} }\pi_{\bm{\theta} }(o\mid q,o_{i,:t})\\
&=\nabla_{\bm{\theta} }\underbrace{\sum_{o}\pi_{\bm{\theta} }(o\mid q,o_{i,:t})}_{=1}=0
\end{align}$$

3、另外不同的强化学习算法归一项 c 的定义是不同的。关键区别是 c 是否依赖 $q,o_{i,:t}$,例如 Dr. GRPO $\mathbb{E}[c\mid q,o_{i,:t}]=\mathbb{E}[c]=\frac{1}{G}$, 而GRPO则是 $\mathbb{E}[c\mid q,o_{i,:t}]=\mathbb{E}\left[\frac{1}{G|\bm{o}_i|}\mid q,o_{i,:t}\right]$。不依赖则可以作为参数提取出来。依赖则加入优势, 对优势塑形。

4、得分函数恒等式

在策略梯度和变分推断里面有一个很重要的技巧是得分函数恒等式,也叫似然比恒等式。

$$\begin{align}
\mathbb{E}_{\pi_{\bm{\theta} }}\left[g(o)\bm{s}\right]=\nabla_{\bm{\theta} }\mathbb{E}_{\pi}\left[g(o)\right]
\end{align}$$

对于策略 $\pi_{\bm{\theta} }(o)$, 样本函数 $g(o)$ 可以是任何依赖于数据的函数,它的职责是定义我们真正想要优化的目标。统计学中得分函数定义为 $\bm{s}(o) = \nabla_{\bm{\theta} } \log \pi_{\bm{\theta} }(o) $, 有:

$$\begin{align}
\mathbb{E}_{\pi_{\bm{\theta} }}[ g(o)\cdot \bm{s}]
&= \int g(o)\pi_{\bm{\theta} }(o) \nabla_{\bm{\theta} } \log \pi_{\bm{\theta} }(o)\mathrm{d}o
= \int \nabla_{\bm{\theta} } g(o)\pi_{\bm{\theta} }(o) \mathrm{d}o\\
&= \nabla_{\bm{\theta} } \int \pi_{\bm{\theta} }(o) g(o) \, \mathrm{d}o = \nabla_{\bm{\theta} }\mathbb{E}_\pi[g(o)]
\end{align}$$

叫似然比恒等式的核心原因是:恒等式里的权重 $s(o)=\nabla_{\bm{\theta} } \log \pi_{\bm{\theta} }(o)$ 本质上是无穷小似然比。它不是两个固定分布之间的普通似然比,而是参数发生微小变动时,新旧分布似然比的一阶展开系数。设有参数化分布 $\pi_{\bm{\theta} }(x)$。两个参数 $\bm{\theta}$ 和 $\bm{\theta}+\bm{\delta}$ 之间的似然比为

$$\begin{align}
L(x)=\frac{\pi_{\bm{\theta}+\bm{\delta} }(x)}{\pi_\bm{\theta}(x)}
\to \log L(x)=\log \pi_{\bm{\theta}+\bm{\delta} }(x)-\log \pi_\bm{\theta}(x)
\end{align}$$

当 $|\bm{\delta}|\to 0$ 时,一阶泰勒展开:

$$\begin{align}
\log L(x)\approx \bm{\delta}^\T\nabla_\bm{\delta} \log \pi_\bm{\delta}(x)
\end{align}$$

所以

$$\begin{align}
\nabla_\bm{\delta} \log \pi_\bm{\delta}(x)
=\lim_{|\bm{\delta}|\to 0}\frac{\log \pi_{\bm{\theta}+\bm{\delta} }(x)-\log \pi_\bm{\theta}(x)}{|\bm{\delta}|}
\end{align}$$
也就是说,得分函数 $\bm{s}(x)=\nabla_\bm{\theta} \log \pi_\bm{\theta}(x)$ 正是对数似然比关于参数的一阶灵敏度。因此它被称为无穷小对数似然比或局部似然比。

2.2、策略梯度参数邻域与曲率

设 $\bm{\delta}$ 是参数的小扰动,$\pi_{\bm{\delta} }(o),\pi_{\bm{\theta}+\bm{\delta} }(o)$,并记得分函数 $\bm{s}(o)=\nabla_{\bm{\theta} }\log\pi_{\bm{\theta} }(o)$。KL 对 $\bm{\theta}$ 求二阶泰勒展开:

$$\begin{align}
\mathbb{KL}\big(\pi_\bm{\theta}\mid\mid\pi_{\bm{\theta}+\bm{\delta} }\big)
=\underbrace{\mathbb{E}_\bm{\theta}\Big[\log\frac{\pi_\bm{\theta} }{\pi_{\bm{\theta}+\bm{\delta} }}\Big]}_{\small\text{一阶项} }
+{\small\text{高阶项} }
\end{align}$$

把 $\log\pi_{\bm{\theta}+\bm{\delta} }(o)$ 在 $\bm{\theta}$ 处展开到二阶:

$$\log\pi_{\bm{\theta}+\bm{\delta} }(o)
=\log\pi_\bm{\theta}(o)
+\bm{\delta}^\T\bm s(o)
+\frac{1}{2}\,\bm{\delta}^\T\nabla^2_{\theta}\log\pi_\bm{\theta}(o)\,\bm{\delta}+O(\mid\mid\bm{\delta}\mid\mid^3)$$

代回:
$$\begin{align}
\mathrm{KL}\big(\pi_\bm{\theta}\,\mid\mid\,\pi_{\bm{\theta}+\bm{\delta} }\big)
=\mathbb{E}_{\pi_\bm{\theta} }\Big[\log\frac{\pi_\bm{\theta} }{\pi_{\bm{\theta}+\bm{\delta} }}\Big]
=-\mathbb{E}_{\pi_\bm{\theta} }\big[\bm{\delta}^\T\bm s\big]+O(\mid\mid\bm{\delta}\mid\mid^2)
=-\bm{\delta}^\T\underbrace{\mathbb{E}_{\pi_\bm{\theta} }[\bm s]}_{= \bm 0}
-\mathbb{E}_{\pi_\bm{\theta} }\Big[\frac12\bm{\delta}^\T\nabla^2\log\pi_\bm{\theta}\,\bm{\delta}\Big]
+O(\mid\mid\bm{\delta}\mid\mid^3)
\end{align}$$

KL 在 $\bm{\delta}=\bm 0$ 处取极小值(值 0),极小值点处一阶导必为零,而极小性本身又由 KL 的非负性(Gibbs 不等式)保证。KL 没有线性项,是距离这类量的定义性特征。麻烦在第二项里的 $\mathbb{E}[\nabla^2\log\pi]$,它不为零。用恒等式 $\nabla^2\log\pi=\dfrac{\nabla^2\pi}{\pi}-\bm s\bm s^\T$(对 $\log\pi$ 求两次导,商法则)拆开:
$$\begin{align}
\mathbb{E}_{\pi_\bm{\theta} }\big[\nabla^2\log\pi_\bm{\theta}\big]
=\mathbb{E}_{\pi_\bm{\theta} }\Big[\frac{\nabla^2\pi_\bm{\theta} }{\pi_\bm{\theta} }\Big]-\mathbb{E}_{\pi_\bm{\theta} }\big[\bm s\bm s^\T\big]
=\int\nabla^2\pi_\bm{\theta}\,\mathrm{d}o- \bm{F}
=\nabla^2_\theta\underbrace{\int\pi_\bm{\theta}\,\mathrm{d}o}_{\equiv 1}- \bm{F}
=\bm 0-\bm{F}=-\bm{F}
\end{align}$$

两个非零的丑东西 $\int\nabla^2\pi$ (测度归一化的二阶变化)与 $\mathbb{E}[\bm s\bm s^\T]$(得分的二阶矩)精确抵消,只因 $\int\pi=1$。 代回:

$$\begin{align}
\mathrm{KL}\big(\pi_\bm{\theta}\,\mid\mid\,\pi_{\bm{\theta}+\bm{\delta} }\big)
=-\bm{\delta}^\T\underbrace{\mathbb{E}[\bm s]}_{\bm 0}
-\frac12\bm{\delta}^\T\big(-\bm{F}\big)\bm{\delta}
+O(\mid\mid\bm{\delta}\mid\mid^3)
=\frac12\,\bm{\delta}^\T\bm{F}\,\bm{\delta}+O(\mid\mid\bm{\delta}\mid\mid^3)
\end{align}$$

$$\begin{align}
\mathrm{KL}\big(\pi_\bm{\theta}\,\mid\mid\,\pi_{\bm{\theta}+\bm{\delta} }\big)
=\frac12\,\bm{\delta}^\T\bm{F}\,\bm{\delta}+O(\mid\mid\bm{\delta}\mid\mid^3)
\end{align}$$

其中 $\bm{F}=\mathbb{E}_{\pi_\bm{\theta} }\big[\bm s\,\bm s^\T\big]$,全部高阶信息被压进 $O(\mid\mid\bm{\delta}\mid\mid^3)$。

2.3、归一化外生情形

归一化参数 c 条件于 $q,o_{i,:t}$ 是常数, 也是外生变量。那么策略梯度可写为 $\nabla_{\bm{\theta} }\mathcal{L}(\bm{\theta})
=c\cdot \sum_{i,t}\mathbb{E}_{q\sim \mathcal{D},\bm{o}\sim \pi_{\bm{\theta} }}
\left[
a_{i,t}\cdot\bm{s}_{i,t}
\right]$ ,为节约符号直接写成:

$$\begin{align}
\nabla_{\bm{\theta} }\mathcal{L}(\bm{\theta})
=c\cdot \sum_{i,t}\mathbb{E}_{\pi_{\bm{\theta} }}
\left[
a\cdot\bm{s}
\right]
\end{align}$$

那么有
$$\begin{align}
\nabla_{\bm{\theta} }\mathcal{L}(\bm{\theta})
&=c\cdot \sum_{i,t}\mathbb{E}_{\pi_{\bm{\theta} }}
\left[a\cdot\bm{s}\right]\\
&=c\sum_{i,t}\bigg[
\underbrace{\mathbb{E}_{\pi_{\bm{\theta} }}\left[a\right]
\mathbb{E}_{\pi_{\bm{\theta} }}\left[\bm{s}\right]}_{\small 漂移项}
+\underbrace{\mathrm{Cov}[a,\bm{s}]}_{\small线性相关项}
\bigg]\\
&=c\sum_{i,t}\mathrm{Cov}[a,\bm{s}]
\end{align}$$

也就是有学习信号 = 得分函数与有效优势的协方差,准确的说其实是

$$\begin{align}
\nabla_{\bm{\theta} }\mathcal{L}(\bm{\theta})=c\sum_{i,t}\mathrm{Cov}[a_{i,t},\bm{s}_{i,t}\mid q,o_{i,:t}]
\end{align}$$

把协方差读成回归,用得分函数当特征去拟合有效优势 $a=\bm{w}^\T \bm{s}+\epsilon$

$$\begin{align}
\bm{w}^*=\arg\min_{\bm{w} }\mathbb{E}\big[(a-\bm{w}^\T\bm{s})^2\big]
\end{align}$$

最小二乘目标函数

$$\begin{align}
J(\bm{w})
&=\mathbb{E}\big[(a-\bm{w}^\T\bm{s})^2\big]\\
&=\mathbb{E}\big[a^2-2a\bm{w}^\T\bm{s}+\bm{w}^\T\bm{s}\bm{s}^\T\bm{w}\big]\\
&=\mathbb{E}[a^2]-2\bm{w}^\T\mathbb{E}[a\cdot\bm{s}]+ \bm{w}^\T\mathbb{E}[\bm{s}\bm{s}^\T]\bm{w}
\end{align}$$

记 Fisher 信息 $\bm{F}=\mathbb{E}[\bm{s}\bm{s}^\T]$, 则有 $J(\bm{w})=\mathbb{E}[a^2]-2\bm{w}^\T\mathbb{E}[a\cdot\bm{s}]+ \bm{w}^\T\bm{F}\bm{w}$ 求梯度

$$\begin{align}
\nabla_{\bm{w} }J(\bm{w})=-2\mathbb{E}[a\cdot\bm{s}]+2\bm{F}\bm{w}=\bm{0}
\end{align}$$

有

$$\begin{align}
\bm{w}^*
= \bm{F}^{-1}\mathbb{E}[a\cdot\bm{s}]
= \bm{F}^{-1}\mathrm{Cov}[a,\bm{s}]
\xrightarrow{\bm F=\bm{E} }\bm{E}^{-1}\mathrm{Cov}[a,\bm{s}]=\mathrm{Cov}[a,\bm{s}]
\end{align}$$

亦有

$$\begin{align}
\nabla_{\bm{\theta} }\mathcal{L}(\bm{\theta})
=c\sum_{i,t}\mathrm{Cov}[a,\bm{s}]
=c\sum_{i,t}\bm{F}\bm{w}^*
\end{align}$$

得分函数零期望导致背景噪声恒为零,这样策略梯度是纯协方差。它不含任何背景噪声,只度量得分方向与调整后优势之间的线性关联。这正是 Sutton 相容函数逼近定理:以 $a({\bm{\theta} })=\bm{w}^\T\nabla_{\bm{\theta} }\log\pi_{\bm{\theta} }$ 线性拟合优势,最小二乘解恰为策略梯度。两种视角在此合流:策略学习 = 以得分函数为特征、以有效优势为目标的线性回归。投影残差 $a-\bm{w}^{*\T}\bm{s}$ 是当前参数化下学不到的部分——无论怎么采样都不会产生梯度。

2.4、曲率的影响
2.4.1、朴素策略梯度

线性回归视角下,朴素策略梯度等于默认了特征已归一化且互不串扰($\mathbb E[\bm{s}\bm{s}^\T]=\bm{E}$), 即需要假设Fisher信息矩阵是单位矩阵,即策略参数得分函数需要完成白化($\bm{F}^{-1/2}$),每个得分方向的波动幅度相同、彼此不相关。实质上是:跳过对特征二阶矩的估计,把除以各方向波动这笔账省了。但省下的账不会消失——它变成瘦方向欠学、肥方向过冲。下面来详细分析:

$\bm{s}_{i,t}$ 是动作 $o_{i,t}$ 通过策略 $\pi_\bm{\theta}$在参数 $\bm{\theta}$ 处的一阶灵敏度特征。它回答的是参数往这个方向微动一步,这个具体动作的概率会怎么变,不含任何好坏评价,纯粹是 $\pi_\bm{\theta}$ 微分结构在 $o_{i,t}$ 处的取值。 $a_{i,t}$ 是回归目标。是环境(价值估计)给这个动作 $o_{i,t}$ 打的分。实际上有 $N=\sum_{i=1}^G|\bm{o}_i|$ 个单点数据集 $\displaystyle \mathcal{D}_{i,t}=\{\bm{s}_{i,t},o_{i,t}\}_{i\times t}$ , 也就是说每个 $(i,t)$ 只有一个采样点,条件于自己的前缀 $o_{i,:t}$。这样协方差计算就变成单样本一阶矩(这是建立在漂移项为零的基础上的)。策略梯度就是 $N$个回归之和,每个单样本数据集贡献一个无偏投票。实际上朴素策略梯度是:

$$\begin{align}
\widehat{\nabla_{\bm{\theta} }\mathcal{L}(\bm{\theta})}
=c\sum_{i,t}a_{i,t}\cdot\bm{s}_{i,t}
\end{align}$$

单样本虽然是无偏的,但是方差是极高的。这与假设的几何结构 $\bm{F}=\bm{E}$ :Fisher信息矩阵是单位矩阵是相矛盾的。由2.2节知:

$$\begin{align}
\mathrm{KL}\big(\pi_\bm{\theta}\,\mid\mid\,\pi_{\bm{\theta}+\bm{\delta} }\big)
\approx \frac{1}{2}\bm{\delta}^\T\bm{F}\,\bm{\delta}=\frac{1}{2}||\bm{\delta}||^2
\end{align}$$

单位矩阵曲率实际上做了两个重要假设:

  • 正交性: 曲率的非对角元为零。往提高动作 A 概率的方向走一步,不会扰动动作 B 的概率。各动作的得分方向彼此正交。
  • 等距性:对角元全为1。在任何参数方向上走同样的 $||\bm{\delta}||$,分布在 KL 意义下挪动同样远。

这样单位矩阵曲率使得策略参数空间的一步等于分布空间的一步、各方向互不串扰且是等距的,参数坐标没有肥瘦之分。而真实策略的两个条件几乎从不满足。拿最简单的 categorical-softmax 看,$\bm{F}=\mathrm{diag}(\bm{\pi})-\bm{\pi}\bm{\pi}^\T$:概率小的动作对应的方向极瘦($F_{ii}=\pi_i(1-\pi_i)\approx \pi_i$ 很小,参数要挪很远才能在分布上挪一点),概率大的方向极肥,且各方向经 $\bm{\pi}\bm{\pi}^\T$ 纠缠。参数空间是各向异性的。下面来详细说明。

2.4.2、一般策略梯度

设 logits 为 $\mu\in\mathbb R^K$,分类分布给出 $\pi_i=\frac{e^{\mu_i} }{\sum_{j=1}^K e^{\mu_j} }$ , 其中 $i=1,\dots,K$。这里参数就是 logits $\mu$。考虑沿第 $i$ 个 logit 方向走一小步 $\mu’ = \mu+t e_i$, 则分布变化为

$$\begin{align}
\pi’_i=\pi_i+\pi_i(1-\pi_i)t+O(t^2),
\end{align}$$

而局部 KL 散度为

$$\begin{align}
D_{\mathrm{KL} }(\pi\mid\mid\pi’)
=
\frac12 t^2 F_{ii}+O(t^3)
=
\frac12 t^2 \pi_i(1-\pi_i)+O(t^3).
\end{align}$$

因此,如果 $\pi_i\ll 1$,则 $F_{ii}\approx \pi_i, \mathbb{KL}(\pi\mid\mid\pi’)\approx \frac12 t^2 \pi_i$, 若想产生固定的分布变化,比如 $\mathbb{KL}(\pi\mid\mid\pi’)=\varepsilon$ 则需要

$$\begin{align}
|t|\approx \sqrt{\frac{2\varepsilon}{\pi_i(1-\pi_i)} }
\sim \sqrt{\frac{2\varepsilon}{\pi_i} },
\qquad \pi_i\to 0.
\end{align}$$

也就是说,当动作 $i$ 的概率很小时,沿它的 logit 方向,参数要移动 $O\left(\frac{1}{\sqrt{\pi_i} }\right)$ 这么大的距离,才能在分布层面产生有限变化。因此该方向在 Fisher 度量下“极瘦”。等价地,看概率本身的一阶变化:

$$\begin{align}
\Delta \pi_i:=\pi’_i-\pi_i
\approx \pi_i(1-\pi_i)t
\approx \pi_i t
\end{align}$$

所以若要把动作 $i$ 的概率改变 $\Delta \pi_i$,大致需要 $t\approx \frac{\Delta \pi_i}{\pi_i}$ 。当 $\pi_i\to 0$ 时,这个步长发散。这就是“概率小的动作对应的方向极瘦,参数要挪很远才能在分布上挪一点”。

反过来,若 $\pi_i$ 不是非常接近 $1$,则 $F_{ii}=\pi_i(1-\pi_i)$ , 相对于小概率动作会大很多。例如二动作情形 $K=2$,非零 Fisher 曲率为 $\pi(1-\pi)$。当 $\pi$ 从很小增加到 $1/2$ 时,曲率增大,说明大概率方向在 Fisher 度量下更“肥”,同样参数步长会造成更大的分布变化。但严格地说,$F_{ii}=\pi_i(1-\pi_i)$ 在 $\pi_i=1/2$ 最大;当 $\pi_i\to 1$ 时, $F_{ii}\approx 1-\pi_i$ 也会变小。因此“大概率方向极肥”应理解为:相对于极小概率方向,在非饱和区域,大概率方向曲率更大。完全接近确定性的动作,其 logit 方向也会变瘦,因为 softmax 对 logits 的局部敏感性降低。

另外,$F=\operatorname{diag}(\pi)-\pi\pi^\T$ 有一个零特征方向 $\mathbf 1$,因为给所有 logits 同时加常数不改变 softmax 分布。因此严格讨论方向时应在商空间或与 $\mathbf 1$ 正交的子空间中进行。

小概率动作 $i$ 满足 $F_{ii}=\pi_i(1-\pi_i)\approx \pi_i\ll 1$ 所以沿该 logit 方向,局部 KL 曲率约为 $\pi_i$,要获得固定分布变化需要参数步长 $O(\pi_i^{-1/2})$,故该方向极瘦。大概率方向在非饱和区相对更肥;但严格最大曲率在 $\pi_i=1/2$,且 $\pi_i\to1$ 时也会变瘦。就是说朴素的策略梯度会在小概率方向放大不足,所以小概率动作学得很慢,出现探索不足的情况。而在大概率方向会过度学习,在几个高概率动作之间来回跳,导致震荡和熵坍缩。

对于更一般情形,设 $\bm{F}$ 的特征分解为 $\sum_k\lambda_k\bm{v}_k\bm{v}_k^\T$,则朴素方向在特征基下是 $\sum_k\bm{v}_k\bm{v}_k^\T\bm{w}^\ast$,而自然方向是 $\sum_k\frac{1}{\lambda_k}\bm{v}_k\bm{v}_k^\T\bm{w}^\ast$:

  • 瘦方向($\lambda_k\ll1$,对应小概率动作/罕见行为):自然梯度把该方向放大 $1/\lambda_k$ 倍——因为在这个方向上“参数走很多、分布挪很少”,需要迈大步;朴素梯度原封不动,于是稀有动作的学习被系统性欠速。这与探索不足互相加强:探索少 → 该方向 $\lambda$ 更小 → 更学不动。
  • 肥方向($\lambda_k\gg1$,大概率动作/主导模式):自然梯度把该方向压回 $1/\lambda_k$;朴素梯度照单全收,于是优势信号在主导方向上过冲——这正是训练后期策略在几个高概率模式间震荡的一个来源。

三、深入理解

3.1、优化器的影响
3.1.1、从曲率视角看优化器引言

在深度强化学习中,代理目标的梯度由各类优化器更新,而优化器对策略参数的曲率信息做了不同假设。回顾第 2.3–2.4 节的核心结论:真实曲率是 Fisher 信息矩阵

$$\begin{align}
\bm F=\mathbb E_{\pi_{\bm\theta} }\big[\bm s\,\bm s^\T\big]
\end{align}$$

自然梯度方向为 $\bm F^{-1}\bm g$。任何优化器都可以看成对 $\bm F$ 做某种近似或替换,从而在参数空间中隐式选定一个约束集与一把尺子。一个关键概念是正交不变性:量 $f$ 称为正交不变的,如果

$$\begin{align}
f(\bm U\bm W\bm V^\T)=f(\bm W)
\qquad
\forall\,\bm U\in O(m), \bm V\in O(n).
\end{align}$$

即不管怎么旋转输入、输出空间的基,这个量都不变。权重初始化从某种分布采样,不同随机种子给出不同初始基。若优化器不是正交不变的,训练结果就部分依赖于初始基的任意选择,而非只依赖初始函数的性质。 正交不变性消除了这种任意性。所有优化器的参数更新共享同一骨架:

$$\begin{align}
\bm\theta_{k+1}=(1-\eta\lambda)\,\bm\theta_k-\eta\,\bm U_k
\end{align}$$

差异只在三点:

  1. 更新方向 $\bm U_k$ 如何从梯度 $\bm g_k$ 构造;
  2. 该构造隐式对应哪个约束集;
  3. 约束集是否正交不变。

下面逐一对 SGD、AdamW、Muon 展开,并在每处标出它对应的有效曲率 $\bm F$。

3.1.2、SGD:朴素策略梯度的化身

SGD 对应朴素策略梯度:

$$\begin{align}
\nabla_{\bm\theta}\mathcal L^{\mathrm{SGD} }
=c\sum_{i,t}a_{i,t}\,\bm s_{i,t}
\end{align}$$

$\bm s_{i,t}$ 只在这里出现一次——这是恒等式 $c\sum_{i,t}\mathrm{Cov}(a,\bm s)$ 的无偏单样本估计,零修饰。

更新方向:
$$\begin{align}
\bm U_k^{\mathrm{SGD} }=\bm g_k
\end{align}$$

参数更新:
$$\begin{align}
\bm\theta_{k+1}=\bm\theta_k-\eta\,\bm g_k
\end{align}$$

隐式约束集(带权重衰减时):$\{\mid\mid\bm\theta\mid\mid_2\le R\}$

有效曲率:
$$\begin{align}
\bm{F}=\bm{E}
\end{align}$$

即 SGD 隐式假设 Fisher 信息是单位矩阵——把参数空间当作欧氏空间。

几何性质:
1、曲率校正 $=\bm{E}$,所有坐标等权;
2、正交不变(Frobenius 球是正交不变的);
3、但对最坏奇异方向无感知——不区分”某个方向增益大”和”增益均匀”。

3.1.3、AdamW:坐标加权 $\ell_2$ 球上的自适应下降

更新方向:
$$\begin{align}
\bm U_k^{\mathrm{AdamW} }
=\frac{\hat{\bm m}_k}{\sqrt{\hat{\bm v}_k}+\epsilon}
\end{align}$$

其中
$$\begin{align}
\bm m_k=\beta_1\bm m_{k-1}+(1-\beta_1)\bm g_k,
\qquad
\bm v_k=\beta_2\bm v_{k-1}+(1-\beta_2)\bm g_k^{\odot2},
\end{align}$$
$$\begin{align}
\hat{\bm m}_k=\frac{\bm m_k}{1-\beta_1^k},
\qquad
\hat{\bm v}_k=\frac{\bm v_k}{1-\beta_2^k}.
\end{align}$$

参数更新(含解耦权重衰减):
$$\begin{align}
\bm\theta_{k+1}=(1-\eta\lambda)\,\bm\theta_k
-\eta\,\frac{\hat{\bm m}_k}{\sqrt{\hat{\bm v}_k}+\epsilon}
\end{align}$$

隐式约束集: $\Big\{\sqrt{\sum_j d_j\,\theta_j^2}\le R\Big\}$, 其中 $d_j=\hat v_j^{1/2}$

有效梯度与曲率:
$$\begin{align}
\nabla_{\bm\theta}\mathcal L^{\mathrm{AdamW} }
\approx\mathrm{diag}(\bm v_k)^{-1/2}\,\mathbb{EMA}\Big[c\sum_{i,t}a_{i,t}\bm s_{i,t}\Big]+\lambda\bm\theta_k
\end{align}$$

$$\begin{align}
\bm{F}=\mathrm{diag}(\bm v_k)^{1/2}
\end{align}$$

这里需注意:有效曲率是 $\mathrm{diag}(\bm v_k)^{1/2}$,因为曲率校正是它的逆 $\mathrm{diag}(\bm v_k)^{-1/2}$。

几何性质:
1、曲率校正 $=\mathrm{diag}(\hat{\bm v}_k)^{-1/2}$,逐坐标不同权;
2、正交不变没有——权重 $d_j$ 依赖坐标基,神经元重编号后 $d_j$ 全变;
3、粒度比 SGD 细(逐坐标),但代价是放弃了基无关性;
4、处在”比 SGD 更细、但比 Muon 更远离函数空间几何”的中间位置。

3.1.4、Muon:谱范数球上的正交化下降

把参数按层切开,对矩阵层 $\bm W^{(\ell)}\in\mathbb R^{m\times n}$:

动量:
$$\begin{align}
\bm M_k^{(\ell)}=\mu\,\bm M_{k-1}^{(\ell)}+(1-\mu)\,\bm g_k^{(\ell)}
\end{align}$$

更新方向:
$$\begin{align}
\bm U_k^{\mathrm{Muon},\,(\ell)}
=\mathrm{Ortho}!\big(\bm M_k^{(\ell)}\big)
\end{align}$$

其中 $\mathrm{Ortho}(\bm M)=\bm U\bm V^\T$(若 $\bm M=\bm U\bm\Sigma\bm V^\T$),即把奇异值全部抹平为 1。

参数更新(含解耦权重衰减):
$$\begin{align}
\bm W_{k+1}^{(\ell)}
=(1-\eta\lambda)\,\bm W_k^{(\ell)}
-\eta\,\mathrm{Ortho}!\big(\bm M_k^{(\ell)}\big)
\end{align}$$

隐式约束集:
$$\begin{align}
\{\mid\mid\bm W^{(\ell)}\mid\mid_{\mathrm{sp} }\le R\}
\end{align}$$

有效曲率:由左极因子视角
$$\begin{align}
\mathrm{Ortho}(\bm M)=\bm U\bm\Sigma^{-1}\bm U^\T\cdot\bm M
\end{align}$$

曲率校正为 $\bm U\bm\Sigma^{-1}\bm U^\T$,故有效曲率为

$$\begin{align}
\bm F=\bm U\bm\Sigma\bm U^\T
\end{align}$$

这是对 Fisher 曲率做谱域近似——用动量的奇异向量作为曲率特征方向的代理,用奇异值作为曲率谱的代理,并把谱压平为等距。

几何性质:
1、曲率校正 $=\bm U\bm\Sigma^{-1}\bm U^\T$(左极因子视角),作用在奇异方向上;
2、正交不变——$\mid\mid\bm W^{(\ell)}\mid\mid_{\mathrm{sp} }$ 是内禀量,不随基旋转而变;
3、粒度最粗(逐矩阵、逐奇异方向),但恰好对准泛化界里出现的量。

3.1.5、三者并排对比

三者共享同一个更新骨架
$$\begin{align}
\bm\theta_{k+1}=(1-\eta\lambda)\bm\theta_k-\eta\bm U_k,
\end{align}$$
区别只在 $\bm U_k$:

优化器 更新方向 $\bm U_k$ 曲率校正 有效曲率 $\bm F$ 约束集 正交不变
SGD $\bm g_k$ $\bm E$ $\bm E$ 欧氏球 √
AdamW $\dfrac{\hat{\bm m}_k}{\sqrt{\hat{\bm v}_k}+\epsilon}$ $\mathrm{diag}(\hat{\bm v}_k)^{-1/2}$ $\mathrm{diag}(\bm v_k)^{1/2}$ 坐标加权 $\ell_2$ 球 ×
Muon $\mathrm{Ortho}(\bm M_k)$ $\bm U\bm\Sigma^{-1}\bm U^\T$ $\bm U\bm\Sigma\bm U^\T$ 谱范数球 √

曲率视角的统一叙述:

  • SGD 假设 $\bm F=\bm E$:把参数空间当欧氏空间,各向同性;
  • AdamW 用 $\mathrm{diag}(\bm v_k)^{1/2}$ 近似 $\bm F$:只取对角,丢弃非对角耦合,破坏正交不变性;
  • Muon 用 $\bm U\bm\Sigma\bm U^\T$ 近似 $\bm F$:在谱域做等距化,恢复正交不变性。

从 SGD 到 Muon 的演进:

  • 预条件粒度从”逐坐标”走向”逐奇异方向”;
  • 约束集从坐标层面回到算子层面;
  • 曲率近似从”各向同性”经”对角截断”最终到达”谱域等距”;
  • 越来越贴近函数空间的真实几何。
3.2、归一化内生情形

回到策略梯度,归一化参数 $c$通常是内生的,由全组回答的长度决定,而长度是随机变量。这就不能移出期望外。
$$\begin{align}
\nabla_{\bm{\theta} }\mathcal{L}(\bm{\theta})
=\sum_{i=1}^G\sum_{t=1}^{T_{max} }\mathbb{E}_{q\sim \mathcal{D},\bm{o}\sim \pi_{\bm{\theta} }}
\Bigg[
\mathbb{I}\big[t\leqslant |\bm{o}_i|\big]\cdot c\cdot a_{i,t}\cdot\bm{s}_{i,t}
\Bigg]
\end{align}$$

归一化内生不改变得分函数期望为零:$\mathbb{E}\big[\bm{s}_{i,t}\,\big|\,q,\bm{o}_{i,:t}\big]
=\sum_{o}\nabla_{\bm{\theta} }\pi_{\bm{\theta} }(o\mid q,\bm{o}_{i,:t})=\nabla_{\bm{\theta} }1=\bm{0}$

$$\begin{align}
\nabla_{\bm{\theta} }\mathcal{L}(\bm{\theta})
=\sum_{i=1}^G\sum_{t=1}^{T_{max} }\mathbb{E}_{\pi_{\bm{\theta} }}
\mathrm{Cov}\big(ca_{i,t},\bm{s}_{i,t}\big)
=\underbrace{\mathbb{E}[c]\sum_{i,t}\mathrm{Cov}\big(\bm{s}_{i,t},a_{i,t}\big)}_{\text{理想协方差信号} }
+\underbrace{\mathrm{Cov}\Big(c, \sum_{i,t}\bm{s}_{i,t}a_{i,t}\Big)}_{\text{内生归一化偏差} }
\end{align}$$

令 $a_{i,t}\equiv a_0$ 为常数(不含任何信息),理想协方差故事预言梯度为零,但造成2个严重后果

$$\begin{align}
\nabla_{\bm{\theta} }\mathcal{L}=a_0\,\mathrm{Cov}\big(\bm{s}_{i,t},\,c\big)\neq\bm{0}
\end{align}$$

1、只要回答长度与 token 内容相关(事实如此:冗长与对错强相关),梯度就朝“压缩或拉伸期望总长”的方向泄漏。数据依赖的归一化自带一条与内容无关的推力。

2、还有一个容易被忽视的后果:baseline 平移不变性被 $c$ 破坏了。 往 $a$ 里减常数 $b$,$\mathbb{E}[c\,\bm{s}\,(a-b)]=\mathbb{E}[c\,\bm{s}\,a]-b\,\mathbb{E}[c\,\bm{s}]$,而 $\mathbb{E}[c\bm{s}]\neq\bm{0}$——于是加基线不改梯度这条策略梯度的经典技巧失效了。GRPO 的组内去均值与长度归一化两个组件,不再可分解地各自引入偏差。

3、考虑 $c$ 之后,回归升级为加权最小二乘:
$$\begin{align}
\bm{w}^{*}=\mathbb{E}\big[c\,\bm{s}\,\bm{s}^{\T}\big]^{-1}\,\mathbb{E}\big[c\,\bm{s}\,a\big]
=\bm{F}_c^{-1}\,\nabla_{\bm{\theta} }\mathcal{L},\qquad
\bm{F}_c=\mathbb{E}\big[c\,\bm{s}\bm{s}^{\T}\big]
\end{align}$$

$c$ 内生时 $\bm{F}_c$ 是长度混合下的加权二阶矩,不再是任何单一分布的 Fisher 信息,策略梯度的白化解释随之弱化。更本质的是,权重本身是被优化的对象的一部分——这是内生加权回归与普通 WLS 的分界线。

不同的归一化方案,对应不同的“加权哲学”:

归一化 有效权重 加权含义 内生性后果
$\frac{1}{\mid\mid\bm{o}_i\mid\mid}$(GRPO) $\frac{1}{G\mid\mid\bm{o}_i\mid\mid}$ 回答级等权 序列级长度偏置 $O(1)$:长回答总贡献被压,Dr.GRPO 指出的长度膨胀
$\frac{1}{\sum_{i=1}^G\mid\mid\bm{o}_i\mid\mid}$(本文) 组内每token均摊 每题等权 组级耦合 $O(1/G)$ 跨序列稀释,冗长题目整组降权
$\frac{1}{\sum_{\text{batch} }\mid\mid\bm{o}\mid\mid}$(DAPO) 全batch每token等权 token等权 跨题耦合,自归一化偏差 $O(1/N)$
$\frac{1}{GL_{\max} }$(Dr. GRPO) 常数 无加权 纯协方差,2.2 的偏差项精确为零

① 组内基线与归一化(GRPO 的 $A_i$)。常数基线 $b$ 在外生权重下是平移回归目标,不改梯度只改方差(最优 $b$ 是加权最小二乘的截距);GRPO 的组内去均值 $\frac{r_i-\bar r}{\mathrm{std} }$ 是数据驱动的基线 + 目标标准化。但 $c$ 内生后,基线的选择开始真正影响期望更新。

② KL 惩罚 $a_{i,t}=A_{i,t}-\beta\log\frac{\pi_{\bm{\theta} }}{\pi_{\bm{\theta}_{ref} }}$。它不是常数基线,而是逐 token 依赖 $\bm{\theta}$ 的量,真正改变协方差:
$$\begin{align}
\mathrm{Cov}(\bm{s},a)
=\mathrm{Cov}(\bm{s},A)-\beta\,\mathrm{Cov}!\Big(\bm{s},\,\log\tfrac{\pi_{\bm{\theta} }}{\pi_{\bm{\theta}_{ref} }}\Big)
\end{align}$$
而 $\log\frac{\pi_{\bm{\theta} }}{\pi_{\bm{\theta}_{ref} }}$ 本身就是一个“指向参考策略的得分方向”。这一项相当于在回归目标里减掉优势在“离开参考策略”方向上的分量。一次朝 $\pi_{\bm{\theta}_{ref} }$ 的收缩,功能上类似回归中的惩罚项把估计拉向先验。

另外把归一化参数 c 换成历史可测的量(常数、EMA 滑动 token 数、或别的样本的长度)都是可以的。训练时可以监控泄漏梯度:把 $a\equiv1$ 代入,直接估 $\mathbb E[\sum c\bm s]$ 的范数及其与总长度的相关。长度漂移监控:EMA 归一后长度分布应平稳;若仍单向漂移,泄漏可能来自别处(组内均值自偏差、reward hacking 等),不要赖在 $c$ 上。

3.3、高方差问题

2.4.1节,线性回归视角是有 $N=\sum_{i=1}^G|\bm{o}_i|$ 个单点数据集 $\displaystyle \mathcal{D}_{i,t}=\{\bm{s}_{i,t},o_{i,t}\}_{i\times t}$ , 也就是说每个 $(i,t)$ 只有一个采样点,条件于自己的前缀 $o_{i,:t}$。这样协方差计算就变成单样本一阶矩(这是建立在漂移项为零的基础上的)。策略梯度就是 $N$个回归之和,每个单样本估计是无偏,但方差巨大。有三个机制修复:

① 跨时间步的平均。 单个 mini-batch 的 $\widehat{\nabla\mathcal{L} }$ 是高噪声估计;但优化器(SGD/AdamW)在训练轨迹上做隐式平均,期望意义下每步都朝 $\sum\mathrm{Cov}$ 的方向,沿优化轨迹用随机逼近(Robbins–Monro 式)渐近完成的,每个数据集贡献一个无偏投票。

② 组内的真重复:GRPO 的 $G$ 条回答。同一道题 $q$,$G$ 条独立回答。组内去均值 $A_i=r_i-\frac{1}{G}\sum_{j=1}^G r_j$ 的估计用 $G$ 个重复把“题目层面的共同噪声”估计出来并减掉来降低方差。

③ KL 惩罚与 clip。前者是给单样本目标加收缩先验,后者直接截断单样本重要性比例的尾部(防止单个离谱样本主导整步更新,稳定更新)。

四、评述

1、从线性回归看强化学习的策略梯度。发现强化学习算法实际上是在做优势和参数得分函数的回归。而这个回归有前提: 参数得分函数的期望是零。而通常情况很难保证(不同精度(bf16 vs int8/fp4)、浮点运算非结合性、不同输入形状 → 不同规约序、异步部署 → 轨迹陈旧性)都会导致训推不匹配 → 分布漂移 → 期望不是零。这正是论文[^1]做的事情,这里就不在叙述,有空在另外介绍。
)
2、在线性回归视角下,发现一个根本性的问题如何处理策略参数的 Fisher信息矩阵是策略更新的关键。因为策略参数不是自动正交的,一方面随机初始化原因,一方面softmax 分布输出就天然不能确保正交。必须通过信息矩阵对策略梯度做白化/重标定,使参数步长与分布位移在 KL 意义下等价
3、归一化参数不是一个可以随便定参数,内生与外生会完全改变线性回归的性质。可以使用各种外生变量(常数、EMA 滑动 token 数、或别的样本的长度)。

参考文献
[^1]: Marek, M., & Ryabinin, M. (2026, September 17). Score centering stabilizes off-policy reinforcement learning. arXiv. https://doi.org/10.48550/arXiv.2609.20807


版权声明
由引线小白创作并维护的柠檬CC博客采用署名-非商业-禁止演绎4.0国际许可证。
本文首发于柠檬CC [ https://www.limoncc.com ] , 版权所有、侵权必究。
本文永久链接https://www.limoncc.com/post/ef608130a2545b76/
如果您需要引用本文,请参考:
引线小白. (Oct. 8, 2026). 《大语言模型研究18——线性回归视角看强化学习策略梯度》[Blog post]. Retrieved from https://www.limoncc.com/post/ef608130a2545b76
@online{limoncc-ef608130a2545b76,
title={大语言模型研究18——线性回归视角看强化学习策略梯度},
author={引线小白},
year={2026},
month={Oct},
date={8},
url={\url{https://www.limoncc.com/post/ef608130a2545b76}},
}

'