块偏置 Mamba 用于长程序列处理

Mamba 通过引入输入依赖动力学扩展了早期的状态空间模型（SSM），并在语言建模、计算机视觉和基础模型等多个领域展现出强大的经验性能。然而，尽管其建立在为长程依赖设计的架构之上，Mamba 在长程序列任务上仍表现不佳。理解和解决这一差距对于提高 Mamba 的通用性和灵活性至关重要。本文从三个角度分析 Mamba 的局限性：表达性、归纳偏置和训练稳定性。我们的理论结果表明，Mamba 在这些方面均不如早期的 SSM 如 S4D。为解决这些问题，我们提出了一种简单的方法，即在 Mamba 的 S6 单元中结合块式选择动力学和通道特定偏置的扩展，即 $\text{B}_2\text{S}_6$。我们证明，这些变更为模型提供了更合适的归纳偏置，同时提升了其表达性和稳定性。经验上，$\text{B}_2\text{S}_6$ 在长程竞技场（LRA）任务上优于 S4 和 S4D，同时保持了 Mamba 在语言建模基准测试上的性能。