1. 离散扩散模型D3PM的核心思想解析
离散去噪扩散概率模型(D3PM)是扩散模型在离散状态空间的自然延伸。与连续空间的DDPM不同,D3PM通过精心设计的转移矩阵来控制离散数据的损坏过程,这为模型性能带来了显著提升。让我们深入剖析其核心机制。
1.1 离散扩散的数学基础
D3PM建立在分类随机变量的马尔可夫链之上。对于具有K个类别的离散变量,前向转移概率可以表示为K×K的矩阵Q_t,其中[Q_t]ij = q(x_t=j|x=i)。这种表示允许我们使用矩阵运算来描述整个扩散过程:
- 单步转移:q(x_t|x_{t-1}) = Cat(x_t; p = x_{t-1}Q_t)
- 多步累积:q(x_t|x_0) = Cat(x_t; p = x_0Q̄_t),其中Q̄_t = Q_1Q_2...Q_t
这种矩阵表示不仅简洁,而且便于计算各种边缘分布和后验分布。特别地,反向过程的后验分布可以表示为:
q(x_{t-1}|x_t,x_0) = Cat(x_{t-1}; p = (x_tQ_t^T ⊙ x_0Q̄_{t-1})/(x_0Q̄_tx_t^T))
1.2 转移矩阵的设计哲学
D3PM的核心创新在于对转移矩阵Q_t的精心设计。与连续空间DDPM通常使用高斯噪声不同,离散空间提供了更丰富的设计可能性:
-
均匀转移矩阵:每个状态以相同概率转移到其他状态,完全忽略状态间的语义关系。这是最基础的设置,但性能通常有限。
-
吸收状态矩阵:引入特殊的[MASK]状态作为吸收态,一旦进入就无法离开。这种设计与BERT等掩码语言模型有密切联系,特别适合文本数据。
-
离散化高斯矩阵:在有序离散数据(如图像像素值)上模拟连续高斯扩散的行为,使相似状态间转移概率更高。
-
嵌入相似度矩阵:利用预训练嵌入空间中token的相似度来指导转移概率,使语义相近的token更可能相互转换。
这些设计使得D3PM能够根据数据类型和任务需求,选择最合适的噪声注入方式,这是其性能优于传统离散扩散模型的关键。
2. D3PM的关键技术实现
2.1 反向过程的参数化技巧
D3PM采用了一种巧妙的参数化方式:不是直接预测p_θ(x_{t-1}|x_t),而是预测p̃_θ(x̃_0|x_t),然后通过以下公式计算反向转移概率:
p_θ(x_{t-1}|x_t) ∝ Σ_{x̃_0} q(x_{t-1},x_t|x̃_0)p̃_θ(x̃_0|x_t)
这种"x_0参数化"有三大优势:
- 自动确保反向转移概率的稀疏模式与前向过程一致
- 允许一次性预测多步的反向转移
- 当预测准确时,KL散度项会精确为0
对于有序离散数据,还可以进一步使用截断逻辑斯蒂分布来参数化p̃_θ(x̃_0|x_t),引入有益的归纳偏置。
2.2 混合损失函数设计
D3PM提出了一种结合变分下界和辅助去噪目标的混合损失函数:
L_λ = L_vb + λE[-log p̃_θ(x_0|x_t)]
其中:
- L_vb是标准的变分下界
- 辅助项鼓励模型在每个时间步都能准确预测原始数据x_0
实验表明,适当的λ值(如0.001)能显著提升模型性能,特别是在图像生成任务上。这是因为辅助项提供了更稳定的梯度信号,缓解了L_vb在不同时间步梯度幅度差异过大的问题。
2.3 噪声调度策略
D3PM针对不同类型的转移矩阵设计了专门的噪声调度方案:
- 均匀转移:采用余弦调度,使累积转移概率随时间呈余弦变化
- 吸收态转移:使用(T-t+1)^{-1}调度,线性增加被吸收的概率
- 离散高斯转移:采用线性增加方差的调度
- 通用情况:基于互信息的调度,使I(x_t;x_0)≈(1-t/T)H(x_0)
这些精心设计的调度策略确保了前向过程能有效地将数据逐渐扰动到平稳分布,同时为反向过程的学习提供合适的难度梯度。
3. D3PM与现有模型的关联
3.1 与BERT的联系
当使用单步吸收态扩散,并以10%概率掩码、5%概率随机替换时,D3PM的损失函数退化为标准的BERT去噪目标。这表明BERT可以视为D3PM的一个特例,这为理解BERT的生成能力提供了新视角。
3.2 与自回归模型的关系
通过设计特定的确定性前向过程(逐步掩码序列中的token),D3PM可以完全重现自回归模型的训练目标。这说明自回归模型实际上是扩散模型的一种特殊形式,这一发现统一了两种看似不同的生成范式。
3.3 与掩码语言模型的对应
生成式掩码语言模型(如SpanBERT)本质上是使用特定调度和吸收态转移矩阵的D3PM。D3PM框架为这类模型提供了更坚实的理论基础,并启发了可能的改进方向。
4. D3PM的实践表现
4.1 文本生成任务
在text8字符级生成任务上,D3PM吸收态模型取得了1.45 bits/char的负对数似然,显著优于均匀扩散的1.61。在LM1B词级任务上,吸收态模型达到76.9的困惑度,远优于均匀扩散的137.9。
关键发现:
- 吸收态设计特别适合文本数据
- 嵌入相似度引导的转移矩阵表现不如预期
- 模型可以高效地使用较少步数(如20步)进行采样
4.2 图像生成任务
在CIFAR-10上,D3PM高斯+逻辑斯蒂模型取得了:
- Inception Score: 8.56
- FID: 7.34
- NLL: 3.435 bits/dim
这甚至超过了原始连续DDPM的性能,证明了离散扩散在图像领域的潜力。
4.3 计算效率考量
D3PM通过矩阵运算的优化,实现了高效的训练和采样。例如,在text8上生成256长度文本仅需0.58秒(256步),比相同规模的Transformer快近20倍。
5. 实际应用建议
5.1 模型选型指南
- 文本数据:优先考虑吸收态D3PM,它自然地与掩码语言模型范式对齐
- 有序离散数据(如图像):使用离散高斯转移矩阵,配合逻辑斯蒂参数化
- 无明确结构的离散数据:可尝试均匀转移或嵌入相似度转移
5.2 实现注意事项
- 对于大词表,需优化Q̄_t的计算,避免内存爆炸
- 损失函数中λ的选择需要验证,通常0.001-0.01效果较好
- 不同的噪声调度对最终性能影响显著,需仔细调整
5.3 未来改进方向
- 探索更复杂的转移矩阵设计,如分层、注意力驱动的转移
- 结合连续时间扩散的理论进一步优化离散噪声调度
- 研究更高效的反向过程参数化方式,减少采样步数
D3PM框架为离散数据生成提供了强大而灵活的新工具,其性能已经展现出超越连续扩散模型的潜力。通过深入理解其设计原理和实现细节,研究者可以在各种离散数据任务上获得显著的性能提升。
