💬 NLP

Improving Sampling for Masked Diffusion Models via Information Gain

本文针对掩码扩散模型（MDM）现有采样器因忽视当前解码选择对后续步骤的累积影响而存在的局限性，提出了一种基于信息增益的 Info-Gain 采样框架，通过平衡即时不确定性与未来信息增益，在推理、编码、创意写作及图像生成等任务中显著提升了生成质量。

原作者： Kaisen Yang, Jayden Teoh, Kaicheng Yang, Yitong Zhang, Alex Lamb

发布于 2026-03-19

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Kaisen Yang, Jayden Teoh, Kaicheng Yang, Yitong Zhang, Alex Lamb

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

这篇论文提出了一种让 AI 写东西（无论是文字、代码还是图片）变得更聪明、更靠谱的新方法。为了让你轻松理解，我们可以把生成内容的过程想象成**“在一个漆黑的房间里拼图”**。

1. 背景：AI 是怎么“写”东西的？

现在的 AI 主要有两种“写作”方式：

** autoregressive（自回归模型，像现在的普通大模型）：** 就像**“从左到右写文章”**。它必须写完第一个字，才能写第二个字，写完第二个字才能写第三个字。这很稳，但一旦开头写错了，后面很难改，而且写长文章时容易“迷路”。
Masked Diffusion Models（掩码扩散模型，MDM，本文的主角）： 就像**“先画满黑块，再慢慢擦除”。一开始，整张纸都是黑色的（全是乱码或占位符），AI 的任务是决定先擦除哪一块**，把它变成正确的字，然后再擦下一块。

问题来了： 面对满纸的黑块，AI 应该先擦哪一块？

2. 旧方法的毛病：只看眼前，不管未来

以前的 AI 在擦除黑块时，用的是一种**“贪心策略”**（Greedy Heuristic）。

它的逻辑是： “我看这块黑块，擦掉后变成‘苹果’的概率是 99%，非常确定！那我就先擦这块。”
比喻： 就像你在解一道复杂的数学题，你发现最后一步的答案（比如“等于 1"）很容易猜，概率很高，于是你先写下了“等于 1"。但是，前面的计算过程（中间步骤）你还没算呢！
后果： 因为前面的步骤没定下来，后面的“等于 1"可能完全是错的。这就导致了**“为了眼前的确定性，牺牲了整体的逻辑”**。结果就是：AI 写出来的东西，局部看挺像样，整体看逻辑不通，甚至完全胡扯。

3. 新方法：Info-Gain Sampler（信息增益采样器）

这篇论文提出了一种新策略，叫**“信息增益采样器”。它的核心思想不再是“谁最确定就先擦谁”，而是“擦掉谁能让剩下的谜题变得最容易解？”**

核心比喻：解开“俄罗斯方块”或“侦探破案”

想象你在玩一个**“俄罗斯方块”游戏，或者当一名“侦探”**：

旧方法（贪心）： 看到地上有个红色的方块（很容易识别），就先把红色的拼上去。结果发现，这个红色方块挡住了关键线索，导致后面根本拼不出完整的图案。
新方法（Info-Gain）： 它会思考：“如果我先把这个最难、最模糊的方块拼好，虽然我现在有点拿不准，但一旦拼好，整个画面的其他部分瞬间就清晰了，剩下的谜题就迎刃而解了。”

具体是怎么做的？

双向视野： 这种模型（MDM）有一个超能力，它能同时看到整张纸。它不需要像旧模型那样只能看前面。
模拟未来： 在决定擦除哪一块之前，它会快速在脑子里“预演”一下：“如果我擦除 A，剩下的黑块会变得多清晰？如果我擦除 B，剩下的黑块会不会变得更乱？”
计算“信息增益”： 它选择那个**“擦除后，能让剩余所有黑块的困惑度下降最多”**的位置。

简单说： 它不再追求“当下最稳”，而是追求**“当下这一步，能让未来变得最稳”**。

4. 实际效果：真的有用吗？

论文在好几个领域做了测试，效果非常惊人：

