前言
新的一个系列，讲讲AI的基础知识，主要是给自己看的。至于标题只能说懂得都懂，去TMD的AI。但话这么说，AI又必须得学习，相恨相杀了属于是。
第一篇打算讲讲Transformer架构中的很重要的概念，QKV，以及对应的优化KV Cache。当然，正文是AI写的（服了）。
正文
什么是QKV?
QKV是Transformer架构中最核心的概念之一,它代表着三个重要的角色:
- Q (Query): 查询向量 - “我在找什么?”
- K (Key): 键向量 - “我有什么特征?”
- V (Value): 值向量 - “我实际的内容是什么?”
一个生活中的例子
想象你在图书馆找书:
- Query(Q): 你的搜索需求(比如”我想找人工智能的书”)
- Key(K): 每本书的书名和标签
- Value(V): 书的实际内容
当你说”我想找AI的书”(Query),系统会把你的需求和每本书的标签(Key)做匹配,匹配度越高,那本书的内容(Value)对你的影响就越大。
QKV在模型中是如何工作的?
在大语言模型中,每个词都会被转换成这三个向量。让我们看一个简单的架构图:
graph TB
A[输入文本] --> B[词嵌入]
B --> C[分割成Q/K/V]
C --> D1[Query向量]
C --> D2[Key向量]
C --> D3[Value向量]
D1 --> E[计算注意力分数]
D2 --> E
E --> F[Softmax归一化]
F --> G[加权求和]
D3 --> G
G --> H[输出结果]
核心流程:
- 输入文本经过Embedding层转换
- 通过三个不同的线性变换,生成Q、K、V
- Q和K做点积,计算注意力分数(Q·K^T)
- 经过Softmax归一化
- 用这些分数对V加权求和,得到最终输出
数学表达式:
$$
\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V
$$
多头注意力:让模型”多管齐下”
单一的QKV就像只从一个角度理解问题,而大模型使用的是多头注意力(Multi-Head Attention):
graph TB
A[输入] --> B1[Head 1: Q1/K1/V1]
A --> B2[Head 2: Q2/K2/V2]
A --> B3[Head 3: Q3/K3/V3]
B1 --> C1[注意力输出1]
B2 --> C2[注意力输出2]
B3 --> C3[注意力输出3]
C1 --> D[拼接]
C2 --> D
C3 --> D
D --> E[线性变换]
E --> F[最终输出]
为什么要多头?
- Head 1: 可能关注语法关系
- Head 2: 可能关注语义相似性
- Head 3: 可能关注实体关系
就像从不同专家的角度分析问题,最后综合判断,效果自然更好。
多头注意力是从什么角度去计算的
Multi-Head Attention 的核心思想是从不同的子空间（subspace）并行地计算注意力，每个 Head 关注不同的特征维度。
具体计算方式：
- 假设原始维度是 512，如果有 8 个 Head
- 每个 Head 处理 512/8 = 64 维
- 每个 Head 有自己独立的 Q、K、V 权重矩阵
为什么要这样设计？
KV Cache:大模型的”记忆加速器”
为什么需要KV Cache?
大模型生成文本是一个字一个字输出的。假设你要生成”人工智能是未来的发展方向”:
1 | 第1步: "人工" → 预测下一个字 |
问题来了: 每次生成新字时,模型都要重新计算前面所有字的K和V。这就像每次考试都要从小学重新学起,太浪费了!
KV Cache的解决方案
核心思想: 把已经计算过的K和V缓存起来,新来的字只需要计算自己的K和V,然后追加到缓存中。
graph LR
A[第1轮计算] --> B[计算所有Token的K/V]
B --> C[存入KV Cache]
C --> D[第2轮计算]
D --> E[只计算新Token的K/V]
E --> F[追加到Cache]
F --> G[使用完整Cache做Attention]
KV Cache工作原理
假设生成序列 “今天天气真好”:
没有KV Cache:
1 | 步骤1: 计算 [今] 的 Q/K/V |
有KV Cache:
1 | 步骤1: 计算 [今] 的 Q/K/V → 缓存K/V |
性能提升有多大?
对于长度为 $n$ 的序列:
- 无Cache: 计算复杂度是 $O(n^2)$
- 有Cache: 计算复杂度降低到 $O(n)$
实际效果:
- 推理速度提升: 2-4倍
- 内存占用: 增加(需要存储KV Cache),但换来速度大幅提升
- 适用场景: 特别适合长文本生成
KV Cache的架构图
graph TB
A[新Token输入] --> B[计算新Token的K/V]
B --> C[KV Cache]
D[历史Token的K/V] --> C
C --> E[拼接完整K/V]
A --> F[计算新Token的Q]
F --> G[Attention计算]
E --> G
G --> H[输出预测]
深入理解:QKV的数学本质
让我们用更直观的方式理解QKV的计算过程:
1. Q和K的点积:计算相似度
$$
\text{Score}_{ij} = Q_i \cdot K_j
$$
这表示第 $i$ 个位置的Query和第 $j$ 个位置的Key的匹配程度。
2. 缩放操作:防止数值过大
$$
\text{Scaled Score}_{ij} = \frac{Q_i \cdot K_j}{\sqrt{d_k}}
$$
除以 $\sqrt{d_k}$ 是为了防止点积结果过大,导致softmax梯度消失。
3. Softmax:转换成概率分布
这样所有权重加起来等于1,可以理解为”关注程度”的分布。
4. 加权求和:得到最终表示
KV Cache的实际应用
在GPT中的应用
1 | # 伪代码示例 |
内存优化技巧
KV Cache虽然快,但占用内存。常见的优化方法:
- PagedAttention: 像操作系统管理内存一样,分页管理KV Cache
- KV Cache量化: 用INT8或FP4存储,减少内存
- 滑动窗口: 只保留最近的KV Cache
graph LR
A[PagedAttention] --> B[将KV Cache分页]
B --> C[按需加载]
C --> D[提高内存利用率]
E[量化] --> F[FP16→INT8]
F --> G[内存减少50%]
H[滑动窗口] --> I[只保留最近N个Token]
I --> J[内存使用恒定]
总结:一图看懂全貌
让我们用一个完整的流程图把所有概念串起来:
graph TB
A[输入序列] --> B[Embedding层]
B --> C[多头注意力]
subgraph "每个Attention Head"
C --> D1[生成Q]
C --> D2[生成K]
C --> D3[生成V]
D1 --> E[Q·K^T计算相似度]
D2 --> E
E --> F[Softmax归一化]
F --> G[加权V得到输出]
D3 --> G
end
G --> H[拼接所有Head]
H --> I[线性变换]
subgraph "KV Cache机制"
J[存储历史K/V] --> K[新Token只计算自己的K/V]
K --> L[追加到Cache]
L --> M[使用完整Cache加速计算]
end
核心要点回顾
| 概念 | 作用 | 比喻 |
|---|---|---|
| Q (Query) | 表示”我在找什么” | 你的搜索关键词 |
| K (Key) | 表示”我有什么特征” | 书籍的标签 |
| V (Value) | 表示”我的实际内容” | 书籍的内容 |
| Attention | 计算关联度并加权求和 | 根据匹配度决定参考哪些书 |
| Multi-Head | 多个角度同时分析 | 多个专家会诊 |
| KV Cache | 缓存历史K/V,加速推理 | 记住学过的知识,不用重学 |
碎碎念
我个人也刷了很多视频还有文章学习这个概念，今天正文感觉有点知识流经大脑的感觉了。很多概念还是很模糊，尤其牵涉到数学。让我多看几遍。
PS.顺带还把博客支持了流程图还有公式的渲染。