minimind/course/notes/24_cispo_notes.md

12 KiB
Raw Blame History

CISPO 这段代码怎么理解

本文只解释第 24 课里 MiniMind 的 CISPO 分支:

trainer/train_grpo.py:135-138

源码:

if args.loss_type == "cispo":
    clamped_ratio = torch.clamp(ratio, max=args.epsilon_high).detach()
    per_token_loss = -(
        clamped_ratio
        * advantages.unsqueeze(1)
        * per_token_logps
        - args.beta * per_token_kl
    )

0. 一句话

CISPO 的想法是:

不要让 clipped ratio 变成一个切断梯度的目标;
把 clipped ratio 当成固定权重;
让真正的策略梯度从 log pi_theta(a|s) 走。

也就是:

ratio 负责告诉模型“这次更新权重多大”
logprob 负责提供“往哪个方向更新”的梯度路径

1. 先看 GRPO 的问题

MiniMind 的 GRPO 分支是:

clipped_ratio = torch.clamp(ratio, 1 - args.epsilon, 1 + args.epsilon)
per_token_loss1 = ratio * advantages.unsqueeze(1)
per_token_loss2 = clipped_ratio * advantages.unsqueeze(1)
per_token_loss = -(torch.min(per_token_loss1, per_token_loss2) - args.beta * per_token_kl)

对应公式:


\mathcal{L}_{\mathrm{GRPO}}
=
-
\mathbb{E}
\left[
\min
\left(
r_t A,
\operatorname{clip}(r_t, 1-\epsilon, 1+\epsilon) A
\right)
-
\beta \operatorname{KL}_t
\right]

其中:


r_t
=
\frac{
\pi_{\theta}(a_t \mid s_t)
}{
\pi_{\mathrm{old}}(a_t \mid s_t)
}

当 ratio 超过 clip 范围时,clipped_ratio 变成常数边界:

1 + epsilon

或者:

1 - epsilon

这能限制更新幅度,但也带来一个副作用:

被 clip 的那一侧,梯度可能直接被截断。

README 里也说了这个问题:

ratio 被 clip 之后,梯度流直接就被硬截断了。

源码证据:

文件:README.md:1156-1170

看它是为了理解:CISPO 为什么不是重新设计一套 RL 算法,而是改 GRPO 的 loss 形式。

代码/文本摘录:

CISPO 的关注点并不是重新设计 group baseline,
而是用非常小的 loss 改动,更直接地修正这个问题。

这段说明:

  • GRPO 的 group reward / advantage 仍然保留。
  • CISPO 主要改策略项。
  • 目标是避免 ratio clip 后梯度路径被截断。

1.1 clip 为什么会让梯度变成 0

先只看 torch.clamp:


y
=
\operatorname{clamp}(x, l, u)

它的导数大致是:


\frac{\partial y}{\partial x}
=
\begin{cases}
1, & l < x < u \\
0, & x < l \text{ or } x > u
\end{cases}

原因很简单:

x 在区间内:clamp(x) = x,输出会跟着 x 变。
x 超过边界:clamp(x) = 边界常数,输出不再跟着 x 变。

所以当 ratio 超过 clip 边界时:

clipped_ratio = torch.clamp(ratio, 1 - epsilon, 1 + epsilon)

如果:

ratio = 1.5
1 + epsilon = 1.2

那么:

clipped_ratio = 1.2

这时把 ratio 从 1.5 改成 1.6,clipped_ratio 仍然是 1.2。也就是说:

clipped_ratio 对 ratio 不敏感了;
梯度就不能从 clipped_ratio 继续传回 ratio。

再放回 GRPO/PPO 的 clipped surrogate:


\min
\left(
r_t A_t,
\operatorname{clip}(r_t, 1-\epsilon, 1+\epsilon)A_t
\right)

这里“梯度被截断”不是所有情况都会发生,而是发生在 min 选中了 clipped 分支、并且 clamp 已经落到边界时。

两个典型情况:

advantage ratio 状态 被限制的方向 结果
A_t > 0 r_t > 1+\epsilon 好 token 概率已经涨太多 clipped 分支是常数边界,继续增大 ratio 没有梯度
A_t < 0 r_t < 1-\epsilon 坏 token 概率已经降太多 clipped 分支是常数边界,继续减小 ratio 没有梯度

所以这句话的准确含义是:

clip 限制了 policy 相对 old policy 变化太大;
但当 clipped 分支被选中时,策略项对 ratio 的梯度可能变成 0。

这就是 CISPO 想缓解的问题:不要把被 clip 的 ratio 本身当作主要梯度路径。

2. CISPO 的公式

README 写的 CISPO 损失是:


\mathcal{L}_{\mathrm{CISPO}}
=
-
\mathbb{E}
\left[
\min(r_t, \epsilon_{\max})
\cdot
A_t
\cdot
\log \pi_{\theta}(a_t \mid s_t)
-
\beta \operatorname{KL}_t
\right]

