🤖 machine learning

Set Diffusion: Interpolating Token Orderings Between Autoregression and Diffusion for Fast and Flexible Decoding

本文介绍了 Set Diffusion，这是一种新型语言模型，它将灵活的任意阶（arbitrary-order）标记集生成与 KV 缓存支持相结合，从而实现比现有的自回归和块扩散方法更快的推理速度以及更优越的速度-质量权衡。

原作者： Marianne Arriola, Volodymyr Kuleshov

发布于 2026-07-03

📖 1 分钟阅读☕ 轻松阅读

原作者： Marianne Arriola, Volodymyr Kuleshov

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

以下是使用简单语言和创意类比对《Set Diffusion》论文进行的解释。

大局观：关于“顺序”的问题

想象你正在尝试写一个故事或解决一道数学题。你主要有两种方式可以做到：

严谨的作家（自回归/Autoregression）： 你一次只写一个词，严格地从左到右。在完成开头之前，你无法写出结尾。这种方式质量很高且准确，但速度很慢，因为你无法同时进行多项工作。
混沌的艺术家（标准扩散/Standard Diffusion）： 你从满篇乱码（噪声）开始，并试图一次性修复所有内容。你可以同时预测许多单词，所以速度很快，但很难保持逻辑连贯。此外，你很难轻松地“记住”之前写了什么来辅助接下来的写作，因此每次做出修改时，你都必须重新阅读整页内容。

问题所在： 之前尝试将这两者结合的方法（称为“块扩散/Block Diffusion”）就像是在以僵硬的块状进行写作。你可以一次写一个句子而不是一个词，但你仍然必须在移动到下一块之前完成整个句子。如果你想修改故事中间的一个词，你必须等待整个块完成后才能进行。

解决方案：Set Diffusion

作者引入了 Set Diffusion。不要把它看作是在一条线上写作，而要把它看作是 用灵活的碎片填充拼图。

不再被强迫逐词写作（太慢），也不再是按块写作（太僵硬），Set Diffusion 允许你选择 任何一组单词 作为下一步生成的对象，只要你遵循一套特定的规则。

“滑动窗口”类比

想象你正在绘制一幅长卷壁画。

旧方法（块扩散/Block Diffusion）： 你画完 4 英尺的部分，等它完全干透，然后再移动到下一个 4 英尺的部分。在完成整个块之前，你无法再次触碰第一个部分。
新方法（Set Diffusion）： 你有一个“滑动窗口”式的油漆刷。你可以画前 4 英尺，但随后你可以将窗口向后滑动，同时绘制第 2、3、4 英尺以及第 5、6、7 英尺的部分。你甚至可以在绘制边缘的同时，跳回中间某个位置进行修补。

核心特性简易说明

1. 灵活的“标记集”（Token Sets，即拼图碎片）
在这个模型中，“标记（token）”就是一个单词或一段代码。

创新点： 模型不仅仅是预测下一个词，它是在预测一个集合（set）。
神奇之处： 你可以告诉模型，“预测接下来的 3 个词”，或者“预测位置 5 和位置 10 的词”。它可以处理不同大小和不同顺序的组合。这使得它既快速（同时预测多个内容），又聪明（能够追踪顺序）。

2. “记忆库”（KV 缓存/KV Caching）
在 AI 领域，“KV 缓存”就像是一个便签本，模型将它已经生成的内容记录在上面，这样它就不必每次都重新计算。

旧问题： 在标准扩散模型中，模型每次进行预测时都必须重新阅读整个文本。这就像为了记住主角的名字而不得不把整本书读一遍。
新方案： Set Diffusion 在每一步之后都会更新它的“便签本”。一旦它预测出一组单词，就会立即将它们存入记忆。这使得它极其迅速，既拥有严谨作家的逻辑，又具备混沌艺术家的速度。

3. “滑动窗口”策略
论文引入了一种特定的选择下一个预测词的方法，称为 滑动窗口（Sliding-Window） 方法。

想象一个在舞台上移动的聚光灯。这个聚光灯可以覆盖一个演员，也可以同时覆盖三个演员。
模型利用这个聚光灯来决定：“我现在要生成聚光灯范围内的单词。”
通过调整聚光灯的大小，你可以控制速度（大聚光灯，一次预测多个词）与准确度（小聚光灯，减少预测数量以求更谨慎）之间的平衡。

