SWIFT: On-the-Fly Self-Speculative Decoding for LLM Inference Acceleration

作者/机构: Heming Xia1, Yongqi Li1∗, Jun Zhang2, Cunxiao Du3, Wenjie Li1 (1Department of Computing, The Hong Kong Polytechnic University; 2College of Computer Science and Technology, Zhejiang University; 3Sea AI Lab)

A1 主要贡献

本文探讨了一种新颖的、基于层稀疏性的即插即用（plug-and-play）推测解码（Speculative Decoding, SD）方案，以解决现有SD方法大多需要额外参数或大量训练的局限性。

核心问题与研究目标:
大型语言模型（LLM）的自回归解码效率低下，而现有的推测解码加速方法（如基于Jacobi解码的方法）存在与LLM预训练目标不一致的问题，或需要额外的模型模块和训练过程，限制了其普适性。本文旨在开发一种无需辅助模型或额外训练的即插即用SD解决方案，以加速不同LLM在多样化输入数据流上的推理。

关键发现与创新点:
研究基于两个关键观察：
1. LLM通过层稀疏性展现出巨大的自加速潜力：研究发现，即使是均匀地跳过模型中的某些层来生成草稿，也能实现显著的1.2倍加速，这为即插即用的SD提供了坚实基础。
2. 层稀疏性具有任务特定性：不同的任务需要不同的最优跳层配置，将为某个任务优化的配置应用于其他任务会导致性能大幅下降（例如，从故事生成任务迁移到推理任务，加速比从1.47倍降至1.01倍）。

基于以上发现，本文提出了 SWIFT，一种在推理过程中动态自适应（on-the-fly）的自推测解码算法。SWIFT的核心贡献如下：
- (i) 对层稀疏性加速LLM的实证分析：揭示了LLM通过层稀疏性进行自加速的潜力及其任务特定性，强调了在推理过程中进行自适应自推测解码的必要性。
- (ii) 提出SWIFT算法：这是首个即插即用的自推测解码算法，它在推理时动态优化目标LLM的跳层集合，从而在不同输入数据流上实现无损加速。SWIFT包含两个关键创新：
- 基于上下文的层集合优化机制：利用LLM生成的上下文（即先前的输出token）来高效地识别与当前输入流相对应的最优跳层集合。
- 置信度感知的推理加速策略：最大化草稿token的利用率，通过置信度分数动态调整草稿生成和验证过程，提高了推测准确性和验证效率。

• (iii) 广泛的实验验证：在LLaMA-2和CodeLLaMA等多种模型和下游任务（如摘要、代码生成、数学推理）上的实验表明，SWIFT能够在不使用任何辅助模型或额外训练的情况下，实现1.3倍至1.6倍的加速，同时理论上保证生成文本的分布与原始模型一致。

下图展示了本文提出的基于稀疏性的草稿生成方法与先前基于Jacobi解码方法的对比。

图 1: 先前解决方案与本文方案的即插即用SD图示。(a) 基于Jacobi的草稿生成方法将多个伪token附加到输入提示中，使目标LLM能够在单一步骤中生成多个token作为草稿。(b) SWIFT采用基于稀疏性的草稿生成方法，利用LLM中固有的稀疏性来促进高效的草稿生成。这项工作是首次使用基于稀疏性的草稿生成方法探索即插即用SD。

下表对比了SWIFT与现有SD方法的特性。

表 1: SWIFT与现有SD方法的比较。“AM”表示该方法是否需要辅助模块，如额外参数或数据存储。“Greedy”、“Sampling”和“Token Tree”分别表示该方法是否支持贪心解码、多项式采样和token树验证。SWIFT是首个即插即用的层跳过SD方法，与诸如Lookahead（Fu et al., 2024）等基于Jacobi的方法是正交的。

A3 背景知识与关键观察

3.1 自推测解码 (SELF-SPECULATIVE DECODING)

自推测解码的定义。与大多数需要额外参数的推测解码（SD）方法不同，自推测解码（Self-SD）首次提出利用LLM的一部分作为紧凑的草稿模型【索引2, Draft& verify: Lossless large language model acceleration via self-speculative decoding+2024+ACL】。在每个解码步骤中，该方法通过跳过LLM的中间层来高效地生成草稿token；然后，这些token由全参数的LLM并行验证，以确保目标LLM的输出分布保持不变。Self-SD的主要挑战在于确定在草稿阶段应跳过哪些层以及跳过多少层——这被称为“跳层集合”，并被形式化为一个优化问题。

Self-SD的优化问题。形式上，给定输入数据$X$和具有$L$层（包括注意力层和MLP层）的目标LLM $M_T$，Self-SD旨在找到一个最优的跳层集合$z$，以最小化每个token的平均推理时间：

$$\boldsymbol{z}^{*}=\arg \min _{\boldsymbol{z}} \frac{\sum_{\boldsymbol{x} \in \mathcal{X}} f\left(\boldsymbol{x} \mid \boldsymbol{z} ; \boldsymbol{\theta}_{\mathcal{M}_{T}}\right)}{\sum_{\boldsymbol{x} \in \mathcal{X}}|\boldsymbol{x}|}, \quad \text { s.t. } \boldsymbol{z} \in\{0,1\}^{L},$$

其中，$f(·)$是一个黑盒函数，返回样本$x$的推理延迟，$z_i \in \{0, 1\}$表示在生成草稿时是否跳过目标LLM的第$i$层，$|x|$表示样本长度。Self-SD通过贝叶斯优化过程【索引24, Efficient global optimization of expensive black-box functions+1998+J. Glob. Optim.】来解决这个问题。在推理之前，该过程基于高斯过程【索引43, Gaussian processes for machine learning+2006+MIT Press】迭代选择新的输入$z$，并在$X$的训练集上评估公式(1)。经过指定次数的迭代后，最佳的$z$被视为$z^*$的近似值，并在推理过程中保持固定。

