苹果提出一种两阶段蒸馏方法，将 Transformer 模型转换为 Mamba 架构，在仅使用 2.7% 训练数据的情况下，使 1B 参数模型的困惑度从 14.89 恢复至 14.11，逼近教师模型的 13.86，实现了推理成本从平方级到线性的转变。
📝 详细摘要
本文介绍了苹果公司最新发表的一篇论文，提出了一种将 Transformer 模型转换为 Mamba 架构的新型跨架构蒸馏方法。核心挑战在于 Transformer 的 Softmax Attention 计算复杂度随序列长度平方增长，而 Mamba 等状态空间模型虽然具有线性复杂度，但性能通常不及 Transformer。苹果的解决方案是两阶段蒸馏：第一阶段将原始的 Softmax Attention 替换为一种通过小型神经网络学习的线性注意力（Hedgehog），并通过余弦相似度蒸馏对齐输出；第二阶段将线性注意力嵌入 Mamba 结构，通过参数映射和结构对齐实现平滑过渡，最后通过微调恢复模型能力。实验表明，在仅使用教师模型 2.7% 训练数据（10B token）的情况下，1B 参数的 HedgeMamba 模型在困惑度上达到 14.11，接近教师模型 Pythia 的 13.86，并在 Arc、PIQA、BoolQ 等多个下游任务上全面超过传统蒸馏基线。消融实验进一步揭示了门控机制的关键作用、两阶段数据分配的最优策略以及该方法的可规模化特性。这项工作为将现有大量 Transformer 模型低成本转换为高效架构提供了新的工程可能性。
💡 主要观点
- 苹果提出两阶段蒸馏方法，将 Transformer 转换为 Mamba 架构。 第一阶段将 Softmax Attention 替换为学习型线性注意力（Hedgehog），通过余弦相似度蒸馏对齐；第二阶段将线性注意力嵌入 Mamba 结构，通过参数映射实现结构对齐，避免直接蒸馏的性能崩塌。
💬 文章金句
- 把又贵又强的 Transformer，改造成又便宜又差不多强的 Mamba。而且，性能基本没怎么掉。
- 过去大家默认一个前提：只要你把 Transformer 换成另一种架构，性能就会明显掉一截。但这篇论文给出的答案是，这个损失可以被大幅追回来。
- 过去几年积累的大量 Transformer 模型，并不需要全部推倒重来，而是有机会通过一套流程，被「转制」为更高效的形态。
📊 文章信息
AI 初评：88
来源：机器之心
作者：机器之心
分类：人工智能
语言：中文
阅读时间：10 分钟
字数：2408
标签： Transformer, Mamba, 模型蒸馏, 线性注意力, 苹果