🍀 Fast Inference from Transformers via Speculative Decoding

大型自回归模型，尤其是大型 Transformer，相较于规模较小的模型具有更强的能力。然而，这些大模型执行单个解码步骤所需的时间显著高于小模型；更为不利的是，这些步骤以串行方式执行——解码 K 个 token 需要串行运行模型 K 次。鉴于大型自回归模型，特别是大型 Transformer 的重要性，已有研究提出多种加速其推理的方法

• 一类方法则源于以下观察：并非所有推理步骤都具有相同的难度——某些步骤需要规模很大的模型，而另一些步骤则可以由效率更高的模型较好地近似。这些自适应计算方法旨在为较简单的推理步骤使用更少的计算资源
• 尽管其中许多方案已被证明在实践中非常有效，但它们通常需要修改模型架构、改变训练流程并重新训练模型，而且无法保持输出完全一致

上述关于推理步骤存在难易之分的关键观察，也是本研究的重要动机。作者进一步观察到，大模型推理的瓶颈往往不在于算术运算，而在于内存带宽和通信，因此仍可能存在可用的额外计算资源

• 基于此，提出将提高并发度作为自适应调整计算量的互补方法。具体而言，能够在不改变模型架构、不改变训练流程或要求重新训练模型、也不改变模型输出分布的前提下，加速模型推理。这一目标通过投机执行得以实现
• 投机执行是一种常见于处理器的优化技术：在执行某项任务的同时，并行验证该任务是否确有必要，从而提高并发度。分支预测是投机执行的一个典型例子。要使投机执行有效，就需要一种高效机制，用于提出很可能实际需要执行的任务
• 在本文中，将投机执行推广至随机情形，即某项任务以一定概率需要执行的情形。将这一思想应用于 Transformer 等自回归模型的解码时，从效率更高的近似模型中采样生成序列，并将其作为速度较慢的目标模型的投机前缀
• 借助一种称为投机采样的新型采样方法，最大化这些投机任务被接受的概率，同时保证系统输出与单独使用目标模型时的输出具有相同的分布

本方法易于应用于实际生产环境，无需训练新模型，也不会改变输出分布。在内存带宽构成瓶颈且仍有可用计算资源的常见场景中，它可能是加速 Transformer 等自回归模型采样的一种合适方案。本文的主要贡献为

• 将投机执行推广至随机情形，并提出一种新的采样方法，称为投机采样
• 提出一种称为投机解码的解码机制，能够在不改变模型架构、训练方式和输出分布的前提下，加速自回归模型的解码

🌿 Speculative Decoding

设为希望加速其推理的目标模型，表示该模型在给定前缀时输出的概率分布。设为针对同一任务、效率更高的近似模型，并用 表示该模型在给定前缀 时输出的概率分布。核心思想是

• 利用效率更高的模型生成个候选 token
• 利用目标模型并行计算各候选位置的概率分布，并与候选在 下的对应概率比较，依次接受能够保持输出分布不变的候选
• 若遇到拒绝，从调整后的分布中采样一个 token 替换首个被拒绝的候选，并丢弃后续候选；若所有候选均被接受，则从目标模型对应的分布中额外采样一个 token

通过这种方式，目标模型每次并行运行都至少产生一个新 token。因此，即使在最坏情况下，目标模型的串行运行次数也不会超过简单自回归方法所需的次数。同时，每次运行也可能产生多个新 token，最多为个，具体取决于对的近似程度

投机采样：为获得服从的样本，即，改为先从中采样，即

• 如果，则保留该样本；如果，则以的概率拒绝该样本，并改为从下面经过调整的分布中重新采样

• 对于任意概率分布 和 ，通过上述方式采样得到的 确实满足 ​
• 给定将应用于条件前缀所得到的分布，可以采样一个 token 。随后，将应用于 ，以计算分布；与此同时，投机地将应用于，并行计算下一个 token 的概率分布
• 待两项计算均完成后，按照上述规则进行处理：如果被拒绝，则丢弃针对 的计算结果，并从调整后的分布中重新采样；如果 被接受，则保留这两个 token。将这一思想推广至多个候选，即可一次采样 1 至 个 token