Self-SD的局限性。虽然Self-SD已被证明是有效的，但它依赖于耗时的贝叶斯优化过程，这带来了一些限制。对于每个任务，Self-SD必须在每次迭代中顺序评估所有选定的训练样本以优化公式(1)；此外，贝叶斯优化的计算负担随着迭代次数的增加而显著增加。因此，在NVIDIA A6000服务器上，仅处理8个CNN/Daily Mail【索引39, Abstractive text summarization using sequence-to-sequence rnns and beyond+2016+CoNLL】样本进行1000次贝叶斯迭代，LLaMA-2-13B需要近7.5小时，LLaMA-2-70B需要20小时。这些计算需求限制了Self-SD在不同模型和任务间的通用性。

3.2 实验观察

本小节深入探讨Self-SD，探索这种基于层跳过的SD范式在无损LLM推理加速方面的即插即用潜力。我们的关键发现如下详述。

图 2: (a) LLM通过层稀疏性具备自加速潜力。通过利用top-k候选词的草稿，我们发现在草稿阶段均匀跳过一半的层可以产生显著的1.2倍加速。(b) 层稀疏性是任务特定的。每个任务都需要其自身最优的跳层集合，将一个任务的跳层配置应用于另一任务可能导致严重的性能下降。“X LS”表示为任务X优化的跳层集合。

3.2.1 LLM通过层稀疏性具备自加速潜力

Top-k预测提升对齐性。我们首先研究目标LLM与其层跳过变体之间的行为对齐潜力。与之前仅关注贪心草稿预测的工作【索引2, Draft& verify: Lossless large language model acceleration via self-speculative decoding+2024+ACL】不同，我们利用了top-k预测中的潜在草稿候选，详见第4.2节。

均匀跳层的有效性。我们在LLaMA-2-13B模型上，针对CNN/Daily Mail【索引39, Abstractive text summarization using sequence-to-sequence rnns and beyond+2016+CoNLL】、GSM8K【索引6, Training verifiers to solve math word problems+2021+CoRR】和TinyStories【索引11, Tinystories: How small can language models be and still speak coherent english?+2023+CoRR】数据集进行了实验。我们应用了统一的层跳过模式，并将k设置为10。实验结果如图2(a)所示，通过利用top-k预测，token接受率平均提高了30%，超过90%的草稿token被目标LLM接受。

挑战传统观念。因此，相较于在此实验设置中最高仅实现1.01倍加速的Self-SD，我们揭示了层跳过SD范式通过统一的层跳过模式，可以比传统的自回归解码实现平均1.22倍的壁钟时间加速。这一发现挑战了普遍认为必须精心策划跳过层的选择的观点，反而表明LLM通过固有的层稀疏性拥有更大的自加速潜力。

3.2.2 层稀疏性是任务特定的

跨任务适用性研究。我们进一步探讨以下研究问题：为一个特定任务优化的跳层集合是否适用于其他任务？为了解决这个问题，我们在LLaMA-2-13B上使用CNN/Daily Mail、GSM8K、TinyStories和WMT16 DE-EN数据集进行了领域转移实验。

实验发现。实验结果如图2(b)所示，揭示了两个关键发现：1) 每个任务都需要其自身最优的跳层集合。如图2(b)所示，最高的加速性能总是由为每个任务专门优化的跳层配置实现的。这些层的详细配置见附录A，表明不同任务的最优配置是不同的。2) 在不同任务间应用静态的跳层配置可能导致效率大幅下降。例如，当将一个故事生成任务的优化跳层集合应用于一个数学推理任务时，加速比从1.47倍下降到1.01倍，这表明为一个特定任务优化的跳层集合不能有效地泛化到其他任务。

为即插即用方案奠定基础。这些发现为我们在层跳过SD中实现即插即用解决方案奠定了基础。第3.2.1节为实时跳层选择提供了坚实的基础，表明可能不需要使用训练数据进行额外优化；第3.2.2节强调了静态层跳过模式在处理跨各种任务的动态输入数据流时的局限性，从而突显了在推理过程中进行自适应层优化的必要性。基于这些见解，我们提出了用于高效和自适应层集优化的动态自推测解码方法。

A2 方法细节

我们引入SWIFT，这是首个即插即用的自推测解码方法，它在推理过程中动态优化目标LLM的跳层集合，从而促进跨不同输入数据流的无损LLM加速。如图3所示，SWIFT将LLM推理分为两个不同阶段：(1) 基于上下文的层集合优化 (§4.1)，旨在给定输入流时识别最优的跳层集合；(2) 置信度感知的推理加速 (§4.2)，该阶段采用确定的配置来加速LLM推理。

图 3: SWIFT推理的时间线。N表示每个实例的最大生成长度。

4.1 基于上下文的层集合优化

层集合优化的挑战与方法。层集合优化是自推测解码中的一个关键挑战，因为它决定了应跳过目标LLM的哪些层来构成草稿模型（见第3.1节）。与依赖于耗时的离线优化的先前方法不同，我们的工作强调在推理过程中进行层集合优化，这对延迟-准确性权衡提出了更大的挑战：优化过程必须足够高效以避免在推理期间造成延迟，同时要确保能准确地草拟后续的解码步骤。为了解决这个问题，我们提出了一种自适应优化机制，该机制平衡了效率与草稿准确性。我们的方法通过每步仅对草稿模型执行一次前向传播来验证潜在的跳层集合候选，从而最小化了开销。核心创新在于使用LLM生成的token（即先前上下文）作为基准（ground truth），从而能够同时验证草稿模型在预测未来解码步骤时的准确性。

优化阶段的详细流程。在接下来的小节中，我们将阐述该优化阶段对每个输入实例的详细过程，包括上下文积累（§4.1.1）和层集合优化（§4.1.2）。

4.1.1 上下文积累

初始化草稿模型并积累上下文。在优化阶段，对于给定的一个输入实例，草稿模型通过在目标LLM中均匀跳过层来初始化。这个初始的层跳过模式会一直保持以加速推理，直到积累了指定数量的LLM生成的token，我们称之为上下文窗口。当达到这个窗口长度时，推理过程便转换到层集合优化阶段。

4.1.2 层集合优化

