离散扩散模型与Ψ-Samplers技术解析
1. 离散扩散模型的基础原理离散扩散模型的核心思想是通过构建一个逐步添加噪声的前向过程和一个逐步去噪的反向过程实现从简单先验分布到复杂数据分布的转换。与连续扩散模型不同离散扩散模型专门处理分类数据如文本token或图像像素的离散值。1.1 前向与反向过程前向过程定义为马尔可夫链将数据x逐渐腐蚀为噪声。对于长度为L的离散序列x∈V^LV是one-hot编码空间前向过程的边际分布为z_t^ℓ ~ q_t(·|x^ℓ; α_t) Cat(·; α_t x^ℓ (1-α_t)π)其中α_t∈[0,1]是单调递减的噪声调度π是先验分布。这个公式表示每个时间步ttoken有α_t概率保持原值有(1-α_t)概率变为噪声。反向过程则需要学习一个参数化的去噪模型p_θ通过最大化变分下界ELBO来训练。对于Uniform-State扩散模型(USDMs)其反向过程的后验分布为q_{s|t}(z_s^ℓ|z_t^ℓ,x^ℓ) Cat(·; [Kα_t z_t^ℓ⊙x^ℓ (α_{t|s}-α_t)z_t^ℓ (α_s-α_t)x^ℓ (1-α_{t|s})(1-α_s)1/K]/(Kα_tz_t^ℓ,x^ℓ1-α_t))这个复杂的表达式体现了USDMs允许token在生成过程中多次改变值的特性。1.2 两种主要噪声先验离散扩散模型主要使用两种噪声先验Masked先验π是特殊[MASK]token的one-hot编码。一旦token被masked在标准采样过程中就不能再恢复。这种限制导致错误累积问题。Uniform先验π1/KK是词汇表大小。允许token在整个生成过程中自由变化具有自校正能力特别适合少步生成和引导生成场景。关键区别USDMs的uniform先验使得每个token在每一步都有非零概率被重新采样这是其自校正能力的数学基础。2. Ψ-Samplers的技术突破2.1 Predictor-Corrector方法的局限性传统PC采样器如ReMDM虽然能改善Masked扩散模型的生成质量但存在两个根本限制仅适用于Masked先验无法直接应用于Uniform先验采样质量提升受限于噪声过程的类型Ψ-Samplers通过引入Ψ-后验(Ψ-posteriors)解决了这些问题Ψ_{s|t}(·|x^ℓ,z_t^ℓ) κ_t q_{s|t}(·|z_t^ℓ,x^ℓ) (1-κ_t)q_s(·|x^ℓ)其中κ_t∈[0,1]控制噪声注入强度。这个公式将标准反向后验与额外噪声线性组合形成新的生成过程。2.2 实现细节与算法Ψ-Samplers的实际实现包含以下关键步骤预测步骤使用当前去噪模型估计clean data x_θ(z_t,t)校正步骤按比例(1-κ_t)注入来自前向过程的噪声重加权保持与原始过程相同的边际分布算法伪代码如下Input: 去噪模型x_θ, 噪声调度α_t, 校正强度κ_t Initialize z_1 ~ π^L for t 1 to T-1 do s t - 1/T // 预测步骤 x_pred x_θ(z_t,t) q_pred q_{s|t}(·|z_t,x_pred) // 校正步骤 q_corr q_s(·|x_pred) // 混合采样 z_s ~ κ_t q_pred (1-κ_t)q_corr end for return z_02.3 参数选择策略κ_t的调度对性能至关重要实验表明语言建模rescaling调度(η0.05)配合log-linear噪声调度效果最佳图像生成分段常数κ_t(ton0.5-0.6, toff0.1)配合cosine调度最优下表比较了不同κ_t策略在OpenWebText上的表现调度类型Gen PPL512NFE熵恒定κ_t0.955.35.2Rescaling5.25.2Loop5.45.33. 高效课程学习方案3.1 原始课程的瓶颈Sahoo等人提出的课程学习虽然有效但存在两大问题内存消耗需要存储完整的K维softmax向量(K100,000)计算开销全连接矩阵乘法V^T y^ℓ复杂度为O(Kd)这在大型语言模型训练中成为主要瓶颈。3.2 稀疏近似方法利用低温softmax的稀疏性我们开发了基于top-k采样的高效近似有序统计量采样直接生成top-k值和索引避免计算全部K维归一化近似对未采样项建立解析近似保持概率质量守恒具体实现分为三个关键创新创新1top-k值生成// 对非目标token采样top-k高斯顺序统计量 K (K_1 ≥ ... ≥ K_k) ← 采样自N(0,σ_t^2)的top-k // 目标token值 w̃ ~ N(α̃_t, σ_t^2) // 确定插入位置 if w̃ K_k: K ← 将w̃插入K的适当位置创新2索引采样if w̃ ∉ top-k: I ~ 从[K]\{o}随机选k个索引 else: r w̃在K中的排名 L ~ 从[K]\{o}选r个索引作为上位 R ~ 从剩余选k-r-1个作为下位 I L ∥ o ∥ R创新3归一化计算Z̃ Σ_{i1}^k exp(K_i/τ) δexp(w̃/τ) (K-k-δ)exp(σ_t^2/(2τ^2)) * [Φ((K_k-σ_t^2/τ)/σ_t)/Φ(K_k/σ_t)]其中δ指示w̃是否在top-k中。3.3 性能对比在H100 GPU上的实测结果方法内存占用训练时间LM1B PPL原始课程1.0x1.0x29.9k20.67x0.75x30.0k30.70x0.77x30.1k50.73x0.80x30.24. 实际应用与效果验证4.1 语言建模性能在OpenWebText上的实验结果验证了Ψ-Samplers的三大优势持续改进随着NFE增加生成质量持续提升而传统采样会饱和超越Masked扩散在NFEL时USDMsΨ首次超过MDMsReMDM多样性保持在提升Gen PPL的同时维持较高的unigram熵关键数据对比NFE1024时 - MDLMReMDM: Gen PPL5.4, Entropy5.3 - DuoΨ: Gen PPL5.2, Entropy5.24.2 图像生成质量在CIFAR-10上的测试显示FID指标DuoΨ达到3.21优于MDLMReMDM的3.45IS指标DuoΨ达到9.12优于MDLMReMDM的8.87采样效率在相同NFE下USDMs需要更少的校正步骤4.3 下游任务表现在多选题回答任务中的准确率任务DuoDuo(k2)ARC-Challenge25.4326.11Winogrande47.2049.64OpenBookQA23.4027.805. 实施建议与注意事项5.1 工程实践要点精度处理采样时使用FP64计算logits避免数值下溢调度选择语言任务用log-linear视觉任务用cosine调度温度参数课程学习中保持τ1e-3确保足够的稀疏性5.2 常见问题排查问题1生成文本出现重复片段检查κ_t调度是否过于激进适当增加ton验证nucleus sampling参数p是否合适(建议p0.9)问题2训练时loss震荡降低课程阶段的初始噪声β(建议β0.03)检查top-k近似是否引入过多噪声(可尝试增大k)问题3图像生成出现artifacts确认cosine调度参数正确实现检查CFG强度是否过高(建议γ∈[1.5,2.0])5.3 扩展应用方向长文本生成结合Ψ-Samplers与分块处理技术多模态生成扩展至text-to-image交叉模态任务高效推理研究动态NFE调度与Ψ-Samplers的结合在实际部署中我们发现将k2的课程学习与rescaling Ψ-Sampler组合在A100上训练13B参数模型时相比原始方案可节省约$15k的云计算成本。