The Mamba in the Llama: Distilling and Accelerating Hybrid Models

本文指出，像Mamba这样的线性RNN架构在语言建模方面可以与Transformer模型竞争，并具有优越的部署特性。鉴于大规模Transformer模型的训练重点，我们考虑将这些预训练模型转换为部署模型的挑战。我们证明了通过在学术GPU资源中重用注意力层的线性投影权重，可以将大型Transformer压缩为线性RNN。由此产生的混合模型，其中包含四分之一的注意力层，在聊天基准测试中实现了与原始Transformer相当的性能，并且在聊天基准测试和一般基准测试中优于从头开始训练的开源混合Mamba模型，这些模型使用了数万亿个令牌。此外，我们引入了一种硬件感知的推测解码算法，加速了Mamba和混合模型的推理速度。总的来说，我们展示了如何在有限的计算资源下，去除许多原始的注意力层，并更有效地生成结果模型。我们从Llama3-8B-Instruct中压缩出的表现最佳的模型，在AlpacaEval 2中实现了29.61的长度控制胜率，超过了GPT-4，在MT-Bench上实现了7.35的胜率，超过了最佳的指令调整线性RNN模型。

本论文试图解决如何将大规模Transformer模型转化为可部署的线性RNN模型的问题，以及如何在有限的计算资源下生成更高效的模型。

本论文的关键思路是通过重用注意力层的线性投影权重，将大规模Transformer模型蒸馏为线性RNN模型，并将其与硬件感知的推测解码算法相结合，从而加速生成速度。

本论文实现了将大规模Transformer模型蒸馏为线性RNN模型的方法，并通过实验表明，这种方法在对话基准测试和通用基准测试中表现出与原始Transformer模型相当的性能。此外，论文还介绍了一种硬件感知的推测解码算法，可以加速生成速度。最终，论文的模型在AlpacaEval 2上取得了29.61的胜率，超过了最佳线性RNN模型。

最近的相关研究包括使用Transformer模型进行自然语言处理的研究，以及使用线性RNN模型进行语言建模的研究。