大模型推理速度飙升3.6倍，「美杜莎」论文来了，贾扬清：最优雅加速推理方案之一

原文来源：机器之心

图片来源：由无界 AI生成

去年，在加速大语言模型推理层面，我们迎来了一个比推测解码更高效的解决方案 —— 普林斯顿、UIUC 等机构提出的 Medusa。如今，关于 Medusa 终于有了完整技术论文，还提供了新的版本。

如你我所知，在大型语言模型（LLM）的运行逻辑中，随着规模大小的增加，语言生成的质量会随着提高。不过，这也导致了推理延迟的增加，从而对实际应用构成了重大挑战。

从系统角度来看，LLM 推理主要受内存限制，主要延迟瓶颈源于加速器的内存带宽而非算术计算。这一瓶颈是自回归解码的顺序性所固有的，其中每次前向传递都需要将完整的模型参数从高带宽内存传输到加速器缓存。该过程仅生成了单个的 token，没有充分利用现代加速器的算术计算潜力，导致了效率低下。

为了解决这一问题，加速 LLM 推理的方法被提出，既可以增加解码过程的算术强度（FLOPs 与总数据移动的比率），也能减少解码步骤数量。这类方法以推测解码（speculative decoding）为代表，使用较小的草稿（draft） 模型在每一步生成 token 序列，然后通过较大的原始模型进行细化以获得可接受的延续。不过获得合适的草稿模型仍然具有挑战性，并且将草稿模型集成到分布式系统中更加困难。

在本文中，来自普林斯顿大学、Together.AI、伊利诺伊大学厄巴纳 - 香槟分校等机构的研究者没有使用单独的草稿模型来顺序生成候选输出，而是重新审视并完善了在主干模型之上使用多个解码头加速推理的概念。他们发现，如果该技术得到有效应用，可以克服推测解码的挑战，从而无缝地集成到现有 LLM 系统中。

具体来讲， 研究者提出了 MEDUSA，一种通过集成额外解码头（能够同时预测多个 tokens）来增强 LLM 推理的方法。这些头以参数高效的方式进行微调，并可以添加到任何现有模型中。至此，不需要任何新模型，MEDUSA 就可以轻松地集成地当前的 LLM 系统中（包括分布式环境），以确保友好用户体验。

值得关注的是，该论文作者之一 Tri Dao 是近来非常火爆的 Transformer 替代架构 Mamba 的两位作者之一。他是 Together.AI 首席科学家，并即将成为普林斯顿大学计算机科学助理教授。

论文地址：https://arxiv.org/pdf/2401.10774.pdf

GitHub 地址：https://arxiv.org/pdf/2401.10774.pdf

在具体实现中，研究者通过两个关键见解进一步增强了 MEDUSA。首先，当前在每个解码步骤生成单个候选延续的方法导致了可接受长度受限和计算资源的低效使用。为了解决这个问题，他们建议使用 MEDUSA 头来生成多个候选延续，并通过对注意力掩码的简单调整来进行验证。其次可以使用类似于推测解码中的拒绝采样方案来生成与原始模型具有相同分布的响应，但对于很多 LLM 应用来说通常不必要。

因此，研究者考虑或许可以引入一种典型的可接受方案，即从 MEDUSA 输出中选择合理的候选者。他们使用温度作为阈值来管理原始模型预测的偏差，为拒绝采样提供了一种有效的替代方案。这种方法有效地解决了拒绝采样的局限性，比如在较高温度下速度降低。

此外，为了给 LLM 配备预测性的 MEDUSA 头，研究者提出了两种针对不同场景量身定制的微调程序。对于计算资源有限或者目标是将 MEDUSA 纳入现有模型而不影响其性能的情况，他们建议使用 MEDUSA-1。该方法需要的内存最少，并且可以使用类似于 QLoRA 中的量化技术来进一步优化，而不会因固定主干模型影响生成质量。

不过，对于 MEDUSA-1，主干模型的全部潜力无法得到充分利用。因此可以进一步进行微调，以提高 MEDUSA 头的预测精度，并直接带来更大加速。因此研究者提出了 MEDUSA - 2，它适用于计算资源充足或从基础模型进行直接监督微调的场景。MEDUSA-2 的关键是一个训练协议，它能够对 MEDUSA 头和主干模型进行联合训练，而不会影响模型下一个 token 的预测能力和输出质量。

在实验部分，研究者主要关注批大小为 1 的场景，这代表了 LLM 本地托管以供个人使用的用例。他们在不同大小和训练设置下测试了 MEDUSA，包括 Vicuna-7B 和 13B（使用公共数据集训练）、Vicuna -33B（使用私有数据集训练）、Zephyr-7B（使用监督微调和对齐训练）。