[论文解读] Mamba-ND: Selective State Space Modeling for Multi-Dimensional Data

Shufan Li, Harkanwar Singh+1 more|arXiv (Cornell University)|Feb 8, 2024

Neural Networks and Applications被引用 4

一句话总结

Mamba-ND 通过在层间交替使用行主序排列方式，将 Mamba 状态空间模型扩展至多维数据，在图像、视频和天气预报任务上实现了最先进性能，参数量显著减少且复杂度为线性，优于 ViT 和 Swin-Transformer 等 Transformer 模型，同时模型大小最多减少 44.5%。

ABSTRACT

In recent years, Transformers have become the de-facto architecture for sequence modeling on text and a variety of multi-dimensional data, such as images and video. However, the use of self-attention layers in a Transformer incurs prohibitive compute and memory complexity that scales quadratically w.r.t. the sequence length. A recent architecture, Mamba, based on state space models has been shown to achieve comparable performance for modeling text sequences, while scaling linearly with the sequence length. In this work, we present Mamba-ND, a generalized design extending the Mamba architecture to arbitrary multi-dimensional data. Our design alternatively unravels the input data across different dimensions following row-major orderings. We provide a systematic comparison of Mamba-ND with several other alternatives, based on prior multi-dimensional extensions such as Bi-directional LSTMs and S4ND. Empirically, we show that Mamba-ND demonstrates performance competitive with the state-of-the-art on a variety of multi-dimensional benchmarks, including ImageNet-1K classification, HMDB-51 action recognition, and ERA5 weather forecasting.

研究动机与目标

• 将原本专为一维序列设计的 Mamba 状态空间模型扩展至图像、视频和科学数据等多维数据。
• 解决将 Mamba 应用于多维输入时，其序列处理依赖特定数据排序所带来的挑战。
• 评估并比较多种 SSM 的多维扩展方法，包括双向、N 方向及多头设计。
• 识别在不引入复杂分解或架构的前提下，将一维 SSM 扩展至高维的最有效且高效的架构。
• 证明仅通过层间简单交替排序策略，即可在性能和效率上超越复杂的多方向设计。

提出的方法

• Mamba-ND 将一维 Mamba 层作为黑箱组件使用，对多维数据进行展平后，在连续层之间采用交替的行主序排列方式处理。
• 模型在不同层之间交替使用序列排序方式——例如行主序、列主序及类似对角线的排序方式——以实现多方向信息流动，而无需修改底层 SSM 机制。
• 通过利用 Mamba 的选择性状态空间机制，保持了与序列长度呈线性关系的复杂度，避免了自注意力机制带来的二次方复杂度。
• 该架构无需分解或专用的二维/三维 SSM 层；相反，它仅在每个一维 SSM 层之前对输入数据进行重排。
• 采用块级设计，每个块应用不同的排序方式，使模型能够跨多个方向捕捉空间和时间依赖性。
• 该方法对输入维度不敏感，通过将数据视为补丁序列，可泛化至二维（图像）、三维（视频）甚至四维（气象数据）数据。

Figure 1: Mamba-ND outperforms Transformers while significantly reducing the number of parameters. On ImageNet-1k, we compare against ViT [ 12 ] . On HMDB-51 [ 22 ] , we compare against Video-Swin [ 28 ] . On ERA5, we compare against Cli-ViT [ 34 ] .

实验结果

研究问题

• RQ1 原本专为一维文本序列设计的 Mamba 架构，能否有效扩展至图像和视频等多维数据？
• RQ2 在建模多维数据时，仅通过层间简单交替序列排序策略，是否优于更复杂的多方向或多头 SSM 设计？
• RQ3 Mamba-ND 的有效感受野与单向和双向基线模型相比，在捕捉空间和时间依赖性方面有何差异？
• RQ4 架构选择如分解、层排序和深度对性能和参数效率有何影响？
• RQ5 仅通过在层间重新排列输入序列这一最小改动，是否能实现与最先进 Transformer 模型相当的性能？

主要发现

• 在 ImageNet-1K 上，Mamba-ND 的 top-1 准确率比 ViT 高 1.5%，同时参数量减少 20.7%。
• 在 HMDB-51 动作识别基准上，Mamba-ND 的准确率比 3D Swin-Transformer 高 0.9%，参数量减少 39.0%。
• 在 ERA5 72 小时天气预报任务中，Mamba-ND 相较于 Cli-ViT 在异常相关系数（ACC）上提升 0.7 分，参数量减少 44.5%。
• 尽管复杂度更高，交替行主序排序策略仍优于 2D+3D 分解、多头 SSM 和双向 SSM 等更复杂设计。
• 有效感受野可视化显示，Mamba-ND 的敏感度分布比单向或双向模型更均匀、更具全局性，解释了其优越性能。
• 消融实验表明，更简单的交替排序设计优于更深或更宽的多方向架构，支持在 SSM 中深度优先于宽度的结论。

Figure 2: Architecture of Mamba-ND: In this figure we visualize Mamba-3D. Given 3D input data, we patchify it into $L$ patches. During this process, we maintain the original 3D structure of the input. This sequence is then passed through $K$ Mamba-ND blocks, each of which consists of a chain of 1-D