优化步骤概述。在此阶段，如图4所示，我们在每个LLM解码步骤之前集成一个优化步骤来精炼跳层集合，该步骤包括两个子步骤：

图 4: SWIFT中的层集合优化过程。在优化阶段，SWIFT在每个LLM解码步骤之前执行一个优化步骤来调整跳层集合，这包括：(a) 高效的层集合建议。SWIFT将随机搜索与区间贝叶斯优化相结合来提出层集合候选；(b) 并行候选评估。SWIFT使用LLM生成的token（即先前上下文）作为基准，从而能够同时验证提出的候选。选择表现最佳的层集合来加速当前的解码步骤。

高效的层集合建议。这个子步骤旨在提出一个潜在的层集合候选。形式上，给定一个有L层的目标LLM $M_T$，我们的目标是识别一个最优的跳层集合 $z \in \{0, 1\}^L$ 来构成紧凑的草稿模型。与完全依赖耗时贝叶斯优化过程的 Zhang 等人（2024）不同，我们引入了一种结合随机搜索和贝叶斯优化的高效策略。在这种方法中，随机采样高效地处理了大部分的探索工作。具体来说，给定一个固定的跳过率 $r$，SWIFT会以 $\beta$ 个优化步骤为固定间隔（例如，$\beta = 25$）应用贝叶斯优化来建议下一个层集合候选，而在其他优化步骤中则采用随机搜索。

$$\begin{aligned} z = \begin{cases} \text{Bayesian\_Optimization}(l) & \text{if } o \ \% \ \beta = 0 \\ \text{Random\_Search}(l) & \text{otherwise} \end{cases}, \end{aligned}$$

其中，$1 \le o \le S$ 是当前的优化步骤；$S$ 表示最大优化步骤数；$l = \binom{L}{Lr}$ 表示输入空间，即所有可以被跳过的层的可能组合。

并行候选评估。SWIFT利用LLM生成的上下文来同时验证候选草稿模型在预测未来解码步骤时的性能。形式上，给定一个输入序列 $x$ 和在上下文窗口内先前生成的token，表示为 $y = \{y_1, \dots, y_\gamma\}$，草稿模型 $M_D$（跳过了目标LLM的指定层 $z$）被用来并行预测这些上下文token：

$$y_{i}^{\prime}=\arg \max _{y} \log P\left(y \mid \boldsymbol{x}, \boldsymbol{y}_{<i} ; \boldsymbol{\theta}_{\mathcal{M}_{D}}\right), 1 \leq i \leq \gamma,$$ <p>其中 $\gamma$ 代表上下文窗口。目标LLM $M_T$ 中缓存的键值对被 $M_D$ 重用，这大概率能使 $M_D$ 的分布与 $M_T$ 对齐，并减少了冗余计算。匹配度分数（matchness score）定义为 $y$ 和 $y'$ 之间的精确匹配率：

$$\text{matchness} = \frac{\sum_{i} \mathbb{I}\left(y_{i}=y_{i}^{\prime}\right)}{\gamma}, 1 \leq i \leq \gamma,$$

其中 $I(·)$ 表示指示函数。该分数在优化过程中作为优化目标，反映了 $M_D$ 在预测未来解码步骤时的准确性。如图4所示，每一步的匹配度分数被整合到高斯过程模型中以指导贝叶斯优化，得分最高的层集合候选被保留下来构成草稿模型。

优化阶段的终止。如图3所示，对每个实例，上下文积累和层集合优化的过程交替进行，直到满足终止条件——要么达到最大优化步骤数，要么最佳候选在多次迭代中保持不变。一旦优化阶段结束，推理过程就转换到置信度感知的推理加速阶段，此时优化的草稿模型被用来加速LLM推理。

4.2 置信度感知的推理加速

加速阶段的解码过程。在加速阶段，优化步骤被移除。SWIFT应用表现最佳的层集合来构成紧凑的草稿模型，并遵循“先草稿后验证”的范式进行解码。具体来说，在每个解码步骤，给定输入 $x$ 和之前的LLM输出 $y$，草稿模型 $M_D$ 以自回归方式预测未来的LLM解码步骤：

$$y_{j}^{\prime}=\arg \max _{y} \log P\left(y \mid \boldsymbol{x}, \boldsymbol{y}, \boldsymbol{y}_{<j}^{\prime} ; \boldsymbol{\theta}_{\mathcal{M}_{D}}\right),$$ <p>
图 5: SWIFT的置信度感知推理过程。(a) 如果置信度分数低于阈值 ϵ，草稿生成会提前终止。(b) 根据置信度动态选择草稿候选，然后由目标LLM并行验证。

其中 $1 \le j \le N_D$ 是当前的草稿步骤，$N_D$ 表示最大草稿长度，$y'_{<j}$ 代表之前的草稿token，$P(·)$ 表示下一个草稿token的概率分布。目标LLM $M_T$ 和先前草稿token $y'_{<j}$ 的KV缓存被重用以降低计算成本。</p>

置信度分数的利用。令 $p_j = \max P(·)$ 表示top-1草稿预测 $y'_j$ 的概率，这可以被视为一个置信度分数。最近的研究【索引31, EAGLE-2: faster inference of language models with dynamic draft trees+2024+CoRR】、【索引9, Glide with a cape: A low-hassle method to accelerate speculative decoding+2024+ICML】表明，这个分数与草稿token $y'_j$ 通过验证的可能性高度相关——更高的置信度分数表明有更大的接受机会。因此，遵循以往的研究【索引2, Draft& verify: Lossless large language model acceleration via self-speculative decoding+2024+ACL】、【索引9, Glide with a cape: A low-hassle method to accelerate speculative decoding+2024+ICML】，我们利用置信度分数来修剪不必要的草稿步骤并选择有价值的草稿候选，从而提高推测准确性和验证效率。