🍀 Accelerating Large Language Model Decoding with Speculative Sampling

将 Transformer 模型的参数规模扩展至 500B 以上，已在自然语言处理、计算机视觉和强化学习等诸多任务上带来显著的性能提升

• 然而，在这一规模下，Transformer 解码仍然是一个开销高昂且效率低下的过程。Transformer 的采样过程通常受限于内存带宽
• 因此，在给定硬件配置下，从一阶近似来看，Transformer 模型生成单个 token 所需的时间，与模型参数及 Transformer 记忆状态的存储规模成正比。语言模型的规模也使得服务部署必须采用模型并行，从而引入通信开销，并使资源需求成倍增加。由于每个新 token 都依赖于此前的内容，因此，采样生成一个新序列需要多次调用 Transformer

本文提出一种面向延迟敏感应用的 Transformer 采样加速算法，称为投机采样。该算法通过以下步骤实现

• 生成长度为的短草稿序列。这既可以通过并行模型完成，也可以通过调用速度更快的自回归模型次来实现。将该模型称为草稿模型​
• 使用希望从中采样的、规模更大且能力更强的模型，对草稿序列进行评分。将该模型称为目标模型​
• 采用一种改进的拒绝采样方案，从左到右接受个草稿 token 中的一部分，并在此过程中恢复目标模型的概率分布

直观而言，在许多序列中，下一个 token 可能是显而易见的。因此，如果草稿模型与目标模型在某个 token 或 token 子序列上的概率分布高度一致，上述机制便允许每次调用目标模型时生成多个 token

🍄 Speculative Sampling

对于投机采样，首先观察到：并行计算一段由个 token 构成的短续写序列的 logits，其延迟与采样生成单个 token 的延迟十分接近。重点关注采用 Megatron 方式进行分片的大型 Transformer 模型。对于这些模型，采样时间主要由以下三个部分构成

• 线性层在小批量下通常受内存带宽限制。在批量大小较小时，每个线性层仅处理少量嵌入向量。这使得前馈层、查询、键、值的计算，以及最终注意力投影中的稠密矩阵乘法均受限于内存带宽。当较小时，这些运算仍然受限于内存带宽，因此所需时间相近
• 注意力计算的内存访问主要来自已有 KV 缓存。 注意力机制同样受限于内存带宽。在采样过程中，我们缓存序列中此前所有 token 的键和值，以避免重复计算。这些 KV 缓存规模较大，占据注意力机制内存带宽消耗的主要部分。然而，由于 KV 缓存的规模并不随 的增大而变化，因此这一部分的开销几乎没有差异，甚至完全相同
• 少量 token 的 all-reduce 通常受通信延迟限制。 随着模型规模增大，其参数需要分布在多个加速器上，从而引入通信开销。在 Megatron 中，这体现为每个前馈层和注意力层之后的一次全归约操作。由于只需传输少量 token 的激活值，因此，无论是采样还是评分，在较小时，该操作通常都受限于延迟，而非吞吐量

需要一种方法，能够利用草稿模型生成的样本，以及两个模型针对这些 token 输出的 logits，恢复目标模型的概率分布。为此，针对草稿 token 提出以下拒绝采样方案。给定 token 序列，以及由生成的个草稿 token ，以下述概率接受

• 其中，和分别表示：在给定当前上下文的条件下，目标模型和草稿模型赋予 token 的概率
• 如果该 token 被接受，则令，并对重复上述过程，直至某个 token 被拒绝，或所有 token 均被接受。如果被拒绝，则从以下分布中重新采样

通过依次应用上述过程，可以在硬件数值精度允许的范围内，为接受的 token 恢复目标模型的分布

• 每轮草稿生成—接受循环始终至少生成一个 token ：即使第一个 token 被拒绝，也会通过重新采样获得一个有效 token
• 由于草稿中的最后一个 token 能够提供下一个 token 的 logits，因此，如果所有草稿 token 均被接受，就可以利用这些 logits 按常规方式再采样一个 token 。这样，每轮循环最多可以生成个 token ，而朴素实现最多只能返回个 token
• 对于核采样、top-k 采样和温度调整等标准采样方法，可以先对概率进行相应变换，再应用上述拒绝采样方案

