arXiv cs.CL· Abhinav Moudgil, Ningyuan Huang, Eeshan Gunesh Dhekane, Pau Rodr\'iguez, Luca Zappella, Federico Danieli·· 11 小时前AI 评分31
Attention to Mamba：一种跨架构知识蒸馏方案
Attention to Mamba: A Recipe for Cross-Architecture Distillation
AI 导读
研究者提出两阶段跨架构蒸馏方案，先将 Transformer 知识蒸馏到线性化 Attention，再蒸馏到不含 Attention 模块的改造版 Mamba。
正文 · AI 翻译
摘要：状态空间模型（SSM），如 Mamba，已成为 Transformer 模型的热门替代方案，因为与基于注意力机制的模型相比，它们的内存消耗更低，生成时的吞吐量更高。另一方面，社区已经积累了相当多的关于如何训练 Transformer 的知识，并且许多预训练的 Transformer 模型可以随时获取。为了在利用现有预训练 Transformer 的同时促进 SSM 的采用，我们旨在确定一种有效的方案，将基于注意力的模型蒸馏为类似 Mamba 的架构。然而，在先前关于跨架构蒸馏的工作中已经表明，从 Transformer 到 Mamba 的朴素蒸馏过程无法保持原始教师模型的性能，这一局限通常通过结合注意力和 SSM 模块的混合方案来克服。我们工作的核心论点是，通过为 Mamba 配备有原则的初始化，我们可以获得一个整体更好的跨架构蒸馏方案。为此，我们提出了一种有原则的两阶段方法：首先，我们使用核技巧的改编版本，将知识从传统 Transformer 蒸馏到注意力的线性化版本中。然后，我们将线性化版本蒸馏到一个不使用任何注意力模块的改编 Mamba 模型中。总体而言，蒸馏后的 Mamba 模型保留了原始 Pythia-1B Transformer 的大部分性能，困惑度达到 14.01，接近教师的 13.86。这在很大程度上延续到了下游任务，尽管在更依赖上下文回忆的任务上出现了一些差距。为了展示我们方案的有效性，我们在 1B 规模、10B token 下进行了全面的消融实验，改变序列混合器架构，对模型大小和总蒸馏 token 进行缩放分析，并对阶段之间的 token 分配进行敏感性分析。
| 评论： | v2：修复了 v1 中影响 SSM 衰减门的实现错误；所有实验均在 Mamba-2 上重新运行并更新了结果；增加了推理效率结果 |
| 主题： | 计算与语言 (cs.CL)；机器学习 (cs.LG) |
| MSC 分类： | 68T07 |
| 引用为： | arXiv:2604.14191 [cs.CL] |
| （或 arXiv:2604.14191v2 [cs.CL] 用于此版本） | |
| https://doi.org/10.48550/arXiv.2604.14191 由 DataCite 颁发的 arXiv DOI |
提交历史
来自：Federico Danieli [查看邮箱]
[v1]
2026年4月1日星期三 09:23:08 UTC (14,935 KB)
[v2]
2026年11月3日星期二 22:13:00 UTC (14,637 KB)
来源：arXiv cs.CL · arxiv.org