大语言模型优化（Large Language Models Optimization） #

为什么 Decoder-only 模型推理时需要 KV Cache？如何优化其内存占用？

Note：在推理阶段，所有的 权重（Weights）已经通过训练学习完毕（不再改变）。输入的句子会按照预定义的规则转换成 Query (Q)、Key (K) 和 Value (V) 向量。无论是训练阶段还是推理阶段，句子中的 每个 token（词）都会有一个唯一对应的 Key 和 Value 向量。在自回归生成过程中，每次输入新的时间步 x_t ，会计算出当前时刻的 Query 向量 Q_t ，同时，新的 Key 和 Value 向量 K_t, V_t 会与之前的 KV 缓存合并，形成当前时刻的完整历史信息。新的 Query 用于查询当前输入和历史信息之间的关系，从而生成有意义的上下文信息，用于生成下一个 token。

在 Decoder-only 模型（如 GPT）中，推理时需要 KV Cache（Key-Value Cache）来提高计算效率和减少内存占用。KV Cache 主要用于存储模型在每个时间步生成的 Key 和 Value 向量，这些向量用于自注意力机制（Self-Attention）中计算当前词与之前所有词之间的依赖关系。具体来说，当模型生成一个新的 token 时，它不仅要计算当前 token 的自注意力，还需要利用已经生成的 tokens 的表示来计算新的 token，因此必须保存这些 Key 和 Value 向量。

Note：假设我们正在进行推理，模型已经生成了序列 “I love deep learning” 中的前四个词 “I love deep learning”（也就是说，当前时间步是第5个词）。在传统推理中，如果没有 KV Cache，模型在计算每个新的 token 时都需要重新计算与之前所有已经生成的 tokens（“I love deep learning”）之间的依赖关系。

• 对于第5个词 “model”，模型首先计算 “model” 和 “I”、“love”、“deep”、“learning” 之间的自注意力（self-attention）。为了计算这个自注意力，模型需要 重新计算每个之前词的 Key 和 Value 向量。

为了避免这种重复计算，使用 KV Cache 的方法是：当我们生成第5个词时，模型会保存 前四个词（“I love deep learning”）的 Key 和 Value 向量。下一次生成新 token 时（比如第6个词），模型只需要利用 缓存中的 Key 和 Value 向量 来计算当前 token 和已经生成的历史 tokens 之间的依赖关系，而无需重新计算历史 tokens 的表示。

• 对于第5个词 “model”，模型首先计算 “model” 和缓存中的 “I love deep learning” 之间的自注意力。
• 在此过程中，模型使用的是 已经缓存的 Key 和 Value 向量，而不是重新计算整个输入序列的 Key 和 Value 向量。
• 当模型生成第6个词时，只需将第5个词 “model” 的 Key 和 Value 向量加入缓存，并计算与缓存中所有其他 tokens 之间的关系。

在传统的推理过程中，模型需要重新计算每个时间步的所有 Key 和 Value 向量，导致计算量和内存占用急剧增加。使用 KV Cache 后，模型只需要保存每一层的 Key 和 Value 向量，从而避免了重复计算，极大地提升了推理效率。

如何优化内存占用：

1. 动态 KV 缓存大小：在一些任务中，并不需要保留所有时间步的 Key 和 Value 向量。例如，对于生成式任务，缓存可以按照一定步长进行清理，或者只保留 前 n 步 的缓存。
2. 分层缓存：根据模型层数和层间依赖，可以在 不同层 采用不同的缓存策略。例如，可以对较低层进行更频繁的缓存清理，对较高层保留更多的缓存信息。
3. 量化（Quantization）：通过降低 Key 和 Value 向量的精度（例如从浮点数精度到低精度存储），减少内存占用，同时尽量保持推理的精度。

Decoder-only 模型如何处理长文本依赖问题？（如稀疏注意力、窗口注意力）

Decoder-only 模型（如 GPT 类模型）通过不同的技术来处理长文本中的依赖问题，尤其是在处理长序列时，传统的 全局注意力（Global Attention） 计算会变得非常消耗资源。为了解决这个问题，Decoder-only 模型采用了 稀疏注意力（Sparse Attention） 和 窗口注意力（Windowed Attention） 等方法，从而有效地减小计算复杂度并增强长文本的建模能力。

• 窗口注意力（Windowed Attention） 是一种将输入序列划分为多个固定大小的窗口（或块），每个窗口内的 token 之间通过注意力进行交互，而窗口之间没有直接的依赖关系。窗口大小是一个超参数，通常会选择较小的窗口以限制每次计算的注意力范围，从而减少计算负担。通过这种方式，模型能够在较低的计算成本下捕捉到长序列中的重要信息，同时避免了全局注意力带来的高昂计算开销。

e.g.：假设我们有一个长度为 6 的序列 [A, B, C, D, E, F]，并且我们选择一个大小为 3 的窗口进行计算。

1. 在窗口注意力中，我们将输入序列分为若干个滑动窗口。例如，窗口大小为 3 的情况下：
   • 第一个窗口：[A, B, C]
   • 第二个窗口：[B, C, D] 每个窗口内部的 token 之间会进行注意力计算，但是不同窗口之间的 token 之间是没有交互的。
2. 对于序列 [A, B, C, D, E, F]，采用窗口大小为 3 的策略，注意力计算矩阵将是：

\[ \begin{array}{|c|c|c|c|c|c|c|} \hline & A & B & C & D & E & F \\ \hline A & 1 & 1 & 1 & 0 & 0 & 0 \\ B & 1 & 1 & 1 & 1 & 0 & 0 \\ C & 1 & 1 & 1 & 1 & 1 & 0 \\ D & 0 & 1 & 1 & 1 & 1 & 1 \\ E & 0 & 0 & 1 & 1 & 1 & 1 \\ F & 0 & 0 & 0 & 1 & 1 & 1 \\ \hline \end{array} \]

• 稀疏注意力（Sparse Attention） 的 核心思想是通过引入局部化注意力机制，使得每个 token 只与部分上下文进行交互，从而减少计算量。具体而言，稀疏注意力只计算一部分的注意力权重而不是全部，这样可以降低模型计算的复杂度。常见的稀疏注意力结构包括 固定模式（Fixed Patterns） 和 学习模式（Learned Patterns），其中一个代表固定的局部上下文窗口，另一个则依赖于模型在训练过程中自适应学习关注哪些位置的关系。稀疏注意力通常通过 Top-k 注意力（Top-k Attention） 或 Block-sparse 格式 来实现。

Note: 稀疏注意力和窗口注意力是非常相似的概念，都属于通过减少计算量来优化自注意力机制的方法。但稀疏注意力不仅仅局限于相邻的 token，还可以是基于某些策略（如全局选择、局部窗口、随机选择等）来选择哪些 token 进行注意力计算。它可以通过灵活的方式选择注意力的稀疏性，可以是全局性策略（如某些重要的 token）或局部性策略（如基于输入特征选择 token）。