跳到正文
原文
arXiv:对齐、欺骗与监督· arXiv:2610.03641· Vladislav Gromadskii, David Li, Samson Gourevitch, Yazid Janati, Eric Moulines, Maxim Panov, Alexander Korotin·本站收录· 原文发表

IDRF:掩码离散扩散模型的逆蒸馏奖励微调

IDRF: Inverse-Distilled Reward Fine-tuning of Masked Discrete Diffusion Models

论文速读

防御

据论文 PDF 整理(AI 生成),以原文为准

IDRF正则少步奖励微调:DNA 8步ATAC-Acc 95.3%,并减轻reward hacking。

问题
Masked diffusion的序列似然要对去噪轨迹边缘化,纯奖励微调会破坏可及性、图像或评论停止等未度量性质。作者要用不依赖参考rollout的序列级KL替代,以便直接微调少步采样器。
方法
从逆KL奖励目标出发,用参考与辅助去噪器的NELBO间隙代替序列KL。少步生成写成终点才给奖励的MDP,在学生自己的轨迹上用leave-one-out策略梯度更新,不需要参考rollout或奖励梯度。
实验与结果
DNA 8步IDRF:ATAC-Acc 95.3%,无正则37.8%(Table 1)。图像4步ImageReward 0.88,对比128步reference的−0.73(Table 2)。文本32步IDRF的EOS 100%,RLOO 9.4%。
局限与可以继续做的
作者指出学生与辅助的min-max可能不稳定、对更新平衡敏感,辅助增加训练成本;实用替代只给一步补全混合分布的KL界,终点分布的控制仍开放。后续是更大模型和更多领域。实验覆盖HepG2 DNA、10类图像和Amazon Fashion评论。
实验设置HepG2 regulatory DNA (200-bp, DRAKES)、ImageNet attribute prompts (10 classes) 等 · Pred-Activity、ATAC-Acc 等
基准
HepG2 regulatory DNA (200-bp, DRAKES)ImageNet attribute prompts (10 classes)Amazon Fashion reviews
指标
Pred-ActivityATAC-Acc3-mer correlationLog-LikEntropyCLIPImageRewardPickScoreHPSv2.1log P(1 star)GenPPLEOS (%)Norm. entropyGM
AI 导读IDRF 是一种面向少步掩码离散扩散生成器的奖励微调框架,用逆蒸馏正则替代难以计算的序列级 KL 惩罚,并证明在最优辅助去噪器下该损失是序列级 KL 散度的上界。它无需参考模型 rollout,将少步生成视为有限时域马尔可夫决策过程,用裁剪策略梯度目标优化奖励。在 DNA、图像和文本生成任务上,IDRF 以最多减少 32 倍去噪步数取得高奖励,同时缓解奖励作弊并保持样本质量。

IDRF 是一种面向少步掩码离散扩散生成器的奖励微调框架,用逆蒸馏正则替代难以计算的序列级 KL 惩罚,并证明在最优辅助去噪器下该损失是序列级 KL 散度的上界。它无需参考模型 rollout,将少步生成视为有限时域马尔可夫决策过程,用裁剪策略梯度目标优化奖励。在 DNA、图像和文本生成任务上,IDRF 以最多减少 32 倍去噪步数取得高奖励,同时缓解奖励作弊并保持样本质量。

深度解读

