《Diffusion Models Without Attention》CVPR2024
摘要
这篇论文探讨了在高保真图像生成领域，去噪扩散概率模型（Denoising Diffusion Probabilistic Models, DDPMs）的重要性。尽管DDPMs在捕捉复杂视觉分布方面表现出色，但在高分辨率图像生成上面临显著的计算挑战。现有的方法，如在U-Net和Transformer架构中采用patchifying（分块处理），虽然加快了处理速度，但牺牲了表示能力。为了解决这一问题，论文提出了一种新的架构——Diffusion State Space Model（DIFFUSSM），它用一个更可扩展的状态空间模型骨干网络替代了注意力机制。这种方法在不进行全局压缩的情况下有效处理更高分辨率的图像，从而在整个扩散过程中保留了详细的图像表示。论文还强调了在扩散训练中对FLOP（浮点运算次数）高效架构的关注，并在ImageNet和LSUN数据集上的评估表明，DiffuSSM在FID（Fréchet Inception Distance）和Inception Score指标上与或超过现有的带有注意力模块的扩散模型，同时显著减少了总FLOP的使用。
概述
拟解决的问题：论文旨在解决DDPMs在高分辨率图像生成中的计算挑战，特别是在不牺牲图像细节和结构完整性的情况下，减少对高成本注意力机制的依赖。
创新之处：
- 架构创新：提出了DIFFUSSM，这是一种无需注意力机制的扩散架构，它使用状态空间模型（SSM）骨干网络来处理长序列，避免了传统注意力机制中的二次复杂度问题。
- 效率提升：通过使用SSM和hourglass（沙漏）架构，DIFFUSSM在保持图像细节的同时，提高了计算效率。
- 性能提升：在高分辨率图像生成任务中，DiffuSSM在FID和Inception Score等指标上达到了与或超过现有模型的性能，同时显著减少了计算资源的使用。
方法
- 输入序列化：将图像数据转换为序列形式，以便通过状态空间模型进行处理。
- 门控双向状态空间模型（Gated Bidirectional SSM）：这是DIFFUSSM块的核心，用于处理序列数据并捕捉长距离依赖关系。
- Hourglass 架构：用于在多层感知机（MLP）层中交替扩展和收缩序列长度，以提高计算效率。
- 输出解码：将处理后的序列数据解码回原始图像空间，生成最终的图像输出。
状态空间模型(SSMs)
状态空间模型通过将系统的状态表示为一个向量，并通过状态转移方程和观测方程来描述系统的动态行为。SSMs通常用于描述线性动态系统，但也可以扩展到非线性情况。
这个表达式描述了状态空间模型中的两个基本方程：状态方程和观测方程。
DIFFUSSM块
Hourglass架构是一种特殊的网络结构，它通过在网络的不同层之间交替扩展和收缩序列长度，来优化计算效率。在DIFFUSSM块中，Hourglass架构的具体实现如下：
- 扩展层（Upscaling Layer）：将序列长度扩展，以捕捉更广泛的上下文信息。
- 收缩层（Downscaling Layer）：将序列长度收缩，以减少计算量并提取关键特征。
- 重复层（Recurrent Layer）：在扩展和收缩层之间交替使用，以实现对序列数据的深度处理。
门控双向SSM是DIFFUSSM块的核心，它结合了状态空间模型和门控机制，以增强模型的处理能力。具体来说：
- 状态空间模型：使用线性或非线性动态系统来描述序列数据的转换过程。
- 门控机制：通过引入额外的控制信号（如类别标签或时间步信息），来调节状态空间模型的行为，从而提高模型的灵活性和适应性。
结论
论文得出结论，DIFFUSSM作为一种不依赖注意力机制的扩散模型架构，能够有效地处理高分辨率图像生成任务，同时在计算效率和图像质量上都取得了显著的改进。此外，该架构的提出为未来在需要长距离扩散的其他领域的应用，如高保真音频、视频或3D建模，提供了新的可能性。尽管论文主要关注图像生成任务，但其提出的方法和架构也可能对其他类型的生成任务有启发作用。