🤖 machine learning

Accelerating Discrete Diffusion Models with Parallel-In-Time Sampling

本文介绍了一种用于离散扩散模型的并行化时间采样方法，该方法利用 $\tau$ -leaping 算法的连续时间随机积分形式和 Picard 迭代，实现了指数阶阶乘级的收敛速度，在保持合成、图像和文本任务生成质量的同时，显著降低了时间复杂度和运行时间。

原作者： Yu Yao, Huanjian Zhou, Andi Han, Wei Huang, Masashi Sugiyama

发布于 2026-07-02

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Yu Yao, Huanjian Zhou, Andi Han, Wei Huang, Masashi Sugiyama

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

想象一下，你正试图重建一份被撕碎的文件，但重建的不是纸张，而是一个句子或一张图像，它们被慢慢变成了一个巨大的问号（即“掩码”/mask）。这就是**离散扩散模型（Discrete Diffusion Models）**的工作原理：它们从一张清晰的图片或一段文本开始，将其转化为噪声（掩码），然后计算机学习如何逆转这个过程来重构原始内容。

问题在于？目前计算机进行这种操作的方式，就像是一个人试图一次只能重建一个词，且必须严格遵守顺序。他们必须先猜出第一个词，然后是第二个，接着是第三个。即使你拥有拥有数千个核心的超级快速计算机（如现代 GPU），这种方法也会迫使计算机在开始下一步之前必须等待上一步完成。这就像一场接力赛，接力棒必须完美传递后，下一位跑者才能出发。

这篇论文介绍了一种名为 Picard $\tau$ -leaping 的新方法，它彻底改变了这场比赛。以下是它的工作原理，我们通过简单的类比来进行说明：

1. 旧方法：单列纵队

把旧的方法想象成排队进入电影院的人群。一个人买完票，走进门，然后下一个人才能买票。即使电影院有 100 个门，一次也只能有一个人使用柜台，因为规则规定要“轮流等待”。在计算机术语中，这被称为顺序采样（sequential sampling）。这种方式很准确，但极其缓慢，因为它无法充分利用计算机的强大性能。

2. 新思路：“时空穿越”小组

作者意识到，与其等待队伍一个接一个地移动，我们可以将一段连续的时间视为一个整体块。想象一下，你想预测下周的天气。与其一个接一个地计算周一、周二、周三，不如说：“让我们一次性猜出整周的天气，然后检查我们的工作，再根据结果进行改进，再次猜测，但这次会更准。”

这就是他们的并行时间（Parallel-in-Time）方法的核心。他们选取一个时间块（例如重建过程中的 10 个步骤），并尝试利用计算机的众多核心同时解决这 10 个步骤。

3. 秘诀：“Picard 迭代”（猜想与校验循环）

如何在不搞砸的情况下同时解决 10 个步骤？作者使用了一个被称为 Picard 迭代（Picard iteration） 的数学技巧。

第一轮（大胆猜测）： 计算机基于起始点，对整个周的天气（或整个图像重建过程）做一个粗略的猜测。
第二轮（修正）： 计算机观察“游戏规则”（数学模型），并找出第一次猜测出错的地方。由于它面前有整周的数据，它可以同时修正所有的错误。
第三轮（精炼）： 它重复这个过程。每一次，猜测都会越来越接近真相。

因为计算机可以在处理“周一到周五”的所有数学运算时处于完全相同的时刻，所以它完成任务的速度比单列纵队快得多。

4. 特殊规则：“首次命中”停止

这里有一个限制。在这种特定类型的游戏（称为吸收扩散/Absorbing Diffusion）中，一旦一个“问号”被变回真实的字母或像素，它就会永远固定在那里。它不会再发生改变。

如果你只是同时猜测整周的情况，你可能会不小心尝试去修改一个在较早步骤中已经被固定的字母。为了解决这个问题，作者添加了一个**“首次命中截断”（First-Hitting Truncation）**。

把它想象成一场带有特殊规则的“抢椅子”游戏：一旦椅子被占用了，它就被锁定了。如果你的“小组猜测”试图移动一个已经坐下的玩家，系统会直接忽略那个动作，让其保持原位。这确保了计算机在追求速度的同时不会破坏规则。

5. 结果：速度与质量并存

该论文声称，通过使用这种“猜想整个区块并进行精炼”的方法：

速度： 在单个计算机芯片（GPU）上，他们生成图像和文本的速度比旧方法快 1.45 到 1.86 倍，同时保持质量完全一致。
效率： 他们需要减少约 50% 的计算步骤（NFE） 即可获得相同的结果。
可扩展性： 从理论上讲，随着问题规模变大（更复杂的图像或更长的文本），这种方法相对于旧方法会变得相对更快。

总结

这篇论文提出了一种运行生成文本和图像 AI 模型的新方法。与其强迫 AI 一个接一个地进行微小且缓慢的步骤，不如让 AI 在时间轴上进行大幅度的并行跨越。他们使用“猜想与校验”的循环来确保这些大幅跨越保持准确，并使用“锁定”规则来确保 AI 不会弄乱已经修复的部分。其结果是一种更快速、更高效的数字内容创作方式。

