关注 Mamba：跨架构蒸馏的配方
计算与语言
2026-04-17 v1 机器学习
摘要
状态空间模型（SSMs）如 Mamba 因其在生成时的内存消耗更低和吞吐量更高，成为 Transformer 模型的备选方案之一。另一方面，社区已积累了相当大量关于如何训练 Transformer 的知识，且许多预训练 Transformer 模型 readily 可用。为促进 SSM 的采用同时利用现有的预训练 Transformer，本文旨在识别一种有效的配方，用于将基于注意力的模型蒸馏到类似 Mamba 的架构。然而，在跨架构蒸馏的先前工作中，已证明将 Transformer 蒸馏到 Mamba 的简单做法无法保留原始教师性能，常需采用结合注意力和 SSM 块的混合解决方案来克服这一限制。我们的工作核心论点是，通过为 Mamba 配备原则性的初始化，我们可以恢复跨架构蒸馏的整体更好配方。为此，我们提出了一种原则性的两阶段方法：首先，采用核技巧的变体将知识从传统 Transformer 蒸馈到注意力的线性化版本；然后，将线性化版本的知识蒸馈到采用了不使用任何注意力块的改进 Mamba 模型。总体而言，蒸馏后的 Mamba 模型能够保留原始 Pythia-1B Transformer 在下游任务中的性能，保持困惑度为 14.11，接近教师的 13.86。为展示我们配方的有效性，我们在 1B 参数规模下进行了彻底的消融实验，变更序列混合器架构、对模型规模和总蒸馏 token 数进行了扩展分析，并对不同阶段的 token 分配进行了敏感性分析。
引用
@article{arxiv.2604.14191,
title = {Attention to Mamba: A Recipe for Cross-Architecture Distillation},
author = {Abhinav Moudgil and Ningyuan Huang and Eeshan Gunesh Dhekane and Pau Rodríguez and Luca Zappella and Federico Danieli},
journal= {arXiv preprint arXiv:2604.14191},
year = {2026}
}