两种置信度感知策略。如图5所示，我们将SWIFT与两种置信度感知的推理策略相结合：1) 提前终止草稿生成。如果置信度 $p_j$ 低于指定的阈值 $\epsilon$，自回归的草稿生成过程将停止，避免后续草稿计算的浪费。2) 动态验证。每个 $y'_j$ 会根据其置信度分数 $p_j$ 动态地扩展其top-k草稿预测以进行并行验证，从而提高推测准确性。具体来说，当 $p$ 分别在 $(0, 0.5]$、$(0.5, 0.8]$、$(0.8, 0.95]$ 和 $(0.95, 1]$ 范围内时，$k$ 分别设为10、5、3和1。所有草稿候选被线性化成一个单一序列，并由目标LLM使用一个特殊的因果注意力掩码（见图5(b)）进行并行验证。

A4 实验环境

• 模型架构:

  • 主要评估模型为LLaMA-2系列（Touvron et al., 2023b）和CodeLLaMA系列（Roziere et al., 2023）。
  • 深入分析中还包括Yi-34B（Young et al., 2024）和DeepSeek-Coder-33B（Guo et al., 2024）及其指令微调版本。

• 数据集与任务:

  • 摘要: CNN/Daily Mail (CNN/DM)，1-shot，最大生成长度64。
  • 数学推理: GSM8K，5-shot，最大生成长度64。
  • 故事生成: TinyStories，1-shot，最大生成长度128。
  • 代码生成: HumanEval，评估pass@1和pass@10，最大生成长度512。
  • 动态输入流分析: 额外使用Alpaca（指令遵循）、WMT14 DE-EN（翻译）和Natural Questions（问答），最大生成长度512。
  • 除HumanEval外，每个数据集均随机采样1000个实例进行测试。

• 硬件配置:

  • 4块NVIDIA RTX A6000 GPU (40GB)。
  • Intel(R) Xeon(R) Platinum 8370C CPU，32核。

• 软件与实现细节:

  • 代码库: Pytorch 2.1.0, Huggingface transformers。

  • 推理设置:

    • 接受策略：推测采样（Speculative Sampling）。
    • 批处理大小：1。
    • 精度：FP16。
    • 优化阶段参数：上下文窗口$\gamma=32$，最大草稿长度$N_D=25$，最大优化步骤$S=1000$，贝叶斯优化间隔$\beta=25$。
    • 优化终止条件：匹配度分数在300步内无提升或超过0.95。
    • 层跳过率$r$：13B模型为0.45，34B和70B模型为0.5。

  • 对比基线:

    • Parallel Decoding (Santilli et al., 2023)。
    • Lookahead Decoding (Fu et al., 2024)。
    • 这两种基线均为即插即用的Jacobi解码方法。

A4 实验结果

主要结果

文本生成任务:
如表2所示，SWIFT在文本生成任务上表现出卓越的效率。

• 性能优越性: SWIFT在各种模型和任务上均实现了1.3倍至1.6倍的稳定加速，优于现有的即插即用方法（Parallel Decoding 和 Lookahead Decoding）。
• 高接受率: SWIFT的效率得益于目标LLM和其层跳过变体之间的高度行为一致性。平均生成长度（M）达到5.01，token接受率（α）在90%到100%之间。特别是在LLaMA-2系列上，接受率稳定在98%-100%，意味着几乎所有草稿token都被接受。
• 模型规模效应: LLaMA-2-70B模型相比13B模型能以更高的层跳过率（0.45 → 0.5）获得更高的加速比，表明更大规模的LLM具有更强的层稀疏性，这预示着SWIFT在更大模型上具有更大的加速潜力。

表 2: SWIFT与先前即插即用方法的比较。报告了贪心解码下的平均生成长度M、加速比和平均解码速度（tokens/s）。†表示token接受率α高于0.98的结果。

代码生成任务:
如表3所示，SWIFT在代码生成任务中同样有效。
- 加速效果: 在贪心解码和随机采样设置下，SWIFT均实现了1.3倍至1.5倍的加速。
- 质量保持: 推测采样理论上保证了SWIFT能维持目标LLM的原始输出分布。实验中的任务性能指标（如pass@10）也验证了这一点，SWIFT的性能与自回归解码几乎完全相同。

表 3: SWIFT在代码生成任务上的实验结果。报告了平均生成长度M、接受率α、准确率（Acc.）和加速比。pass@1使用贪心解码，pass@10使用温度为0.6的随机采样。

深入分析

推理过程图解与计算分解:
- 推理过程 (图6左): SWIFT的推理过程分为优化和加速两个阶段。在优化阶段，匹配度分数迅速从0.45提升至0.98，触发优化终止。随后进入加速阶段，实例加速比稳定在平均1.53倍。整体加速比随着生成token数的增加而稳步提升，显示出SWIFT的效率随输入长度和实例数量增加而提高。
- 计算分解 (图6右): 在LLaMA-2-13B上对1000个CNN/DM样本的分析表明，优化步骤仅占总推理延迟的0.8%，证明了该策略的高效性。与Self-SD（需要7.5小时优化）相比，SWIFT实现了近180倍的优化时间缩减。草稿生成阶段占据了大部分延迟（64.4%），与高平均生成长度的结果一致，说明约80%的输出token由高效的草稿模型生成。

图 6: SWIFT推理过程图解与延迟分解。左图显示，在基于上下文的层集合优化阶段后，SWIFT的整体加速比稳步增加，在加速阶段达到平均实例加速比。右图显示，额外的优化步骤仅占总推理延迟的0.8%。

动态输入数据流处理 (图7):
- 对比Self-SD: 在处理由摘要、推理、指令遵循、翻译和问答任务组成的动态数据流时，Self-SD对领域漂移非常敏感，平均token接受率从92%降至68%，加速比从1.33倍降至平均1.05倍。
- SWIFT的适应性: SWIFT展现出强大的领域自适应能力，平均token接受率保持在96%，实现了1.3倍至1.6倍的稳定加速。

图 7: SWIFT与Self-SD在处理动态数据输入流时的比较。与Self-SD在分布变化时效率下降不同，SWIFT保持了稳定的加速性能，接受率超过0.9。