技术摘要：通过并行化时间采样加速离散扩散模型

问题陈述

离散扩散模型已成为类别数据生成建模的核心，广泛应用于分子设计、蛋白质工程以及文本/图像生成。一种主流的方法是吸收式离散扩散（Absorbing Discrete Diffusion），它通过逐步将标记（tokens）替换为特殊的掩码状态（mask state）并学习逆转这一过程。然而，生成过程本质上是序列化的。虽然存在针对连续扩散模型的并行化时间算法，但针对离散扩散模型的对应方法在很大程度上仍未得到充分探索。目前的序列采样器（如 $\tau$ -leaping）无法充分利用当代并行计算资源（如 GPU），从而造成了高效推理的瓶颈。本文探讨的核心问题是：能否在不牺牲生成质量的前提下，从根本上加速离散扩散模型的采样过程。

方法论：Picard $\tau$ -Leaping

作者提出了 Picard $\tau$ -leaping 方法，这是首个针对吸收式离散扩散的并行化时间算法。该方法将通常用于连续扩散的 Picard 迭代法扩展到了离散状态空间中的 $\tau$ -leaping 随机积分形式。

核心算法

时间离散化： 将反向采样区间 $[\eta, T]$ 分为若干个大块（blocks）。每个块进一步细分为精细的时间单元（time cells）。
并行块处理： 算法不再是按步顺序计算轨迹，而是同时求解整个块内的转换过程。
Picard 迭代：
- 初始化： 块从一个初始猜测值（即块起始时的状态）开始。
- 并行评估： 在每个 Picard 轮次中，块内的所有微步（microsteps）都使用来自前一次迭代的轨迹进行并行评估。
- 首次命中截断（First-Hitting Truncation）： 这是针对吸收式扩散的关键组件。由于掩码标记只能被取消掩码（unmask）一次，算法在块内应用了“首次命中截断”。该操作确保对于每个坐标，仅应用最早提出的取消掩码事件，从而在实现并行前缀操作的同时，保持了吸收式的结构特性。
- 收敛： 过程重复进行，直到收敛。理论分析表明，误差随 Picard 迭代次数呈指数阶乘级收缩。

理论框架

该方法依赖于通过泊松随机测度（Poisson random measures）来表示离散扩散。作者引入了关于得分估计误差、有界得分估计以及“归一化事件切换”（类似于 Lipschitz 连续性）的具体假设，以限制误差传播。他们证明了 Picard 迭代收敛于序列化 $\tau$ -leaping 轨迹。

核心贡献

1. 理论进展

并行化时间算法： 引入了针对吸收式离散扩散的 Picard $\tau$ -leaping 方法（算法 1）。
复杂度降低： 证明了其时间复杂度相对于函数评估次数（NFE）从主流序列化方法的 $O(d \log S)$ 降低至 $O(\log(d \log S) \cdot \log d)$ 。
收敛性证明： 提供了在离散设置下 Picard 迭代具有指数阶乘收敛性的证明，这一特性与典型的连续扩散中的指数收缩有所不同。
空间复杂度： 该方法通过存储块轨迹，引入了额外的 $\tilde{O}(dS)$ 空间复杂度，以空间换取并行性。

2. 实验结果

作者在合成数据和真实任务上对该方法进行了评估：

合成数据： 在二维分布（棋盘格和圆环）上，并行方法在保持与序列化基准方法相当的 KL 散度的同时，实现了显著的运行时间加速（在特定合成缩放场景下高达 7–9 倍）。
图像生成： 使用基于 MaskGiT 的模型在 ImageNet (256×256) 上进行测试，并行采样器达到了与序列化 $\tau$ -leaping 基准相当的 Fréchet Inception Distance (FID) 分数，但 NFE 减少了 50%。在单块 RTX 4090 GPU 上，这转化为 1.86 倍的墙钟时间（wall-clock）加速。
文本生成： 使用基于 RADD 的模型在 OpenWebText（GPT-2 规模）上进行测试，并行方法在固定 NFE 下实现了更好的困惑度（perplexity），并在使用一半 NFE 的情况下达到了序列化求解器的困惑度水平。在单块 H100 GPU 上，这带来了 1.45 倍的加速。

重要性与主张

本文声称，这项工作扩展了离散扩散模型进行高效并行推理的潜力。通过利用吸收式扩散的特定结构属性（特别是取消掩码的首次命中特性），作者展示了如何打破序列依赖，从而实现沿时间维度的并行计算。

作者指出，这实现了从坐标级并行（每步取消多个标记）到时间级并行（同时求解多个时间步）的根本转变。他们提到，虽然当前的单 GPU 加速受限于内存带宽，但多项式对数级的复杂度预示着在多 GPU 扩展方面具有更大的潜力。这项工作为加速离散生成模型提供了理论基础，对于需要快速采样的应用（如分子结构生成和大规模语言建模）具有重要意义。

作者对未来的工作保持谦逊，承认目前的分析侧重于全变分（Total Variation）误差，并指出未来的研究应致力于更精确的 KL 级分析以及更具内存效率的变体（例如滑动窗口更新），以进一步增强可扩展性。