Provably adaptive sampling with uniform and remasking discrete diffusion models

在人工智能领域，存在着一场关于教计算机如何创造新事物的持续竞赛，从撰写连贯的故事到生成真实的蛋白质结构。多年来，利用文本或数据序列进行此类操作的主导方法是一种逐步推进的方法，即模型根据之前出现的所有词汇来预测下一个词，就像人类逐词阅读句子一样。虽然这种顺序方法行之有效，但由于它无法同时处理句子的多个部分，因此速度较慢。一种被称为离散扩散（discrete diffusion）的新型、更快速的替代方案已经出现。这种方法不是从头开始构建序列，而是从一团混乱的随机数据开始，并逐渐清理它，将噪声精炼成清晰、有意义的模式。这种方法的妙处在于，它可以同时更新数据的许多部分，为更快的生成提供了一条路径。然而，为了使这种方法在现实世界中发挥作用，它必须是高效的。如果清理噪声的过程需要过多的步骤，速度优势就会消失，使该模型在处理大规模任务时变得不切实际。

这些扩散模型的核心挑战在于它们如何处理引入数据的“噪声”。想象一个系统，它将清晰的句子随机替换为无意义的内容或将其遮掩（masking）。为了生成新文本，模型必须学习逆转这一过程，从受损的数据中推测出原始词汇。长期以来，研究人员认为这种逆转的速度很大程度上取决于系统中单词或符号的总数，即所谓的“维度”。如果一个句子有一千个位置，旧理论认为模型需要大约一千步才能完成清理工作，无论实际句子是简单还是复杂。这种对规模的线性依赖意味着，即使对于高度结构化、可预测的数据，计算机也必须像处理完全随机的噪声一样努力工作，从而抵消了并行处理带来的好处。

宾夕法尼亚大学的一个研究小组现在挑战了这一假设，证明了缓慢并非离散扩散方法本身的根本缺陷，而是由于清理过程执行方式的结果。他们开发了一种新的采样策略，允许模型在过程中纠正自己的错误，而不是被锁定在早期可能错误的决策中。他们的工作表明，生成样本所需的步骤数并不取决于词汇表的规模或序列的长度，而是取决于所创建数据的内部结构。如果数据具有简单的、可预测的模式，即部分内容相互依赖，那么模型生成它的步骤可以比此前认为的要少得多。

研究人员专注于两种特定类型的噪声过程：一种是将标记（tokens）随机替换为任何其他有效标记，另一种是将标记遮掩，且当模型不确定时可以进行取消遮掩或重新遮掩。在过去，用于逆转这些过程的标准算法（例如广泛采用的“tau-leaping”方法）被发现对于均匀过程效率低下。这些旧方法通常会对数据进行单次处理，在不检查这些变化是否与序列其余部分一致的情况下同时更新许多位置。如果模型在早期犯错，这个错误就会持续并影响后续的所有步骤，导致高错误率，从而需要更多步骤来修复。本文引入的新方法使用了一种“留一法”（leave-one-out）策略。模型不再通过观察整个序列来预测单个标记，而是考虑如果移除该特定标记，序列的其他部分会是什么样子。这使得模型能够对每个位置进行更明智、更独立的并行更新，并且至关重要的是，它允许模型在后续更新显示之前的预测不正确时，修正其选择。

通过使用这种改进的方法，研究人员表明，生成样本的计算成本受控于数据各部分之间相互依赖程度的度量。在技术层面，他们将效率与一个称为“对偶总相关性”（dual total correlation）的概念联系起来，该概念量化了整个序列中的共享信息量。对于高度结构化的数据集，例如具有清晰语法结构的句子或具有特定折叠模式的蛋白质，由于序列各部分受到彼此的严格约束，该度量值很小。新的分析证明，对于此类数据，生成样本所需的步骤数随其结构复杂度而变化，而不是随总位置数变化。这意味着，对于一个遵循严格语法规则的长而复杂的句子，只要底层结构简单，模型生成它的速度几乎可以与短句一样快。论文提供了数学证明，证明这种效率提升是真实的，而非仅仅是偶然的观察结果，从而确立了之前的局限性是由于清洗算法的选择，而非扩散过程本身。

为了验证这些理论发现，研究人员在旨在模拟现实世界结构的合成数据上进行了数值实验。他们在遵循马尔可夫链模式（即下一个比特取决于前一个比特）的二进制序列上，将新采样器与旧的标准方法进行了对比测试。在这些测试中，新方法表现始终优于传统方法，即使在保持极低步骤数的情况下也能维持较低的错误率。结果显示，当数据的维度增加时，旧方法表现挣扎，而新方法依然稳健，其性能取决于数据的内在可预测性而非其规模。他们还在混合二进制字符串的情景下测试了该方法，在这种情景下，数据来自一组特定的模式。在此，新采样器同样展示了它能够适应底层分布的低维特性，以比旧理论预测的最坏情况更少的计算步骤实现了高准确度。

这项工作的意义不仅在于一种更快的算法，它从根本上改变了我们对离散扩散模型极限的理解。通过证明这种不利的维度依赖性是一个可以解决的算法设计问题，而非内在障碍，研究人员为更高效的大规模生成模型打开了大门。这对于自然语言处理和蛋白质设计等应用场景尤为重要，因为这些领域的数据具有高维但高度结构化的特点。能够并行生成复杂序列且不被庞大的标记数量所拖累，意味着离散扩散有望在速度和质量上追平甚至超越自回归模型。该研究还强调了允许模型修正中间决策的重要性，这一特征模仿了人类在写作或思考时使用的迭代精炼过程，而非旧模型那种僵化的、单向的生成方式。

最终，这项研究为提高生成式人工智能的效率提供了一条清晰的路径。它证实了离散扩散并行生成数据的潜力不仅是一个理论上的承诺，而且是一个现实，前提是使用正确的工具来应对噪声。这项工作将数学近似过程引入的误差与模型学习引入的误差区分开来，表明前者可以由数据本身的结构进行严密控制。随着领域向更大、更复杂的模型迈进，这些见解对于确保计算成本不会随着问题规模的扩大而失控至关重要。研究结果表明，离散生成的未来不在于暴力计算，而在于更聪明、更具适应性的策略，这些策略能够利用数据中天然存在的顺序和依赖关系。