灵活的优化策略与规模法则 (图8):
- 优化灵活性 (图8a): 通过调整最大优化步数S和贝叶斯间隔β，SWIFT可以灵活适应不同类型的输入数据。对于少量实例，可以减少S和增加β以提前进入加速阶段；对于充足数据，可以探索更多优化路径以提升整体加速比。
- 规模法则 (图8b): 随着模型尺寸的增加，最优的层跳过率和整体加速比均得到提升，表明更大的LLM表现出更强的层稀疏性。这突显了SWIFT在加速更大规模模型（如175B）方面的潜力。

图 8: SWIFT的深入分析，包括：(a) 灵活的优化策略。最大优化迭代S和贝叶斯间隔β可以灵活调整以适应不同的输入数据类型。(b) 规模法则。SWIFT的加速比和最优层跳过率随着模型尺寸的增加而增加，表明更大的LLM表现出更强的层稀疏性。

其他LLM主干 (图9):
- 通用性: SWIFT在LLaMA系列之外的LLM主干上也表现出色。在Yi-34B和DeepSeek-Coder-33B及其指令微调版本上，SWIFT分别实现了26%至54%的效率提升，证明了其作为通用加速方案的兼容性。

图 9: SWIFT在不同LLM主干及其指令微调变体上的加速比。

A5 结论

本文介绍了SWIFT，一种在推理过程中动态自适应地选择并跳过LLM中间层的自推测解码算法。该方法无需额外训练或辅助模型，使其成为一个即插即用的解决方案，能够加速处理多样化的输入数据流。在多种LLM和任务上的广泛实验证明，SWIFT能够实现超过1.3倍至1.6倍的加速，同时保持生成文本的原始分布。此外，深入分析突出了SWIFT在处理动态输入数据流方面的有效性及其与各种LLM主干的无缝集成能力，展示了该范式在实际LLM推理加速中的巨大潜力。

A6 附录

A 初步细节

Self-SD跳层配置的可视化。我们在图10中展示了Self-SD在四个任务领域上的详细配置，证明了最优的跳层配置因具体任务而异。

图 10: 在不同任务领域上由Self-SD（Zhang et al., 2024）优化的LLaMA-2-13B跳层集合配置可视化。灰色方块表示保留的层，红色方块表示跳过的注意力层，蓝色方块表示跳过的MLP层。

B 实验设置

B.1 模型和数据集

实验模型与任务。我们的实验主要评估SWIFT在LLaMA-2【索引57, Llama 2: Open foundation and fine-tuned chat models+2023+CoRR】和CodeLLaMA系列【索引45, Code llama: Open foundation models for code+2023+CoRR】上的效果。我们在多种生成任务上进行了实证验证。对于摘要、数学推理、故事生成和代码生成任务，我们分别选择了CNN/Daily Mail (CNN/DM)【索引39, Abstractive text summarization using sequence-to-sequence rnns and beyond+2016+CoNLL】、GSM8K【索引6, Training verifiers to solve math word problems+2021+CoRR】、TinyStories【索引11, Tinystories: How small can language models be and still speak coherent english?+2023+CoRR】和HumanEval【索引5, Evaluating large language models trained on code+2021+CoRR】数据集。

具体设置。我们对CNN/DM和TinyStories进行1-shot评估，对GSM8K进行5-shot评估。CNN/DM、GSM8K和TinyStories的最大生成长度分别设置为64、64和128。我们比较了HumanEval的pass@1和pass@10。在进一步的分析中，我们增加了三个数据集来验证SWIFT处理动态输入数据流的能力。具体来说，我们选择了Alpaca【索引52, Stanford alpaca: An instruction-following llama model+2023+https://github.com/tatsu-lab/stanford_alpaca】 、WMT14 DE-EN和Natural Questions【索引28, Natural questions: A benchmark for question answering research+2019+Transactions of the Association for Computational Linguistics】分别用于指令遵循、翻译和问答任务。HumanEval和所有分析的最大生成长度设置为512。除HumanEval外，我们从每个数据集的测试集中随机抽取1000个实例。

B.2 推理设置

优化与推理参数。在优化阶段，我们同时采用随机搜索和贝叶斯优化来建议潜在的跳层集合候选，以平衡优化性能和效率。上下文窗口 $\gamma$ 设置为32。最大草稿长度 $N_D$ 设置为25。对于代码生成任务中的随机采样，我们应用温度为0.6，top p = 0.95。层集合优化的最大步骤数 $S$ 设置为1000，每 $\beta = 25$ 步执行一次贝叶斯优化。如果匹配度分数在300步内没有提高或超过0.95，则提前停止优化阶段。层跳过率 $r$ 对13B模型固定为0.45，对34B和70B模型固定为0.5。

软硬件环境。所有实验均在4×NVIDIA RTX A6000 GPU (40GB)上使用Pytorch 2.1.0进行，CUDA版本为12.1，CPU为具有32核的Intel(R) Xeon(R) Platinum 8370C。我们的方法和所有基线的推理均使用Huggingface transformers包执行。遵循先前的工作，我们采用推测性采样【索引29, Fast inference from transformers via speculative decoding+2023+ICML】作为我们的接受策略，批大小设置为1。

B.3 评估指标

指标定义。本小节详细说明了我们的评估指标，即平均生成长度 $M$ 和token接受率 $\alpha$。具体来说，平均生成长度 $M$ 指的是目标LLM每次前向传播产生的平均输出token数；接受率 $\alpha$ 定义为接受的token数与总草稿步数的比率。换句话说，它代表了目标LLM在一次草稿模型前向传播中接受一个潜在token的期望概率。这两个指标与计算硬件无关，因此被认为是更客观的指标。

预期加速比公式。给定平均生成长度 $M$、接受率 $\alpha$ 和层跳过率 $r$，加速阶段的预期壁钟时间加速比的数学公式推导如下：

$$\mathbb{E}(\text{Spd.}) = \frac{M}{(M-1) \times \frac{c}{\alpha} + 1} = \frac{M\alpha}{(M-1)c + \alpha}, \quad c = 1 - r,$$

