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

Together AI 将 Llama 3 蒸馏到 Mamba，性能不减甚至更优，推理速度最高可提升 1.6 倍。该项目由提出 Mamba 架构的 Tri Dao 参与。在蒸馏过程中，先进行从 Transformer 到线性 RNN 的初始化，利用 Transformer 注意力机制与 RNN 计算的相似性建立联系并复制参数，后采用三阶段蒸馏流程提升 Mamba 性能，包括基于伪标签的蒸馏、在指令数据集上监督微调及用人类反馈数据基于奖励模型优化。

在 8 块 80G A100 GPU 上，整个蒸馏过程不到五天。之后提出推测解码算法加速推理过程，该算法使用轻量级 Draft 模型预测多个 token，再用验证模型验证。测试结果表明混合模型在聊天对话任务上与 Llama – 3 相当甚至更优，不同混合比例中 1:1 比例混合模型表现最佳，在零样本通用 NLP 任务评测中平均成绩优于同等规模 RNN 模型，在少样本榜单上与最好开源 RNN 模型相当并在部分任务上超过 Instruct 模型。此外，对纯 Mamba 模型和蒸馏的混合模型测试，推理速度分别提升 1.7 – 2.6 倍和 1.6 倍左右。