最近我把一套偏好对齐流程从 PPO 换到 DPO，跑完一轮之后发现一个挺微妙的问题：模型在整体胜率上看着还行，但只要去抠细节——比如指令里明确写了“不要输出解释”“必须包含编号”，它总在这些关键位置翻车。我一开始以为是数据量不够，后来看到一篇被顶级会议接收为 Oral 的 DPO 改进工作，才反应过来问题可能出在 DPO 的“评分方式”上：它看的是整条回复的总分，而不是具体哪个词贡献了这个总分。这篇工作把 DPO 的句子级优化拆到 token 级，用 Token 重要性去重塑大模型对齐目标。我随后在自己手头的数据上复现了一遍，踩了不少坑，也摸清了它的适用边界。这篇文章就把原理、复现步骤、超参细节和选型建议一次性说清楚。
1. 从“总分制”到“关键分制”：DPO 的隐式奖励到底丢了什么
1.1 DPO 的原貌：一个标量是怎么把偏好压进去的
DPO 的核心思想说起来很直接：不用额外训练一个奖励模型，而是直接用当前策略模型和参考模型的概率比，构造一个“隐式奖励”。对于一条回复 y，隐式奖励可以写成
r(y) = β·[ log π_θ(y|x) - log π_ref(y|x) ]
其中 π_θ 是正在训练的策略模型，π_ref 是冻结的参考模型（一般是 SFT 后的模型），β 是控制偏好强度的温度系数。在偏好对 (y_w, y_l) 上，DPO 的损失函数是
L = -log σ( β·( log π_θ(y_w|x) - log π_ref(y_w|x) - log π_θ(y_l|x) + log π_ref(y_l|x) ) )
意思是：好回复 y_w 的隐式奖励要比坏回复 y_l 的隐式奖励高出越多越好。整个过程里，模型拿到的“监督信号”就是括号里这一个标量——好回复和坏回复各自的整句对数概率差。
我常用一个生活类比来解释这件事：这就像老师批改一份主观题试卷，只看总分，不给每道题的得分。学生知道自己总分提高了还是降低了，但不知道到底是哪道题做对了、哪道题做错了。DPO 也是这样，它只知道“整句话更好”或“整句话更差”，缺少“哪个 token 导致了这个结果”的细粒度信息。
1.2 只看总分意味着什么：关键 token 被稀释
问题在于，整句对数概率是一个“累加”的结果。假设一条回复有 50 个 token，其中某个关键 token 决定了回复是否合规，比如“需要带身份证”里的“身份”两个字，而其余 48 个 token 都是“的”“了”“，”“。”这类对语义影响极小的成分。好回复和坏回复在那个关键 token 上的 logprob 差异可能高达 3-4，但无关 token 的差异接近 0。累加到句级之后，关键差异被大量无关 token 平均稀释，最终 DPO 的梯度在关键 token 上的信号占比可能不足 10%。
这带来的实际后果是：模型学到的不是“把关键位置改对”，而是“整体概率往某个方向平移”。它可能在第一轮训练后开始习惯性缩短句子、改变语气、调整标点，但真正决定好坏的那个词还是经常错。我在复盘自己实验时遇到过非常典型的情况：模型回答的流畅度明显提升，但“必须拒绝该请求”类指令中的“拒绝”二字，会在部分生成样本里被替换成“考虑接受”，而标准 DPO 的 loss 并没有给这种错误足够的惩罚。
1.3 为什么 PPO 能拿到 token 级信号，而 DPO 拿不到
PPO 的对齐流程是：训练一个奖励模型，用奖励模型给整条回复打分，再用某个基线（critic）估计每个 token 的 advantage。虽然奖励模型打分本身也是句级标量，但 PPO 会把“这个 token 之后的总收益预期”和“当前 token 的实际收益”做差，得到 per-token 的 advantage 信号。所以 PPO 天然有 token 级的信用分配能力，只是代价极其高昂：要维护奖励模型、critic、GAE、KL 调度，训练稳定性还差。
DPO 之所以受欢迎，就是因为它把这一套全砍了，直接用参考模型概率比构造隐式奖励。但这个封闭解同时把 token 级的信息也丢掉了——隐式奖励最终表现为一个总和标量，模型只能感知“整句概率比”的变化，无法区分哪个 token 更重要。这其实不是数学上的缺陷，而是目标函数设计层面的取舍。TI-DPO 想做的事情，就是在保留 DPO 简单性的前提下，把“重要性”这个概念重新塞回目标函数。
2. TI-DPO 的解题思路：重要性从哪来，权重用在哪
2.1 定义 Token 重要性：先回答“哪个词决定了胜负”
TI-DPO 最关键的改动，是引入了一个 per-token 的重要性权重 w_t。但权重不是凭空来的，它需要回答一个问题：在这个偏好对里，哪个 token 的好坏决定了整条回复的好坏？目前主流实现里有几种估计重要性的思路。
第一种是 per-token log-ratio 绝对值法。因为句子级隐式奖励本质上是每个 token 的 log π_θ(y_t) - log π_ref(y_t) 的总和，那么自然可以把每一项拆出来，作为“这个 token 有多重要”的度量。我用的公式是
Δ_t = β·( log π_θ(y_w_t|x,y_<t) - log π_ref(y_w_t|x,y_<t) - log π_θ(y_l_t|x,y_<t) + log π_ref(y_l_t|x,y_<t) )
然后定义权重 w_t = softmax( |Δ_t| / τ )，τ 是温度。|Δ_t| 越大，说明好回复和坏回复在这个 token 位置的概率差异越大，也就越可能是决定偏好反转的关键位置。
第二种是“替换扰动法”。把好回复中某个 token 替换成坏回复对应位置的 token，观察句子级隐式奖励变化了多少。变化越大，重要性越高。这个思路更贴近“反事实”直觉，但计算量随序列长度线性增长，训练时基本不现实，只适合离线分析小批量样本。
第三种是影响函数法。用 Fisher 信息或者梯度范数近似估计“如果去掉这个 token，模型参数更新方向会怎么变”。这个理论上很干净，实现起来却繁琐，我日常复现时没有采用。
实际项目里，第一种方法最值得用：它只需要一次前向计算，就能同时拿到 θ 和 ref 的 per-token logprob，没有任何额外开销。我自己的观测是，那些真正决定输出正确性的 token——比如“拒绝”“不能”“必须”“身份证”——在偏好对中的 |Δ_t| 通常会显著高于相邻的虚词。这个直觉和 TI-DPO 的动机是一致的。
2.2 损失函数变形：加权求和，而不是加权平均
拿到权重之后，有三种常见的损失改造方案，我分别说一下利弊。
第一种是软加权隐式奖励版：把句子级差异加权汇总，公式为
L = -log σ( β·Σ_t w_t·Δ_t )
这里 w_t 经过归一化，总和为 1。当 w_t ≡ 1/T 时，它退化成标准 DPO（按长度平均）。当 w_t 集中在少数 token 上时，模型主要优化那些关键位置。这个版本的优点是数学上和标准 DPO 最接近，训练稳定性最好；缺点是 w_t 和 Δ_t 耦合在同一个 sigmoid 内部，如果权重更新太频繁，目标函数会跟着抖。
第二种是逐 token 独立 sigmoid 版：每个 token 单独算一个 sigmoid 损失，再按权重求和，公式为
L = -Σ_t w_t·log σ(β·Δ_t)
这个版本的好处是每个 token 都有独立的偏好信号，坏处是它改变了语义：原来 DPO 比较的是“整句谁更好”，现在变成了“每个位置谁更好”。很多 token 位置本身无所谓好坏，强行给它们单独的 sigmoid 压力，反而会让模型在句法层面产生奇怪的偏移。
第三种是硬掩码版：只对 top-k 重要性的 token 计算标准 DPO 损失，其余 token 完全不参与。这个实现最简单，效果也最粗暴。我试过按 top-20% token 做掩码，关键任务上效果不错，但训练初期容易因为权重估计不稳定而崩。
我最终采用的是第一种，也就是软加权隐式奖励版。它既保留了 DPO 的整体偏好语义，又把训练信号的重心引向关键 token。需要强调的是，w_t 是“权重”而不是“标签”，它只在聚合时改变 token 的贡献比例，并不改变 DPO 本身的对比逻辑。
2.3 为什么不用 PPO 或 RLAIF 替代
很多人会问：既然要 token 级信号，直接用 PPO 不就行了？答案是可以，但不划算。PPO 需要额外训练奖励模型，需要为每个 batch 维护 critic 网络，还需要精心调节 KL 惩罚系数和 GAE 参数。在中小规模的团队或单卡实验环境里，PPO 的调试成本非常高，稍微一个参数没调好，策略就崩掉。
TI-DPO 的价值恰恰在于：它只改损失函数，不改训练框架。参考模型还是冻结那一个，前向计算还是两次——一次当前模型，一次参考模型——只是把原来的“总额比较”变成“加权总和比较”。这意味着你现有的 DPO 数据管道、混洗逻辑、评估流程全部可以复用，只需要替换 loss 计算部分。这种“低成本升级”的思路，在实际工程里远比再搭一套 RLHF 管线现实得多。
3. 复现 TI-DPO：从公式到训练脚本的关键细节
3.1 数据准备与前向计算：logprob 提取最容易出错
复现 TI-DPO 的第一步，仍然是准备标准偏好对数据：每个样本包含 prompt、chosen（好回复）、rejected（坏回复）。数据质量这里不多讲，直接说技术细节。
计算 per-token logprob 时，我踩过的第一个坑是 padding 对齐。很多开源实现里，chosen 和 rejected 的长度不一样，batch 内会做右 padding。但右 padding 的情况下，模型在生成位置看到的 attention mask 没问题，计算 logprob 时却容易把 padding 位置的“无意义 token”也纳入计算。正确做法有两种：要么在 tokenizer 侧统一左 padding，要么在 loss 计算时显式构造 label mask，只保留回复部分的真实 token。
第二个坑是 logprob 的提取方式。对 decoder-only 模型，每个 token 的 logprob 不是直接拿最后一层 hidden state 的 softmax 值，而是用交叉熵损失函数在 label 位置上的输出。具体的做法是：让模型输出 logits，形状为 [batch, seq_len, vocab_size]，然后对 label 序列计算交叉熵，关闭 reduction，得到每个 token 的负对数似然。我建议用模型自带的 loss 函数接口，而不是手动做 log_softmax，因为很多模型在 embedding 或 loss 里做了 label smoothing 之类的小处理，手动实现容易不一致。
第三个坑是 prompt 部分要排除。我们只关心回复部分的 logprob，prompt 里的 token 不应该影响隐式奖励。实现时，把 prompt 长度记为 prompt_len，在 label mask 里把前 prompt_len 个位置全部置为 -100，只保留回复部分的标签。此外，chosen 和 rejected 共享同一个 prompt，所以 prompt 部分的 logprob 理论上是一样的，DPO 公式里也会被减掉，但在 TI-DPO 的逐 token 加权里，如果不 mask 掉 prompt，权重计算会混入 prompt 位置，导致权重分布偏向 prompt 前半段。
3.2 计算重要性权重：温度、归一化和冻结策略
拿到两个回复的 per-token logprob 后，就能计算 Δ_t 和权重 w_t。我复现时的具体流程分四步。
第一步，先冻结参考模型，用参考模型分别计算 chosen 和 rejected 的 per-token logprob。第二步，用当前策略模型（可以是初始的 SFT 模型，也可以是训练中的模型）计算同样的 per-token logprob。第三步，按公式求 Δ_t，注意 chosen 和 rejected 在 token 位置上是对齐的——如果两者长度不同，一般会在数据构造阶段把较短的一侧 padding 到相同长度，或者按字符对齐切分。第四步，用 softmax(|Δ_t| / τ) 得到权重，并按序列长度做归一化，确保长样本不会天然获得更大总权重。
关于 τ 的取值，我建议从 1.0 起步。τ 越大，权重分布越均匀，TI-DPO 退化得越接近标准 DPO；τ 越小，权重越尖锐，优化信号越集中。我自己的实验里，τ=0.7~1.0 是个比较稳的区间。如果 τ 小于 0.3，权重几乎只落在 1-2 个 token 上，训练曲线会剧烈震荡。
还有一个容易踩的细节：权重到底应该由哪个模型来计算？如果每一轮都用当前策略模型重新算权重，那么权重和目标函数会互相影响，形成正反馈循环——模型把某些无关 token 的 Δ_t 放大，权重又进一步聚焦到这些 token 上，最后训练发散。稳妥做法是：权重只由参考模型或初始模型的 logprob 计算一次，训练过程中固定下来。如果希望更精细，可以每 200~500 步用当前模型重估一次权重，但要配合学习率衰减，避免目标函数来回跳。
3.3 核心训练脚本：一个可用的伪代码版本
下面给出一段我实验里使用的 PyTorch 风格伪代码。它不依赖任何特定框架，只展示 TI-DPO 的 loss 计算逻辑。
def compute_logprobs(model, input_ids, attention_mask, labels): # labels 已经将 prompt 部分设为 -100 outputs = model(input_ids=input_ids, attention_mask=attention_mask) logits = outputs.logits # [batch, seq_len, vocab_size] loss_fct = torch.nn.CrossEntropyLoss(reduction="none") per_token_loss = loss_fct( logits.view(-1, logits.size(-1)), labels.view(-1) ).view(labels.size()) # [batch, seq_len] # 将 padding / prompt 位置的 loss 置零 per_token_loss = per_token_loss * (labels != -100).float() return -per_token_loss # 转为 per-token logprob def weighted_dpo_loss(policy_logprob_w, ref_logprob_w, policy_logprob_l, ref_logprob_l, beta=0.2, tau=1.0): # 计算 per-token 差异 delta_w = beta * (policy_logprob_w - ref_logprob_w) delta_l = beta * (policy_logprob_l - ref_logprob_l) delta = delta_w - delta_l # [seq_len] # 计算重要性权重，只使用绝对值 with torch.no_grad(): weights = torch.softmax(delta.abs() / tau, dim=-1) # TI-DPO 损失：加权隐式奖励版 implicit_reward = (weights * delta).sum(dim=-1) loss = -torch.nn.functional.logsigmoid(implicit_reward).mean() return loss这段代码里有一个关键点：weights 的计算被包在torch.no_grad()里。这是为了保证权重不会回传梯度。因为权重的作用是“重新分配不同 token 的贡献”，而不是“被优化对象”；如果权重也参与梯度，模型会想办法通过操纵权重来降低 loss，而不是通过真正提升关键 token 的概率。
实际训练时，我会把 label 里 prompt 部分设置为 -100，padding 设置为 -100，只保留回复部分的 token。这样 per-token logprob 的形状和权重形状完全对齐，不会出现维度错位。
3.4 超参与稳定性：β、τ、权重上限
复现 TI-DPO 时，我总结出三个影响训练成败的超参。
第一个是 β。标准 DPO 的 β 常用区间是 0.1~0.5，但加权版本里，因为权重集中会放大有效梯度，β 往往需要适当调低。我一般从标准 β 的 70% 开始试。比如原来 β=0.3，TI-DPO 可以试试 β=0.2。如果 β 太大，模型会对关键 token 的变化反应过度，生成结果出现“局部的、不自然的重复修正”，比如某个关键词被反复强调，语气变得机械化。
第二个是 τ。这个前面说过，从 1.0 开始通常比较安全，然后再根据权重分布的尖锐程度微调。判断标准是：打印一批样本的 top-5 权重 token，如果它们确实是语义上决定性的词（否定词、指令动词、专有名词），说明 τ 合理；如果 top-5 全是“的”“了”“，”之类，说明权重分布太均匀，需要调小 τ。
第三个是权重上限。softmax 出的权重本身严格小于 1，但如果序列很长，归一化后单个 token 的权重可能仍达到 0.3 以上。我建议给权重设一个硬上限，比如 0.3，防止某个 token 主导整个样本的优化方向。这种做法很像在 loss 里做梯度裁剪，本质上都是为了稳定。
4. 效果观察与避坑：对齐质量、训练稳定性和我踩过的坑
4.1 我在复现中看到的提升模式
我在一组内部偏好数据上做了对比实验，对照组是标准 DPO，实验组是 TI-DPO（τ=0.85，β 调低 20%，权重由初始 SFT 模型预计算并固定）。数据主要来自指令遵从类场景，比如“必须用列表回答”“不要提及某个专有词”“输出必须包含数字编号”。
最明显的差异出现在“关键位置修正”上。标准 DPO 训练出的模型，经常把“不要提及”理解成“不要过多提及”，在边界场景下会输出敏感词；TI-DPO 训练出的模型，在同样的生成样本里，对否定词的遵循率明显更高。我抽查生成结果时还发现一个细节：TI-DPO 模型在遵循“必须包含编号”这类硬约束时，编号格式基本不会乱；而标准 DPO 模型偶尔会把编号漏掉，或者从“1.”突然跳到“4.”。
这不是说 TI-DPO 在所有任务上都更强。在“让回复更友好”“让回复更幽默”这类整体风格偏好上，TI-DPO 和标准 DPO 的差异很小，甚至偶尔更差。原因也很简单：风格类偏好散布在几乎所有 token 上，重要性权重被平均化，加权操作基本退化成标准 DPO，反而多了一层计算噪声。
4.2 坑一：权重抖动导致训练崩溃
我第一次复现时犯了个错误：权重在每个 step 都用当前策略模型重新计算。那几天训练曲线像过山车，前一步 loss 还在下降，下一步突然飙升，生成结果开始出现大量重复 token。后来排查发现，问题出在权重和目标函数的耦合上。
具体来说，随着训练推进，模型会在某些无关 token 上产生较大概率波动。这些波动让 |Δ_t| 变大，softmax 后的权重进一步聚焦到这些 token 上，模型又被逼着调整这些 token 的概率，形成自激震荡。修复方式就是我前文说的：让权重来自参考模型或初始模型的前向结果，并且计算权重时关闭梯度。如果确实需要定期更新权重，我建议用一个滑动平均版本：w_t = α·w_t_old + (1-α)·w_t_new，其中 α 取 0.9 左右，可以明显抑制震荡。
4.3 坑二：正负样本长度差异导致权重失真
第二个坑和数据的真实分布有关。在不少开源偏好数据集里，chosen 通常比 rejected 长，因为标注者倾向于认为“更详细、更完整的回复”更好。但 TI-DPO 的序列级对齐依赖 chosen 和 rejected 在 token 位置上对应，如果长度差异过大，后面多出来的一截 token 会让权重归一化产生偏差。
我遇到的典型问题是：chosen 有 80 个 token，rejected 只有 30 个 token，padding 到 80 之后，rejected 端有 50 个位置是 padding。这些 padding 位置如果没被 mask 干净，Δ_t 会变成一个很大的负值，softmax 权重会大量集中在 padding 位置，真实 token 反而得不到权重。解决办法是：数据预处理时强制将 chosen 和 rejected 截断到相同长度，或者用“回复尾部对齐”的方式拼接——即让两条回复的最后一个真实 token 对齐，而不是开头对齐。尾部对齐在事实型任务里比头部对齐更合理，因为关键信息往往分布在回复中后段。
4.4 评估建议：别只盯着胜率
如果你复现 TI-DPO，我强烈建议不要只用 pairwise win rate 做唯一指标。因为 win rate 是句级指标，它和标准 DPO 的优化目标同构，不容易体现 token 级加权的优势。更好的做法是在指标集里加两类：一类是关键位置准确率，另一类是约束违反率。
具体操作是：写一个小脚本，针对每条测试样本定义关键断言。比如“输出里是否包含编号列表”“输出里是否出现禁用词”“数字格式是否是两位小数”。然后分别统计标准 DPO 和 TI-DPO 在这些断言上的通过率。我自己的实验里，TI-DPO 在“禁用词出现率”上比标准 DPO 低大概一个量级，在“编号完整性”上高了十几个百分点。这些差异如果用整体胜率看，可能只有两三个点，容易被噪声淹没。
5. 它适合你的任务吗：选型建议与适用边界
5.1 优先推荐 TI-DPO 的场景
根据我的复现经验，TI-DPO 在四类任务上优势明显。
第一类是“硬约束指令遵循”。比如系统提示词里写了“必须拒绝任何涉及隐私的请求”，一条回复的好坏几乎完全取决于模型有没有在关键位置正确输出“拒绝”。第二类是“事实性修正”，尤其是数字、专有名词、日期这类信息，一个 token 错了整句就错了，重要性天然集中。第三类是“安全对齐中的拒绝行为”。这类任务里，好回复和坏回复往往只在开头几句话分叉，后面内容差异不大，句级 DPO 容易被长尾内容稀释，TI-DPO 能把信号拉回关键分叉点。第四类是“结构化输出”，例如 JSON 格式里的{、}、"key"这些格式化 token，它们数量少但决定了解析成败，权重集中非常有益。
5.2 不建议硬上的场景
如果你的任务本质是“全局风格迁移”——比如让回复更温柔、更幽默、更像某个特定作者——TI-DPO 大概率帮不上忙。这类偏好的信号散布在全句，重要性权重算出来会非常平坦，加权操作退化成 DPO 之后还会引入额外的权重估计噪声。
另一个不建议的场景是数据量很小的项目。假设你只有几百条偏好对，权重估计本身方差很大，再叠加温度 τ 的调节，训练很容易不稳定。这种情况下，老老实实跑标准 DPO，再把精力花在清洗数据上，收益通常更高。我自己在只有 600 条数据的实验里试过 TI-DPO，loss 曲线明显比标准 DPO 更毛糙，最终生成质量也没有稳定的提升。
5.3 后续可以这样扩展
TI-DPO 的“Token 重要性加权”思路，完全可以和其他对齐技巧组合。比如和长度控制结合：对过短回复里的填充词降权，对必要信息 token 加权；比如和 chunk 级对齐结合：不只看单个 token，而是把回复切成语义块，在块级别做重要性加权，这样可以降低逐 token 权重估计的方差；再比如和奖励模型蒸馏结合：先用奖励模型对样本打分，把分数差值作为重要性的辅助信号，弥补纯概率比估计的不足。
这些方向我没有全部跑完，但至少从 TI-DPO 的结果看，把“整句有多少差异”拆解成“差异具体在哪里”，是偏好对齐里一条值得继续挖的路。
最后再分享一个我实验中总结的小技巧：如果你训练时发现 TI-DPO 和标准 DPO 结果差不多，先不要急着怀疑方法，去打印几个 batch 的权重分布看看。如果 top-5 权重 token 全是虚词，说明重要性度量没有抓到关键位置，多半是 τ 太大或数据本身偏好差异不集中；如果 top-5 权重 token 确实是关键的指令词、否定词、专有名词，但效果还是没提升，那才需要去检查优化器、学习率这些通用环节。重要性权重是一把放大镜，它的作用是让你更清楚地看到该优化的地方，但如果你原本就看错了位置，放大只会让错误更明显。