🌿 MEDUSA: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads

大语言模型的自回归解码通常受限于内存带宽：每次前向传播需要从高带宽内存读取模型参数，却仅生成一个 token，难以充分利用加速器的算术计算能力

• 为解决这一问题，一种加速大语言模型推理的思路是，提高解码过程的算术强度，同时减少解码步骤数。基于这一思路，研究者提出投机解码。该方法使用较小的草稿模型生成 token 序列，再由原始的大模型对其进行修正，从而得到可接受的续写
• 然而，获取合适的草稿模型仍然具有挑战性，而将草稿模型集成到分布式系统中则更加困难

本文不再使用独立的草稿模型顺序生成候选输出，而是重新审视并改进在主干模型之上增加多个解码头以加速推理的思路。具体而言，提出 MEDUSA，通过增加额外的解码头，并行预测多个 token ，从而提高大语言模型的推理效率

• 这些解码头采用参数高效的方式进行微调，并且可以添加到任意现有模型之上。由于无需独立的草稿模型，MEDUSA 可以方便地集成到现有的大语言模型系统中，包括分布式环境中的系统，从而提供便捷的使用体验
• 现有方法在每个解码步骤中只生成一条候选续写，导致计算资源利用不充分。为此，提出利用 MEDUSA 头生成多条候选续写，并通过对注意力掩码进行简单调整，对这些候选进行并行验证
• 可以复用投机采样中的拒绝采样规则，使输出分布与原始模型一致。为进一步提高接受率，MEDUSA 引入典型性接受方案，根据目标分布的概率与熵设定阈值，以相近的生成质量换取更高的解码速度；这种方案不保证严格保持原始输出分布
• 实验结果表明，所提出的典型接受方案能够在保持相近生成质量的同时，进一步提高解码速度

为使大语言模型具备能够进行预测的 MEDUSA 头，针对不同应用场景提出两种微调流程

• 对于计算资源有限，或希望保持主干模型参数不变的场景，可以采用 MEDUSA-1。该方法冻结主干模型，仅训练新增预测头，并可结合类似 QLoRA 的量化技术降低显存需求；输出分布是否严格保持一致，仍取决于推理时采用的接受规则
• MEDUSA-1 尚未充分发挥主干模型潜力。可以进一步微调主干模型，提高 MEDUSA 头的预测准确率，从而直接获得更高的加速比。提出 MEDUSA-2，适用于计算资源充足，或直接对基础模型进行监督微调的场景。MEDUSA-2 的关键在于一套训练方案，使 MEDUSA 头与主干模型能够在不损害模型下一 token 预测能力和输出质量的前提下进行联合训练。根据模型的训练方案及数据集的可用性，提出不同的训练数据获取策略

🌊 Key Components

MEDUSA 遵循与投机解码相同的框架，每个解码步骤主要包含三个子步骤：（1）生成候选；（2）处理候选；（3）接受候选

• 在 MEDUSA 中，步骤（1）由 MEDUSA 头完成，步骤（2）通过树形注意力实现。由于 MEDUSA 头位于原始模型之上，步骤（2）中计算的 logits 可以用于下一解码步骤的候选生成。最后的步骤（3）既可以通过拒绝采样实现，也可以采用典型性接受方案
• 在投机解码中，后续 token 由辅助草稿模型进行预测。该草稿模型必须足够小，同时又要具备足够的预测能力，以生成能够被原始模型接受的续写。满足这些要求并不容易，现有方法通常采用单独预训练一个较小模型的方式。这个预训练过程需要大量额外的计算资源。此外，单独预训练还可能导致草稿模型与原始模型之间出现分布偏移，使得生成的续写不符合原始模型的偏好

MEDUSA 头是在原始模型最后一层隐藏状态之上附加的额外解码头。给定原始模型在位置上的最后一层隐藏状态，在其上增加个解码头。第个头用于预测位置上的 token ，而原始语言模型输出头用于预测位置上的 token

• 将第个头的预测记为，它表示词表上的概率分布；原始模型的预测则记为。每个头均采用一个带残差连接的单层前馈网络

