KOALA：提升大语言模型的推测解码效率

在当前人工智能的浪潮中，大语言模型（LLMs）以其卓越的推理和生成能力吸引了广泛的关注。然而，随着模型规模的不断扩大，推理延迟的问题日益凸显。本文将介绍一种新颖的架构——KOALA（K-layer Optimized Adversarial Learning Architecture），旨在通过多层草稿头和对抗学习的结合，显著提升推测解码的效率。

引言

大语言模型如GPT-4、Llama 2和PaLM 2在多种任务中展现出卓越的性能，然而其自回归解码的特性使得推理过程面临着高延迟的挑战。推测解码作为一种新兴的加速方法，利用草稿模型提高目标LLM的推理效率，实现了草稿-再验证的工作流程。尽管草稿头在此过程中发挥了重要作用，其潜力仍未被充分挖掘。

KOALA提出了一种将传统单层草稿头转变为多层结构的方案，并将对抗学习引入传统的监督训练中。这一创新显著提升了草稿头在预测后续标记时的准确性，使其功能更接近于目标LLM。尽管这种改进稍微增加了草稿负担，但KOALA有效地释放了草稿头的潜力，极大地增强了推测解码的效率。

理论基础

自回归解码

自回归解码是一种序列生成的基本技术，模型按照从左到右的顺序生成标记。对于输入序列 $x_1,x_2,\dots ,x_n$，LLM ${\mathcal{M}}_q$生成下一个标记 $x_{n+1}$，遵循公式：

$$x_{n+1}\sim q_{n+1}\leftarrow {\mathcal{M}}_q(x|x_{\le n})$$

推测解码

推测解码通过草稿模型加速目标LLM的推理，同时确保采样结果与目标LLM一致。它遵循草稿-再验证的范式。在每次解码迭代中，草稿模型首先高效地预测多个未来标记，然后由目标LLM并行验证。

$${\bar{x}}_1,{\bar{x}}_2,\dots ,{\bar{x}}_t\sim d_1,d_2,\dots ,d_t\leftarrow {\mathcal{M}}_d(x|x_{\le n})$$

目标LLM随后计算并验证这些预测标记。

对抗学习

对抗学习（Adversarial Learning）是一种机器学习范式，主要涉及生成器（$G$）和判别器（$D$）两个组成部分。通过竞争和共同进化的过程，这一框架提高了生成数据的真实感。其目标是训练生成器生成的数据能够以至于判别器无法区分。

KOALA：架构与训练

多层草稿头

在传统的草稿头结构中，通常只有一个层级。KOALA通过扩展这一结构至多层次，显著缩小了草稿头与目标LLM之间的性能差距。多层草稿头的设计允许模型更好地模拟目标LLM的功能，从而提高整体协作效率。

例如，在Medusa和EAGLE草稿头中，KOALA将单层结构扩展为K层结构。此举不仅增强了草稿头的功能，还提升了其预测准确性。

训练过程

KOALA的训练过程结合了对抗学习和传统的监督学习。生成器（草稿头）和判别器的能力相当，以确保训练效果最优。训练的主要目标是使草稿头更好地捕捉目标LLM的功能，从而提高其预测准确性。

训练过程的核心公式如下：

$$L_{\mathcal{G}}=-\lambda {\mathbb{E}}_{\tilde{x}\sim p_d(d)}[\log (D(\tilde{x}))]+L_{\text{Distill}}(d,q)$$

其中，$L_{\mathcal{G}}$是生成器的损失，$L_{\text{Distill}}$是监督学习的损失。

实验与结果

实验设置

为了评估KOALA的效率，我们选择了Medusa和EAGLE作为非自回归和自回归草稿头的代表，使用Vicuna模型（7B，13B，33B）作为目标LLM。实验结果显示，KOALA在推理延迟的加速比率上实现了0.24x至0.41x的提升，速度比原始草稿头快10.57%至14.09%。

主要结果

实验表明，KOALA优化后，草稿头的接受率显著提高，生成的标记数量增加，推理效率得到了提升。具体而言，草稿头生成的标记平均长度增加了0.26至0.45，减少了推测解码算法所需的迭代次数。

消融研究

在消融研究中，我们探讨了多层结构的优势。结果显示，多层架构相较于单层结构在平均接受长度和速度提升比率上均表现出显著改进。这表明，多层架构使草稿头能更好地模拟目标LLM的功能。

结论

KOALA通过将传统的单层草稿头转变为多层结构，并结合对抗学习的训练方法，显著提升了推测解码的效率。尽管草稿负担有所增加，但这一创新方法有效缩小了草稿头与目标LLM之间的性能差距，从而加速了LLM的推理过程。

参考文献

1. Achiam et al. (2023). GPT-4 technical report. arXiv preprint arXiv:2303.08774.
2. Ainslie et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. arXiv preprint arXiv:2305.13245.
3. Anil et al. (2023). PaLM 2 Technical Report. arXiv preprint arXiv:2305.10403.
4. Cai et al. (2024). Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads. In Forty-first International Conference on Machine Learning.
5. Zhou et al. (2023). Distillspec: Improving Speculative Decoding via Knowledge Distillation. arXiv preprint arXiv:2310.08461.