做数学题/写代码（逻辑推理）：
- 旧方法： 经常先写个简单的结论，结果逻辑全崩，算不对。
- 新方法： 先解决最难的逻辑步骤（比如先确定乘数，再算乘积），最后得出的答案准确率提升了 3.6%。
- 比喻： 就像先搭好房子的地基和梁柱，而不是先刷最漂亮的墙漆。
写小说/创意写作：
- 旧方法： 容易写偏，或者前后矛盾。
- 新方法： 写出来的故事更连贯，更有创意。在“谁写得好”的比赛中，新方法赢了 63.1% 次。
- 比喻： 就像一位老练的编剧，知道哪里埋下伏笔能让后面的剧情更精彩，而不是只顾着把第一句台词说得最响亮。
画图片：
- 画出来的图更符合描述，位置关系更准确（比如“左边有一只猫”真的在左边，而不是乱跑）。

5. 总结：为什么这很重要？

这篇论文的核心贡献在于，它指出了 AI 在“非顺序生成”（即可以乱序填空）时的一个盲点：不要只看眼前，要看全局。

以前的 AI： 像个短视的工人，哪里好干先干哪里，结果最后发现路走不通。
现在的 AI（用了 Info-Gain）： 像个高明的总指挥。它知道，有时候先攻克最难的堡垒，虽然当下风险大一点，但能瞬间打通整个战局，让后续工作变得无比顺畅。

一句话总结：
这就好比在黑暗中拼图，以前的 AI 喜欢先拼那些颜色最鲜艳、最容易认的碎片；而这篇论文教给 AI 的新技能是：先拼那些虽然难认、但能帮你把整幅图轮廓定下来的关键碎片，这样剩下的拼图就会变得简单又准确。

1. 研究背景与问题 (Problem)

背景：
掩码扩散模型（Masked Diffusion Models, MDMs）作为自回归模型（ARMs）的替代方案，在离散序列生成任务中展现出巨大潜力。利用双向注意力机制，MDMs 打破了严格的从左到右生成限制，在解码路径上具有前所未有的灵活性，适用于代码填充、生物序列设计和长程规划等任务。

核心问题：
尽管 MDMs 训练时采用随机掩码模式，但在推理阶段，生成过程是一个多步且对顺序敏感的解码过程。现有的采样器（Samplers）主要依赖贪婪启发式策略（Greedy Heuristics），即优先选择模型置信度最高（局部不确定性最低）的位置进行解码。

现有方法的局限性：

短视（Myopia）： 贪婪策略仅关注当前的局部确定性，忽略了当前解码决策对后续步骤的下游影响。
累积不确定性未最小化： 这种策略往往优先解码看似语法自信但语义次优的 token，导致错误传播，无法最小化整个生成轨迹的累积不确定性。
未充分利用双向性： 现有的贪婪方法未能利用 MDMs 的双向注意力优势，即无法在单步前向传播中评估一个解码决策如何重塑所有剩余掩码位置的 token 概率/不确定性。

典型案例：
论文通过“单向乘法”实验（ $a \times b = c$ ）指出，贪婪采样器倾向于先解码低不确定性的二进制乘积 $c$ ，而忽略了高不确定性的十进制因子 $a$ 和 $b$ 。一旦 $c$ 被错误固定，后续无法正确推导 $a$ 和 $b$ ，导致生成失败。最优策略应是先解决高不确定性的因子。

2. 方法论 (Methodology)

为了解决上述问题，作者提出了 Info-Gain Sampler（信息增益采样器）。这是一个基于原则的解码框架，旨在平衡“即时不确定性”与“未来信息增益”。

2.1 核心概念

状态不确定性 (State Uncertainty)： 定义为当前状态 $z_t$ 下所有掩码位置边际熵的平均值。
信息增益 (Information Gain, IG)： 定义为执行解码动作 $a_t$ 后，状态不确定性的减少量（即剩余掩码位置边际熵的降低量）。
$IG(a_t; z_t) := H(z_t) - H(z_{t-1})$
目标函数： 采样器的目标是最大化以下目标函数 $J_{IG}$ ：
$J_{IG}(a_t | z_t) = \underbrace{IG(a_t; z_t)}_{\text{信息增益}} - \underbrace{C(a_t | z_t)}_{\text{即时成本}}$
其中 $C(a_t | z_t)$ 是当前步骤解码 token 的边际熵之和（即时成本）。

2.2 算法流程