MiniMind 源码里的变量对应:

min(r_t, epsilon_max)      -> clamped_ratio
A_t                       -> advantages.unsqueeze(1)
log pi_theta(a_t | s_t)   -> per_token_logps
beta * KL_t               -> args.beta * per_token_kl

所以源码:

per_token_loss = -(
    clamped_ratio
    * advantages.unsqueeze(1)
    * per_token_logps
    - args.beta * per_token_kl
)

就是:


\operatorname{loss}_{t}
=
-
\left[
\bar{r}_t
A
\ell^{\theta}_{t}
-
\beta \operatorname{KL}_t
\right]

其中:


\bar{r}_t
=
\min(r_t, \epsilon_{\max})

\ell^{\theta}_{t}
=
\log \pi_{\theta}(a_t \mid s_t)

3. 为什么要乘 logprob

policy gradient 的基础形式是:


\nabla_{\theta} J(\theta)
\approx
A_t
\nabla_{\theta}
\log \pi_{\theta}(a_t \mid s_t)

所以如果要让模型:

A_t > 0 时提高 token 概率
A_t < 0 时降低 token 概率

最直接的可微路径就是:

log pi_theta(a_t | s_t)

CISPO 让损失里显式出现:

advantages * per_token_logps

这样梯度就从 per_token_logps 走。

直观上:

如果 advantage > 0:
  loss = -positive_weight * logprob
  最小化 loss 会提高 logprob。

如果 advantage < 0:
  loss = -negative_weight * logprob
  最小化 loss 会降低 logprob。

4. ratio 在 CISPO 里做什么

ratio 仍然有用。

源码:

ratio = torch.exp(per_token_logps - old_per_token_logps)
clamped_ratio = torch.clamp(ratio, max=args.epsilon_high).detach()

它对应:


r_t
=
\exp
\left(
\log \pi_{\theta}(a_t \mid s_t)
-
\log \pi_{\mathrm{old}}(a_t \mid s_t)
\right)

也就是:


r_t
=
\frac{
\pi_{\theta}(a_t \mid s_t)
}{
\pi_{\mathrm{old}}(a_t \mid s_t)
}

CISPO 不把 ratio 当成主要梯度路径,而是把它当作权重:


\bar{r}_t
=
\min(r_t, \epsilon_{\max})

如果某个 token 的当前概率相对 old policy 变大,ratio 大,那么这个 token 的策略项权重也变大。

但上限是:

epsilon_high

避免权重爆炸。

5. 为什么要 detach

这行最关键:

clamped_ratio = torch.clamp(ratio, max=args.epsilon_high).detach()

detach() 的意思是:

clamped_ratio 参与前向计算;
但不从 clamped_ratio 反传梯度。

也就是说,CISPO 的策略项:

clamped_ratio * advantages * per_token_logps

反向传播时:

clamped_ratio 被当成常数权重
advantages 被当成常数权重
per_token_logps 提供梯度

数学上近似看:


\nabla_{\theta}
\left(
\bar{r}_t^{\mathrm{detach}}
A_t
\log \pi_{\theta}(a_t \mid s_t)
\right)
=
\bar{r}_t^{\mathrm{detach}}
A_t
\nabla_{\theta}
\log \pi_{\theta}(a_t \mid s_t)

这就是 CISPO 的核心:

用 ratio 控制权重;
用 logprob 保留梯度;
不要让梯度穿过 ratio 本身。

如果不 detach,梯度还会穿过:

ratio = exp(new_logp - old_logp)

那策略项会同时从 ratio 和 logprob 两条路径反传,更新会更复杂、更容易偏离 CISPO 想表达的“裁剪权重 × logprob”。

6. 和 GRPO 的差别

GRPO 策略项:


\min(r_t A_t, \bar{r}_t A_t)

它的核心是:

ratio 自己是优化目标的一部分。

CISPO 策略项:


\bar{r}_t^{\mathrm{detach}}
A_t
\log \pi_{\theta}(a_t \mid s_t)

它的核心是:

ratio 只是权重;
logprob 才是梯度路径。

对比:

项目 GRPO CISPO
advantage 来源 group reward group reward
是否需要 critic 不需要 不需要
是否使用 old logprob 使用 使用
ratio 作用 进入 clipped surrogate 作为 detached weight
梯度主路径 ratio / clipped surrogate logprob
KL 约束 token KL token KL

6.1 PPO/GRPO 更新 policy 到底靠谁

PPO/GRPO 更新 policy 不是“只靠 ratio”,也不是“只靠 advantage”。更准确的说法是:

ratio/logprob 提供梯度路径;
advantage 提供方向和权重。

先看没有 clip 的策略目标:


J_t
=
r_t A_t

其中:


r_t
=
\exp
\left(
\ell_t^\theta
-
\ell_t^{\mathrm{old}}
\right)

