Block-Biased Mamba for Long-Range Sequence Processing

📄 arXiv: 2505.09022v1 📥 PDF

作者: Annan Yu, N. Benjamin Erichson

分类: cs.LG, cs.AI, stat.ML

发布日期: 2025-05-13

💡 一句话要点

提出Block-Biased Mamba（B2S6）以提升Mamba在长序列任务上的性能。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 长序列建模 状态空间模型 Mamba 长程依赖 归纳偏置 分块处理 选择性扫描

📋 核心要点

  1. Mamba在长序列建模中表现出潜力，但在长程依赖任务上存在性能瓶颈，限制了其通用性。
  2. 论文提出Block-Biased Mamba（B2S6），通过分块选择性动态和通道偏置增强模型的表达能力和稳定性。
  3. 实验结果表明，B2S6在LRA任务上超越S4和S4D，同时保持了Mamba在语言建模上的优势。

📝 摘要（中文）

Mamba通过引入输入依赖的动态机制扩展了早期的状态空间模型（SSM），并在包括语言建模、计算机视觉和基础模型在内的多个领域表现出强大的经验性能。然而，一个令人惊讶的弱点仍然存在：尽管建立在为长程依赖设计的架构之上，但Mamba在长程序列任务上的表现不佳。理解和解决这一差距对于提高Mamba的通用性和多功能性非常重要。在这项工作中，我们从表达能力、归纳偏置和训练稳定性三个角度分析了Mamba的局限性。我们的理论结果表明，与早期的SSM（如S4D）相比，Mamba在这些方面存在不足。为了解决这些问题，我们提出了$ ext{B}_2 ext{S}_6$，这是Mamba的S6单元的一个简单扩展，它结合了分块选择性动态和通道特定的偏置。我们证明这些改变使模型具有更适合的归纳偏置，并提高了其表达能力和稳定性。在经验上，$ ext{B}_2 ext{S}_6$在长程竞技场（LRA）任务上优于S4和S4D，同时保持了Mamba在语言建模基准上的性能。

🔬 方法详解

问题定义：Mamba虽然在许多领域表现出色，但在处理长程依赖序列时性能下降。现有方法，如S4和S4D，在长程任务上表现更好，但可能在其他任务上不如Mamba。Mamba的痛点在于其在表达能力、归纳偏置和训练稳定性方面存在不足，导致无法有效捕捉长程依赖关系。

核心思路：论文的核心思路是通过改进Mamba的S6单元，使其更适合处理长程序列。具体来说，引入分块选择性动态和通道特定的偏置。分块选择性动态允许模型在不同的序列块上应用不同的动态，从而更好地捕捉局部和全局的依赖关系。通道特定的偏置则有助于模型更好地学习不同通道之间的关系，提高表达能力。

技术框架：B2S6建立在Mamba的S6单元之上。主要修改包括：1) 将选择性扫描机制应用于序列的分块；2) 为每个通道引入一个可学习的偏置项。整体流程与Mamba类似，包括输入嵌入、S6单元处理和输出预测。关键在于S6单元内部的改进，使其能够更好地处理长程依赖。

关键创新：最重要的技术创新点是分块选择性动态和通道特定偏置的结合。与原始Mamba相比，B2S6能够更有效地捕捉长程依赖关系，同时保持模型的效率和可扩展性。与S4和S4D相比，B2S6在长程任务上表现更好，同时在其他任务上也能保持竞争力。

关键设计：分块大小是一个关键参数，需要根据具体的任务进行调整。通道特定偏置的初始化也很重要，可以采用一些启发式方法来提高训练的稳定性。损失函数与原始Mamba相同，通常采用交叉熵损失函数。

🖼️ 关键图片

fig_0

📊 实验亮点

实验结果表明，B2S6在长程竞技场（LRA）任务上显著优于S4和S4D，证明了其在长程依赖建模方面的优势。同时，B2S6在语言建模基准上保持了与Mamba相当的性能，表明其在提高长程建模能力的同时，没有牺牲在其他任务上的表现。具体的性能提升幅度在论文中进行了详细的量化。

🎯 应用场景

该研究成果可应用于需要处理长序列数据的各种领域，如基因组学、视频分析、语音识别和自然语言处理。特别是在需要捕捉长期依赖关系的任务中，B2S6有望提供更强大的建模能力。例如，在长篇文档摘要、视频内容理解和长时间语音识别等场景中具有潜在的应用价值。

📄 摘要（原文）

Mamba extends earlier state space models (SSMs) by introducing input-dependent dynamics, and has demonstrated strong empirical performance across a range of domains, including language modeling, computer vision, and foundation models. However, a surprising weakness remains: despite being built on architectures designed for long-range dependencies, Mamba performs poorly on long-range sequential tasks. Understanding and addressing this gap is important for improving Mamba's universality and versatility. In this work, we analyze Mamba's limitations through three perspectives: expressiveness, inductive bias, and training stability. Our theoretical results show how Mamba falls short in each of these aspects compared to earlier SSMs such as S4D. To address these issues, we propose $\text{B}_2\text{S}_6$, a simple extension of Mamba's S6 unit that combines block-wise selective dynamics with a channel-specific bias. We prove that these changes equip the model with a better-suited inductive bias and improve its expressiveness and stability. Empirically, $\text{B}_2\text{S}_6$ outperforms S4 and S4D on Long-Range Arena (LRA) tasks while maintaining Mamba's performance on language modeling benchmarks.