在人工智能的广阔天地中，大语言模型（LLM）以其非凡的生成能力，成为探索未知语言世界的先锋。想象一下，当你向一个智能体提问，它能够逐字逐句地为你编排答案，这背后隐藏着怎样的秘密？本文将带你走进 LLM 的自回归解码世界，揭示其推理过程中的两个关键阶段：预填和解码。让我们一同揭开这层神秘的面纱，探索智能生成的奥秘。
仅解码器大模型
如果使用过 ChatGPT 或其他大型语言模型（LLM），可能会注意到模型在逐字输出其响应之前存在短暂延迟。这种行为与传统语言模型有很大不同，LLM 只有一个编码，但有多个解码。向 GPT 提交问题后的初始延迟是编码时间（预填阶段），而逐字生成答案的时间则是解码时间（解码阶段）。
最流行的仅解码器大模型，例如 Llama 3 等，预先训练以实现因果建模的目标。它们本质上充当下一个词的预测器。这些 LLM 以一系列标记作为输入，并以自回归方式生成后续标记。生成过程会持续进行，直到满足停止条件，例如达到生成标记数的限制或遇到预定义列表中的停止词。或者，生成一个特殊的 <end> 标记以指示生成结束。
本文主要内容是关于大模型自回归解码的两个主要阶段推理：
- 预填阶段（Prefill stage）：此阶段采用提示序列，并为 LLM 的每个 Transformer 层生成 KV 缓存（KV cache）。
- 解码阶段（Decoding stage）：此阶段利用并更新 KV 缓存逐步生成 token，每个令牌的生成取决于先前生成的 token。
使用带有 KV 缓存的 LLM 生成推断的过程如下图所示：
预填阶段：处理提示
在预填阶段，如上图所示，LLM 处理输入 token 以计算用于生成 “第一个” 新 token 的中间状态（K 和 V）。换句话说，该阶段计算和缓存了每一层的密钥和值，而其他阶段不需要缓存，这个缓存被称为 KV 缓存，在后续解码过程中起着关键作用。
下面是这个阶段的正式描述（只考虑单个头的情况），包括参数、权重矩阵和中间变量如下：
图中列出了几个与Transformer模型相关的参数和变量的定义，包括：
- b：批次大小
- s：输入序列长度
- n：输出序列长度
- h_1：Transformer 隐藏层的维度
- h_2：第二个全连接网络（FFN）的隐藏维度
- l：Transformer 层数
- W_Q, W_K, W_V, W_O ∈ ℝ^{h_1×h_1}：查询、键、值和输出的权重矩阵
- W_1 ∈ ℝ^{h_1×h_2}：第一个 FFN 的权重
- W_2 ∈ ℝ^{h_2×h_1}：第二个 FFN 的权重
- x_i ∈ ℝ^{b×s×h_1}：第 i 层的输入
- Q, K, V, Out^i ∈ ℝ^{b×s×h_1}：注意力层中的键、值、查询和输出。
第 i 层的缓存键和值可以通过以下方式计算：
第 i 层中的剩余计算如下：
由于已知输入的整个范围，在高层次上，这是一个高度并行化的 “矩阵-矩阵” 运算。它有效地利用了 GPU 的效率。
解码阶段：生成输出
在解码阶段，如上图所示，LLM 逐个生成输出 token，这个阶段是顺序的。它利用先前生成的 token 来生成下一个 token，直到满足停止条件。每个顺序输出的 token 需要了解所有先前迭代的输出状态（K 和 V）。
为了避免在每次迭代中重新计算所有标记的 K 和 V ，大多数实现都将这些值存储在 KV 缓存中。
在解码阶段，第 i 层中当前生成的 token 的嵌入是：
在这个阶段，需要做两件事情：
- 更新 KV 缓存：
- 计算当前层的输出：
从上述过程可以看出，解码阶段执行与预填充阶段类似的操作。但是，解码阶段的输入张量形状 [b, 1, h1] 和预填充阶段的输入张量形状 [b, s, h1] 存在差异。这可以类比为 “矩阵-向量” 运算。
与预填阶段不同，解码阶段并未充分利用 GPU 的计算能力。
结构优势
将其分为两个阶段的好处是显而易见的。预填阶段只需填充一次键值缓存，而解码阶段只涉及更新和查找键值缓存。这种方法有助于防止不必要的计算：
- 减少重复计算：预填阶段只需要一次计算并缓存所有层的键和值，这避免了在解码阶段对每个新生成的令牌重复进行相同的计算，从而提高了整体的计算效率。
- 提高并行性：预填阶段的计算是高度并行化的，可以充分利用 GPU 的并行处理能力，这在处理大量数据时尤为重要，因为它可以显著加快处理速度。
- 优化资源使用：通过将计算和缓存分开，系统可以更有效地管理内存和计算资源，减少在解码阶段的内存占用和计算负担。
- 提升响应速度：预填阶段完成后，解码阶段可以快速生成输出，因为所需的键值对已经预先计算并存储在缓存中，这减少了生成每个令牌所需的时间。
- 增强灵活性：这种分阶段的方法提供了更大的灵活性，允许在预填阶段进行优化，例如调整缓存大小或使用不同的算法来提高性能，而不会影响解码阶段。
- 支持动态调整：在解码阶段，可以根据生成的上下文动态调整生成策略，例如通过改变停止条件或使用不同的停止词列表，这为生成过程提供了更多的控制。
- 提高可扩展性：通过将生成过程分为两个阶段，系统更容易扩展以处理更大规模的语言模型或更长的文本序列，因为可以独立地优化每个阶段。
结语
通过本文的深入剖析，我们得以窥见大型语言模型在自回归解码过程中的精妙设计。预填阶段的高效缓存与解码阶段的逐步生成，共同构成了 LLM 推理的双翼。这不仅是对技术深度的一次探索，更是对人工智能如何更智能地服务于人类的一次思考。随着技术的不断进步，我们有理由相信，LLM 将在未来的智能世界中扮演更加重要的角色，为我们带来更多惊喜与可能。