\ell_t^\theta
=
\log \pi_\theta(a_t \mid s_t)

old_per_token_logps 是 rollout 时保存下来的旧 policy logprob,在训练当前 step 时当常数用。所以:


\nabla_\theta J_t
=
A_t r_t
\nabla_\theta \ell_t^\theta

也就是:


\nabla_\theta J_t
=
A_t r_t
\nabla_\theta
\log \pi_\theta(a_t \mid s_t)

这条公式说明三件事:

1. 真正连到 policy net 参数的是 current logprob。
2. ratio 里面包含 current logprob,所以 ratio 是一条梯度路径。
3. advantage 自己通常不反传梯度,但它乘在梯度前面,决定方向和强度。

所以 advantage 不是“只和 reward model / critic 有关,然后和 policy 更新没关系”。它的来源确实可能来自 reward model、critic 或 group reward,但最后会直接进入 policy loss。

PPO 里:

reward / reward_model -> reward
critic_model -> value
reward + value -> GAE advantage
policy_loss -> ratio * advantage 的 clipped 形式

GRPO 里:

reward_model / rule -> reward
同一个 prompt 的多条 rewards -> group mean/std
group 标准化 -> advantage
policy_loss -> ratio * advantage 的 clipped 形式

两者的共同点:

advantage 不一定接收梯度;
但 advantage 一定控制 policy gradient。

如果:


A_t > 0

优化会倾向于提高这个 token 的概率。

如果:


A_t < 0

优化会倾向于降低这个 token 的概率。

CISPO 的变化是:

GRPO/PPO:
  ratio 既是权重,也是梯度路径的一部分。

CISPO:
  clamped_ratio.detach() 只当固定权重;
  per_token_logps 明确承担策略梯度路径。

7. 用正负 advantage 看方向

假设忽略 KL,只看 CISPO 策略项:


\mathcal{L}
=
-
\bar{r}_t
A_t
\log \pi_{\theta}(a_t \mid s_t)

7.1 advantage 大于 0

如果:


A_t > 0

那么:


\mathcal{L}
=
-
\text{positive}
\cdot
\log \pi_{\theta}(a_t \mid s_t)

最小化 loss 会让:

log pi_theta(a_t | s_t) 变大

也就是提高这个 token 的概率。

7.2 advantage 小于 0

如果:


A_t < 0

那么:


\mathcal{L}
=
+
\text{positive}
\cdot
\log \pi_{\theta}(a_t \mid s_t)

最小化 loss 会让:

log pi_theta(a_t | s_t) 变小

也就是降低这个 token 的概率。

这和 policy gradient 的直觉一致。

8. 源码证据:MiniMind 的 CISPO 分支

文件:trainer/train_grpo.py:132-143

看它是为了理解:GRPO 和 CISPO 共用 group advantage、ratio、reference KL,只切换 policy loss 形式。

代码摘录:

kl_div = ref_per_token_logps - per_token_logps
per_token_kl = torch.exp(kl_div) - kl_div - 1
ratio = torch.exp(per_token_logps - old_per_token_logps)
if args.loss_type == "cispo":
    clamped_ratio = torch.clamp(ratio, max=args.epsilon_high).detach()
    per_token_loss = -(clamped_ratio * advantages.unsqueeze(1) * per_token_logps - args.beta * per_token_kl)
else:
    clipped_ratio = torch.clamp(ratio, 1 - args.epsilon, 1 + args.epsilon)
    per_token_loss1 = ratio * advantages.unsqueeze(1)
    per_token_loss2 = clipped_ratio * advantages.unsqueeze(1)
    per_token_loss = -(torch.min(per_token_loss1, per_token_loss2) - args.beta * per_token_kl)
policy_loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1).clamp(min=1)).mean()

这段代码说明:

  • per_token_kl 两种 loss 共用。
  • ratio 两种 loss 都要计算。
  • GRPO 使用 torch.min(ratio*A, clipped_ratio*A)。
  • CISPO 使用 clamped_ratio.detach() * A * per_token_logps。
  • completion_mask 只让有效 response token 参与 loss。

9. 记忆方式

GRPO:
  clipped ratio objective
  ratio 是优化目标

CISPO:
  clipped importance weight * logprob
  ratio 是权重
  logprob 是梯度路径

一句话记:

CISPO = GRPO 的 group advantage + reference KL + detached clipped ratio weighted logprob。

10. 检查问题

  1. CISPO 为什么还需要 old logprob?
  2. clamped_ratio.detach() 的作用是什么?
  3. CISPO 的策略梯度主要从哪个变量走?
  4. advantages.unsqueeze(1) 为什么可以乘到每个 token 上?
  5. CISPO 和 GRPO 哪些部分是共用的?
  6. 为什么 torch.clamp 到边界以后,对 ratio 的梯度可能是 0?
  7. PPO/GRPO 里 advantage 不反传梯度,为什么仍然会影响 policy 更新?