MambaMixer: Efficient Selective State Space Models with Dual Token and Channel Selection

📄 arXiv: 2403.19888v4 📥 PDF

作者: Ali Behrouz, Michele Santacatterina, Ramin Zabih

分类: cs.LG, cs.AI, cs.CV

发布日期: 2024-03-29 (更新: 2024-07-23)

💡 一句话要点

提出MambaMixer以解决长序列建模的效率问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 长序列建模 选择性状态空间模型 计算机视觉 时间序列预测 深度学习 Transformer

📋 核心要点

  1. 现有的Transformer架构在长序列建模中存在时间和空间复杂度高的问题，限制了其可扩展性。
  2. MambaMixer通过双重选择机制在令牌和通道之间进行选择，采用加权平均连接选择性混合器，提升了特征访问效率。
  3. ViM2在图像分类、目标检测和语义分割任务中表现优异，TSM2在时间序列预测中超越了现有最优方法，且计算成本显著降低。

📝 摘要（中文）

近年来，深度学习的进展主要依赖于Transformer架构，但其注意力模块在输入大小上表现出平方级的时间和空间复杂度，限制了其在长序列建模中的可扩展性。尽管已有一些针对多维数据的高效架构设计，但现有模型要么与数据无关，要么未能实现维度间和维度内的有效通信。为此，本文提出了MambaMixer，一种具有数据依赖权重的新架构，采用双重选择机制进行令牌和通道的选择。通过加权平均机制连接选择性混合器，使得层能够直接访问早期特征。我们设计了基于MambaMixer模块的视觉和时间序列架构ViM2和TSM2，并在多个任务中验证了其性能。实验结果表明，选择性混合在令牌和通道之间的重要性，ViM2在多个视觉任务中表现出色，TSM2在时间序列预测中也取得了显著的性能提升。

🔬 方法详解

问题定义：本文旨在解决现有Transformer在长序列建模中的效率问题，尤其是其注意力机制导致的高时间和空间复杂度。现有模型在多维数据处理上存在数据独立性或缺乏维度间通信的不足。

核心思路：MambaMixer的核心思路是通过引入数据依赖的权重和双重选择机制，优化令牌和通道的选择，从而提高模型在长序列任务中的性能和效率。这样的设计使得模型能够更好地利用输入数据的特征。

技术框架：MambaMixer的整体架构包括选择性令牌和通道混合器，通过加权平均机制连接不同层，使得后续层能够直接访问早期特征。该架构分为多个模块，分别处理不同维度的数据。

关键创新：MambaMixer的主要创新在于其双重选择机制和加权平均连接方式，这与传统的Transformer架构在处理长序列时的全局注意力机制形成鲜明对比，显著降低了计算复杂度。

关键设计：在设计中，MambaMixer采用了特定的参数设置以优化选择性混合过程，损失函数经过调整以适应不同任务的需求，网络结构则通过模块化设计增强了灵活性和可扩展性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

在实验中，ViM2在ImageNet分类、目标检测和语义分割任务中表现出色，达到了与现有视觉模型相当的性能，并超越了基于SSM的视觉模型。TSM2在时间序列预测中表现优异，相较于最先进的方法显著降低了计算成本，展示了其在实际应用中的优势。

🎯 应用场景

MambaMixer的研究成果在多个领域具有广泛的应用潜力，包括计算机视觉、时间序列预测和多模态数据分析。其高效的长序列建模能力可以为实时数据处理、智能监控和自动化决策系统提供支持，未来可能在工业、医疗和金融等领域产生深远影响。

📄 摘要（原文）

Recent advances in deep learning have mainly relied on Transformers due to their data dependency and ability to learn at scale. The attention module in these architectures, however, exhibits quadratic time and space in input size, limiting their scalability for long-sequence modeling. Despite recent attempts to design efficient and effective architecture backbone for multi-dimensional data, such as images and multivariate time series, existing models are either data independent, or fail to allow inter- and intra-dimension communication. Recently, State Space Models (SSMs), and more specifically Selective State Space Models, with efficient hardware-aware implementation, have shown promising potential for long sequence modeling. Motivated by the success of SSMs, we present MambaMixer, a new architecture with data-dependent weights that uses a dual selection mechanism across tokens and channels, called Selective Token and Channel Mixer. MambaMixer connects selective mixers using a weighted averaging mechanism, allowing layers to have direct access to early features. As a proof of concept, we design Vision MambaMixer (ViM2) and Time Series MambaMixer (TSM2) architectures based on the MambaMixer block and explore their performance in various vision and time series forecasting tasks. Our results underline the importance of selective mixing across both tokens and channels. In ImageNet classification, object detection, and semantic segmentation tasks, ViM2 achieves competitive performance with well-established vision models and outperforms SSM-based vision models. In time series forecasting, TSM2 achieves outstanding performance compared to state-of-the-art methods while demonstrating significantly improved computational cost. These results show that while Transformers, cross-channel attention, and MLPs are sufficient for good performance in time series forecasting, neither is necessary.