【第614期】从Attention到Mamba：跨架构蒸馏方案

Seventy3：借助NotebookLM的能力进行论文解读，专注人工智能、大模型、机器人算法、crypto方向，让大家跟着AI一起进步。

如果你想要解读自己的论文，获得更多曝光度。请联系小助手微信：seventy3_podcast 加群。
合作邮箱：zhiwudazhanjiangshi#gmail.com

今天的主题是：

Attention to Mamba: A Recipe for Cross-Architecture Distillation

Summary

状态空间模型（State Space Models, SSMs），例如 Mamba，由于相比基于注意力机制（Attention-based）的模型在生成阶段具有更低的内存消耗和更高的吞吐效率，近年来成为 Transformer 的一种流行替代方案。

另一方面，社区已经积累了大量关于 Transformer 训练的经验，并且有许多预训练 Transformer 模型可以直接使用。为了在利用这些预训练资源的同时促进 SSM 的采用，我们希望找到一种有效的配方（recipe），将基于 Attention 的模型蒸馏到类似 Mamba 的架构中。

然而，先前关于跨架构蒸馏（cross-architecture distillation）的研究表明，从 Transformer 到 Mamba 的朴素蒸馏方法无法很好地保留教师模型的性能，这一问题通常需要通过结合 Attention 与 SSM 模块的混合架构来缓解。

我们工作的核心观点是：如果能够为 Mamba 提供一种合理的初始化方式，就可以恢复一种更优的跨架构蒸馏方案。

为此，我们提出一种原则性的两阶段方法：

将传统 Transformer 的知识蒸馏到一种“线性化注意力（linearized Attention）”版本中，并通过核技巧（kernel trick）的改造实现该过程。

再将该线性化注意力模型进一步蒸馏到一个经过适配的 Mamba 模型中，该模型完全不使用 Attention 模块。

总体而言，该蒸馏得到的 Mamba 模型能够在下游任务中保持原始 Pythia-1B Transformer 的性能，其困惑度（perplexity）为 14.11，接近教师模型的 13.86。

为了验证该方法的有效性，我们在 1B 规模模型与 10B tokens 的设置下进行了系统性实验，包括：