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

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