动作采样 (Action Sampling)： 生成一组候选动作集合 $\mathcal{C}$ 。通过两阶段采样（Token 采样 + 位置采样）探索巨大的动作空间。
并行评估 (Parallel Evaluation)： 利用 MDMs 的双向注意力特性，在单次批处理前向传播中，同时计算所有候选动作的目标函数 $J_{IG}$ 。这是该方法高效的关键，避免了自回归模型中昂贵的多步前瞻（Look-ahead）计算。
状态转移 (Transition)： 选择 $J_{IG}$ 最大的动作 $a^*_t$ 执行，更新状态，重复直到生成完整序列。

2.3 高效实现优化

KV 缓存优化： 将信息增益计算限制在当前活跃块（Active Block）内，以利用 KV 缓存。
高置信度旁路 (High-Confidence Bypass)： 如果某些位置的 token 概率超过阈值 $\gamma$ ，直接固定这些位置，减少不必要的计算。由于 Info-Gain 本身降低了不确定性，该机制触发更频繁，显著降低了推理延迟。

3. 主要贡献 (Key Contributions)

问题识别： 通过失败案例分析，实证揭示了现有基于贪婪确定性的采样器在 MDMs 中的根本局限性（短视导致累积不确定性高）。
提出 Info-Gain Sampler： 引入了一种新的解码框架，通过简单的目标函数平衡即时成本与未来信息增益。提出了计算高效的实现方案，利用 MDM 的双向特性实现单次前向传播的评估。
广泛验证： 在多种架构（全注意力、半自回归）和任务（数学推理、代码生成、规划、创意写作、图像生成）上进行了广泛评估。结果表明，Info-Gain Sampler 在几乎所有基准测试中均优于现有的贪婪采样器及并发的前沿方法（如 LookUM）。

4. 实验结果 (Results)

实验涵盖了全注意力 MDM（Dream-7B）、半自回归 MDM（SDAR-8B, TraDo-8B）以及多模态生成任务。

推理任务 (Reasoning)：
- 在 GSM8K, MATH-500, HumanEval 等基准上，Info-Gain Sampler 相比最佳贪婪基线（如 PC-Sampler）平均准确率提升了 3.6%。
- 累积不确定性显著降低： 在推理任务中，累积熵从贪婪基线的 78.4 降至 48.6，降幅巨大，证明了其能发现更全局优化的解码轨迹。
创意写作 (Creative Writing)：
- 在 AlpacaEval 上，Info-Gain Sampler 取得了 63.1% 的平均胜率（Win-rate）。
- 在高温度（ $\tau=1.5$ ）设置下，表现出极强的鲁棒性，峰值胜率达 80.3%，有效平衡了创造力与连贯性。
图像生成 (Image Generation)：
- 在 GenEval 和 ImageNet-512 上，Info-Gain Sampler 在位置（Pos）和属性（Attr）对齐指标上显著优于基线。
- ImageNet-512 的 FID 从 43.3 提升至 38.1，IS 从 53.3 提升至 63.0。
效率与开销：
- 相比贪婪采样，推理时间仅增加约 24%，显存占用增加约 20%，但性能提升显著。
- 消融实验表明，Info-Gain 采样器（Beam Size=1）在帕累托前沿附近，优于简单的 Best-of-N 策略。

5. 意义与影响 (Significance)

理论突破： 首次系统性地证明了在 MDMs 中，利用双向注意力进行信息增益评估是可行的且高效的，打破了自回归模型中“前瞻计算昂贵”的固有认知。
性能提升： 显著提升了 MDMs 在复杂推理和长序列生成任务中的性能上限，解决了 MDMs 长期以来因采样策略不当导致的“训练 - 推理不匹配”问题。
通用性： 该方法不仅适用于全注意力模型，也适用于半自回归（Semi-AR）架构，且对温度参数不敏感，具有广泛的适用性。
未来方向： 为从局部启发式向全局规划的非自回归生成提供了原则性的桥梁，为后续更高效的采样算法设计奠定了基础。

总结：
该论文提出了一种名为 Info-Gain Sampler 的创新采样策略，通过量化解码动作对未来不确定性的减少作用（信息增益），克服了传统贪婪采样的短视缺陷。实验证明，该方法能以极小的计算代价换取生成质量的显著提升，特别是在需要复杂逻辑推理和全局一致性的任务中表现卓越。