编辑：杜伟、小舟
去年，在加速大措辞模型推理层面，我们迎来了一个比推测解码更高效的办理方案 —— 普林斯顿、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.pdfGitHub 地址：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（利用监督微调和对齐演习）。
结果表明，MEDUSA 在不影响天生质量的情形下，可以在不同的 promt 类型中实现 2.3 至 3.6 的推理加速。如下动图为 Vicuna-7b 上有无 Medusa-1 时推理速率比较。
论文共同一作 Tianle Cai 表示，自 Medusa 项目推出以来，它在 TensorRT、TGI 以及浩瀚开源项目和公司中得到采取。在新的技能论文中，我们推出了用于全模型调优的 Medusa-2 方案、用于将 Medusa 集成到任何微调 LLM 的自蒸馏以及其他更多加速技能。
对付这项研究，Lepton AI 创始人贾扬清表示，Medusa 可能是他们见过的最优雅的加速推理解决方案之一，能够与 int8/fp8、编译等互补，在实践中实现 2 倍性能增益。
并且，他们已将 Medusa 与很多现有优化方法、稠浊加速方案进行集成，结果在合理的并发下，加速保持正值，并在 A100 和 H100 等卡中尤其有效。此外，他们还已经为 Llama 模型演习了通用 Medusa 头。
方法概览
MEDUSA 遵照推测解码框架，个中每个解码步骤紧张由三个子步骤组成：(1) 天生候选者，(2) 处理候选者， (3) 接管候选者。对付 MEDUSA，(1) 是通过 MEDUSA 头（head）实现的，(2) 是通过树把稳力（tree attention）实现的，并且由于 MEDUSA 头位于原始主干模型之上，因此 (2) 中计算的 logits 可以用于子步骤 (1) 的下一个解码步骤。末了一步 (3) 可以通过谢绝采样（rejection sampling）或范例接管（typical acceptance）来实现。MEDUSA 的整体流程如下图 1 所示。
关键组件
MEDUSA 的关键组件紧张包括 MEDUSA 头和树把稳力。
首先，MEDUSA 头与原始主干模型一起进行演习。个中，原始主干模型可以在演习期间保持冻结状态 (MEDUSA-1) 或一起演习 (MEDUSA-2)。这种方法乃至可以在单个 GPU 上微调大模型，利用强大的根本模型学得的表征。
此外，MEDUSA 头的分布确保与原始模型的分布同等，从而缓解了分布偏移问题，并且 MEDUSA 不会增加做事系统设计的繁芜性，对分布式设置很友好。
由于候选者增加会提高打算需求，该研究采取树状构造的把稳力机制来同时处理多个候选者。这种把稳力机制不同于传统的因果把稳力范式。在其框架内，只有来自同一 continuation 的 token 才被视为历史数据。受图神经网络领域提出的将图构造嵌入把稳力的启示，研究团队还将树构造合并到把稳力掩码中，如下图 2 所示。
演习策略
冻结主干模型来演习 MEDUSA 头的方法很大略，并且须要的打算资源很少，但是将主干网络与 MEDUSA 头结合演习可以显著提高 MEDUSA 头的准确性。因此，根据打算资源和用例的详细哀求，研究团队为 MEDUSA 头提出了两个级别的演习策略，即 MEDUSA-1：冻结主干网络，MEDUSA-2：联合演习。
末了，该研究提出了 MEDUSA 的两个扩展，包括自蒸馏（self-distillation）和范例接管（typical acceptance），分别用于处理 MEDUSA 没有可用演习数据的情形和提高解码过程的效率。
实验
为了证明 MEDUSA 在不同设置下的有效性，该研究进行了两组实验：首先，在 Vicuna-7B/13B 模型上评估 MEDUSA，以展示 MEDUSA-1 和 MEDUSA-2 的性能；其次，在 Vicuna-33B 和 Zephyr-7B 模型上评估 MEDUSA，以研究自蒸馏的有效性，由于 Vicuna-33B 模型的演习数据集不公开，而 Zephyr-7B 模型利用 RLHF 进行演习。
用例研究 1：在 Vicuna-7B/13B 模型上评估 MEDUSA
在 Vicuna-7B/13B 模型上评估 MEDUSA-1、MEDUSA-2 的结果如下图 4 所示。
用例研究 2：在 Vicuna-33B 和 Zephyr-7B 利用自蒸馏演习
研究者关注了须要自蒸馏的情形，利用 Vicuna-33B 和 Zephyr-7B 作为示例。他们首先利用一些种子 prompt 来天生数据集，然后将 ShareGPT 和 UltraChat 作为种子数据集，并为以上两个示例网络了包含大约 100k 样本的数据集。
下表 1 展示了不同 MEDUSA-2 模型在 MT-Bench 基准下的加速比、开销和质量。
下图 5 为利用 MEDUSA-2 时不同模型的加速情形。
溶解实验
下图 6a 比较了随机采样密集树设置（蓝点）和优化稀疏树设置（红星）的加速率。6b 比较了密集和稀疏树设置的速率。
下图 7 展示了不同采样设置下，模型性能的比较剖析。
两阶段微调的有效性。研究者针对 Vicuna-7B 模型，评估了两种微调策略下的性能差异。