丢掉注意力的集中模型:Mamba带火的SSM被苹果、康奈尔盯上了

替代注意力体制,SSM 真的大有可为?为了用更少的算力让集中模型生成高分辨率图象,注意力体制可以不要,这是康奈尔大学和苹果的一项最新研究所给出的结论。众所周知,注意力体制是 Transformer 架构的核心组件,对于高质量的文本、图象生成都至关重要。但它的缺陷也很明显,即算计复杂度会随着序列长度的增加呈现二次方增长。这在长文本、高分辨率的图象处理中都是一个令人头疼的问题。为了解决这个问题,这项新研究用一个可扩展性更强的状态空间模型(SSM)主干替代了传统架构中的注意力体制,开发出了一个名为 Diffusion S

替代注意力体制,SSM 真的大有可为?

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

丢掉注意力的集中模型:Mamba带火的SSM被苹果、康奈尔盯上了

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

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

丢掉注意力的集中模型:Mamba带火的SSM被苹果、康奈尔盯上了

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

丢掉注意力的集中模型:Mamba带火的SSM被苹果、康奈尔盯上了

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

丢掉注意力的集中模型:Mamba带火的SSM被苹果、康奈尔盯上了

论文概览

图象生成领域的迅速进展得益于去噪集中概率模型(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 更少。

丢掉注意力的集中模型:Mamba带火的SSM被苹果、康奈尔盯上了

论文链接:https://arxiv.org/pdf/2311.18257.pdf

DIFFUSSM 框架

作家的目标是设计一种集中架构,能够在高分辨率下学习长程相互作用,而无需像分块化那样进行「长度缩减」。与 DiT 类似,该方法通过展平图象并将其视为序列建模问题来实现。然而,与 Transformer 不同,这种方法在这个序列的长度上运用次二次(sub-quadratic)算计。

DIFFUSSM 的核心组件是优化处理长序列的门控双向 SSM。为了提高效率,作家在 MLP 层中引入沙漏架构。这种设计在双向 SSM 周围交替扩展和收缩序列长度,同时在 MLP 中特定地减少序列长度。完整的模型架构如图 2 所示。

丢掉注意力的集中模型:Mamba带火的SSM被苹果、康奈尔盯上了

具体来说,每个沙漏层接收经过缩短并展平的输入序列 I ∈ R^(J×D),其中 M = L/J 是缩小和放大的比例。同时,整个块,包括双向 SSM,在原始长度上进行算计,充分利用全局上下文。文中运用 σ 表示激活函数。对于 l ∈ {1 . . . L},其中 j = ⌊l/M⌋,m = l mod M,D_m = 2D/M,算计方程如下所示:

丢掉注意力的集中模型:Mamba带火的SSM被苹果、康奈尔盯上了

作家在每个层中运用跳跃连接集成门控 SSM 块。作家在每个位置集成了类标签 y ∈ R^(L×1) 和时间步 t ∈ R^(L×1) 的组合,如图 2 所示。

参数:DIFFUSSM 块中参数的数量主要由线性变换 W 决定,其中包含 9D^2 + 2MD^2 个参数。当 M = 2 时,这产生了 13D^2 个参数。DiT 变换块在其核心变换层中有 12D^2 个参数;然而,DiT 架构在其他层组件(自适应层归一化)中具有更多的参数。研究者在实验中通过运用额外的 DIFFUSSM 层来匹配参数。

FLOPs:图 3 比较了 DiT 和 DIFFUSSM 之间的 Gflops。DIFFUSSM 一层的总 Flops 为丢掉注意力的集中模型:Mamba带火的SSM被苹果、康奈尔盯上了,其中 α 代表 FFT 实现的常数。当 M = 2 且线性层主导算计时,这大约产生 7.5LD^2 Gflops。相比之下,如果在这个沙漏架构中运用全长的自注意力而不是 SSM,会有额外的 2DL^2 Flops。

丢掉注意力的集中模型:Mamba带火的SSM被苹果、康奈尔盯上了

考虑两种实验场景:1) D ≈ L = 1024,这将带来额外的 2LD^2 Flops,2) 4D ≈ L = 4096,这将产生 8LD^2 Flops 并显著增加成本。由于双向 SSM 的核心成本相对于运用注意力的成本较小,因此运用沙漏架构对基于注意力的模型不起作用。正如前面讨论的,DiT 通过运用分块化来避免这些问题,以代价是压缩表示。

实验结果

类别条件图象生成

表 1 是 DIFFUSSM 与目前所有的最先进的类别条件生成模型的比较结果。

丢掉注意力的集中模型:Mamba带火的SSM被苹果、康奈尔盯上了

当没有运用无分类器辅导时,DIFFUSSM 在 FID 和 sFID 两方面均优于其他集中模型,将之前非无分类器辅导潜在集中模型的最佳分数从 9.62 降至 9.07,同时运用的训练步骤减少到原来的 1/3 左右。在训练的总 Gflops 方面,未压缩模型相较于 DiT 减少了 20% 的总 Gflops。当引入无分类器辅导时,模型在所有基于 DDPM 的模型中获得了最佳的 sFID 分数,超过了其他最先进的策略,表明 DIFFUSSM 生成的图象对于空间失真更具鲁棒性。

DIFFUSSM 在运用无分类器辅导时的 FID 分数超越了所有模型,并在与 DiT 相比时保持了相当小的差距(0.01)。需要注意的是,在没有应用无分类器辅导的情况下,以减少 30% 的总 Gflops 训练的 DIFFUSSM 已经超过了 DiT。U-ViT 是另一种基于 Transformer 的架构,但采用了基于 UNet 的架构,块之间有长跳连接。U-ViT 在 256×256 分辨率下运用较少的 FLOPs,并在性能上表现更好,但在 512×512 数据集中情况并非如此。作家主要与 DiT 进行比较,为了公平,没有采用这种长跳连接,作家认为采用 U-Vit 的思想可能对 DiT 和 DIFFUSSM 都有益处。

作家进一步在更高分辨率的基准上运用无分类器辅导进行比较。DIFFUSSM 的结果相对强劲,并接近最先进的高分辨率模型,仅在 sFID 上不及 DiT,并获得了可比较的 FID 分数。DIFFUSSM 在 302M 张图象上进行了训练,观察了 40% 的图象,运用的 Gflops 比 DiT 少了 25%。

无条件图象生成

作家将模型的无条件图象生成能力与现有基线进行比较。结果显示在表 2 中。作家的研究发现,DIFFUSSM 在与 LDM 相当的训练预算下取得了可比较的 FID 分数(差距为 – 0.08 和 0.07)。这个结果突显了 DIFFUSSM 在不同基准和不同任务中的适用性。与 LDM 类似,由于只运用 ADM 总训练预算的 25%,因此在 LSUN-Bedrooms 任务中,该方法并未超过 ADM。对于这个任务,最佳 GAN 模型在模型类别上胜过集中模型。

丢掉注意力的集中模型:Mamba带火的SSM被苹果、康奈尔盯上了

更多内容请参考原论文。

给TA打赏
共{{data.count}}人
人已打赏
工程

HumanGaussian开源:鉴于Gaussian Splatting,高质量 3D 人体生成新框架

2023-12-11 15:28:00

工程

当GPT-4V充当机器人大脑,可能你都没AI会方案

2023-12-12 14:43:00

0 条回复 A文章作者 M管理员
    暂无讨论,说说你的看法吧
个人中心
今日签到
搜索