从 Transformer 到 SSM：将二次知识蒸馏至次二次模型
机器学习
2025-02-11 v2 人工智能
摘要
Transformer 架构已成为语言建模等领域的主导范式，但由于其二次时间的自注意力机制，在许多推理场景中表现不佳。最近提出的次二次架构（如 Mamba）已显示出前景，但其预训练所消耗的计算资源远少于最强的 Transformer 模型。在本工作中，我们提出了一种方法，能够将预训练的 Transformer 架构蒸馏为替代架构，如状态空间模型（SSM）。我们方法的核心思想是，可以将 Transformer 和 SSM 都视为在 token 序列上应用不同形式的混合矩阵。因此，我们可以通过在 SSM 中匹配不同程度粒度的信息来逐步蒸馏 Transformer 架构：首先匹配混合矩阵本身，然后匹配每个块中的隐藏单元，最后匹配端到端预测。我们的方法称为 MOHAWK，能够仅使用 3B 个 token 将基于 Phi-1.5 架构的 Mamba-2 变体（Phi-Mamba）进行蒸馏，并使用 5B 个 token 蒸馏混合版本（Hybrid Phi-Mamba）。尽管使用的训练数据不到从头训练模型通常所需数据的 1%，Phi-Mamba 的性能仍显著优于所有过去的开源非 Transformer 模型。MOHAWK 使 SSM 等模型能够利用在训练 Transformer 架构上投入的计算资源，为构建此类模型开辟了新途径。
引用
@article{arxiv.2408.10189,
title = {Transformers to SSMs: Distilling Quadratic Knowledge to Subquadratic Models},
author = {Aviv Bick and Kevin Y. Li and Eric P. Xing and J. Zico Kolter and Albert Gu},
journal= {arXiv preprint arXiv:2408.10189},
year = {2025}
}