1. 为什么Token级策略梯度会进入视野我第一次做生成式模型的强化学习微调时用的是最朴素的写法模型把整条回复生成完奖励模型给一个总分数然后把这个分数乘到整条序列所有token的log概率梯度上更新一次参数。刚开始看起来没问题跑起来就发现训练过程非常别扭——loss降得飞快但采样出来的文本质量并没有同步变好甚至有时候loss越低模型越平庸。后来才意识到问题出在策略梯度损失是拿一整条序列当一个“时间步”在算中间那些本该提供方向的token-level信号全被平均掉了。这是整个话题的入口token级策略梯度损失本质上就是把强化学习里最基础的那套“逐时间步策略梯度”原原本本套到生成任务上而不是把生成过程压缩成一个单步bandit。这篇文章不打算罗列概念我会从为什么需要它、怎么重新定义生成过程的MDP、损失函数怎么推、代码怎么落地以及我在实际项目中踩过的坑这几个角度讲透。适合读这篇文章的人有两类一类是已经跑通过序列级RL训练、但总觉得梯度更新太粗糙、效果不稳定的工程师另一类是刚接触生成式模型对齐、想搞清楚“每个token都能拿到一条梯度”到底意味着什么的同学。看完之后你应该能独立实现一套可运行的token级策略梯度训练管线。1.1 序列级训练的三个致命痛点序列级策略梯度最常见的损失形式是L - E_τ [ R(τ) · Σ_t log π_θ(a_t | s_t) ]也就是把整条轨迹的总回报R(τ)当作每个token动作的共享权重。这个形式足够简单但在生成式模型上至少有三个硬伤。第一是奖励稀疏。模型生成了几百个token最后只拿到一个标量。中间哪些行为导致了高分、哪些动作在拖后腿模型只能靠猜。尤其当一个回复前段在铺垫、后段才进入正题时最后的分数其实是后段的功劳前段所有token也跟着被放大或打压。第二是credit assignment混乱。一句话里前面用了错误的语气词、后面给出了精准答案总分数可能因为后段优秀而偏高前段的错误词也被当作好样本被抬高反过来也一样整体低分会让所有token共享惩罚包括那些本来正确的部分。序列级公式剥夺了模型分辨“哪些token该奖、哪些该罚”的能力。第三是梯度方差大。生成式采样本身随机性很大同一个prompt可能生成质量悬殊的多个回复。序列级目标把整条序列的回报作为一个整体乘以所有token的梯度不同序列的回报波动会直接传导成很大的梯度抖动。为了压方差只能调小学习率训练速度自然慢。1.2 从“序列反馈”到“逐步反馈”的动机转变其实经典策略梯度定理一直是逐时间步展开的。问题不在于强化学习理论缺了什么而是在生成式模型的应用里很多实现把整个生成过程当成了一步式的“上下文-bandit”输入是prompt输出是整段回复奖励是最终得分。当回复只有几个token时这种近似还凑合一旦序列长度飙到几百、上千单步bandit的近似误差就会被放大。Token级策略梯度的动机就是恢复“每一步都有一个状态、一个动作、一个反馈”的原始设定。对于生成任务状态是“prompt 已经生成的token前缀”动作是“下一个要填的token”反馈可以是每个token的即时奖励也可以是通过价值函数拆解出来的优势估计。切换到这个视角之后训练信号从“整条序列一个数字”变成了“每个token都有自己的优势”。这带来的好处不是理论上的锦上添花而是实打实训练稳定性的提升模型能更精准地调整真正该调整的位置梯度的方差也会因为有了baseline和逐位置分解而降低。下面我们就从MDP建模开始一步步把这个损失构造出来。2. 把生成过程还原成Token级MDP很多做NLP的开发者对强化学习里的MDP有点距离感总觉得要写环境、写状态转移、写奖励函数很麻烦。其实生成式模型天然就是一个MDP只是大家平时用监督学习视角看习惯了没有往这个方向想。这一节会把这层窗户纸捅破。2.1 状态、动作与转移的重新定义定义符号之前先说结论一次生成过程完全可以看作一个确定性转移的马尔可夫决策过程。状态s_t (x, a_{t})其中 x 是输入prompta_{t} 是已经生成的前缀token序列。动作a_t 是第t步从词表 V 中采样得到的token对应MDP中的动作。转移函数是确定性的给定状态 s_t 和动作 a_t下一个状态就是 s_{t1} (x, a_{≤t})。这里不存在环境随机性因为语言模型的自回归生成本质上就是“写下一个token就拼接到前缀后面”。终止状态是生成结束符的时刻也就是完整的回复序列 τ (a_1, a_2, ..., a_T)。折扣因子γ 按需取0.95到1.0之间。强调“确定性转移”这点很重要因为很多经典的策略梯度推导假设转移是随机的而生成任务中转移其实是确定的这会让Bellman方程的计算和value function的估计都变得非常干净。你不必担心模型自身的采样随机性——那已经包含在策略 π_θ(a_t | s_t) 里了。2.2 奖励拆解只有最终奖励时中间token怎么算在上面的MDP里完整轨迹回报 R(τ) 可以拆成每个token即时奖励的和R(τ) Σ_t r(s_t, a_t)但问题是大多数生成式对齐场景下奖励模型只对完整回复打分并没有给每个token标注奖励。这时可以有两种处理方式。第一种是稀疏奖励设 r_T R(τ)中间每一步 r_t 0。这样做简单但优势估计还是需要靠value function在中间位置提供baseline否则策略梯度退化成序列级。第二种是把最终奖励按某种规则分摊到每个token。比如用每个token的KL散度作为即时正则奖励或者用一个逐位置的奖励模型/奖励分解函数来提供dense reward。实际项目中经常做的是混合最终得分R(τ)保留在最后一个token上同时每一步都叠加一个负的KL惩罚项相当于每个token都拿到“即时的小奖励/小惩罚”。也就是说token级策略梯度不一定要求你有一个完美的逐token奖励模型。只要奖励定义拆成了r_t r_dense(s_t, a_t) r_terminal · 1[t T]就可以进入标准强化学习的框架剩下的交给value function去估计优势。2.3 Token级MDP与经典RL环境的对应关系为了直观我列一张对应表经典RL概念生成式任务中的对应环境词表、上下文与解码规则状态 s_tprompt x 与已生成前缀 a_{t} 的拼接动作 a_t下一个token在词表中的下标策略 π_θ(a_t | s_t)语言模型当前step的softmax分布即时奖励 r_ttoken级奖励或KL惩罚终态奖励 R(τ)奖励模型对完整回复的评分折扣因子 γ通常取0.95~1.0采样轨迹从prompt开始自回归采样直到结束符这张表说明一个事实语言生成不是强化学习的特殊案例它就是强化学习的标准案例。很多人觉得序列级策略梯度才是常规做法其实那只是因为实现起来省事。如果把生成过程真正当成一个多步MDPtoken级策略梯度反而更贴近理论原貌。有了这套MDP定义下一步就是把它代入策略梯度定理得到最终的token级损失形式。3. Token级策略梯度损失的数学推导与最终损失形式3.1 从总回报期望推导策略梯度策略梯度的起点是最大化期望总回报J(θ) E_{τ ~ π_θ} [ R(τ) ]对参数求梯度用log-derivative trick可以得到∇_θ J(θ) E_τ [ Σ_t (G_t - b(s_t)) · ∇_θ log π_θ(a_t | s_t) ]其中 G_t 是从第t步到序列末尾的折扣累计奖励G_t Σ_{kt}^{T} γ^{k-t} r_kb(s_t) 是baseline用来降低方差当取 b(s_t) V^π(s_t)状态价值函数时G_t - b(s_t) 就变成了优势函数 A_t。注意这个式子天然是token级的每个时间步 t 都有属于自己的 log π_θ(a_t | s_t) 和权重 A_t。序列级损失其实就是把 A_t 强行换成整条序列总回报 R(τ)等于丢掉了中间所有优势信息。Token级做法的核心只是把公式里省略掉的 A_t 老老实实算出来。3.2 优势函数的Token级估计MC、TD与GAE既然上面出现了 A_t下一个问题就是A_t怎么算。这里有三条路线复杂度和方差特性各不同。Monte Carlo估计直接用后续真实奖励的总和 G_t 减去baseline。无偏但方差很大因为未来若干步的采样噪声全部堆在一起。TD估计用 A_t r_t γV(s_{t1}) - V(s_t)方差小但有偏差而且对value function的准确性非常敏感。GAEGeneralized Advantage Estimation在两者之间取折中。GAE的定义是δ_t r_t γ V(s_{t1}) - V(s_t)A_t^{GAE} Σ_{l0}^{T-t-1} (γλ)^l δ_{tl}其中 λ ∈ [0,1] 控制偏差与方差的平衡。λ越接近1越像MC估计偏差小方差大λ越接近0越像TD估计方差小偏差大。在生成式模型的token级训练里我常用 λ0.95、γ0.99这个组合在大多数对话生成场景下表现都比较稳如果value网络训练得比较好可以适当把 λ 提到0.97~0.99。3.3 常用损失展开式与简化变体有了优势估计token级策略梯度的原始损失就是逐token的加权负对数似然L_PG - E_τ [ Σ_t A_t · log π_θ(a_t | s_t) ]但在实际训练中直接用这个损失容易因为单次更新步长过大而崩掉所以通常会和PPO中的importance ratio结合。假设旧策略是 π_old当前策略是 π_θ定义ρ_t π_θ(a_t | s_t) / π_old(a_t | s_t)那么带clip的token级PPO损失为L_PPO - E_τ [ Σ_t min( ρ_t · A_t, clip(ρ_t, 1-ε, 1ε) · A_t ) ]这里每个token都有自己独立的 ρ_t 和 A_t。和序列级PPO的唯一区别还是在于A_t不再是同一个标量而是每个前缀位置上单独估计的值。更复杂的变体还包括在奖励里混入每一项的KL惩罚和熵正则r_t r_t - β · KL(π_θ(a_t|s_t) || π_ref(a_t|s_t)) η · H(π_θ(·|s_t))KL惩罚让策略不要偏离参考策略太远熵正则防止策略过早熵坍塌。很多开源实现里会把这些项写进损失而不是写进奖励效果上两者基本等价只是对数尺度不同。4. Token级损失落地的实现细节mask、优势估计与归一化数学推导只有一页纸真正让token级策略梯度变得能用的是工程细节。这一节我把实现时最容易出错的几个地方拆开讲内容包括数据组织、mask处理、GAE伪代码、归一化和分布式训练。4.1 数据组织和padding mask训练时一个batch里会有长度不一的序列。通常的做法是在尾部padding到相同长度然后训练时用attention mask和loss mask把padding位置屏蔽掉。在token级策略梯度里需要为每个token准备的数据包括token id、log_probs、values、rewards、advantages、以及一个0/1 mask。下表是一个典型的数据组织方式字段形状说明tokens[B, T]序列token idlog_probs[B, T]旧策略下每个token的log概率用于importance ratiovalues[B, T]value网络对每个前缀状态的估计rewards[B, T]每个token的即时奖励只有最后token有终态奖励时其他位置为0mask[B, T]有效token为1padding为0advantages[B, T]按GAE计算得到的优势关键点rewards在padding位置必须为0values在padding位置也建议置0。否则GAE递推时padding位置会参与δ_t的计算导致同一个batch里短序列的最后一个有效token被padding处的value污染产生了一个“伪动作价值”。4.2 Token级优势与损失伪代码直接给一套PyTorch风格的实现它会在batch维度上并行做GAE递推同时用mask隔离padding。def compute_gae(rewards, values, mask, gamma0.99, lam0.95): rewards: [B, T] float values: [B, T] float mask: [B, T] float, 1表示有效token, 0表示padding B, T rewards.shape # 序列末尾的下一个状态value视为0 next_values torch.cat([values[:, 1:], torch.zeros(B, 1)], dim1) deltas rewards gamma * next_values - values deltas deltas * mask # padding位置不参与优势递推 advantages torch.zeros_like(deltas) running torch.zeros(B) for t in reversed(range(T)): running deltas[:, t] gamma * lam * running * mask[:, t] advantages[:, t] running return advantages注意递推时running也要乘mask否则前一个序列的累计优势会跨过padding边界“漏”到下一个序列。这个bug非常隐蔽因为loss表面上看起来很正常但实际上相邻序列在互相影响。策略梯度的损失计算如下def token_pg_loss(log_probs_old, log_probs_new, advantages, mask, clip_range0.2): # log_probs_old/new: [B, T]都是log_softmax后按token id gather得到的 ratio torch.exp(log_probs_new - log_probs_old) surr1 ratio * advantages surr2 torch.clamp(ratio, 1.0 - clip_range, 1.0 clip_range) * advantages loss -torch.min(surr1, surr2) * mask return loss.sum() / mask.sum()这里的advantage建议在Enter loss之前做一次归一化见下一节。4.3 Whitening、reward scaling和KL惩罚实战中直接用原始GAE计算出的advantage做训练效果往往不如先归一化。常用三种方式全局whitening所有有效token的advantage减去均值、除以标准差得到一个均值为0、方差为1的分布。绝对值均值归一化除以mean(|A|)保留符号方向但不强行把均值拉到0。适合advantage全为正或全为负的情况。固定常数缩放直接用常数 c 控制更新幅度。我自己的经验是如果batch里序列长度差异大尽量用全局running统计做whitening而不是每个batch内独立重新计算均值和方差。因为每个batch的样本分布可能漂移batch内独立标准化会让同一个优势值在不同batch里被放大或缩小训练容易震荡。reward scaling同样需要注意。最终奖励乘以一个系数 c会等价地放大或缩小所有advantage。很多实现里把reward缩放到0到1之间再配合 β0.01~0.1 的KL惩罚。调参顺序建议是先固定reward scaling再调β最后动了reward scaling后必须回头重调β两者耦合很强。每次改动记录per-token KL的变化如果KL在一两千步内突然塌到0附近多半是β过大或reward scaling过小。4.4 多卡训练中的全局统计与同步当训练从单卡切换到多卡DDP时token级策略梯度会多出两个不大不小的问题。第一个问题是advantage归一化的统计量。如果你用的是batch内whitening那么单卡basebatch的均值和全局大batch的均值会不一样global batch越大样本方差越小导致advantage的绝对值被动态压小策略更新幅度也跟着变。稳妥的做法是跨卡做统计量的同步或者干脆使用固定常数缩放例如把advantage除以一个固定值如5.0或10.0不再每次动态计算。第二个问题是数据采样与生成的一致性。如果生成rollout是在多卡上并行做的同一批prompt在不同卡上采样得到的回复可能分布在不同的batch里。这本身没问题但KL惩罚的期望统计会因此波动更频繁——某个batch里恰好全是高KL的样本另一个batch里全是低KL样本训练loss看起来忽高忽低。解决办法是观察滑动窗口的均值不要因为某一步的剧烈波动就去调学习率。这块容易误导人很多新人会误以为是学习率问题实际是采样分布的方差在作祟。5. 训练工程中反复踩到的问题与调参心得理论再清晰落到工程上总是有一堆“文档不会写”的细节。下面这些坑我基本都在项目里亲手踩过有些花了很长时间才定位到根因。5.1 时间步错位与长度偏置生成式模型训练时batch内序列被padding到同一长度所以“第t时间步”在不同的序列里对应的绝对位置通常是相同的。但如果你的训练框架支持序列打包packing多个短序列被拼接在一条样本里这时“token下标”和“MDP时间步”就不再是一回事了。打包场景里GAE递推必须按原序列边界切开不能把两条序列头尾相接。否则第二条序列的第一个token会被当成第一条序列的延续value和advantage全混掉。实现上我建议打包样本额外维护一个segment_id并在计算GAE时按segment_id对mask做分段处理。如果只是为了对齐浪费时间不如老老实实用padding mask。另一个相关问题是长度偏置。长序列天然有更多token在梯度中占的weight更大如果每条序列的loss按token数平均倒是无所谓但如果按sum reduction长序列会主导训练方向。建议loss reduction用有效token数的batch平均值同时关注长度统计直方图确保训练集里没有少数“超长样本”把损失撬偏。5.2 优势约等于零时的“假稳定”token级策略梯度实现到位之后你会看到training loss曲线很好看地下降但生成样本质量提升非常慢。这时候要小心“假稳定”。原因是当value网络拟合得足够好时GAE优势的平均绝对值会越来越小loss随之变小但策略的更新幅度可能远低于预期。判断模型到底有没有在学不要只看loss还要看两个指标一是advantage绝对值的滑动均值如果长期低于0.01说明更新步长被“无优势可乘”的处境限制了二是生成样本端reward的滑动均值必须持续监控。遇到这种情况可以尝试增大reward scaling、降低GAE的λ让优势更接近MC或者让value网络别更新得太快人为制造一些“优势修正空间”。这里有一个trick如果发现优势全部挤在一个很小的区间但方向是对的可以不用whitening直接用sign-preserving归一化比如 A_hat A / mean(|A|)。这样既保留每个token间的相对大小又把整体量级放到1附近策略更新幅度更可控。5.3 策略崩溃与熵保护生成式策略梯度最怕的一件事是策略过快坍缩成一个窄分布。Token级策略会一条一条地更新如果某个token位置偶然获得了正优势模型可能把该位置的分布拉得很尖哪怕这个尖只是噪声造成的。为了对抗这个我一般会在损失里加熵正则项系数从0.01起步。另外一个更直观的做法是设定一个“熵下限”如果某个batch里平均熵低于参考策略熵的60%就暂时把策略梯度幅度降一半。实现上不复杂维护一个滑动平均熵低于阈值时报警或调整系数即可。KL惩罚同样能帮助防坍缩但注意不要惩罚过头。KL惩罚是逐token累加的长序列累计KL惩罚会非常大如果β沿用序列级实验里的数值长文本生成任务的策略基本被压死。我的经验是token级训练里β要比序列级小一个数量级。举个例子序列级里 β0.02 还好token级可能要先从0.002开始再根据per-token KL反馈调整。5.4 一些少有人写但很实际的细节最后补充几个零碎但会影响成败的点都是我实际踩过的采样temperature要和部署一致。训练时用temperature0.8采样生成的rollout算出的advantage如果线上推理用greedy策略优化方向和部署行为可能错位。最好在rollout阶段就固定一个和部署接近的采样策略。value网络不要和策略同步更新。经验上value网络的更新频率可以比策略低或者让value网络使用EMA参数计算target否则advantage会一直在移动靶子上训练很难收敛。定期把value optimizer的步长设成策略的一半效果更稳。log_probs要用float32甚至更高精度保存。生成rollout时保存的旧策略log_probs如果在半精度下和当前策略算ratio时误差会被放大clip边界附近的token尤其敏感。一个很小的精度误差可能导致几千token的梯度被错误clip。离线buffer不要屯太久。Token级策略梯度优势是基于旧value和旧策略的如果rollout buffer在内存里放太久新旧策略差异变大importance ratio普遍偏离1clip后的梯度信息量下降。我通常每个rollout epoch只保留一小部分历史buffer不超过当前batch的20%。说回我自己的训练习惯gamma取0.99GAE的lambda取0.95advantage做全局running whiteningreward scaling控制在0.05~0.1RL训练一个epoch后重新采样一版rollout喂给模型。这套组合在我的多个生成式对齐项目中都跑得比较稳不一定适合所有场景但至少是个不容易翻车的起点。真正适合自己的参数还是要在始终监控生成质量reward和per-token KL的前提下慢慢调出来。 
长河君 · 2026/10/21 17:21:40 · 阅读 1.2k Token级策略梯度：生成式模型强化学习的逐步优化指南
1. 为什么Token级策略梯度会进入视野我第一次做生成式模型的强化学习微调时用的是最朴素的写法模型把整条回复生成完奖励模型给一个总分数然后把这个分数乘到整条序列所有token的log概率梯度上更新一次参数。刚开始看起来没问题跑起来就发现训练过程非常别扭——loss降得飞快但采样出来的文本质量并没有同步变好甚至有时候loss越低模型越平庸。后来才意识到问题出在策略梯度损失是拿一整条序列当一个“时间步”在算中间那些本该提供方向的token-level信号全被平均掉了。这是整个话题的入口token级策略梯度损失本质上就是把强化学习里最基础的那套“逐时间步策略梯度”原原本本套到生成任务上而不是把生成过程压缩成一个单步bandit。这篇文章不打算罗列概念我会从为什么需要它、怎么重新定义生成过程的MDP、损失函数怎么推、代码怎么落地以及我在实际项目中踩过的坑这几个角度讲透。适合读这篇文章的人有两类一类是已经跑通过序列级RL训练、但总觉得梯度更新太粗糙、效果不稳定的工程师另一类是刚接触生成式模型对齐、想搞清楚“每个token都能拿到一条梯度”到底意味着什么的同学。看完之后你应该能独立实现一套可运行的token级策略梯度训练管线。1.1 序列级训练的三个致命痛点序列级策略梯度最常见的损失形式是L - E_τ [ R(τ) · Σ_t log π_θ(a_t | s_t) ]也就是把整条轨迹的总回报R(τ)当作每个token动作的共享权重。这个形式足够简单但在生成式模型上至少有三个硬伤。第一是奖励稀疏。模型生成了几百个token最后只拿到一个标量。中间哪些行为导致了高分、哪些动作在拖后腿模型只能靠猜。尤其当一个回复前段在铺垫、后段才进入正题时最后的分数其实是后段的功劳前段所有token也跟着被放大或打压。第二是credit assignment混乱。一句话里前面用了错误的语气词、后面给出了精准答案总分数可能因为后段优秀而偏高前段的错误词也被当作好样本被抬高反过来也一样整体低分会让所有token共享惩罚包括那些本来正确的部分。序列级公式剥夺了模型分辨“哪些token该奖、哪些该罚”的能力。第三是梯度方差大。生成式采样本身随机性很大同一个prompt可能生成质量悬殊的多个回复。序列级目标把整条序列的回报作为一个整体乘以所有token的梯度不同序列的回报波动会直接传导成很大的梯度抖动。为了压方差只能调小学习率训练速度自然慢。1.2 从“序列反馈”到“逐步反馈”的动机转变其实经典策略梯度定理一直是逐时间步展开的。问题不在于强化学习理论缺了什么而是在生成式模型的应用里很多实现把整个生成过程当成了一步式的“上下文-bandit”输入是prompt输出是整段回复奖励是最终得分。当回复只有几个token时这种近似还凑合一旦序列长度飙到几百、上千单步bandit的近似误差就会被放大。Token级策略梯度的动机就是恢复“每一步都有一个状态、一个动作、一个反馈”的原始设定。对于生成任务状态是“prompt 已经生成的token前缀”动作是“下一个要填的token”反馈可以是每个token的即时奖励也可以是通过价值函数拆解出来的优势估计。切换到这个视角之后训练信号从“整条序列一个数字”变成了“每个token都有自己的优势”。这带来的好处不是理论上的锦上添花而是实打实训练稳定性的提升模型能更精准地调整真正该调整的位置梯度的方差也会因为有了baseline和逐位置分解而降低。下面我们就从MDP建模开始一步步把这个损失构造出来。2. 把生成过程还原成Token级MDP很多做NLP的开发者对强化学习里的MDP有点距离感总觉得要写环境、写状态转移、写奖励函数很麻烦。其实生成式模型天然就是一个MDP只是大家平时用监督学习视角看习惯了没有往这个方向想。这一节会把这层窗户纸捅破。2.1 状态、动作与转移的重新定义定义符号之前先说结论一次生成过程完全可以看作一个确定性转移的马尔可夫决策过程。状态s_t (x, a_{t})其中 x 是输入prompta_{t} 是已经生成的前缀token序列。动作a_t 是第t步从词表 V 中采样得到的token对应MDP中的动作。转移函数是确定性的给定状态 s_t 和动作 a_t下一个状态就是 s_{t1} (x, a_{≤t})。这里不存在环境随机性因为语言模型的自回归生成本质上就是“写下一个token就拼接到前缀后面”。终止状态是生成结束符的时刻也就是完整的回复序列 τ (a_1, a_2, ..., a_T)。折扣因子γ 按需取0.95到1.0之间。强调“确定性转移”这点很重要因为很多经典的策略梯度推导假设转移是随机的而生成任务中转移其实是确定的这会让Bellman方程的计算和value function的估计都变得非常干净。你不必担心模型自身的采样随机性——那已经包含在策略 π_θ(a_t | s_t) 里了。2.2 奖励拆解只有最终奖励时中间token怎么算在上面的MDP里完整轨迹回报 R(τ) 可以拆成每个token即时奖励的和R(τ) Σ_t r(s_t, a_t)但问题是大多数生成式对齐场景下奖励模型只对完整回复打分并没有给每个token标注奖励。这时可以有两种处理方式。第一种是稀疏奖励设 r_T R(τ)中间每一步 r_t 0。这样做简单但优势估计还是需要靠value function在中间位置提供baseline否则策略梯度退化成序列级。第二种是把最终奖励按某种规则分摊到每个token。比如用每个token的KL散度作为即时正则奖励或者用一个逐位置的奖励模型/奖励分解函数来提供dense reward。实际项目中经常做的是混合最终得分R(τ)保留在最后一个token上同时每一步都叠加一个负的KL惩罚项相当于每个token都拿到“即时的小奖励/小惩罚”。也就是说token级策略梯度不一定要求你有一个完美的逐token奖励模型。只要奖励定义拆成了r_t r_dense(s_t, a_t) r_terminal · 1[t T]就可以进入标准强化学习的框架剩下的交给value function去估计优势。2.3 Token级MDP与经典RL环境的对应关系为了直观我列一张对应表经典RL概念生成式任务中的对应环境词表、上下文与解码规则状态 s_tprompt x 与已生成前缀 a_{t} 的拼接动作 a_t下一个token在词表中的下标策略 π_θ(a_t | s_t)语言模型当前step的softmax分布即时奖励 r_ttoken级奖励或KL惩罚终态奖励 R(τ)奖励模型对完整回复的评分折扣因子 γ通常取0.95~1.0采样轨迹从prompt开始自回归采样直到结束符这张表说明一个事实语言生成不是强化学习的特殊案例它就是强化学习的标准案例。很多人觉得序列级策略梯度才是常规做法其实那只是因为实现起来省事。如果把生成过程真正当成一个多步MDPtoken级策略梯度反而更贴近理论原貌。有了这套MDP定义下一步就是把它代入策略梯度定理得到最终的token级损失形式。3. Token级策略梯度损失的数学推导与最终损失形式3.1 从总回报期望推导策略梯度策略梯度的起点是最大化期望总回报J(θ) E_{τ ~ π_θ} [ R(τ) ]对参数求梯度用log-derivative trick可以得到∇_θ J(θ) E_τ [ Σ_t (G_t - b(s_t)) · ∇_θ log π_θ(a_t | s_t) ]其中 G_t 是从第t步到序列末尾的折扣累计奖励G_t Σ_{kt}^{T} γ^{k-t} r_kb(s_t) 是baseline用来降低方差当取 b(s_t) V^π(s_t)状态价值函数时G_t - b(s_t) 就变成了优势函数 A_t。注意这个式子天然是token级的每个时间步 t 都有属于自己的 log π_θ(a_t | s_t) 和权重 A_t。序列级损失其实就是把 A_t 强行换成整条序列总回报 R(τ)等于丢掉了中间所有优势信息。Token级做法的核心只是把公式里省略掉的 A_t 老老实实算出来。3.2 优势函数的Token级估计MC、TD与GAE既然上面出现了 A_t下一个问题就是A_t怎么算。这里有三条路线复杂度和方差特性各不同。Monte Carlo估计直接用后续真实奖励的总和 G_t 减去baseline。无偏但方差很大因为未来若干步的采样噪声全部堆在一起。TD估计用 A_t r_t γV(s_{t1}) - V(s_t)方差小但有偏差而且对value function的准确性非常敏感。GAEGeneralized Advantage Estimation在两者之间取折中。GAE的定义是δ_t r_t γ V(s_{t1}) - V(s_t)A_t^{GAE} Σ_{l0}^{T-t-1} (γλ)^l δ_{tl}其中 λ ∈ [0,1] 控制偏差与方差的平衡。λ越接近1越像MC估计偏差小方差大λ越接近0越像TD估计方差小偏差大。在生成式模型的token级训练里我常用 λ0.95、γ0.99这个组合在大多数对话生成场景下表现都比较稳如果value网络训练得比较好可以适当把 λ 提到0.97~0.99。3.3 常用损失展开式与简化变体有了优势估计token级策略梯度的原始损失就是逐token的加权负对数似然L_PG - E_τ [ Σ_t A_t · log π_θ(a_t | s_t) ]但在实际训练中直接用这个损失容易因为单次更新步长过大而崩掉所以通常会和PPO中的importance ratio结合。假设旧策略是 π_old当前策略是 π_θ定义ρ_t π_θ(a_t | s_t) / π_old(a_t | s_t)那么带clip的token级PPO损失为L_PPO - E_τ [ Σ_t min( ρ_t · A_t, clip(ρ_t, 1-ε, 1ε) · A_t ) ]这里每个token都有自己独立的 ρ_t 和 A_t。和序列级PPO的唯一区别还是在于A_t不再是同一个标量而是每个前缀位置上单独估计的值。更复杂的变体还包括在奖励里混入每一项的KL惩罚和熵正则r_t r_t - β · KL(π_θ(a_t|s_t) || π_ref(a_t|s_t)) η · H(π_θ(·|s_t))KL惩罚让策略不要偏离参考策略太远熵正则防止策略过早熵坍塌。很多开源实现里会把这些项写进损失而不是写进奖励效果上两者基本等价只是对数尺度不同。4. Token级损失落地的实现细节mask、优势估计与归一化数学推导只有一页纸真正让token级策略梯度变得能用的是工程细节。这一节我把实现时最容易出错的几个地方拆开讲内容包括数据组织、mask处理、GAE伪代码、归一化和分布式训练。4.1 数据组织和padding mask训练时一个batch里会有长度不一的序列。通常的做法是在尾部padding到相同长度然后训练时用attention mask和loss mask把padding位置屏蔽掉。在token级策略梯度里需要为每个token准备的数据包括token id、log_probs、values、rewards、advantages、以及一个0/1 mask。下表是一个典型的数据组织方式字段形状说明tokens[B, T]序列token idlog_probs[B, T]旧策略下每个token的log概率用于importance ratiovalues[B, T]value网络对每个前缀状态的估计rewards[B, T]每个token的即时奖励只有最后token有终态奖励时其他位置为0mask[B, T]有效token为1padding为0advantages[B, T]按GAE计算得到的优势关键点rewards在padding位置必须为0values在padding位置也建议置0。否则GAE递推时padding位置会参与δ_t的计算导致同一个batch里短序列的最后一个有效token被padding处的value污染产生了一个“伪动作价值”。4.2 Token级优势与损失伪代码直接给一套PyTorch风格的实现它会在batch维度上并行做GAE递推同时用mask隔离padding。def compute_gae(rewards, values, mask, gamma0.99, lam0.95): rewards: [B, T] float values: [B, T] float mask: [B, T] float, 1表示有效token, 0表示padding B, T rewards.shape # 序列末尾的下一个状态value视为0 next_values torch.cat([values[:, 1:], torch.zeros(B, 1)], dim1) deltas rewards gamma * next_values - values deltas deltas * mask # padding位置不参与优势递推 advantages torch.zeros_like(deltas) running torch.zeros(B) for t in reversed(range(T)): running deltas[:, t] gamma * lam * running * mask[:, t] advantages[:, t] running return advantages注意递推时running也要乘mask否则前一个序列的累计优势会跨过padding边界“漏”到下一个序列。这个bug非常隐蔽因为loss表面上看起来很正常但实际上相邻序列在互相影响。策略梯度的损失计算如下def token_pg_loss(log_probs_old, log_probs_new, advantages, mask, clip_range0.2): # log_probs_old/new: [B, T]都是log_softmax后按token id gather得到的 ratio torch.exp(log_probs_new - log_probs_old) surr1 ratio * advantages surr2 torch.clamp(ratio, 1.0 - clip_range, 1.0 clip_range) * advantages loss -torch.min(surr1, surr2) * mask return loss.sum() / mask.sum()这里的advantage建议在Enter loss之前做一次归一化见下一节。4.3 Whitening、reward scaling和KL惩罚实战中直接用原始GAE计算出的advantage做训练效果往往不如先归一化。常用三种方式全局whitening所有有效token的advantage减去均值、除以标准差得到一个均值为0、方差为1的分布。绝对值均值归一化除以mean(|A|)保留符号方向但不强行把均值拉到0。适合advantage全为正或全为负的情况。固定常数缩放直接用常数 c 控制更新幅度。我自己的经验是如果batch里序列长度差异大尽量用全局running统计做whitening而不是每个batch内独立重新计算均值和方差。因为每个batch的样本分布可能漂移batch内独立标准化会让同一个优势值在不同batch里被放大或缩小训练容易震荡。reward scaling同样需要注意。最终奖励乘以一个系数 c会等价地放大或缩小所有advantage。很多实现里把reward缩放到0到1之间再配合 β0.01~0.1 的KL惩罚。调参顺序建议是先固定reward scaling再调β最后动了reward scaling后必须回头重调β两者耦合很强。每次改动记录per-token KL的变化如果KL在一两千步内突然塌到0附近多半是β过大或reward scaling过小。4.4 多卡训练中的全局统计与同步当训练从单卡切换到多卡DDP时token级策略梯度会多出两个不大不小的问题。第一个问题是advantage归一化的统计量。如果你用的是batch内whitening那么单卡basebatch的均值和全局大batch的均值会不一样global batch越大样本方差越小导致advantage的绝对值被动态压小策略更新幅度也跟着变。稳妥的做法是跨卡做统计量的同步或者干脆使用固定常数缩放例如把advantage除以一个固定值如5.0或10.0不再每次动态计算。第二个问题是数据采样与生成的一致性。如果生成rollout是在多卡上并行做的同一批prompt在不同卡上采样得到的回复可能分布在不同的batch里。这本身没问题但KL惩罚的期望统计会因此波动更频繁——某个batch里恰好全是高KL的样本另一个batch里全是低KL样本训练loss看起来忽高忽低。解决办法是观察滑动窗口的均值不要因为某一步的剧烈波动就去调学习率。这块容易误导人很多新人会误以为是学习率问题实际是采样分布的方差在作祟。5. 训练工程中反复踩到的问题与调参心得理论再清晰落到工程上总是有一堆“文档不会写”的细节。下面这些坑我基本都在项目里亲手踩过有些花了很长时间才定位到根因。5.1 时间步错位与长度偏置生成式模型训练时batch内序列被padding到同一长度所以“第t时间步”在不同的序列里对应的绝对位置通常是相同的。但如果你的训练框架支持序列打包packing多个短序列被拼接在一条样本里这时“token下标”和“MDP时间步”就不再是一回事了。打包场景里GAE递推必须按原序列边界切开不能把两条序列头尾相接。否则第二条序列的第一个token会被当成第一条序列的延续value和advantage全混掉。实现上我建议打包样本额外维护一个segment_id并在计算GAE时按segment_id对mask做分段处理。如果只是为了对齐浪费时间不如老老实实用padding mask。另一个相关问题是长度偏置。长序列天然有更多token在梯度中占的weight更大如果每条序列的loss按token数平均倒是无所谓但如果按sum reduction长序列会主导训练方向。建议loss reduction用有效token数的batch平均值同时关注长度统计直方图确保训练集里没有少数“超长样本”把损失撬偏。5.2 优势约等于零时的“假稳定”token级策略梯度实现到位之后你会看到training loss曲线很好看地下降但生成样本质量提升非常慢。这时候要小心“假稳定”。原因是当value网络拟合得足够好时GAE优势的平均绝对值会越来越小loss随之变小但策略的更新幅度可能远低于预期。判断模型到底有没有在学不要只看loss还要看两个指标一是advantage绝对值的滑动均值如果长期低于0.01说明更新步长被“无优势可乘”的处境限制了二是生成样本端reward的滑动均值必须持续监控。遇到这种情况可以尝试增大reward scaling、降低GAE的λ让优势更接近MC或者让value网络别更新得太快人为制造一些“优势修正空间”。这里有一个trick如果发现优势全部挤在一个很小的区间但方向是对的可以不用whitening直接用sign-preserving归一化比如 A_hat A / mean(|A|)。这样既保留每个token间的相对大小又把整体量级放到1附近策略更新幅度更可控。5.3 策略崩溃与熵保护生成式策略梯度最怕的一件事是策略过快坍缩成一个窄分布。Token级策略会一条一条地更新如果某个token位置偶然获得了正优势模型可能把该位置的分布拉得很尖哪怕这个尖只是噪声造成的。为了对抗这个我一般会在损失里加熵正则项系数从0.01起步。另外一个更直观的做法是设定一个“熵下限”如果某个batch里平均熵低于参考策略熵的60%就暂时把策略梯度幅度降一半。实现上不复杂维护一个滑动平均熵低于阈值时报警或调整系数即可。KL惩罚同样能帮助防坍缩但注意不要惩罚过头。KL惩罚是逐token累加的长序列累计KL惩罚会非常大如果β沿用序列级实验里的数值长文本生成任务的策略基本被压死。我的经验是token级训练里β要比序列级小一个数量级。举个例子序列级里 β0.02 还好token级可能要先从0.002开始再根据per-token KL反馈调整。5.4 一些少有人写但很实际的细节最后补充几个零碎但会影响成败的点都是我实际踩过的采样temperature要和部署一致。训练时用temperature0.8采样生成的rollout算出的advantage如果线上推理用greedy策略优化方向和部署行为可能错位。最好在rollout阶段就固定一个和部署接近的采样策略。value网络不要和策略同步更新。经验上value网络的更新频率可以比策略低或者让value网络使用EMA参数计算target否则advantage会一直在移动靶子上训练很难收敛。定期把value optimizer的步长设成策略的一半效果更稳。log_probs要用float32甚至更高精度保存。生成rollout时保存的旧策略log_probs如果在半精度下和当前策略算ratio时误差会被放大clip边界附近的token尤其敏感。一个很小的精度误差可能导致几千token的梯度被错误clip。离线buffer不要屯太久。Token级策略梯度优势是基于旧value和旧策略的如果rollout buffer在内存里放太久新旧策略差异变大importance ratio普遍偏离1clip后的梯度信息量下降。我通常每个rollout epoch只保留一小部分历史buffer不超过当前batch的20%。说回我自己的训练习惯gamma取0.99GAE的lambda取0.95advantage做全局running whiteningreward scaling控制在0.05~0.1RL训练一个epoch后重新采样一版rollout喂给模型。这套组合在我的多个生成式对齐项目中都跑得比较稳不一定适合所有场景但至少是个不容易翻车的起点。真正适合自己的参数还是要在始终监控生成质量reward和per-token KL的前提下慢慢调出来。