论文

选择性状态空间模型中状态使用的精确工具,以及它揭示的输入驱动迁移

An Exact Instrument for State Usage in Selective State-Space Models, and the Input-Driven Migration It Reveals

摘要

选择性状态空间模型(例如 Mamba)通过一组一阶模式传递信息,这些一阶模式的输入耦合由学习选择机制设置。我们提供了一种精确的工具来测量经过训练的模型如何使用这些模式。因为状态矩阵是对角矩阵,所以每个通道的输出精确地分解为每个模式的贡献,并且每个(层、通道、窗口)Gram 张量产生在任何预算下离线丢弃任何模式子集的精确输出误差。在 Mamba-1 系列上,根据参考实现的相对误差 $2.3\times10^{-7}$ 进行验证,该工具在 $4{,}464$ 配置上预测层的部署修剪误差为 $5\times10^{-7}$ 的中值相对偏差,其下限由重建设置。将仪器应用于 Mamba-1 系列 (130M--2.8B)、部署的 7B Falcon-Mamba 和 Mamba-2,我们发现经过训练的模型会根据输入重新分配其状态空间:哪些模式携带信号在上下文之间迁移,并且在受影响最严重的层,每个输入预言机将固定模式集的输出误差大致减半。冻结信号反事实将迁移主要归因于依赖于输入的写入映射 $B_t$;通常与选择性相关的时间步长几乎没有任何选择性。在从 130M 到部署的 7B Falcon-Mamba 的每个规模上,此测量的输入调度模式修剪都优于静态、基于 Hankel 和层自适应排名,并且在国家预算的一半下,它与未修剪模型相匹配。因为调度程序从第一遍读取每个窗口的模式使用情况,所以这展示了可实现的空间;我们声称没有部署计算或内存节省。