• 将初始化为与原始语言模型输出头相同的权重，并将初始化为零。这样，MEDUSA 头的初始预测便与原始模型保持一致
• 与独立的草稿模型不同，MEDUSA 头依托原始主干模型进行训练；主干模型既可以在训练过程中保持冻结，即 MEDUSA-1，也可以与新增预测头共同训练，即 MEDUSA-2。该方法利用强大基础模型已经学习到的表示，使得即使在单张 GPU 上也能够完成大模型的微调。此外，它使 MEDUSA 头的分布与原始模型的分布对齐，从而缓解分布偏移问题

通过 MEDUSA 头，可以获得后续个 token 的概率预测，并据此构造长度为的候选续写。尽管投机解码研究提出仅采样一条续写作为候选，但在解码过程中使用多个候选，可以提高单个解码步骤内的期望接受长度。然而，增加候选数量也会提高计算需求。为在两者之间取得平衡，采用树结构的注意力机制，对多个候选进行并行处理

• 这种注意力机制不同于传统的因果注意力：在该框架中，只有属于同一条续写路径的 token 才被视为历史上下文。借助 MEDUSA 头所生成候选的结构，采用自顶向下的方式构建候选树。对于第个预测头，其概率最高的个预测结果作为构造候选的基础。对各头选出的这个预测结果求笛卡尔积，即可得到候选续写
• 在这棵树中，每个 token 仅关注自身、祖先节点及已验证前缀，不读取其他分支。配合按树深度设置的位置索引，可以在不增加批量大小的情况下并行处理多个候选。除原始语言模型输出头对应的根节点外，新增节点数量为

在投机解码中，作者采用拒绝采样，生成与原始模型分布一致的多样化输出。然而，后续实现表明，随着采样温度升高，这种采样策略的效率会下降

• 若草稿模型与目标模型分别独立采样，再以 token 是否相同判断匹配，即使两者分布相同，也可能产生不一致的结果。前文基于概率比的投机采样接受规则则不同：当两者分布完全一致时，候选的接受概率为 1
• 更高的温度应当为原始模型接受草稿模型的输出提供更多机会。通常没有必要严格匹配原始模型的分布。因此，提出采用典型性接受方案来选择合理的候选，而不是使用拒绝采样。目标是选择具有典型性的候选，即这些候选由原始模型生成的概率不应低得不合理。为此，使用原始模型给出的预测概率作为自然的衡量标准，并根据预测分布设定阈值，以决定是否接受候选
• 具体而言，给定上下文，在评估候选序列时，该序列由原始语言模型输出头和 MEDUSA 头的高概率预测结果组成，考虑以下条件

• 其中，表示熵函数，和分别为硬阈值和熵相关阈值参数。基于两项观察：概率相对较高的 token 具有合理性；当分布的熵较高时，多种续写都可能被认为是合理的。在解码过程中，使用这一判据评估每个候选；如果候选的某个前缀满足条件，就接受该前缀。为保证每个步骤至少生成一个 token ，对第一个 token 采用贪心解码，并无条件接受它；对于后续 token ，则采用典型性接受。当前步骤的最终预测结果，由所有候选中被接受的最长前缀确定
• 首先，当温度设为 0 时，它退化为贪心解码，因为只有概率最高的 token 具有非零概率。当温度大于 0 时，在适当设置和的情况下，贪心解码的结果始终会被接受，因为这些 token 具有最高概率，从而获得最大的加速效果。同样，在一般情况下，提高温度会相应地增加被接受序列的长度

📡 Training Strategies

在冻结主干模型的情况下，使用 MEDUSA 头的预测结果与真实标签之间的交叉熵损失进行训练。具体而言，给定位置上的真实 token ，第个预测头的损失为，其中，表示第个预测头为 token 分配的概率

• 观察到，随着增大，也会增大。这是合理的，因为预测位置越远，第个预测头所面临的不确定性就越大。因此，可以为引入权重，以平衡不同预测头的损失。MEDUSA 的总损失为

• 在实际应用中，将设置为某个常数的次幂。由于主干模型仅用于提供隐藏状态，因此可以使用量化后的主干模型来降低显存占用