他们证明了什么？

作者在三个主要任务上测试了这种新方法：

数学推理： 解决数学应用题（例如 GSM8K 数据集）。
摘要生成： 将长篇文章浓缩为短摘要。
填空（Infilling）： 填充故事中缺失的部分（类似于“填词游戏/Mad Libs”）。

结果：

速度 vs. 质量： Set Diffusion 找到了以往模型错过的“甜点位（sweet spot）”。它比僵硬的块模型更快，比混沌的扩散模型更准确。
填空任务： 它在填充故事缺失部分方面明显优于之前的“块扩散（Block Diffusion）”方法。
灵活性： 它可以生成任何长度和任何顺序的文本，这对于编辑文档或修改文件中间的代码至关重要。

总结隐喻

如果说 自回归（Autoregression） 是一场 接力赛（一名跑者将接力棒传给下一位，严格遵守顺序），而 标准扩散（Standard Diffusion） 是一场 自由竞争（所有人同时奔跑，但场面混乱），那么 Set Diffusion 就是一支 配合默契的舞团。

舞者们（标记/tokens）可以成组（集合）移动。他们可以从左向右移动，也可以跳跃式地填补空白，但由于有了共享记忆（KV 缓存），他们始终清楚谁在身边以及之前的动作是什么。这使得他们能比接力队更快、更灵活地完成复杂的编舞（如数学、故事），同时又不会像自由竞争那样混乱。

技术摘要：集合扩散 (Set Diffusion)

问题陈述

离散扩散模型已成为自回归 (AR) 语言模型的一个具有竞争力的替代方案，提供了更快的生成速度和更好的可控性。然而，现有的扩散方法面临着显著的局限性：

固定长度与双向约束： 大多数离散扩散模型依赖于需要完整双向上下文的架构，这阻碍了键值 (KV) 缓存的使用，并限制了生成长度为固定长度。
块扩散 (Block Diffusion) 的僵化性： 虽然块扩散 (BD3LM) 通过从左到右生成固定大小的顺序块，部分弥补了这一差距，但它存在缺乏灵活性的问题。它无法支持任意位置的插入（填空）或块内的变长生成，且 KV 缓存的更新被延迟到整个块解码完成后才进行。
权衡限制： 当前的方法主要通过改变固定的块大小来在自回归和扩散之间进行插值，这限制了解码的灵活性和并行性。

核心挑战在于开发一种语言模型，不仅通过块大小，而且通过改变生成顺序来在自回归和扩散之间进行插值，从而实现灵活长度的生成、KV 缓存以及任意顺序的解码，同时保持高质量的样本。

方法论

作者引入了 集合扩散 (Set Diffusion)，这是一个将生成过程分解在灵活位置、灵活长度的标记集合（token sets）之上，而非固定块之上的框架。

1. 概率模型

集合扩散定义了一个关于一系列标记集合 $\sigma = (\sigma_1, \dots, \sigma_N)$ 的自回归概率分布，其中每个 $\sigma_n$ 是标记位置的一个子集。似然函数分解为：
$p_\theta(x) = \sum_{\sigma} \prod_{n=1}^N \pi(\sigma_n | x_{<\sigma_n}) \cdot p_\theta(x_{\sigma_n} | x_{<\sigma_n})$
这里， $\pi$ 是一个选择下一个位置集合的顺序策略，而 $p_\theta$ 使用扩散模型对该集合中标记的条件似然进行建模。该模型优化了一个负证据下界 (NELBO)，该下界是这些集合上扩散损失的总和。

2. 集合因果架构与 KV 缓存

一个关键的架构创新是集合因果 Transformer (set-causal transformer)。与标准的双向扩散模型不同，该架构支持在每次推理步骤后进行 KV 缓存。

机制： 当生成一个标记集合 $x_{\sigma_n}$ 时，其键 (Key) 和值 (Value) 向量会被缓存并重复使用，以对后续预测进行条件化。
训练效率： 作者提出了一种向量化训练算法，通过使用专门的注意力掩码 (attention mask)，在单次前向传递中同时处理干净的上下文和噪声查询。该掩码确保集合中的噪声标记仅关注同一集合内的其他噪声标记以及先前集合中的所有干净标记。