其中 $c$ 定义为 Leviathan 等人（2023）中的成本系数。它计算为草稿模型的单次前向传播时间与目标LLM的单次前向传播时间之比。总之，理想的加速比会随着 $M$ 和 $\alpha$ 的增大以及 $c$ 的减小而提高。

C 实验细节

C.1 主要结果的细节

详细统计数据。我们在表4中展示了我们主要实验结果的详细统计数据。SWIFT在所有评估设置中始终实现了超过90%的token接受率 $\alpha$，平均生成长度 $M$ 从2.99到8.21不等。这些统计数据表明目标LLM与其层跳过草稿变体之间存在强烈的行为对齐，正如第5.2节所讨论的。此外，我们报告了使用公式(6)计算的预期加速比E(Spd.)，表明SWIFT的当前实现有很大的潜力可以通过进一步优化来提升其效率。

表 4: SWIFT在LLaMA-2系列上进行文本生成任务的详细结果。我们报告了在FP16精度下贪心解码设置的平均生成长度M、token接受率α和由公式(6)计算的预期加速比E(Spd.)。

C.2 LLaMA-70B模型的附加结果

在70B模型上的表现。除了表2中呈现的主要结果外，我们还提供了SWIFT在LLaMA-70B模型上的进一步实验评估，包括LLaMA-2-70B和LLaMA-3-70B，以及它们的指令调优变体，实验设置相同。结果表明，SWIFT在LLaMA-2和LLaMA-3系列上均实现了1.4倍∼1.5倍的壁钟时间加速。值得注意的是，SWIFT在各种评估设置中实现了超过85%的token接受率 $\alpha$，平均生成长度 $M$ 从3.43到7.80不等。尽管不同模型之间观察到层冗余度的差异（例如，LLaMA-2-70B与LLaMA-3-70B的跳过率 $r$ 不同），SWIFT表现出强大的适应能力，无论模型版本如何，都能保持一致的加速性能。

表 5: SWIFT在LLaMA-70B系列上进行文本生成任务的实验结果。我们报告了贪心解码设置下的平均生成长度M、token接受率α和加速比。LLaMA-2模型的跳过率r设为0.5，LLaMA-3模型设为0.4。

C.3 LLM主干的详细结果

在其他LLM上的有效性。为了进一步验证SWIFT的有效性，我们使用了除LLaMA系列之外的其他LLM主干进行了实验。具体来说，我们选择了两个最近具有代表性的LLM：用于文本生成的Yi-34B和用于代码生成的DeepSeek-Coder-33B。实验结果如表6和表7所示，证明了SWIFT在这些LLM主干上的功效。SWIFT在Yi-34B系列上实现了1.2倍∼1.3倍的稳定壁钟时间加速，在DeepSeek-Coder-33B系列上实现了1.3倍∼1.5倍的加速。值得注意的是，对于DeepSeek-Coder-33B系列，SWIFT的平均生成长度 $M$ 在3.16到4.17之间，同时token接受率 $\alpha$ 超过83%。这些发现证实了SWIFT作为一种通用的、即插即用的SD方法的实用性，为不同的LLM主干提供了有希望的推理加速。

表 6: SWIFT在Yi-34B系列上进行文本生成任务的实验结果。我们报告了在FP16精度下贪心解码设置的平均生成长度M、token接受率α和加速比。跳过率r设为0.45。

表 7: SWIFT在DeepSeek-Coder-33B系列上进行代码生成任务的实验结果。跳过率r设为0.5。我们对pass@1使用贪心解码，对pass@10使用温度为0.6的随机采样。“DS”是DeepSeek的缩写。

D 进一步分析与讨论

D.1 消融研究

各组件贡献分析。表8展示了在CNN/DM上使用LLaMA-2-13B进行的SWIFT消融研究。实验结果表明，SWIFT的每个组件都对其整体加速做出了贡献。具体来说，提前停止草稿生成有效地减少了无效的草稿步骤，使token接受率 $\alpha$ 提高了55%。动态验证通过根据置信度分数从top-k预测中选择合适的草稿候选，进一步提高了效率；移除该组件会导致平均生成长度（M）和整体加速比的下降。此外，优化阶段精炼了跳层集合，与初始的统一层跳过策略相比，加速比提高了34%。总之，这些结果证实了SWIFT中每项创新策略的有效性。

表 8: SWIFT的消融研究。“ver.”是verification的缩写。

D.2 上下文窗口

上下文窗口大小的影响。在表9中，我们展示了对上下文窗口 $\gamma$ 的详细分析，该窗口决定了在层集优化过程中使用的LLM生成token的数量。较小的 $\gamma$ 会在匹配度分数计算中引入更大的随机性，导致性能次优，而较大的 $\gamma$ 则会增加优化步骤的计算开销。结果表明，$\gamma = 32$ 在优化性能和计算开销之间提供了最佳平衡。

表 9: SWIFT在不同上下文窗口γ下的加速比。报告了优化步骤的延迟以说明相关的开销。

D.3 与先前层跳过方法的比较

与LayerSkip和Self-SD的比较。本小节中，我们将SWIFT与两种代表性的层跳过推测解码（SD）方法进行比较：LayerSkip【索引12, LayerSkip: Enabling early exit inference and self-speculative decoding+2024+ACL】和Self-SD【索引2, Draft& verify: Lossless large language model acceleration via self-speculative decoding+2024+ACL】。具体来说，LayerSkip通过实现提前退出草稿生成来创新自推测解码，其中LLM仅使用其较早的层来生成草稿。然而，这种方法需要耗时的预训练或微调过程，这会改变目标LLM的原始输出分布，可能损害生成输出的可靠性；Self-SD提出通过跳过中间层来构建紧凑的草稿模型，在推理前使用广泛的贝叶斯优化过程来确定目标LLM内的最优跳过层。如第3.1节所述，虽然有效，但Self-SD存在显著的优化延迟（LLaMA-2-13B近7.5小时，LLaMA-2-70B近20小时）。这种漫长的优化过程限制了其在不同模型和任务间的实用性和通用性。

