蒸馏后的 Mamba 模型可以达到 14.01 的困惑度，几乎与 Pythia-1B Transformer 教师模型的 13.86 困惑度相匹配。这一由 Apple Machine Learning Research 发布的研究结果，展示了一条将预训练 Transformer 的庞大知识库迁移到状态空间模型 (SSMs) 的可行路径，SSMs 在生成过程中具有更优越的效率。
,
弥合架构差距
像 Mamba 这样的状态空间模型因在生成过程中提供更低的内存消耗和更高的吞吐量，作为 Transformer 的替代方案而受到青睐。然而，机器学习社区已经投入大量资源来训练和完善 Transformer 架构，从而建立了一个庞大的预训练知识库。
采用 SSMs 的一个重大障碍是“朴素”蒸馏——即直接从 Transformer 转移到 Mamba 模型——往往无法保留教师模型的性能。以往解决此问题的尝试通常需要结合 Attention 和 SSM 块的混合模型，而不是依赖纯粹的 Mamba 架构。
— Attention to Mamba: A Recipe for Cross-Architecture Distillation
> “总体而言，蒸馏后的 Mamba 模型保留了原始 Pythia-1B Transformer 的大部分性能，达到了 14.01 的困惑度，接近教师模型的 13.86。”
两阶段蒸馏配方
为了解决这个问题，研究人员提出了一种原则性的两阶段蒸馏方法。在第一阶段，知识从传统的 Transformer 蒸馏到线性化版本的 Attention 中。这一步骤利用了“核技巧”的变体，将注意力机制简化为与 SSM 结构更兼容的形式。
在第二阶段，这个线性化 Attention 模型随后被蒸馏到包含任何 Attention 块的适配 Mamba 模型中。通过使用这个中间步骤，研究人员能够为 Mamba 提供“原则性的初始化”，使其能够恢复在直接蒸馏过程中损失的性能。
性能与局限性
HN Digest 根据 machinelearning.apple.com 报告的数据绘制的图表；数值按来源所示。
结果显示，蒸馏后的 Mamba 模型保留了 Pythia-1B Transformer 的大部分性能。虽然困惑度得分非常接近，但研究人员指出，在高度依赖上下文召回的特定任务中会出现一些性能差距。该研究在 1B 规模和 10B Token 上进行了彻底的消融实验，以测试不同的序列混合架构。他们还对模型大小进行了缩放分析，并对两个蒸馏阶段之间的 Token 分配敏感性进行了分析，以优化整个过程。