稀疏化状态空间模型是高效的高速公路网络

状态空间模型（SSMs）为序列建模提供了一种很有前途的架构，通过用线性递归取代昂贵的自注意力机制，为 Transformer 提供了一种替代方案。在本文中，我们提出了一种简单而有效的技巧，通过对 SSMs 进行稀疏化，以在给定计算预算内增强其性能。我们的直觉是，由于渐进的递归更新，SSMs 中的令牌高度冗余，并且密集的递归操作会阻碍过去信息的传递。特别是，我们观察到 SSMs 的上层倾向于更冗余，因为它们编码全局信息，而下层编码局部信息。受此启发，我们引入了 Simba，一种基于令牌剪枝的 SSMs 分层稀疏化方法。Simba 对上层比下层进行更多的稀疏化，鼓励上层表现得像高速公路。为了实现这一点，我们提出了一种新颖的 SSMs 令牌剪枝准则，通过累积局部递归来衡量令牌对最终输出的全局影响。我们证明，Simba 在各种自然语言任务中以相同的 FLOPS 优于基线模型 Mamba。此外，我们阐释了高速公路的效果，表明 Simba 不仅提高了效率，而且改善了长序列的信息流。代码可在 https://github.com/woominsong/Simba 获取。