SWIFT的优势。表10和11分别总结了在加速性能和训练/优化成本方面的比较结果。下面，我们详细说明SWIFT相对于这些方法的优势：
- 与LayerSkip的比较：LayerSkip实现了激进的跳过率（r = 0.8），平均生成长度为2.42，token接受率为0.64。然而，其对预训练或微调的依赖改变了目标LLM的原始分布，可能降低可靠性。相比之下，SWIFT在不需额外训练的情况下，保持了目标LLM的原始分布，同时提供了可比的1.56倍加速。
- 与Self-SD的比较：Self-SD依赖于耗时的贝叶斯优化过程，在推理前产生大量延迟。SWIFT通过在线优化策略消除了这一瓶颈，实现了约200倍的优化延迟降低，同时保持了相同的1.56倍加速。我们进一步用我们的置信度感知推理加速策略增强了Self-SD（Self-SD w/ dynamic ver.）。即使与这个增强版本相比，SWIFT也实现了有竞争力的加速。

结论。这些发现突显了SWIFT相对于以前的层跳过SD方法的效率和实用性。作为首个即插即用的层跳过SD方法，我们希望SWIFT能为该领域提供有价值的见解并激发进一步的研究。

表 10: SWIFT与先前层跳过SD方法的比较。我们报告了贪心解码下的跳过率r、平均生成长度M、token接受率α和加速比。结果是在CNN/DM上使用LLaMA-2-13B获得的。“ver.”是verification的缩写。

表 11: SWIFT与先前层跳过SD方法在LLaMA-2-13B上的训练成本和优化延迟比较。训练成本来源于原始论文，而优化延迟是在A6000 GPU上我们重新实现时测量的。SWIFT在无需额外训练的情况下，与先前方法相比，优化延迟减少了约200倍，确立了其作为高效即插即用SD方法的地位。

D.4 与SELF-SD的详细比较

优化延迟与加速比对比。本小节我们提供了SWIFT与Self-SD【索引2, Draft& verify: Lossless large language model acceleration via self-speculative decoding+2024+ACL】的详细比较。图11展示了Self-SD在不同优化延迟下的加速比，反映了贝叶斯优化迭代次数的增加。如图所示，当贝叶斯迭代次数少于50次（对应优化延迟低于1474秒）时，Self-SD的加速改进微乎其微——几乎等同于统一跳过。在100次贝叶斯迭代时，Self-SD实现了1.19倍的加速；然而，其优化延迟几乎是SWIFT的25倍（2898秒 vs. 116秒）。

图 11: SWIFT与Self-SD在优化延迟和加速比方面的比较。SWIFT以116秒的优化延迟实现了1.56倍的加速。

相似延迟下的性能对比。表12在相似的优化延迟下比较了SWIFT和Self-SD（前两行）。结果凸显了SWIFT在优化效率（116秒 vs. 155秒）和加速比（1.56倍 vs. 0.97倍）上的优越性。即使与增强版的Self-SD（w/ dynamic verification）相比，SWIFT在加速比上仍有30%的相对提升。下面我们分析导致这一优势的因素（在第3.1节中已阐述）：
- 优化目标粒度：Self-SD在多样本级别计算其优化目标，要求每次迭代顺序解码所有选定的训练样本（例如，8个样本，每个32个token）来优化公式1。相比之下，SWIFT采用步级优化目标，在每个解码步骤动态优化层集。
- 贝叶斯优化复杂度：贝叶斯优化的计算复杂度随迭代次数显著增加。SWIFT通过将随机搜索与区间贝叶斯优化相结合来减轻这一负担，加速了优化过程的收敛，同时降低了计算开销。

进一步的优化权衡分析。为了进一步检验优化权衡，我们将Self-SD的顺序优化要求减少到单个样本8个token，从而在可比的延迟内实现更多的贝叶斯优化迭代。相应的结果，表示为Self-SDc（第3-4行），呈现在表12中。即使在这些优化设置下，SWIFT仍然表现出显著优越的加速和效率，突显了我们所提出策略的有效性。

表 12: SWIFT与Self-SD在相似优化延迟下的比较。我们报告了贪心解码下的跳过率r、平均生成长度M、token接受率α和加速比。结果是在CNN/DM上使用LLaMA-2-13B获得的，“ver.”表示verification。

D.5 即插即用SD方法的必要性

SD研究的两个方向。最近对推测解码（SD）的兴趣激增，导致该领域发展出许多有前景的策略，大致可分为两个方向：
- 需要训练的SD。这些方法需要额外的预训练或微调来提高推测准确性，通常涉及集成额外的参数。例如，Medusa【索引3, Medusa: Simple LLM inference acceleration framework with multiple decoding heads+2024+ICML】和Eagle【索引30, Eagle: Speculative sampling requires rethinking feature uncertainty+2024+】将轻量级草稿头集成到目标LLM中并进行微调，实现了3倍∼4倍的加速。
- 即插即用的SD。这些方法无需依赖辅助模型或额外训练即可立即加速LLM推理。著名的例子包括Parallel Decoding【索引46, Accelerating transformer inference for translation via parallel decoding+2023+ACL】和Lookahead【索引14, Break the sequential dependency of LLM inference using lookahead decoding+2024+ICML】，它们利用基于Jacobi的草稿生成，在各种LLM上实现了1.2倍∼1.4倍的加速。