6个问题,点开问题读完整回答

  1. 这篇论文试图解决什么问题?

    少步masked diffusion的序列KL不可算,纯奖励微调会破坏未度量性质。

    少步 masked diffusion 的序列似然不可计算,只优化奖励会破坏未度量性质;IDRF 用 inverse distillation 代替序列级 KL。

    • 场景:MDM 并行去掩码,用于语言、生物序列和离散图像。奖励可对准 DNA 活性、图像属性或评论情感,但作者称只优化奖励会毁掉奖励没测的性质(Figure 1)。
    • 计算障碍:序列似然要对全部去噪轨迹边缘化,因而 intractable。向预训练参考做 KL 正则是常见约束,但作者称它因此难以直接用在 MDM 上。作者还称迭代采样代价高。
    • 现有不足:已有工作用去噪似然近似或证据界,或惩罚反向转移之间的 KL。作者称后者要沿 rollout 评估参考,并把正则绑在参考采样路径上。
    • 本文做法:IDRF 把少步生成写成有限时域 MDP,在学生自己的轨迹上做策略梯度,用 inverse distillation 间隙作正则。不需要参考似然、参考 rollout、预训练数据或奖励梯度。
  2. 有哪些相关研究?

    蒸馏、奖励微调和联合微调之外,IDRF把inverse distillation用作少步奖励正则。

    少步蒸馏

    • SDTT(Deschenaux & Gulcehre, 2025)把多步教师蒸馏成更少学生步;Duo(Sahoo et al., 2025)把 consistency distillation 用于 uniform-state discrete diffusion。
    • DiDi-Instruct(Zheng et al., 2026)用积分 KL 做分布匹配;IDLM(Li et al., 2026)用 inverse distillation 训练快速离散生成器。IDLM 在单 token 和联合反向过程上给出匹配界,并把实践中的逐 token 去噪损失留给后续。
    • 作者称这些方法复现教师,而 IDRF 用 inverse distillation 在优化奖励时保持接近参考。

    奖励微调

    • 引言把指向预训练参考的 KL 正则视为标准约束(Ziegler et al., 2019; Korbak et al., 2022),并指出 MDM 序列似然不可计算。
    • d1 与 wd1 用去噪似然近似;SPG 用证据上下界;TUBE 另给出 masked diffusion 证据的上界。AGRPO 把去噪写成逐步动作概率可计算的 MDP;TraceRL 沿记录的推理轨迹更新。DMPO 与 Discrete Tilt Matching 学习奖励倾斜分布。DRAKES 在有奖励梯度时用可微松弛做生物序列设计。
    • 作者只把“要沿 rollout 评估参考、正则绑在采样路径上”说在反向转移 KL 这一路上,没有把这句批评写到似然近似或证据界上。

    联合蒸馏与微调

    • 连续扩散中,Diff-Instruct++ 用积分 KL 正则一步奖励优化;DMDR 结合分布匹配与强化学习;RTDMD 用奖励倾斜的分布匹配训练少步 flow。
    • Soft-Di[M]O 用软嵌入细化 masked image diffusion 的一步生成器。作者称其对抗细化使用真实图像,奖励微调需要可微奖励。
    • 作者称 IDRF 把 inverse distillation 用作轨迹策略梯度的正则,直接在少步机制做奖励微调,且不需要预训练数据或奖励梯度。

    基线与基准 DNA 对比 CG、SMC、TDS、CFG、DRAKES(有/无 KL)、SEPO+GF,以及 Han et al. (2026) 报告的 D2-DPO、VRPO、DDPP-IS、diffu-GRPO、SDPO;任务是 200-bp HepG2 regulatory DNA。图像对比 reference、ReDGE 与去掉 LID 的 RLOO;文本对比 MDLM reference 与同一消融。

    作者把 IDRF 放在“复现教师的蒸馏”和“绑在参考 rollout 上的正则”之间:正则来自 inverse distillation,优化发生在学生自己的少步采样器上。

  3. 论文如何解决这个问题?

    用inverse distillation间隙上界序列KL,并在学生少步轨迹上做策略梯度。

    整体是在标准逆 KL 奖励目标上,用 inverse distillation 换成可计算正则,再在学生自己的少步轨迹上做策略梯度。

    目标与序列级界

    起点是终端奖励减去学生终点分布相对冻结参考的 KL。作者写成 maxθ 𝔼[R(x0)]−λ DKL(pθ|pref) 即提高完整序列奖励,并用 λ>0 拉向 pref。终点律要对全部去噪轨迹边缘化,对数密度一般不可计算。

    替代量是参考去噪器与在学生分布上重拟合的去噪器的 NELBO 之差。作者称,辅助去噪器最优、参考 NELBO 有限、前向掩码独立且日程严格下降时, LID(θ)≥ DKL(pθ|pref) 对任意序列长度和任意学生采样器成立;若参考是其自身终点律的 Bayes 去噪器,则等号成立(命题1,附录A)。作者称,用该群体损失替换 KL,得到 KL 正则目标的一个下界。实践中的损失打在 rollout 状态的软预测上,而不是直接把界写成终点分布本身。

    少步生成器作为 MDP

    • 学生:Gθ 从参考初始化,把部分掩码序列映射为各掩码位置上干净 token 的类别分布,不接收时间输入。本文实验的预训练模型也都省略时间条件。
    • 转移:固定网格上做 SUBS。第 k 步每个仍被掩码的位置以概率 bk 揭示,数值从 Gθ 抽取,已揭示位置复制;b1=1,故 K 次网络评估后得到完整序列。Max-confidence 改为揭示最有信心的掩码位置。
    • MDP:状态是步数和当前序列,动作是下一步部分序列,奖励只在最后一步等于 R。轨迹概率是各步策略之积,优化的是 K 步采样器的期望回报。

    奖励项与正则

    • 策略梯度:bk 不依赖 θ,梯度只涉及该步新揭示的 token,不穿过离散抽样,也不穿过 R。损失是裁剪的 leave-one-out 替代:优势为该条奖励减去同组其余样本均值。Ancestral 在 θ=θold 时,作者称该损失梯度的负号是回报梯度的无偏估计;max-confidence 则固定记录的揭示集合。
    • 软目标:在轨迹上均匀抽深度,学生预测当作单纯形上的目标,再新鲜掩码,由冻结参考和可训练辅助打 NELBO。已揭示位置的目标是复制 token 的 one-hot。
    • 界的对象:作者称,辅助最优且与正则使用同一状态律和同一掩码时,该期望等于混合分布 q̄θ 的群体损失,并上界 q̄θ 到 pref 的 KL;对式(9)采样器,最后一档等于终点律 pθ(推论1)。梯度是 semi-gradient:状态和掩码上下文固定,只穿过学生软概率。前向代价是一次学生评估,加上参考和辅助各一次。

    训练循环

    算法1把两者都从参考初始化,使初始间隙为 0。每轮先在部分 rollout 上拟合辅助,学生停止梯度;再由旧策略采完整 K 步轨迹,记录新揭示位置和 leave-one-out 优势;然后冻结辅助,学生最小化奖励损失加 λ 乘 NELBO 间隙。辅助在结束后丢弃。数值若停在正的噪声地板,剩余掩码由最后一次类别抽样填上,并计入轨迹。任务相关的 K、λ 和更新次数在第4问。

  4. 论文做了哪些实验?

    三项任务上IDRF以8、4、16–32步取得高奖励,纯奖励微调出现可及性或EOS退化。

    实验设置

    • DNA:学生与辅助从 DRAKES 预训练模型初始化,生成 200-bp。训练奖励是其 fine-tuning oracle;Pred-Activity 用 held-out oracle。λID=0.2,K=8,ancestral 或 max-confidence。每采样种子 640 条、三种子,报告均值和样本 SD。
    • DNA指标:ATAC-Acc 是 HepG2 ATAC oracle 的 output 1 超过 0.5 的百分比。3-mer 相对 Gosai et al. 活性高于 99.9 分位的 736 条增强子。Log-Lik 是预训练模型 NELBO 似然估计的中位数,128 评估步、1 个 MC 样本。Entropy 是序列内碱基熵。
    • 图像:类条件 Halton-MaskGIT-L,ImageNet 384×384,24×24 VQ,码本 16384。奖励是 CLIP ViT-L/14 的 100×CLIPScore;类标签进模型,10 个属性提示只进奖励。每类 100 张、共 1000 张,均值±95% CI(类内方差合并后用 Student-t)。K=4 另加一次 unmasking;IDRF 行 w=0。held-out 为 ImageReward、PickScore v1、HPSv2.1(表中 ×100)。
    • 文本:MDLM 的 OpenWebText 检查点经 50000 步 EOS-aware SFT 后,初始化 reference、学生和辅助。Amazon Fashion,序列长 128。奖励是冻结 BERT 分类器的 log P(1 star)。生成行 1000 条,评测种子 1,10000 次更新后。held-out 含 GPT-2-large 的 GenPPL、长度、EOS、Norm. entropy 和 GM。
    • 基线:DNA 用发布实现复测 CG、SMC、TDS、CFG、DRAKES 有/无 KL,并评测 SEPO 检查点。带 * 的行来自 Han et al. (2026),其发布评测为 128 步。图像另有 ReDGE。文本对照是去掉 LID 的 RLOO。
    • 优化:DNA 为 AdamW、学习率 10^{-4}、每轮 256 条 rollout、辅助/学生更新 4/2、500 轮、ε=0.2、训练种子 0(Table 4)。图像为学习率 3×10^{-5}、64 条 rollout(16 组×4)、80 轮、λID=0.5、种子 20(Table 8)。文本为 8 张 A100、全局 batch 64、学习率 10^{-6}、λID=0.01;每批只更新一次,裁剪不生效。

    主结果

    表格较宽,可左右滑动

    设置步数奖励相关未优化性质出处
    DNA ancestral IDRF8Pred-Activity 7.22 (0.01)ATAC-Acc 95.3% (0.7)Table 1
    DNA max-conf IDRF8Pred-Activity 7.31 (0.03)ATAC-Acc 92.3% (0.5)Table 1
    DNA RLOO 无 LID8Pred-Activity 7.68 (<0.01)ATAC-Acc 37.8% (5.5)Table 1
    DNA reference128Pred-Activity 0.18 (0.02)ATAC-Acc 2.0% (0.4)Table 1
    图像 IDRF,w=04CLIP 27.74±0.06ImageReward 0.88±0.02Table 2
    图像 reference,w=3128CLIP 23.11±0.13ImageReward −0.73±0.05Table 2
    文本 IDRF32log P(1⋆) −0.044EOS 100.0%Table 3
    文本 RLOO 无 LID32log P(1⋆) −0.008EOS 9.4%Table 3

    DNA 括号是三种子的样本 SD;图像是 95% CI。表内 bold 排名排除无 LID 消融,且不排 * 行和 Log-Lik。

    • 作者称 ancestral 的 ATAC-Acc 为 Table 1 最高,并与 SDPO* 的 94.8% 相当;max-conf 活性为次高。更高的 SEPO+GF 为 7.66,ATAC-Acc 47.4%。不参与排名的 RLOO 无 LID 在 8 步为 7.68、128 步为 7.43。作者称基线都是 128 步,为 IDRF 的 16×。摘要另称相对 reference 最多少 32× 去噪步。
    • 作者称去掉 LID 后 oracle 被利用:Table 1 中 8 步 3-mer 为 0.378、Entropy 为 1.036,128 步 ATAC-Acc 为 0.1%。Log-Lik 升到 −204 与 −163,高于本文其他测量,但作者把这和塌缩放在一起写。
    • 图像上作者称 4 步且无 guidance 的 IDRF 在每个 held-out 指标上都超过 w=3 的 128 步 reference。文本上作者称 16 与 32 步都保住评论格式,相对 128 步少 4–8× 采样步;32 步奖励略低于 RLOO,对真实一星评论的 GM 作者称低 2.6×(文中 10.9 对 28.8;Table 3 为 10.86 对 28.81)。

    DNA可及性与SEPO复测

    • 组成:Table 7 在 1920 条序列上统计。RLOO 无 LID 在 K=8 只有 13 个 T、K=128 只有 5 个 T,IDRF 为 56344 个;无 T 序列占 99.32% 与 99.74%,IDRF 为 0。CA/CG 相关二核苷酸从 44.46% 升到 72.66% 与 79.61%。
    • 3-mer 不是表中最高:ancestral IDRF 为 0.683,max-conf 为 0.718;DRAKES 为 0.887,SDPO* 为 0.900,D2-DPO* 为 0.944(Table 1)。作者没有称 IDRF 在 3-mer 上最高。
    • SEPO:附录 C.3 称发布笔记本的 ATAC 图用 output 0(GM12878),本文与 DRAKES 用 output 1(HepG2)。复测 SEPO 为 8.70%(1.49),SEPO+GF 为 47.40%(1.82);论文引用的发表值为 99.5%(0.2)与 99.9%(0.09)。作者称这指出公开笔记本的输出索引不一致,并不确定每条发表结果的运行历史。Plain SEPO 的 Pred-Activity 为 7.60(0.01),128 步。

    图像held-out

    • ReDGE:w=3、20 步,另加每图 400 次奖励梯度更新。CLIP 37.77±0.10,为 Table 2 最高;ImageReward 0.30±0.04、PickScore 21.38±0.03、HPSv2.1(×100)25.98±0.10,低于 IDRF 的 22.11±0.02 与 27.72±0.04。作者称它在推理时逐图优化,held-out 仍低于 IDRF。
    • 无 LID 与 4 步 reference:RLOO 的 CLIP 26.12±0.10、ImageReward −0.49±0.03,低于 IDRF;PickScore 20.61±0.02、HPSv2.1 22.57±0.08,低于 128 步 reference 的 21.10±0.04 与 25.05±0.12。4 步 reference 的 CLIP 为 14.00±0.18、ImageReward −2.17±0.01。
    • 图:Figure 1 图注称相同步数预算下 reward-only 图像变模糊。Figure 4 只标明各类第一张样本的排布。正文没有给出可回查的模糊度数值。

    文本停止与GM

    • 不终止:RLOO 有 906/1000 条从不发出 EOS,并达到 127 token 上限(Figure 5;与 Table 3 的 9.4% 一致)。Norm. entropy 为 0.719,真实一星为 0.948。GenPPL 为 76.9,低于 reference 的 91.7;作者称 GenPPL 单独抓不住这一失败。
    • IDRF 的步数差:16 步与 32 步 EOS 都是 100.0%,Norm. entropy 为 0.922 与 0.902。32 步 GenPPL 65.4,低于 128 步 reference 的 91.7;16 步 GenPPL 107.3,高于该 reference。长度:IDRF 32 步 44.1、16 步 37.4,RLOO 126.8,真实一星 32.2,reference 27.8。真实一星的 log P(1⋆) 为 −0.950,reference 为 −5.070。
    • GM:Table 10 相对全部评分的真实评论,32 步 IDRF 为 12.13±0.56,RLOO 为 29.53±0.47。作者称对两个参考,32 步 IDRF 的 GM 比纯奖励微调低 2.4–2.7×,并认为正则减轻了分布偏移。同分布对照为 0.03±0.26 与 0.06±0.24。附录 E 的词重复中位数:RLOO 66%,IDRF 14%,真实一星 9%;论文未按 16/32 步拆开 IDRF。

    其他消融与分析

    • Figure 3,K=8:活性随 λID 增强单调下降;到 0.2,ATAC-Acc 从 37.8% 到 95.3%,3-mer 从 0.38 到 0.68。作者称更大权重再换活性,但不再有进一步收益。
    • Table 6,K=8:forward-masked 完成样本相对 on-policy,ancestral 的 ATAC-Acc 从 95.3% 到 52.5%、3-mer 从 0.683 到 0.273、Pred-Activity 从 7.22 到 5.89;max-conf 的 ATAC-Acc 从 92.3% 到 42.8%、3-mer 从 0.718 到 0.408、Pred-Activity 从 7.31 到 6.23。
    • Table 1,RLOO 无 LID、K=128:Pred-Activity 7.43(<0.01),ATAC-Acc 0.1%(0.1),3-mer 0.240,Log-Lik −163,Entropy 1.072。
    • 附录 B:ancestral 时间网格末端,DNA 为 10^{-5},图像与文本为 10^{-3}。DNA max-confidence 每步揭示 25 个位置。MaskGIT 奖励损失从轨迹中无放回抽 4 个转移,和按长度/4 缩放。LID 状态在 DNA 与 MaskGIT 排除 cleanup,文本包含。
    • 附录 E:文本 K 次反向更新另加一次随机终端步,去噪器评估为 K+1;Table 3 的 Steps 仍写 16 与 32。DNA 评测按 DRAKES 实现,在剩余掩码上取最可能 token;训练的最后一步用类别抽样。
    • 附录 E.1:GM 用两组各 50 条评论的 GPT-2-large token NLL 梯度差做内积,报告 10 次估计的均值和标准误。一星参考是 2748 条合格评论中的 1000 条,全评分参考是 21448 条中的 1000 条。
  5. 有什么可以进一步探索的点?

    作者称min-max可能不稳定,终点采样分布的KL控制仍开放。

    作者指出的局限与后续方向

    • 对抗训练:实践中学生与辅助去噪器是 min–max。作者称,和其他对抗训练一样,这可能带来优化不稳定,并对双方更新的平衡敏感(第6节 Limitations)。
    • 训练成本:辅助去噪器增加训练成本(第6节 Limitations)。论文没有给出墙钟时间或额外 FLOPs。
    • 界的对象:实用替代对一步补全分布的混合有相关 KL 保证;控制最终采样分布仍是开放问题(第6节 Limitations)。
    • 后续:作者将把 IDRF 扩展到更大模型,并延伸到更多领域(第6节)。

    实验覆盖范围

    • 生成模型为 DRAKES 预训练 MDM、类条件 Halton-MaskGIT-L,以及经 EOS-aware SFT 的 MDLM。任务是 HepG2 200-bp DNA、ImageNet 上 10 个类的属性提示、Amazon Fashion 评论(第5节)。
    • DNA 未标 * 的行是 3 个采样种子、每种子 640 条的均值和样本 SD(附录 C.1)。图像每方法 1000 张并给 95% CI(Table 2)。文本生成行 1000 条;GM 是 10 次估计的均值和标准误(附录 E.1)。
    • DNA 对比 CG、SMC、TDS、CFG、DRAKES 有/无 KL、SEPO+GF,以及 Han et al. 报告的五行 * 结果。图像对比 reference、ReDGE 与 RLOO 无 LID。文本对比 reference 与 RLOO 无 LID(Table 1–3)。
    • 论文给出 DNA 的 λID 扫描(Figure 3)和 on-policy 对 forward-masked 的状态对比(Table 6),并在附录 C.3 对照 SEPO 发表 ATAC-Acc 与 HepG2 output 1 复测。
    • 论文写明训练种子为 DNA 0、图像 20、文本训练 0、文本评测 1,未报告多个训练种子下的结果,也未报告墙钟时间。文本与图像在每批一次 on-policy 更新下裁剪不生效;DNA 的 ε=0.2(Table 4、附录 E)。
  6. 总结一下论文的主要内容

    IDRF用inverse distillation正则少步奖励微调,在DNA、图像和文本上减轻reward hacking。

    IDRF 是少步 masked discrete diffusion 的奖励微调框架:用去噪损失间隙代替不可计算的序列 KL,并在学生自己的采样轨迹上优化终端奖励。

    • 问题:MDM 的序列似然要对去噪轨迹边缘化。只优化奖励会丢掉奖励没测的性质。Figure 1 图注把这些失败写成 DNA 缺少可及性、图像变模糊、评论重复且很少以 EOS 结束。
    • 方法:学生从冻结参考初始化,K 步 SUBS 采样写成只在终点给奖的 MDP。辅助去噪器在 on-policy 状态上拟合学生;正则是参考与辅助的 NELBO 差,不需要参考 rollout 或奖励梯度。作者称,辅助最优时群体损失上界到参考的序列级 KL;实用目标上界的是一步补全律的混合。
    • DNA:8 步 ancestral 的 Pred-Activity 为 7.22、ATAC-Acc 为 95.3%;同预算去掉 LID 为 7.68 与 37.8%(Table 1,每种子 640 条、三种子)。作者称基线为 128 步。3-mer 为 0.683,低于 DRAKES 的 0.887。
    • 图像与文本:4 步、w=0 的 IDRF,CLIP 为 27.74、ImageReward 为 0.88,高于 w=3 的 128 步 reference 的 23.11 与 −0.73(Table 2,1000 张)。32 步文本的 log P(1⋆) 为 −0.044、EOS 为 100.0%、GM 为 10.86;RLOO 为 −0.008、9.4%、28.81(Table 3)。16 步 GenPPL 为 107.3,高于 128 步 reference 的 91.7。
    • 作者结论:IDRF 能以更少去噪步取得高奖励,并减轻 reward hacking。摘要写最多少 32×,正文分别写 DNA 的 16× 和文本的 4–8×。他们同时称 min-max 可能不稳定,终点采样分布的控制仍是开放问题。
阅读原文arxiv.org