研究者提出将Transformer蒸馏为Mamba架构的有效方法

Attention to Mamba：一种跨架构蒸馏方案

作者：Abhinav Moudgil†、Ningyuan Huang‡、Eeshan Gunesh Dhekane、Pau Rodríguez Lopez、Luca Zappella、Federico Danieli

状态空间模型（SSM）如 Mamba 已成为 Transformer 模型的热门替代方案，原因在于与基于 Attention 的模型相比，它们在生成时内存消耗更低、吞吐量更高。另一方面，学术界已经积累了相当丰富的关于如何训练 Transformer 的知识，并且许多预训练的 Transformer 模型可以直接获取。为了在利用现有预训练 Transformer 的同时促进 SSM 的采用，我们旨在找到一种有效的方案，将基于 Attention 的模型蒸馏为类 Mamba 架构。然而，在先前关于跨架构蒸馏的工作中，研究表明从 Transformer 到 Mamba 的朴素蒸馏过程无法保持原始教师模型的性能，这一局限通常通过结合 Attention 和 SSM 模块的混合方案来克服。我们工作的核心论点是：通过为 Mamba 配备有原则的初始化，我们可以获得一个整体更优的跨架构蒸馏方案。为此，我们提出了一种有原则的两阶段方法：首先，我们使用核技巧的改编版本，将传统 Transformer 中的知识蒸馏到 Attention 的线性化版本中。然后，我们将该线性化版本蒸馏到一个不使用任何 Attention 模块的适配 Mamba 模型中。总体而言，蒸馏后的 Mamba 模型保留了原始 Pythia-1B Transformer 的大部分性能，困惑度达到 14.01，接近教师模型的 13.86。这一效果在很大程度上延续到了下游任务，尽管在更依赖上下文回忆的任务上出现了一些差距。为了展示我们方案的有效性，我们在 1B 规模、10B token 的条件下进行了全面的消融实验，变化序列混合器架构，对模型尺寸和总蒸馏 token 数进行了缩放分析，并对两阶段之间的 token 分配进行了敏感性分析。

2025年7月4日 研究领域：方法与算法、研究领域：语音与自然语言处理 会议：ICML

状态空间模型（SSM），特别是 Mamba，最近作为 Transformer 的有前景替代方案崭露头角。Mamba 在其 SSM 层（S6）中引入了输入选择性，并在其模块定义中融入了卷积和门控。虽然这些修改确实提升了 Mamba 相对于其 SSM 前代模型的性能，但 Mamba 如何利用输入选择性提供的额外功能，以及这些功能如何……目前仍在很大程度上不清楚。

2025年7月1日 研究领域：方法与算法、研究领域：语音与自然语言处理 会议：ICML

我们提出了一种蒸馏缩放定律，根据计算预算及其在学生模型和教师模型之间的分配来估计蒸馏模型的性能。我们的发现通过为学生模型和教师模型实现计算最优分配以最大化学生性能，降低了大规模蒸馏相关的风险。我们为两个关键场景提供了计算最优蒸馏方案：当教师模型已存在时，以及当……

我们在机器学习领域的研究每天都在开辟新的天地。