丢掉注意力的扩散模型：Mamba带火的SSM被苹果、康奈尔盯上了

2023.12.11

来源：机器之心

替代注意力机制，SSM 真的大有可为？

为了用更少的算力让扩散模型生成高分辨率图像，注意力机制可以不要，这是康奈尔大学和苹果的一项最新研究所给出的结论。

众所周知，注意力机制是 Transformer 架构的核心组件，对于高质量的文本、图像生成都至关重要。但它的缺陷也很明显，即计算复杂度会随着序列长度的增加呈现二次方增长。这在长文本、高分辨率的图像处理中都是一个令人头疼的问题。

为了解决这个问题，这项新研究用一个可扩展性更强的状态空间模型（SSM）主干替代了传统架构中的注意力机制，开发出了一个名为 Diffusion State Space Model（DIFFUSSM）的新架构。这种新架构可以使用更少的算力，媲美或超越具有注意力模块的现有扩散模型的图像生成效果，出色地生成高分辨率图像。

得益于上周「Mamba」的发布，状态空间模型 SSM 正受到越来越多的关注。Mamba 的核心在于引入了一种新的架构 ——「选择性状态空间模型（ selective state space model）」，这使得 Mamba 在语言建模方面可以媲美甚至击败 Transformer。当时，论文作者 Albert Gu 表示，Mamba 的成功让他对 SSM 的未来充满了信心。如今，康奈尔大学和苹果的这篇论文似乎又给 SSM 的应用前景增加了新的例证。

微软首席研究工程师 Shital Shah 提醒说，注意力机制可能要从坐了很久的王座上被拉下来了。

论文概览

图像生成领域的迅速进展得益于去噪扩散概率模型（DDPMs）。这类模型将生成过程建模为迭代去噪潜变量，当执行足够的去噪步骤时，它们能够产生高保真度的样本。DDPMs 捕捉复杂视觉分布的能力使其在推动高分辨率、照片级合成方面具有潜在的优势。

在将 DDPMs 扩展到更高分辨率方面仍然存在重要的计算挑战。主要瓶颈是在实现高保真生成时依赖自注意力。在 U-Nets 架构中，这个瓶颈来自将 ResNet 与注意力层相结合。DDPMs 超越了生成对抗网络 (GANs)，但需要多头注意力层。在 Transformer 架构中，注意力是中心组件，因此对于实现最新的图像合成结果至关重要。在这两种架构中，注意力的复杂性，与序列长度成二次方关系，所以当处理高分辨率图像时将变得不可行。

计算成本促使以往的研究者们使用表示压缩方法。高分辨率架构通常采用分块化（patchifying）或多尺度分辨率。通过分块化可以创建粗粒度表示，降低计算成本，但代价是牺牲关键的高频空间信息和结构完整性。多尺度分辨率虽然可以减少注意层的计算，但也会通过降采样减少空间细节并在应用上采样时引入伪影。

扩散状态空间模型（DIFFUSSM）是一种不使用注意力机制的扩散架构，它旨在解决在高分辨率图像合成中应用注意力机制时出现的问题。DIFFUSSM 在扩散过程中采用了门控状态空间模型 (SSM)。之前的研究表明，基于 SSM 的序列模型是一种有效而且高效的通用神经序列模型。通过使用这种架构，可以使 SSM 核心处理更细粒度的图像表示，消除全局分块化或多尺度层。为进一步提高效率，DIFFUSSM 在网络的密集组件中采用沙漏 (hourglass) 架构。

作者在不同分辨率下验证了 DIFFUSSM 的性能。在 ImageNet 上的实验证明，在各种分辨率下，DIFFUSSM 在 FID、sFID 和 Inception Score 上都取得了一致的改进，并且总 Gflops 更少。