点击下方卡片，关注 「AI视界引擎」 公众号
( 添加时备注：方向+学校/公司+昵称/姓名 )
大型语言模型（LLMs）通过将任务统一为文本生成，彻底改变了自然语言处理领域。然而，由于其庞大的参数规模和自动回归性质，推理速度受到了限制。
SAM-Decoding通过引入一种新颖的基于检索的推测解码方法来解决这个问题，该方法使用后缀自动机进行高效的草稿生成。
与现有方法使用的n-gram匹配不同，SAM-Decoding在生成文本和文本语料库时找到最长的后缀匹配，实现了每个生成步骤的平均时间复杂度为O(1)。SAM-Decoding为文本语料库和输入 Prompt 分别构建了静态和动态后缀自动机，实现了快速、精确的草稿生成。
同时，SAM-Decoding设计为可以与现有方法相结合的方法，允许根据匹配长度自适应选择草稿生成策略，从而增加LLM的推理速度。
与Token Recycling相结合时，评估结果显示SAM-Decoding在Spec-Bench上的速度比现有无模型方法快2.27倍，比自动回归解码快2.49倍，超越了所有现有方法。作者的代码已在作者的仓库中发布。
1 Introduction
基于Transformer的大型语言模型已成为自然语言处理领域的基础模型。通过大型语言模型（LLM），所有自然语言处理任务都被统一为文本生成任务，这大大提高了解决自然语言处理任务的可便利性。尽管LLM在各种任务上表现良好且易于使用，但包含在LLM中的大规模参数和文本生成范式，在实际部署中带来了挑战，特别是在LLM的推理速度方面。LLM以自回归方式逐步生成 Token 。在每个生成轮次中，LLM根据已经生成的 Token 生成一个新的 Token 。此外，用于部署LLM的GPU，这个过程需要从高带宽内存（HBM）读取包含在LLM中的大量参数到片上内存，但这些参数只作用于对应新 Token 的少量参数，导致无法充分利用现代GPU的并行能力，从而限制了LLM的推理速度。
针对上述问题，一种有效的策略是增强大型语言模型（LLM）生成过程的算术强度，这涉及在减少生成轮次的同时，增加每轮生成的 Token 数量。在这个概念基础上，推测解码方法已经被引入，并引发了一系列后续研究。这些方法并不是每步生成一个 Token ，而是采用草案模型预测 Token 序列作为草案。然后，预测的序列被验证与LLM，只采用准确的 Token 。
大多数猜测性解码技术都采用了各种神经网络，例如较小的语言模型（LLMs）或嵌入层，作为草案模型，作者将其归类为基于模型的方法。虽然基于模型的猜测性解码算法有效地提高了处理速度，但它们并非没有局限性。首先，应用基于模型的猜测性解码算法的必要条件是训练草案模型，而这个过程通常需要特定的领域知识，这增加了其实际部署的复杂性。其次，这些初步模型与主要的LLMs竞争GPU内存资源。此外，使用Transformer架构作为初步模型引入了额外的键值缓存需求，进一步复杂了LLM推理系统中的GPU内存管理（Kwon等人，2023；Zheng等人，2024）。
除了基于模型的预测解码方法外，还有一种替代方法称为检索式方法。这些方法在每次生成步骤中，从之前生成的内容或文本数据库中检索潜在的后续 Token 。与基于模型的方法不同，检索式方法不需要训练草稿模型或使用GPU进行推理。它们在总结和代码重写等任务中特别有利。
在审视了当前的检索方法后，作者发现它们存在几个局限性。首先，这些方法通常依赖于单一的检索源。例如，PLD（Saxena，2023）仅关注 Prompt ，而REST（He等人，2024）局限于文本数据库。其次，所采用的检索技术往往效率低下；对n-gram匹配的依赖限制了匹配的最大长度，可能导致生成较不理想的草案。最后，尽管检索方法在摘要等任务上表现出色，但与基于模型的方法在其他领域相比通常表现较差，限制了其非法使用。
为了应对之前提到的基于检索的方法的局限性，本文引入了一种名为SAM-Decoding的创新性基于检索的推测解码技术。与传统基于检索的技术相比，该方法具有更高效、更精确的检索能力。与传统的基于n-gram匹配的方法不同，SAM-Decoding利用后缀自动机（SAM）解决生成文本在历史生成文本和文本数据库中的最长后缀匹配问题。此过程得到的草案比n-gram匹配更准确。此外，在文本生成阶段，后缀自动机有效地捕获了相邻后缀之间的关系，确保了每一步识别最长后缀匹配的平均时间复杂度保持在，这比n-gram匹配的计算效率有显著提高。同时，最长后缀匹配的长度作为生成文本与检索源之间对齐的指示器。通过利用这个指标，作者的方法可以将现有的推测解码方法作为辅助方法，并动态决定是否根据匹配内容生成草案，或者利用辅助解码方法。这种适应性使得文本生成过程的整体速度显著提高。
具体而言，SAM-Decoding首先为现有文本语料库构建一个静态后缀自动机。对于每个请求，SAM-Decoding然后根据输入 Prompt 构建一个动态后缀自动机。后缀自动机包含一系列节点，每个节点对应于检索源中的多个子串。作者在检索源中记录每个节点的最早对应位置。这些信息可以通过直接将自动机中的节点映射到生成文本以及文本语料库中的索引来促进快速生成草稿。作者在静态和动态自动机中维护生成文本的相应节点。因此，在每次生成过程中，作者可以根据节点中记录的位置从检索源中直接检索草稿，并根据匹配长度过滤草稿。生成完成后，随着新 Token 的产生，对于静态自动机，作者根据这些 Token 进行节点转移。对于动态自动机，作者首先扩展结构以容纳新 Token ，然后根据需要相应地转移节点。
广泛的评估表明，SAM-Decoding可以在各种任务上都实现与现有最先进（SOTA）方法竞争的结果。尤其是在适用于检索方法的任务中，SAM-Decoding 优于所有现有方法。在 Spec-Bench 最新的基准测试中，SAM-Decoding 与 Token Recycling（Luo等人，2024）相结合，与自回归解码相比，速度提高了2.26倍，超越了所有模型无（模型无）方法。此外，当 SAM-Decoding 与 EAGLE2（Li等人，2024）结合使用时，SOTA基于模型的 Speculative Decoding 方法，与自回归解码相比，速度提高了2.49倍，超越了所有现有基于模型和模型无方法。
2 Background
在本节中，作者将介绍与本工作相关的背景信息。
Suffix Automaton
后缀自动机是一种高效的数据结构，用于表示给定字符串的所有子串的子串索引，从而实现对所有子串的存储、处理和检索压缩信息。构建后缀自动机的时间复杂度为,其中为字符串的长度，并且可以分步构建。
一个后缀自动机包含一系列节点和两种类型的状态转移边，分别是next和link。在自动机中的一个节点表示一个状态，对应于字符串中具有相同结束位置的多个子串。同时，使用两种类型的转移边来将新的 Token 添加到当前状态（next），并获取与当前状态对应的子串的后缀（link），从而转移到新的状态。
基于两种类型的转移边，对于逐步生成的 Token 序列，作者可以在生成过程中的每个步骤上，用平均的时间复杂度找到一个匹配后缀自动机序列的最长后缀。
Speculative Decoding
LLM 通过自回归方式生成 Token 。假设模型输入为 ，则在每一代中，LLM 生成一个新的 Token 如下：
推测解码通过增加每个步骤生成的 Token 数量和减少生成步骤的总数，提高了LLM的生成速度。在每次生成步骤中，推测解码方法首先根据草案模型和输入生成一个包含多个 Token 的草案序列。草案随后由LLM验证，并接受正确的 Token 。
通过检查草案中的每个 Token 是否是LLM中前一个 Token 的预测结果，作者可以过滤出接受的 Token 。具体而言，作者找到连续 Token 的最大数量，如下所示：
在上述方法中，作者假设草案是一系列 Token 。最近关于预测解码的研究也提出了使用树状草案的方法。具体来说，这些工作在注意力模块中引入了一个树状草案，以在不改变现有框架的情况下计算树状草案。使用树状草案允许LLM同时验证多个可能的 Token 序列，从而增加了草案的接受率。
3 SAM-Decoding
在本节中，作者介绍了作者提出的SAM-Decoding方法。SAM-Decoding是一种基于检索的推测解码方法，旨在解决现有基于检索的推测方法中发现的三个关键限制：
-
使用较少的全面检索源。
-
使用较不高效的检索方法和匹配长度的限制。
-
在专业领域表现不佳。
为了克服前两个局限性，SAM-Decoding利用了后缀自动机，这显著提高了检索过程的全面性和效率，同时允许更灵活的匹配长度。接下来，作者详细介绍如何将SAM-Decoding与无模型和有模型方法相结合。通过利用后缀自动机提供的精确匹配信息，作者的方法不仅克服了第三个局限性，而且在各种任务上确保了有效的性能改进。
Suffix Automaton Construction
SAM解码通过利用后缀自动机来促进快速草稿生成。与依赖单一检索源的传统方法不同，作者的方法利用了文本语料库和生成文本本身作为检索源。为了适应这两种来源，作者构建了两种不同的后缀自动机：静态后缀自动机和动态后缀自动机。对于文本语料库，作者离线预先构建了一个静态后缀自动机。在推理阶段，仅在此静态自动机上执行状态匹配。相反，对于生成文本，作者创建了一个随文本生成逐步扩充的动态后缀自动机，同时进行状态匹配。
可以使用Blumer's算法（Blumer等，1984年）在线性时间内构建一个后缀自动机。由于后缀自动机是为单个引用字符串设计的，在处理文本语料库中的多个字符串时，作者使用特殊符号（如End-of-Sentence（EOS） Token ）将它们拼接成一个单独的长字符串。然后，作者为这个拼接字符串构建一个静态后缀自动机。
同时，作者还对后缀自动机进行了几项修改，以更好地适应草稿生成。具体来说，在每一代步骤中，生成的文本对应于自动机中的一个节点，表示最长的后缀匹配。作者需要在文本语料库或生成文本中确定匹配位置。基于这个位置，作者可以使用随后的 Token 作为潜在的草稿。为了实现这一目标，在后缀自动机中的每个节点，作者记录了该节点在参考字符串中对应的所有子串的最小结束位置。作者将这个位置定义为 min_endpos。通过使用min_endpos，作者可以直接将节点映射到参考字符串中的位置。后缀自动机的构建过程详细地描述在附录A.1中的算法3中。
### Drafting with Suffix Automaton
在构建后缀自动机之后，它可以作为生成过程每个步骤的草案模型来生成草案。具体而言，假设当前输入到LLM的 Token 为,其中预测的下一个 Token 为。作者将后缀自动机定义为，并将其相关参考文本定义为。对应于序列的后缀自动机中的状态。状态转换基于新生成的 Token 和当前状态：
随后，作者从参考文本中提取个连续的 Token ，形成一个草案，使用对应状态的节点中的min_endpos值。如果中的min_endpos值是，那么草案被定义为：
生成草稿 的长度为 。
在实际应用中，作者使用两种类型的自动机：一个静态自动机用于处理文本语料库，一个动态自动机用于生成文本。为了区分这两种自动机，作者明确记录状态的属性。这个属性表示生成文本在自动机中的最长后缀匹配长度。这有助于作者确定是否应该使用静态后缀自动机或动态后缀自动机生成草稿。作者的实验结果表明，从动态自动机生成的草稿通常优于从静态文本语料库生成的草稿。因此，作者优先使用来自动态自动机的草稿。具体地，令和分别表示静态自动机和动态自动机的匹配长度。如果，则仅使用静态自动机提供的草稿；否则，使用动态自动机提供的草稿。这里，是一个预定义的常数，大于0。
作者说明了算法1中后缀自动机的状态转移过程。使用摊销分析，作者可以证明状态转移的平均复杂度为，最坏情况的时间复杂度为，其中表示生成文本的长度。详细的证明见附录A.2。回顾现有方法，PLD使用-grams并在 Prompt 中进行蛮力匹配搜索，导致时间复杂度为。REST也使用-grams，但在使用后缀数组在文本中进行匹配搜索，导致时间复杂度为。在这里，代表预定义的最大匹配长度，分别代表 Prompt 的长度和语料库中的总文本长度。相比之下，作者提出的模型（SAM-Decoding）具有较低的时间复杂度。此外，SAM-Decoding可以在不限制匹配长度的情况下找到确切的最长后缀匹配。这意味着SAM-Decoding可以更快、更准确地生成草案。
### Update of Suffix Automaton
在生成草稿后，作者使用大型语言模型（LLM）验证草稿，并接受草稿的正确前缀。将接受的 Token 表示为 。然后，根据这些接受的 Token 更新后缀自动机的状态。对于静态后缀自动机，作者按照算法1简单地转移匹配状态：
对于动态后缀自动机，作者首先根据接受的 Token 扩展其状态，然后转移匹配状态。令 表示与生成的文本 对应的动态后缀自动机。这个过程如下：
在附录A.1中的算法3详细介绍了扩展后缀自动机的流程。
Adaptive Draft Selection
一种直接的想法是，后缀匹配的长度可以反映出自动机生成的草稿的质量。更长的匹配意味着草稿中可能会有更多可接受的 Token ，而较短的匹配则暗示草稿包含的信息较少。除了基于检索的推测解码方法外，还存在各种其他推测解码技术。因此，如果基于检索的方法无法生成令人满意的草稿，作者可以采用这些替代方法来提高草稿的有效性。
为了实现这一目标，作者同时采用了一种额外的推测解码技术，并与后缀自动机并行工作。在生成的每个迭代过程中，作者根据生成文本与自动机中的匹配长度，在自动机提供的草案和补充推测解码方法提供的草案之间进行自适应选择。自适应草案选择过程详细说明在算法2中，其中表示一个预定的常数，用于确定是否使用自动机或辅助推测解码进行草案生成。在作者的研究中，作者将两种最先进的推测解码方法作为辅助选项：模型无关的Token Recycling和基于模型的EAGLE。其中，Token Recycling为每个 Token 维护一个相邻列表，其中包含每个 Token 的可能下一个 Token 。在每次生成的每个步骤中，通过广度优先搜索使用这个相邻列表构建一个草案树，该列表会根据最新的生成的 Token 不断更新。相反，EAGLE在每个生成的周期中利用一个集成LLM输出的Transformer解码层构建一个草案树。
4 Experiments
在本节中，作者首先介绍实验设置，然后呈现实验结果，最后进行消融实验。
模型与任务作者在Vicuna-7B-v1.3 上进行了实验。作者在Spec-Bench（Xia等人，2024年）、HumanEval 、MBPP（Austin等人，2021年）和HARGID 上评估了SAM-Decoding。Spec-Bench是一个综合基准，旨在评估跨多种场景的Speculative Decoding方法。它基于六个常用的数据集：MT-Bench（Zheng等人，2023年）、WMT14 DE-EN、CNN/Daily Mail（Nallapati等人，2016年）、Natural Question（Kwiatkowski等人，2019年）、GSM8K（Cobbeet等人，2021年）和DPR（Karpukhin等人，2020年），包括六个方面：多轮对话、翻译、总结、问答、数学推理和检索增强生成。此外，MBPP、HumanEval和HARGID分别用于评估代码生成任务和上下文问答任务中解码方法的性能。
Baseline 方法 作者考虑了以下 Baseline 方法，包括基于模型的方法 EAGLE (Li等人，2024；b)，基于模型的自由方法 Token Recycling (Luo等人，2024) 和 Lookahead Decoding (Fu等人，2024)，以及基于检索的方法 PLD (Saxena，2023) 和 REST (He等人，2024)。
指标。 作者从以下几个方面评估了猜测解码方法（Li等人，2024）
加速比：与自回归生成方法相比，猜测解码方法的墙时加速比。
平均接受的 Token 数：每代平均接受的 Token 数。
吞吐量：每秒生成的 Token 平均数量。
实验设置。作者在配备20核CPU和单NVIDIA RTX A6000 GPU（48GB）的服务器上进行了实验。实验基于PyTorch 2.3.0和CUDA 12.1实现。对于模型，作者使用了float16数据类型，并应用了贪心解码，批处理大小为1。关于超参数，除非另有说明，否则和分别设置为5，自动机生成的草稿大小设置为40个 Token 。对于辅助预测解码算法，作者使用各自原始论文中描述的默认配置。唯一例外是 Token 回收： instead of the draft tree structure used in the original paper，作者采用了深度为6且包含60个节点的草稿树。
对于SAM解码，作者基于斯坦福-alpaca、python-code-instruction-18k和GSK8k构建了一个静态后缀自动机。为了增强作者的模型，作者融入了两种辅助方法：无模型方法Token Recycling和有模型方法EAGLE。在这里，SAM-Decoding[T]，SAM-Decoding[E]和SAM-Decoding[E2]分别表示作者的基础模型与Token Recycling、EAGLE和EAGLE2的组合。
实验结果 表1显示了在Spec-Bench上的实验结果。Spec-Bench包括六个任务，即多轮对话、总结和检索增强生成，其中检索技术特别适用于这些任务。结果显示将SAM解码集成到这些任务中导致了显著的速度提升。具体来说，对于Token Recycling，三个任务的速度提升比分别从1.92倍、1.96倍和1.68倍增加到2.48倍、2.86倍和2.14倍。在EAGLE中，速度提升比分别从2.63倍、2.28倍和1.95倍增加到2.78倍、2.65倍和2.10倍。对于EAGLE2，速度提升因子分别从2.87倍、2.33倍和2.03倍增加到3.02倍、2.76倍和2.23倍。然而，对于剩余的任务，检索方法产生的草案并没有显著帮助推理解码，反而引入了额外的计算成本，导致性能略有降低。尽管如此，当SAM解码与这些现有方法结合时，速度提升因子总体上有所增加：对于Token Recycling，从2.83倍增加到2.98倍；对于EAGLE，从3.57倍增加到3.77倍；对于EAGLE2，从2.38倍增加到2.49倍。
值得注意的是，在摘要任务中，模型无关方法（SAM-Decoding[T]）超过了基于模型的方法，实现了2.86倍的速度提升。这一发现强调了将SAM-Decoding与不同的辅助策略相结合的潜在优势，尤其是在检索可以有效支持解码过程的任务中。
表2展示了在HumanEval、MBPP和HAGRID数据集上的实验结果。对于代码数据集HumanEval和MBPP，作者基于python-code-instructions-18k数据集构建了静态后缀自动机。对于上下文问答数据集HAGRID，作者使用其自身的训练集构建了一个静态后缀自动机。
在编码任务中，SAM-Decoding显著加速了无模型方法。当SAM-Decoding与Token Recycling相结合时，加速比从HumanEval和MBPP数据集上的和分别提高到和。值得注意的是，SAM-Decoding在HumanEval上表现出更显著的加速效果。这可以归因于HumanEval通常具有较长的、语义丰富的变量名称，并且在输入 Prompt 中包括额外的标注。这些条件有利于检索型方法。相反，对于模型型方法如EAGLE，集成SAM-Decoding往往会导致接受的 Token 数平均减少和生成吞吐量降低。
转向上下文问答数据集HAGRID，这个任务对检索方法非常有益，SAM-Decoding有效地提高了推理速度。当与Token Recycling、EAGLE和EAGLE2相结合时，它们在HAGRID数据集上的推理速度分别从、和增加到、和。此外，可以明显看出，模型自由方法SAM-Decoding[T]在该项目上的速度比最先进的基于模型方法EAGLE2更高。
消融实验 为了进一步理解SAM-Decoding中各个组件的贡献以及不同超参数对推理速度的影响，作者对一系列超参数进行了消融实验。除非另有说明，这些实验是在Spec-Bench上针对vicuna-7B-v1.3模型进行的，其中SAM-Decoding[T]作为解码策略。
首先，作者通过网格搜索来研究和对推理速度的影响。这些参数控制了在检索过程中生成文本的偏好，以及在创建草稿时使用检索结果的偏好。结果总结如图2所示。如图所示，这两个指标在和达到一定值时，均有所增加。然而，在两个参数都超过5之后，这些指标开始下降。
此外，作者还研究了SAM-Decoding使用的草案大小如何影响推理速度。图3说明了在不同草案大小下，SAM-Decoding[T]的平均接受 Token 和推理速度。数据表明，随着草案大小的增长，平均接受 Token 和推理速度最初上升。然而，一旦草案大小超过40，这些指标开始下降，当草案大小达到70时，下降变得更加明显。
最后，作者研究了SAM-Decoding中不同模块对其推理速度的影响。SAM-Decoding包括三个草案生成模块：静态后缀自动机、动态后缀自动机和辅助推测解码模块。作者分别移除了这三个模块中的每一个，然后测量了SAM-Decoding的推理速度。结果如表3所示。从实验结果可以看出，每个模块都对解码过程的加速做出了贡献。值得注意的是，动态后缀自动机的影响远大于静态后缀自动机。这表明在很多情况下，从动态上下文中生成草案比从预先存在的文本语料库中检索草案更有效。此外，由于并非所有场景都适合检索方法，因此辅助推测解码模块在提高SAM-Decoding的整体性能方面发挥着关键作用。
5 Related Work
预测解码是一种可以显著加快大型语言模型（LLM）速度，同时不牺牲其输出质量的方法。通过在每个推理步骤中增加生成的 Token 数量，同时减少所需的推理步骤总数来实现。大多数预测解码技术依赖于在推理过程中使用较小的神经网络创建初始草稿。这些技术被称为基于模型的预测解码方法。基于模型的预测解码的早期实现，如Speculative Decoding（Leviathan等人，2023年）中所描述的那样，主要关注使用现有较小规模的LLM生成草稿序列。后来，Medusa（Cai等人，2024年）和SpecInfer（Miao等人，2024年）引入了基于树的预测方法，并开始了专门针对预测解码的专用草稿模型的开发。例如，Medusa开发了一组专门的解码头，可作为草稿模型。最近，EAGLE（Li等人，2024年；20）强调了将模型输出的隐藏状态信息整合到预测解码过程中的重要性。遵循这一洞察，EAGLE训练了一个集成模型输出隐藏状态信息的草稿模型，从而在各种任务上实现了更高效的加速。
与基于模型的方法相比，某些方法侧重于通过利用模型的近期预测结果来生成草稿，例如解码头的输出。这类方法的示例包括Lookahead Decoding（Fu等人，2024年）和Token Recycling（Luo等人，2024年）。Lookahead Decoding 跟踪生成的文本的 -gram，而 Token Recycling 则维护每个 Token 的顶 个最可能的下一个 Token 。
此外，一些策略通过检索先前生成的文本或语料库，直接生成草稿，采用 - gram 匹配。其中 PLD（Saxena，2023）和 REST（He 等人，2024）尤为突出。作者提出的 SAM-Decoding 方法也采用检索进行草稿生成。然而，由于其具有优越的速度和准确性，并且可以无缝集成其他推测解码技术，因此它脱颖而出。
此外，除了上述方法之外，研究还开展了一种依赖于模型的推测解码方法（Kou等人，2024年），以及依赖于大型架构内的子模型的推测解码方法（Elhoushi等人，2024年）。
高效LLM架构。还有从模型结构角度提高模型推理速度的工作。这部分工作包括模型蒸馏、量化与剪枝。模型蒸馏将大型模型的知识蒸馏到小型模型中，从而在保持模型性能的同时加速推理。量化减少了存储参数所需的比特数，并减少了在推理过程中从HBM传输到芯片内存的数据传输时间，从而实现有效的推理加速。剪枝用于移除模型中的冗余参数。对于结构剪枝，可以与模型蒸馏相结合训练高效的小模型，而半结构剪枝可以通过结合特殊硬件（稀疏张量核心）来减少模型的内存访问和计算开销，从而通过结合专用硬件提高推理速度。
6 Conclusion
在本工作中，作者提出了一种名为SAM解码的推测解码方法，该方法通过构建由生成文本和文本语料库生成的后缀自动机来实现。
SAM解码可以有效地从检索源中检索草稿，从而加速推理。SAM解码还设计得可以与现有方法无缝集成。因此，在检索不可行的情况下，SAM解码可以适应性切换到其他草稿生成方法。
实验结果表明，结合最先进的技术，SAM解码可以在多轮对话、摘要和检索增强生成任务中显著提高性能。
[0]. SAM Decoding: Speculative Decoding via Suffix Automaton.
点击上方卡片，关注 「AI视界引擎」 公众号