即插即用方法的价值。尽管需要训练的SD方法通常能提供更高的加速比，但它们对额外训练和参数的依赖限制了其通用性和实用性。这引发了学术界关于即插即用SD方法价值的争论。为了解决这些担忧，我们下面提供详细分析，以突显即插即用SD方法的必要性，并强调我们提出的SWIFT的贡献：
1. 需要训练的SD方法的训练成本通常令人望而却步。像Medusa和Eagle这样的方法虽然实现了更高的加速，但产生了巨大的训练成本。尽管努力减少训练开销，这些方法仍然需要大量的计算资源（如GPU时间和数据集）才能提供有效的加速性能。例如，Eagle对于LLaMA-33B需要使用8个RTX 3090 GPU训练1-2天，或者对于LLaMA-2-Chat-70B在4个A100 (40G) GPU上训练长达2天，使用的数据集包含来自ShareGPT的70k对话。这样的计算负担在几种情况下带来了挑战：
- 用户必须为不受支持的目标LLM训练新的草稿模型。例如，如果用户的目标LLM不在发布的检查点中，或者基础模型更新了（如LLaMA-3.x），用户被迫训练一个新的草稿模型，这可能超出他们可用的GPU资源（如GPU时间）。
- 小规模加速需求的用户面临效率低下的问题。例如，一个需要评估一小组样本（如10小时的评估）的研究人员会发现1-2天的训练要求不成比例，并妨碍整体研究效率。

1. 即插即用SD填补了需要训练方法无法解决的关键空白。包括SWIFT在内的即插即用SD方法是模型无关和免训练的，无需额外的计算开销即可提供即时加速。这些特性对于大型模型（70B–340B）和需要快速集成的用例尤其关键。即插即用SD方法（如Lookahead）日益增长的采用，进一步凸显了它们的重要性。这些方法满足了易用性和计算效率至上的场景，验证了它们的研究意义。
2. SWIFT开创了基于层跳过草稿的即插即用SD。SWIFT代表了首个将层跳过草稿纳入其中的即插即用SD方法。它在不同模型和任务上始终比原生自回归解码实现1.3倍∼1.6倍的加速。此外，它比Lookahead表现出10%∼20%的更高效率。尽管其有效，SWIFT为现有的即插即用SD引入了一个互补的研究方向。其方法与Lookahead解码正交，两者结合可能会进一步放大它们的集体效率。我们相信这项研究提供了宝贵的见解，并为未来的SD进步铺平了道路，特别是在实用和成本效益高的LLM加速方面。

总结。总而言之，虽然需要训练的SD方法实现了更高的加速比，但其高计算成本和有限的灵活性降低了实用性。像SWIFT这样的即插即用SD方法提供了免训练、模型无关的加速，使其成为各种场景的理想选择。我们希望这一澄清能促进对即插即用SD研究价值的更大认识和认可。

D.6 与相关工作的额外讨论

层跳过与提前退出。在这项工作中，我们通过层跳过利用LLM固有的层稀疏性，即选择性地绕过目标LLM内的中间层来构建紧凑的草稿模型。除了层跳过，SD中还有另一个研究方向专注于提前退出，即在较早的层停止推理以提高计算效率【索引61, Predictive pipelined decoding: A compute-latency trade-off for exact LLM decoding+2023+CoRR】、【索引20, SPEED: speculative pipelined execution for efficient decoding+2023+CoRR】、【索引2, Fast and robust early-exiting framework for autoregressive language models with synchronized parallel decoding+2023+EMNLP】、【索引12, LayerSkip: Enabling early exit inference and self-speculative decoding+2024+ACL】。特别是，LayerSkip【索引12, LayerSkip: Enabling early exit inference and self-speculative decoding+2024+ACL】通过仅使用目标LLM的较早层生成草稿，然后用全参数模型进行验证，探索了提前退出草稿。这种方法需要涉及层丢弃和提前退出损失的训练。类似地，PPD【索引61, Predictive pipelined decoding: A compute-latency trade-off for exact LLM decoding+2023+CoRR】采用提前退出，但为每个层训练单独的分类器，而不是依赖于单个最终层分类器。尽管有效，这些方法依赖于广泛的微调来启用提前退出能力，产生了巨大的计算成本。此外，训练过程改变了目标LLM的原始输出分布，可能损害生成输出的可靠性。相比之下，我们提出的SWIFT不需要辅助模型或额外训练，保留了目标LLM的原始输出分布，同时提供了可比的加速效益。

自回归与非自回归草稿。有一条平行的需要训练的SD研究路线，专注于非自回归草稿策略【索引50, Blockwise parallel decoding for deep autoregressive models+2018+NeurIPS】、【索引3, Medusa: Simple LLM inference acceleration framework with multiple decoding heads+2024+ICML】、【索引15, Better & faster large language models via multi-token prediction+2024+CoRR】、【索引26, Accelerating blockwise parallel language models with draft refinement+2024+NeurIPS】。这些方法将多个草稿头集成到目标LLM中，从而在每个解码步骤中能够并行生成草稿token。值得注意的是，Kim等人（2024）【索引26, Accelerating blockwise parallel language models with draft refinement+2024+NeurIPS】建立在Stern等人（2018）【索引50, Blockwise parallel decoding for deep autoregressive models+2018+NeurIPS】引入的块并行解码范式之上，通过使用任务无关的n-gram和使用较小LM的轻量级重排序器来精炼块草稿，从而加速推理。虽然这些方法取得了显著的加速，但它们也需要对草稿模型进行广泛的训练。SWIFT通过开创即插即用的SD来补充这些努力，消除了对辅助模型或额外训练的需求，为不同的用例提供了更灵活和实用的解决方案。

D.7 优化步骤

优化过程中的跳层集变化。我们在图12中展示了SWIFT在不同优化步骤中的详细配置。随着过程的继续，跳层集合逐渐被精炼至最优配置。

图 12: SWIFT在不同优化步骤中优化的LLaMA-2-13B跳层集合配置的可视化。灰色方块表示保留的层，红色方块表示跳过的注意力层，蓝色方块表示跳过的MLP层。

D.8 案例研究

与Lookahead的比较案例。我们分别在表13和表14中展示了来自CNN/DM和Humaneval的示例，比较了SWIFT与先前的即插即用方法Lookahead【索引14, Break the sequential dependency of LLM inference using lookahead decoding+2024+ICML】之间的token接受情况。结果表明，SWIFT接受了更多的草稿，这些草稿均匀分布在输出中，突显了我们方法的有效性。

表 13: 使用LLaMA-2-13B的CNN/DM示例；红色token表示接受的草稿。

表 14: 使用CodeLLaMA-13B的Humaneval示例；红色token表示接受的草稿。