Attention to Mamba：跨架构蒸馏新配方
Attention to Mamba: A Recipe for Cross-Architecture Distillation
Mamba 等状态空间模型（SSMs）凭借低内存和高吞吐优势，正成为Transformer的热门替代方案。但如何有效利用现有的Transformer知识来训练Mamba，一直是个难题。以往直接蒸馏往往效果不佳，而混合架构又增加了复杂度。我们提出了一种原则性的两阶段蒸馏方案：首先将传统Transformer的知识蒸馏到线性化的Attention中，利用kernel trick进行适配；随后再将线性化版本蒸馏到纯Mamba模型中，完全不依赖Attention模块。实验表明，经过蒸馏的Mamba模型在Pythia-1B基准上表现优异，困惑度从13.86降至14.01，几乎保留了教师模型的性能。这一方法在下游任务中也展现出强大潜力，尽管在依赖上下文记忆的任务上仍有提升空间。
通过为Mamba配备原则性的初始化，我们可以恢复一种更优的跨架构蒸馏方案。
- buildbot
在我看来，关于所有状态空间注意力类型类似想法的更好、更全面的解读，是这篇来自 AWS 的论文：https://arxiv.org/abs/2605.08301
- kooi
很酷，我喜欢状态空间模型，因为它们“感觉”上是对线性动力系统控制与建模力量的恰当致敬。（Mamba 本质上是一串由非线性激活函数连接的时间变线性系统）。
我现在正在读其中一篇原始论文，这里有一些有趣的片段。我很好奇在数据中心中，有多少推理是使用传统 Transformer 进行的，又有多少是 SSM，因为效率提升非常显著。
> 我们提出了一类新的选择性状态空间模型，它在多个维度上改进了先前的工作，在实现与 Transformer 相当的建模能力的同时，序列长度的扩展呈线性增长。
> 我们的 Mamba 语言模型相比同等规模的 Transformer，生成吞吐量高出 5 倍，而 Mamba-3B 的质量则与两倍规模的 Transformer 相当。
接着，关于训练，有趣的是：
> 通常，该模型 [SSM] 使用卷积模式（3）进行高效的并行化训练（此时整个输入序列是预先可见的），然后切换到循环模式（2）以进行高效的自回归推理（此时输入是逐个时间步看到的）。
- vlovich123
仅针对 1B 模型。没有迹象表明这是否能扩展到更大的模型，或者从 Mamba 的角度看那会是什么样子。
> 这很大程度上延续到了下游任务中，尽管在一些更依赖上下文回忆的任务上出现了一些差距。
没错，这就是 SSM 著名的阿喀琉斯之踵——回忆能力差。