3. 插值标记排序

为了控制样本质量（边界的紧密程度）与推理效率（并行性）之间的权衡，作者引入了位置依赖的揭示时间调度 (position-dependent reveal-time schedules)。

位置偏移调度 (Position-Offset Schedules)： 每个标记位置 $\ell$ 被分配一个随机揭示时间 $R_\ell$ ，该时间是从由累积分布函数 (CDF) $\alpha^\ell_\tau$ 定义的分布中抽取的。
滑动窗口实例化 (SW-SetDLM)： 作者提出了一个特定的实例化，其中标记具有一个活跃生成区间 $[a_\ell, a_\ell + w]$ $[a_{ℓ}, a_{ℓ} + w]$ 。窗口大小 $w$ $w$ 控制“从左到右”的偏置：
- 小 $w$ （趋近于 $1/L$ ）使顺序集中在从左到右的生成上，恢复 AR 行为并收紧 ELBO。
- 大 $w$ （趋近于 1）允许无序生成，恢复标准的掩码扩散。
- 中间值则实现了滑动窗口解码，允许模型在任意位置预测灵活长度的标记集合，支持如填空等任务。

4. 推理

在推理期间，模型根据当前排序时间和位置偏移调度采样候选位置。然后，它使用 KV 缓存并行地对这些位置进行去噪。被接受的标记被提交，并且缓存立即更新，从而实现变长生成和高效解码。

核心贡献

集合扩散框架： 一类新的语言模型，它将生成分解在灵活位置、灵活长度的标记集合之上，推广了自回归、无序扩散和块扩散。
集合因果架构： 一种支持在每一步推理后更新 KV 缓存的 Transformer 架构，实现了高效的变长生成。
顺序插值： 一种通过调整位置依赖的排序调度（特别是滑动窗口调度）而非固定块大小，在 AR 和扩散之间进行插值的原则性方法。
最先进的性能： 证明了集合扩散在数学推理、摘要生成和无条件生成方面，相比之前的扩散模型，实现了更优的速-质权衡；并且在填空任务上显著优于块扩散。

实验结果

作者在多个基准测试上评估了集合扩散（特别是 SW-SetDLM）：

数学推理 (GSM8K)： 集合扩散实现了比块扩散更优的“速度-准确率”帕累托前沿。例如，当解码窗口大小 $S \le 8$ 时，它在吞吐量 $\ge 60.42$ tokens/sec 的情况下达到了 66.41% 的 pass@1 准确率，优于块扩散（ $S=8$ ）在 $\ge 55.39$ tokens/sec 下实现的 63.53% 准确率。
填空 (ROCStories)： 集合扩散在填空任务（填充中间 1 个或 3 个句子）中显著优于块扩散，实现了更高的 ROUGE 分数，且解码速度快了约 25%。块扩散由于其僵化的从左到右块结构，在此类任务中表现挣扎。
摘要生成 (CNN/DailyMail)： 集合扩散实现了与扩散基准模型相当的 ROUGE 分数，且比块扩散快了高达 10% 的解码速度。它也优于 AR 基准模型，作者指出 AR 模型在该数据集上容易过拟合。
无条件生成与似然度： 在 OpenWebText 和 LM1B 上，集合扩散在提高 MAUVE 分数的同时，比块扩散解码更快。在 LM1B 上，它在保持比以往扩散模型更高吞吐量的同时，实现了扩散模型中最好的困惑度 (27.83 PPL 对于 $S \le 8$ )。

意义与主张

本文声称，集合扩散通过将自回归与扩散之间的插值从僵化的块结构中解耦，代表了一项重大进展。通过利用标记排序进行插值，该框架提供了：

更大的灵活性： 支持任意位置的插入和变长生成，这对于填空等任务至关重要。
高效性： 能够在每一步之后更新 KV 缓存，这是此前扩散模型无法实现的功能，从而实现了更快的推理。
高性能： 与现有的扩散语言模型相比，在样本质量与推理速度之间提供了更优的权衡。

作者将集合扩散定位为一个统一的框架，它将 AR、无序扩散和块扩散视为特例，提供了一种可调机制，可以根据特定推理任务的需求，在并行性和样本质量之间进行平衡。