离散扩散模型的ODE采样器：指数残差流加速生成

扩散模型在连续状态空间（如图像生成）中已取得显著成功，其采样过程可通过概率流常微分方程（ODE）实现确定性轨迹，从而支持加速采样与一致性蒸馏。然而，离散扩散模型（如文本、代码生成）是否也存在类似的确定性ODE形式，一直是一个开放问题。最近一篇论文提出了指数残差流（Exponential Residual Flows, ERFs），正面回答了这个问题，并为离散扩散模型的快速采样提供了新思路。

离散扩散的确定性采样难题

连续扩散模型中的概率流ODE能够精确复现模型边缘分布，且其确定性轨迹允许使用黑箱ODE求解器进行快速采样，并可通过一致性蒸馏进一步压缩采样步数。但在离散状态空间（如词表中的token），由于状态跳跃的离散性，直接定义ODE面临挑战：如何在不破坏边缘分布的前提下，构建一个确定性的连续演化过程？

传统离散扩散采样通常依赖随机采样方法（如祖先采样），虽然能保证分布正确，但需要较多步数且轨迹随机，难以应用一致性蒸馏等加速技术。因此，探索离散扩散的ODE形式具有重要的实用价值。

概率流常微分方程（Probability Flow ODE） 是连续扩散模型中的核心概念。在基于分数匹配的扩散模型中，前向加噪过程可以描述为一个随机微分方程（SDE），而反向生成过程同样是一个SDE。概率流ODE是与此SDE具有相同边缘分布的一个确定性ODE：它构造一个向量场，使得沿该向量场演化的粒子轨迹在任意时刻的边缘分布与SDE完全相同，但轨迹本身是确定性的。这一性质带来两个关键好处：一是可以使用高阶黑箱ODE求解器（如Runge-Kutta）以较少步数获得高精度采样；二是确定性映射允许一致性蒸馏（consistency distillation）——训练一个学生模型直接将任意噪声水平映射到最终样本，从而将采样步数压缩到1-4步。连续扩散中的这一成功范式正是离散扩散研究者试图复现的目标。

祖先采样（Ancestral Sampling） 是离散扩散模型最常用的采样方法。它从完全噪声状态出发，按照反向扩散过程的转移概率逐步采样：每一步根据当前token和扩散模型的反向速率，从类别分布中随机抽取下一个token。祖先采样能严格保证生成样本服从模型分布，但存在两个固有缺陷：一是每步都引入随机性，导致采样轨迹不可复现且方差较大；二是通常需要与训练时扩散步数相当（数百到数千步）的迭代次数，采样效率低。更重要的是，由于轨迹是随机的，无法像连续扩散那样通过一致性蒸馏将多步采样压缩为少步确定性映射，这严重限制了离散扩散在实际部署中的竞争力。ERFs的动机正是绕开祖先采样的这些限制。

指数残差流（ERFs）的核心思想

ERFs的核心创新在于为每个token和位置引入一个连续的残差变量，从而将离散状态嵌入到连续空间中。具体而言，模型状态由离散的token序列和一个对应的残差向量组成。残差按照分段定义的ODE进行确定性演化，其向量场取决于当前的离散状态以及扩散模型的前向和反向速率矩阵。

当某个残差变量触及零时，对应的离散状态发生跳跃（即token改变）。这种机制类似于连续空间中的“事件触发”更新，但严格保证了边缘分布与原始离散扩散模型一致。作者从理论上证明，该构造能够精确保持原扩散模型的边缘分布，从而为确定性采样奠定了基石。

速率矩阵（Rate Matrix） 是离散扩散模型的数学基础。在连续时间马尔可夫链（CTMC）框架下，离散扩散的前向过程由一个速率矩阵 R_t 描述：对于任意两个不同的token类别 i 和 j，R_t(i,j) 表示从 i 跳到 j 的瞬时速率，对角线元素由行和为负的条件确定。前向过程使数据分布逐渐趋于均匀或全掩码分布；反向过程则需要知道反向速率矩阵，它可以通过前向速率和模型学习的分数（即去噪分布）解析计算得到。ERFs的向量场正是由这些前向和反向速率矩阵共同决定的：残差的连续演化速率取决于当前离散状态对应的速率矩阵行，当残差归零时触发按速率矩阵比例确定的跳跃。因此，理解速率矩阵是理解ERFs构造的前提。

针对掩码与均匀扩散的离散化实现

为了使ERFs实用化，作者推导了一种实用的离散化方案，并将其专门应用于两种常见的离散扩散类型：掩码扩散（masked diffusion）和均匀扩散（uniform diffusion）。掩码扩散中，token以一定概率被掩码替换；均匀扩散则是在所有token类别上均匀跳转。针对这两种情况，ERFs的向量场可以解析表达，从而允许高效的数值积分。

离散化方案的关键在于如何处理残差触及零的事件。由于残差演化是确定性的，可以通过事件驱动或固定步长积分来模拟，论文中给出了具体的算法流程。这种确定性轨迹不仅支持标准的ODE求解器，还为一致性蒸馏铺平了道路——通过蒸馏可以训练一个少步采样器，大幅降低采样成本。

掩码扩散与均匀扩散 是离散扩散中最常用的两种具体形式，区别在于前向过程的破坏方式。掩码扩散（Masked Diffusion）在前向过程中以一定速率将任意token替换为特殊的[MASK]标记，最终状态为全掩码序列；其反向过程即逐步预测并填充被掩码的位置，这与BERT式掩码语言模型的训练目标高度契合，因此在文本和代码生成中尤为流行。均匀扩散（Uniform Diffusion）则在前向过程中以均匀速率将token跳转到词表中任意其他类别（包括随机token），最终趋于均匀分布；它不引入特殊标记，理论上更对称，但反向去噪难度更大。两种扩散的速率矩阵结构不同，导致ERFs残差向量场的解析表达式也不同，论文分别针对二者推导了可高效计算的离散化方案，使ERFs能直接应用于现有主流离散扩散模型。

实验验证：代码与文本生成上的提升

作者在多种基准上评估了ERFs，包括代码生成和文本生成任务。实验结果表明，与之前的离散扩散采样方法相比，ERFs在生成质量上取得了经验性的改进。由于确定性轨迹支持一致性蒸馏，ERFs可以进一步加速采样，在保持生成质量的同时显著减少采样步数。

虽然论文未披露具体的量化指标，但作者声称“在多种基准上优于先前方法”，这为离散扩散模型的快速采样提供了新的技术路径。考虑到离散扩散在自然语言处理、代码生成等领域的潜力，ERFs有望推动这些应用的实际部署。

意义与展望

ERFs的提出填补了离散扩散模型确定性采样理论的空白，将连续扩散中的概率流ODE思想推广到离散空间。其理论保证和实用离散化方案为后续研究奠定了基础。未来可能的方向包括：将ERFs扩展到更复杂的离散结构（如图、序列决策），结合更高效的一致性蒸馏方法，以及在大型语言模型上的应用探索。对于从事生成模型研究的从业者，ERFs提供了一个值得关注的新工具。