arXiv:2408.10189v2 宣布类型: replace-cross 摘要：Transformer架构已成为语言建模等领域的主要范式，但由于它们在自注意力上的二次时间复杂度，在许多推理设置中表现出色。最近提出的亚二次架构，如Mamba，显示出一定的前景，但其预训练使用的计算资源远少于最强的Transformer模型。在本工作中，我们提出了一种方法，能够将预训练的Transformer架构蒸馏为替代架构，如状态空间模型（SSMs）。我们方法的关键思想是，我们可以将Transformer和SSMs视作在令牌序列上应用不同形式的混合矩阵。因此，我们可以逐步将Transformer架构蒸馏为SSMs的不同粒度：首先匹配混合矩阵本身，然后是每个块中的隐藏单元，最后是端到端的预测。我们的方法称为MOHAWK，能够使用仅30亿个令牌蒸馏基于Phi-1.5架构的Mamba-2变体（Phi-Mamba），以及使用50亿个令牌的混合版本（Hybrid Phi-Mamba）。尽管仅使用了通常用于从头训练模型的训练数据不到1%的数据，Phi-Mamba的表现仍然远远优于所有之前的开源非Transformer模型。MOHAWK允许像SSMs这样的模型利用在训练基于Transformer的架构中投资的计算资源，突显出一种新的构建此类模型的途径。