扩散对偶性第二章：Ψ采样器与高效课程学习 The Diffusion Duality, Chapter II: Ψ-Samplers and Efficient Curriculum
提出Ψ-采样器和高效课程学习，显著提升离散扩散模型的生成质量和训练效率
前置知识
离散扩散模型
离散扩散模型是一类生成模型，通过逐步向离散数据添加噪声然后反向去噪来生成样本。主要有两种噪声分布：Masked Diffusion使用特殊 token作为prior，Uniform-State Diffusion使用均匀分布prior。模型训练通过最小化负证据下界（NELBO）来学习去噪网络，其中关键概念包括前向过程、反向过程和噪声调度α_t。
本文聚焦于离散扩散模型的采样效率问题，需要理解离散扩散的前向过程定义、边缘分布q_t(z|x) = Cat(·; α_t x + (1-α_t)π)以及NELBO优化目标，才能理解Ψ-posteriors如何通过修改反向过程来改进采样质量。
预测-校正采样器
预测-校正采样器是一类改进扩散模型采样质量的方法。预测步骤使用模型预测的去噪方向从一个时间步转移到另一个时间步，校正步骤注入额外的随机噪声来修正预测误差。在连续扩散中，这类似于DDIM/ODE采样加Langevin校正；在离散扩散中，关键挑战是如何在保持边缘分布正确的前提下实现噪声注入。ReMDM是应用于Masked Diffusion的PC方法，允许已经unmask的token重新被mask。
本文的核心贡献Ψ-samplers就是一种预测-校正方法，通过参数κ_t控制预测和校正的混合比例。理解PC方法的基本原理有助于理解为什么Ψ-samplers能够在保持正确边缘分布的同时实现自我纠错。
扩散变换算子T
扩散变换算子T: [0,1] → [0,1]建立了连续高斯扩散与离散扩散之间的联系。具体来说，如果连续潜在量w_t ~ N(·; ˜α_t x, (1-˜α_t^2)I)，那么离散潜在量z_t = argmax(w_t)的边缘分布满足z_t ~ Cat(·; T(˜α_t)x + (1-T(˜α_t))/K·1)。T的定义涉及高斯CDF和CDF的K-1次幂的积分：T(˜α_t) = (K-1)∫_{-∞}^{∞} Φ(z - ˜α_t/√{1-˜α_t^2})Φ_{K-1}(z)dz - 1。这个算子使得离散扩散模型可以从高斯扩散推导而来，也为课程学习提供了理论基础。
课程学习方法依赖于在Gaussian diffusion上使用softmax替代argmax将Gaussian latents转换为离散token，这需要准确计算T(˜α_t)来维持正确的训练目标。本文利用T的泰勒展开和多项式近似来高效计算，避免了预计算10万个值对的繁琐做法。
课程学习
课程学习是一种训练策略，通过从简单任务逐步过渡到困难任务来加速模型训练。在离散扩散的上下文中，Sahoo et al. (2025a)提出在训练前半阶段使用低温度softmax替代硬argmax将Gaussian latents转换为离散token。这给模型提供了更温和的信号：输入是多个embedding的线性组合而非完全噪声的embedding，使去噪任务更容易。这种curriculum需要复杂的计算来近似softmax对embedding矩阵的加权求和。
本文的高效curriculum是对原有方法的优化，利用低温度softmax的稀疏性（大部分权重接近零）来只计算top-k个embedding的加权组合，显著减少内存和计算开销。理解curriculum的基本动机有助于理解为什么稀疏近似仍然有效。
研究动机
离散扩散模型分为两大类：Masked Diffusion Models使用 token作为prior，Uniform-State Diffusion Models使用均匀分布prior。USDMs在少步采样和引导生成任务中表现优异，因为它们允许token在生成过程中多次修正，具备自我纠错能力。然而，USDMs存在两个关键问题：当使用ancestral采样器时，随着采样步数增加，生成质量会出现plateau，无法继续改善；在likelihood建模能力上，USDMs仍然落后于MDLMs。此外，原有的课程学习方法虽然能缩小likelihood差距，但计算成本极高：需要显式构造K维向量（K是词汇表大小，通常超过10万）并在训练期间预计算10万个T函数值对，导致内存占用过大。这些问题限制了USDMs在高采样步场景下的应用。
本文的目标是本文的目标有三个：为离散扩散模型设计一种高效的预测-校正采样器，能够在任意噪声prior（uniform或masked）下工作，并随着采样步数增加持续提升质量；证明USDMs配合这种采样器可以匹配甚至超越MDLMs在高NFE regime的性能，挑战Masked diffusion是扩散语言模型必然未来的观点；开发一种内存高效的课程学习方法，减少训练时间和内存使用，同时保持与原有curriculum相当的训练效果。
与已有工作不同的是，本文的切入点是重新思考离散扩散的反向过程。标准离散扩散使用马尔可夫反向后验q_{s|t}，但多个联合分布可以具有相同的边缘分布。作者提出Ψ-posteriors作为标准反向后验和正向过程的线性叠加：Ψ_{s|t} = κ_t q_{s|t} + (1-κ_t) q_s。这种非马尔可夫后验允许在保持边缘分布正确的条件下注入额外噪声，实现了真正的预测-校正机制。与ReMDM仅适用于Masked diffusion不同，Ψ框架泛化到任意噪声prior，为USDMs带来了预测-校正能力。此外，作者观察到课程学习中使用的低温度softmax具有高度稀疏性，利用这一特性设计了只计算top-k个embedding权重的高效近似，彻底改变了原有curriculum的计算效率。
核心方法
本文的核心方法包含两个部分：Ψ-采样器和高效课程学习。Ψ-采样器的核心思想是通过线性叠加标准反向后验和正向过程来构造非马尔可夫后验分布。具体来说，Ψ-posteriors定义为Ψ_{s|t}(·|x_ℓ, z_ℓ^t) = κ_t q_{s|t}(·|z_ℓ^t, x_ℓ) + (1-κ_t) q_s(·|x_ℓ)，其中κ_t ∈ [0,1]控制预测和校正的混合比例。当κ_t = 1时，恢复标准的ancestral sampler；当κ_t < 1时，额外的(1-κ_t)q_s项注入噪声，实现了预测-校正机制。高效课程学习利用低温度softmax的稀疏性：在温度τ = 10^{-3}时，softmax权重高度集中在少数几个最大的entries上。通过只保留top-k个entries并使用order statistics直接采样它们的值（无需构造完整的K维向量），再加上对未采样entries的均值近似，可以高效计算softmax-weighted embedding组合，将内存使用从O(K)降低到O(k)。
Ψ-samplers的关键创新在于认识到边缘分布相同的联合分布可以有不同的动态。通过线性组合q_{s|t}和q_s，Ψ-posteriors保持了与标准扩散相同的边缘分布（通过归纳法证明），但在生成过程中允许更灵活的噪声注入。对于Masked diffusion，当π = m时，Ψ-posteriors允许已经unmask的token重新回到masked状态，这是ancestral sampler无法做到的。对于Uniform-state diffusion，当π = 1/K时，(1-κ_t)(1-α_s)π项确保每个token都有非零采样概率，即使去噪器赋予正确token极小的概率。这实现了真正的错误修正能力。高效课程学习的关键洞察是低温度softmax的稀疏性：虽然理论上涉及K个embedding，但实际上只有极少数权重显著，因此可以用top-k近似而不损失训练效果。
方法步骤详情
Ψ-samplers的工作步骤如下：从时间t的latent z_t开始，使用去噪网络x_θ预测clean token x̂_θ = x_θ(z_t, t)；计算预测步骤的分布q_{s|t}(·|z_t, x̂_θ)，这是标准ancestral sampler使用的后验；计算校正步骤的分布，这是前向过程在时间s的边缘分布，但x被替换为x̂_θ；将两个分布以κ_t和(1-κ_t)的比例混合，得到Ψ后验；从Ψ后验采样得到z_s。高效课程学习的实现步骤是：使用order statistics算法直接采样K-1个零均值高斯变量的top-k个值，无需生成完整的K维向量；采样对应clean token的特殊值˜w ~ N(˜α_t, ˜σ_t^2)；将˜w与top-k值比较，确定clean token是否在top-k中；使用Floyd算法采样k个不重复的token索引；计算top-k项的指数和，并使用闭式表达式估计未采样entries的均值贡献；归一化得到softmax权重并与对应的embedding rows相乘。
技术新颖性
Ψ-samplers的技术新颖性在于它提供了一个统一的框架来理解和设计离散扩散的预测-校正方法。传统方法如ReMDM是为Masked diffusion设计的，而Ψ框架自然扩展到任意噪声prior，包括uniform prior。通过参数κ_t的适当选择，Ψ-samplers可以恢复ReMDM（κ_t = 1 - σ_t/(1-α_s)）以及其他早期PC方法。更重要的是，Ψ框架证明了保持边缘分布正确的非马尔可夫后验的存在性，这为未来的研究开辟了新方向。高效课程学习的新颖性在于将order statistics理论应用于softmax稀疏化：通过直接采样高斯order statistics，避免了构造完整向量的开销，同时利用闭式表达式估计未采样项的贡献，保证了近似的准确性。此外，作者提出了T函数的泰勒展开和多项式近似，进一步简化了计算。
实验结果
实验结果显示，Ψ-samplers在语言和图像建模任务上均显著提升了生成质量。在OpenWebText语言建模上，Duo++配合Ψ-samplers在32到4096个NFE的范围内超越了MDLM+ReMDM和ancestral sampling。最重要的是，当NFE超过序列长度（1024）后，Ψ-samplers和ReMDM继续改善样本质量，而ancestral sampling出现plateau。在CIFAR-10图像生成上，Duo++配合Ψ-samplers达到了最佳FID（15.05在4096步）和IS（8.46在1024步），超越了MDLM的ancestral sampling（FID 23.93）和ReMDM（FID 23.71）。高效课程学习方面，Duo++在保持与Duo相当的perplexity（OpenWebText上25.2 vs 25.2，LM1B上30.0 vs 29.9）和下游任务准确性的同时，将峰值内存使用降低了33%，训练速度提升了25%（25%更少的GPU-hours）。Zero-shot perplexity在7个数据集上也表现相似，证明高效curriculum没有牺牲泛化能力。
![FID on CIFAR-10 with Ψ-samplers, where Ψ-samplers are activated for steps with t ∈ [toff, ton], when κt is kept constant.](/papers/2602.21185/images/table_9.jpg)
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 语言建模 (OpenWebText) | 生成困惑度 (Gen. PPL, 越低越好) | 29.19 (4096步, Ψ-sampler, rescale schedule) | 36.76 (MDLM, 4096步, ancestral sampling) | 约20.6%降低 |
| 图像生成 (CIFAR-10) | Fréchet Inception Distance (FID, 越低越好) | 15.05 (Duo++, 4096步, Ψ-sampler) | 23.93 (MDLM, 4096步, ancestral sampling) | 约37.1%降低 |
| 图像生成 (CIFAR-10) | Inception Score (IS, 越高越好) | 8.46 (Duo++, 1024步, Ψ-sampler) | 7.97 (MDLM, 1024步, ReMDM) | 约6.1%提升 |
| 训练效率 (内存使用) | 峰值内存 (GiB, 越低越好) | 63.4 (Duo++) | 94.3 (Duo) | 约33%降低 |
| 训练效率 (GPU小时) | 训练时间 (GPU-hours, 越低越好) | 比Duo少25% | Duo baseline | 25%节省 |
局限与改进
作者承认的局限性包括：Ψ-samplers的性能依赖于κ_t schedule的选择，不当的κ_t可能导致性能劣于ancestral sampling；虽然Duo++在多项指标上超越了MDLM，但在某些下游任务（如HellaSwag）上仍然落后于MDLM，这与较高的困惑度一致；高效课程学习虽然减少了内存使用，但当k较大时（如k = 5），性能开始略有下降；本文主要研究了138M参数规模的模型，更大规模的行为需要进一步验证。此外，我的观察是：论文中的最佳κ_t schedule在不同任务和数据集上有所不同（语言推荐rescale schedule with η = 0.05，图像推荐cosine schedule with κ_t = 0.95），这意味着实际应用时需要超参数调优；论文没有探讨Ψ-samplers在非常低NFE regime（如32步以下）的表现，这可能对实时应用很重要；高效curriculum的理论保证依赖于低温度softmax的稀疏性，但对于更高的温度或不同的噪声水平，稀疏性可能减弱。
独立分析的弱点
基于独立分析，我识别出以下弱点：κ_t schedule需要针对不同任务手动调优，缺乏自适应或学习机制。改进方向：可以设计自适应κ_t schedule，根据中间latent的置信度（如预测概率分布的熵）动态调整κ_t；在某些场景下，过度的噪声注入（κ_t过小）可能导致采样质量下降。改进方向：可以探索非均匀的κ_t schedule，在不同token位置使用不同的κ_t值，或者根据上下文重要性调整；高效课程学习依赖于低温度的强稀疏性假设，对于更温和的温度或不同的diffusion schedule，稀疏性可能减弱。改进方向：可以开发自适应k选择策略，根据当前噪声水平动态调整保留的entries数量；论文主要关注训练效率和采样质量，但对推理速度的详细分析较少。改进方向：可以研究Ψ-samplers的推理复杂度优化，如并行化多token的top-k采样；论文没有探讨Ψ-samplers与条件生成或引导（如CFG）的交互。改进方向：可以研究Ψ-samplers在强引导设置下的表现，探索引导强度与κ_t的关系。
未来方向
作者提出的未来方向包括：扩展到更一般的离散扩散过程，包括masked和uniform prior的组合；与其他采样方法结合，如高阶采样器或自适应步长调度；将Ψ框架应用于更大的模型规模以验证scaling behavior。基于本文成果，可以延伸的研究方向包括：学习κ_t schedule：将κ_t参数化为可学习函数，通过强化学习或元学习自动找到最优schedule；稀疏curriculum的理论分析：深入研究top-k近似的误差界限，提供理论保证；多模态应用：将Ψ-samplers应用于视频、音频等多模态离散数据；与其他生成范式结合：探索Ψ框架与flow matching或一致性模型的联系；条件生成优化：研究Ψ-samplers在分类器引导、文本到图像等条件生成任务中的应用，探索如何平衡引导强度和噪声注入；采样步数压缩：研究Ψ-samplers是否支持更激进的步数压缩，如配合蒸馏技术实现极低步数的高质量生成。
复现评估
论文的可复现性较好。作者已开源代码、checkpoint和视频教程（https://s-sahoo.com/duo-ch2）。实验在16个H100 GPU上以bfloat16精度训练，使用batch size 512，训练100万步。对于语言建模，使用预训练的MDLM和Duo checkpoint，可以直接复现采样实验；对于图像建模，使用35M参数U-Net训练150万步，架构细节在表3中完整描述。超参数选择（如κ_t schedule）在附录中有详细消融研究，为复现提供了指导。潜在的复现挑战包括：语言建模评估需要GPT-2 Large计算生成困惑度，这需要额外的资源；CIFAR-10评估需要50k生成样本计算FID，计算成本较高；高效curriculum的实现涉及order statistics和多项式近似，需要仔细实现以保证数值稳定性。总体而言，论文提供了足够的细节和开源资源，中等实验室应该能够复现主要结果。
论文图表
