Consistency LLM 推理加速方法实现 3.5 倍提速

Consistency LLM 推理加速方法实现 3.5 倍提速 | 前端进阶之旅

Full translation

完整中文译文

背景：Jacobi 解码

Jacobi 解码的局限性

然而，在实际应用中，用于 LLM 的普通 Jacobi 解码相比 AR（自回归）解码只能取得边际加速，例如平均 $1.05\times$ 的加速比。这是因为一个经过 AR 训练的 LLM 在前面的令牌中存在错误时，很难生成正确的令牌。因此，大多数 Jacobi 迭代对于 $n$ 令牌序列而言只能获得一次纠正，导致轨迹更长，如图 3 左侧所示。

前瞻解码（Lookahead decoding）和推测解码（speculative decoding）方法试图减轻 Jacobi 解码和传统 AR 解码的低效率，但在推理时会产生额外的内存成本。而 CLLM 则无需额外成本。

一致性 LLM（CLLM）

Jacobi 解码初步

用 Jacobi 轨迹训练

Jacobi 轨迹准备：对于每个提示词，我们按顺序对 $n$ 令牌的每个截断执行 Jacobi 解码，直到整个响应序列 $\mathbf l$ 生成完毕，这相当于所有连续不动点的串联。沿轨迹生成的每个序列都计为一个数据条目。需要注意的是，对于长度为 $N$（$N ≫ n$）个令牌的较长响应 $\mathbf l$，这种截断方法避免了对长输入的缓慢模型评估。

用一致性损失和 AR 损失训练：我们联合优化两个损失函数来调整 CLLM，一致性损失保证了一次预测多个令牌，AR 损失防止 CLLM 偏离目标 LLM，从而保持生成质量。

一致性和 AR 损失

令 $p$ 表示目标 LLM。令 $q_\theta(\cdot| \mathbf x)$ 表示参数初始化自 $p$ 的 CLLM，其参数为 $\theta$。对于提示词 $\mathbf x$ 及其对应的 Jacobi 轨迹 $\mathcal{J}$，令 $\mathbf y$ 和 $\mathbf y^*$ 分别表示轨迹上的随机状态和不动点。

我们可以通过最小化以下损失（称为全局一致性（GC）损失）来促使 CLLM 以 $\mathbf y$ 作为输入输出 $\mathbf y^*$：

其中 $\theta^{-} = \text{stopgrad}(\theta)$，我们滥用记号来表示从数据集中的均匀采样。$D(\cdot||\cdot)$ 表示两个分布之间的距离，具体选择在 GKD 方法中讨论，本文主要采用前向 KL。

或者，采用遵循一致性模型形式的局部一致性（LC）损失，其中 Jacobi 轨迹 $\mathcal{J}$ 中的相邻状态 $(\mathbf y^{(j)}, \mathbf y^{(j+1)})$ 被驱动产生相同的输出：

特定领域：从图 5 可以看出，与其他基线（包括原始目标模型、Medusa2 和推测解码）相比，CLLM 实现了最显著的加速。

开放域对话挑战（MT-bench）：从 LLaMA2-7B 使用 ShareGPT 数据集训练的 CLLM，与前瞻解码相结合时，可以实现与 Medusa2 大致相同的加速，在 MT-bench 上的得分相当。但是，CLLM 提供了更高的适应性和内存效率，因为它不需要对目标模型的原始架构进行修改，也不需要辅助组件。

CLLM 的微调成本是适度的，例如，对于 LLaMA-7B，在 Spider 数据集上只需传递约 100 万个令牌就可以实现 $3.4\times$ 的加速。在数据集规模较大的情况下，例如对于 CodeSearchNet-Python，只需要数据集的 10% 来生成 Jacobi 轨迹用于训练 CLLM 就可以获得约 $2.5\times$ 的加速。总令牌数可以通过以下公式估计：

$$N = \text{平均每个提示词的轨迹数} \times \text{平均轨迹长度} \times \text{提示词数量}$$

快速前进和静止令牌

图 6 左侧显示，目标 LLM 通常在一次迭代中只生成一个正确的令牌。相比之下，在 CLLM 中，我们识别出了快速前进现象，即在单个 Jacobi 迭代中正确预测多个连续令牌。

此外，提前正确生成的令牌（例如图 6 左侧索引 6 和 7 处的"country"和"H"）在目标 LLM 的后续迭代中常常被不准确地替换。另一方面，CLLM 即使在前面有错误令牌的情况下，也能够抢先预测正确的令牌，同时保持这些令牌不变。我们将这样的令牌称为静止令牌。这两种现象都有助于 CLLM 的 Jacobi 解码快速收敛，从而导致显著的生成加速。

我们观察到，CLLM 通过训练获得了一个关键的语言学概念——搭配词（collocation）：一系列词或术语以比随机机会预期的频率更高的频率共同出现。语言不仅由孤立的单词组成，还严重依赖于特定的词对搭配。搭配词在自然语言和编程语言中都大量存在。它们包括动词 + 介词组合（例如"talk to"、"remind … of …"）、动词 + 名词结构（例如"make a decision"、"catch a cold"）以及许多领域特定的句法结构（例如"SELECT … FROM …"、"if … else"用于编程）。一致性生成目标允许 CLLM 从 Jacobi 轨迹中的任何点推断这些结构，促使 CLLM 精通众多搭配词，从而能够同时预测多个单词以最小化迭代步骤。

我们感谢 Yang Song、Canwen Xu、Yonghao Zhuang、Dacheng Li 和 Yichao Fu 提供的富有见地的反馈。

@misc{kou2024cllms,
      title={CLLMs: Consistency Large Language Models}, 
      author={Siqi Kou and Lanxiang Hu and Zhezhi He and Zhijie Deng and Hao Zhang},
      year={2024},
      eprint={2403.00835},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}