计算自注意力时,Q(查询)向量在每次解码步骤中都是全新的,而 K(键)和 V(值)向量大部分是重复的,缓存 K 和 V 可以避免大量的重复计算。
下面通过一个具体的例子来详细解释为什么。
自注意力机制
在解码(生成)过程中,对于每一个新生成的 token,都有:
-
Q (Query):来自当前新生成的 token。它负责去“询问”历史上下文中哪些部分是重要的。
-
K (Key):来自所有已生成的 token(包括当前新的这个)。它作为“标识”,被 Q 查询。
-
V (Value):来自所有已生成的 token(包括当前新的这个)。它包含每个 token 的实际信息内容。
注意力分数的计算方式是:Softmax( (Q * K^T) / sqrt(d_k) ) * V
具体例子
假设我们正在生成句子 “I love you”。
第 1 步:生成 “I”
-
输入:
[BEGIN]
(假设起始符) -
输出:
"I"
-
此时,我们计算出了
Q1, K1, V1
(对应 token “I”)。 -
KV Cache 状态:缓存
K1, V1
。
第 2 步:生成 “love”
-
输入:
[BEGIN], I
-
输出:
"love"
-
在这个步骤中:
-
我们只需要为新 token “love” 计算
Q2_new
。 -
K 和 V 从哪里来?
-
一部分来自新 token “love” 自己计算出的
K2_new, V2_new
。 -
另一部分来自 KV Cache 中读取的上一步的
K1, V1
。
-
-
-
所以,计算注意力时,完整的 K 和 V 是:
-
K = [K1, K2_new]
(拼接操作) -
V = [V1, V2_new]
-
-
Q 是什么? 只有
Q2_new
。 -
计算:
Attention = Softmax( (Q2_new * K^T) / sqrt(d_k) ) * V
-
更新 KV Cache:将新的
K2_new, V2_new
拼接到缓存中,现在缓存里有[K1, K2_new]
和[V1, V2_new]
。
第 3 步:生成 “you”
-
输入:
[BEGIN], I, love
-
输出:
"you"
-
在这个步骤中:
-
我们只需要为新 token “you” 计算
Q3_new
。 -
K 和 V 从哪里来?
-
一部分来自新 token “you” 自己计算出的
K3_new, V3_new
。 -
绝大部分来自 KV Cache 中读取的上一步的
[K1, K2_new], [V1, V2_new]
。
-
-
-
所以,计算注意力时,完整的 K 和 V 是:
-
K = [K1, K2_new, K3_new]
-
V = [V1, V2_new, V3_new]
-
-
Q 是什么? 只有
Q3_new
。 -
计算:
Attention = Softmax( (Q3_new * K^T) / sqrt(d_k) ) * V
-
更新 KV Cache:将新的
K3_new, V3_new
拼接到缓存中。
Q 不进行缓存
从上面的过程可以清晰地看出:
-
Q 的生命周期很短:每个解码步骤中,当前的 Q 向量只在这个步骤中被使用一次,用于和所有已缓存的 K 向量计算注意力分数。一旦这个步骤结束,这个 Q 向量就完成了它的使命,在后续的步骤中再也用不到了。
-
K 和 V 是累积的:历史所有步骤的 K 和 V 向量在后续的每一个解码步骤中都会被重复使用。如果不缓存,我们在生成第 N 个 token 时,就需要为前 N-1 个 token 重新计算它们的 K 和 V,这造成了巨大的计算浪费。
向量 | 是否缓存 | 原因 |
---|---|---|
K (Key) | 是 | 历史所有 token 的 K 向量在后续每一步都会被新的 Q 查询。不缓存会导致大量重复计算。 |
V (Value) | 是 | 历史所有 token 的 V 向量在后续每一步的注意力加权求和中都会被使用。 |
Q (Query) | 否 | 每个 token 的 Q 向量仅在当前解码步骤中使用一次,之后便失效。缓存它只会浪费显存,毫无益处。 |