LIama 3+Mamba强强联手！蒸馏到线性RNN，推理速度提升1.6倍

文章预览

克雷西 发自 凹非寺 量子位 | 公众号 QbitAI 把Llama 3蒸馏到Mamba，推理速度最高可提升1.6倍！ 而且性能不减，甚至表现比原始模型还要优异。 这是来自Together AI的新作，通过蒸馏将Transformer和Mamba模型结合到了一起，同时还为混合模型涉及了推理加速算法 提出Mamba架构的大神、FlashAttention作者Tri Dao，也参与了这一项目。 Together AI创始人兼CEO表示，Transformer和Mamba的混合，是未来大模型的一大发展方向。 将Transformer蒸馏进Mamba 在蒸馏正式开始之前，需要先进行从Transformer到线性RNN的初始化。 作者观察到，Transformer的注意力机制与RNN的计算之间存在一定的相似性。 因此可以将Transformer的注意力线性化，从而建立二者的联系。 利用这种对应关系，可以将预训练的Transformer模型的参数复制到Mamba模型中。 在完成参数初始化后，作者采用了一个三阶段的蒸馏流程进 ………………………………