2024 年 8 月 22 日，Kevin Li 在社交媒体上发布了一条消息：通过一种名为 MOHAWK 的新蒸馏方法，可以将预训练的 Transformer 模型蒸馏到 Mamba 或混合架构中，从而在仅使用 30 亿个 token 的情况下，蒸馏出一个 1.5B 的模型。这一消息迅速引起了广泛关注和讨论。
MOHAWK 方法的核心在于将 Transformer 模型中的知识迁移到 Mamba 等替代架构中，使得模型在保持较低计算成本的同时，性能提升。通过这种方法，研究团队将 Phi-1.5 架构的 Mamba-2 变体（Phi-Mamba）进行蒸馏，并使用 5B 个 token 进行混合版本（Hybrid Phi-Mamba）的蒸馏。尽管使用的训练数据不到从头训练模型所需数据的 1%，Phi-Mamba 的性能增强。
MOHAWK 方法的提出是为了应对 Transformer 模型在处理长序列时计算量巨大的问题。Transformer 模型依赖于二次自注意力机制，这使得其在处理序列数据时能够捕捉序列内部的长距离依赖关系，但也导致了计算成本的增加。MOHAWK 方法通过将注意力和 SSMs 视为可以对齐的序列转换，允许跨架构蒸馏，从而降低了计算成本。
具体来说，MOHAWK 方法分为三个阶段：
- ** 矩阵对齐（Matrix Orientation）**：在这一阶段，研究人员对齐序列变换矩阵本身，使得不同架构之间的转换更加顺畅。
- ** 隐藏状态对齐（Hidden-State Alignment）**：这一阶段的目标是对齐网络每个单独层的隐藏状态表示，确保不牺牲预先学习的表示。实验发现，隐藏状态对齐可以提高后续阶段的性能。
- ** 权重转移和知识蒸馏（Weight-Transfer and Knowledge Distillation）**：在这一阶段，通过一个端到端训练过程，将权重转移到目标模型中，最终使用只有一小部分训练数据来蒸馏网络的最终输出。
实验结果显示，Phi-Mamba 模型在多个基准测试中表现良好，达到了开源非 Transformer 架构中的 SOTA（State of the Art）。例如，在语言建模任务中，Phi-Mamba 模型的性能超过了许多传统的 Transformer 模型，且推理速度提升。
Mamba 架构由 FlashAttention 作者 Tri Dao 和 CMU 助理教授、Cartesia AI 联合创始人及首席科学家 Albert Gu 在去年年底提出。Mamba 是第一个实现匹配 Transformer 性能的线性时间序列模型，也是状态空间模型（SSM，State Space Model）。Mamba 建立在更现代的适用于深度学习的结构化 SSM（S4, Structured SSM）基础上，与经典架构 RNN 有相似之处。
Mamba 架构的主要创新点包括对输入信息的选择性处理、硬件感知的算法和更简单的架构。与 Transformer 架构相比，Mamba 在处理长序列时具有优势。Transformer 模型依赖于二次自注意力机制，这使得其在处理长序列时计算成本增加。而 Mamba 架构通过线性时间复杂度的设计，能够处理长序列数据。
在多个领域的应用中，Mamba 架构表现良好。例如，在语言建模任务中，Mamba 模型不仅在预训练质量上达到了 Transformer 的水平，而且在下游评估中表现良好。研究者发现，Mamba 模型在处理长文本时的推理速度是 Transformer 模型的五倍。
此外，Mamba 架构在音频和基因组学等领域也表现良好。在音频波形和 DNA 序列建模方面，Mamba 在预训练质量和下游指标方面都优于 SaShiMi、Hyena、Transformer 等先前的 SOTA 模型。例如，在具有挑战性的语音生成数据集上，Mamba 模型将 FID（Fréchet Inception Distance）降低了一半以上。