关注 Mamba：跨架构的秘诀 蒸馏

诸如 Mamba 之类的状态空间模型 (SSM) 已成为 Transformer 模型的流行替代品，因为与基于注意力的模型相比，它们减少了内存消耗，并且生成时的吞吐量更高。另一方面，社区已经建立了大量关于如何训练 Transformer 的知识体系，并且许多预训练的 Transformer 模型都很容易获得。为了促进 SSM 的采用，同时利用现有的预训练 Transformer，我们的目标是找到一种有效的方法，将基于注意力的模型提炼成类似 Mamba 的架构。然而，在跨架构 蒸馏 的先前工作中，已经表明从 Transformer 到 Mamba 的简单 蒸馏 过程无法保留原始的教师性能，而结合注意力和 SSM 块的混合解决方案通常可以克服这一限制。我们工作的关键论点是，通过为 Mamba 配备有原则的初始化，我们可以为跨架构 蒸馏 恢复整体更好的配方。为此，我们提出了一种有原则的两阶段方法：首先，我们使用内核技巧的改编，将传统 Transformer 中的知识提炼成注意力的线性化版本。然后，我们将线性化版本提炼为不使用任何 Attention 块的改编 Mamba 模型。总体而言，蒸馏后的 Mamba 模型能够在下游任务中保留原始 Pythia-1B Transformer 的性能，保持 14.11 的困惑度接近老师的 13.86。为了展示我们的方案的有效性，我们使用 10B 词元不同的序列混合器架构进行了 1B 规模的彻底消融，对模型大小和总 蒸馏 词元进行了缩放分析，并对阶段之间的词元分配进行了敏感性分析。