提示工程

提示原则

原则一:编写清晰、具体的指令

  • 使用分隔符实现不同文本部分的隔断,比如````,""",< >, `等,可以防止提示词注入
  • 结构化输出,json形式
  • 要求模型检查是否满足条件
  • 提供少量示例

原则二:给模型时间去思考

  • 指定完成任务所需的步骤
  • 指导模型在下结论之前找出一个自己的解法

直接给模型输入一份解答,模型可能默认会判断输入是对的,需要让它自己思考一次再对比

提示词设计

语言类prompt检查是否满足四要素:

  • 立角色
  • 述问题
  • 定目标
  • 补要求

统一模板:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
你是一名具有【领域】经验的【角色】

## 任务背景
...

## 任务说明
子任务1:...
子任务2:...
...

## 任务示例

## 输入[不一定都有输入]
<INPUT>
{{输入内容}}
</INPUT>

## 任务约束
1. 【具体要求1】
2. 【具体要求2】
3. 【具体要求3】
4. 不得添加输入中不存在的事实
5. 信息不足时明确指出,不要自行假设

## 输出结构
请严格按照以下结构输出:
一般以json形式输出,且要规定好输出的json结构

最重要的是:任务、约束、输出格式

提示词优化方法:

图像类prompt一般不止包含正向提示词,还需要负向提示词来限制画面输出,正向提示词包含:

  • 主体
  • 细节
  • 风格参考

多模态prompt需要明确模态分工,减少不同模态之间的对应关系歧义

  • 图片:核心的语义负载,表达“是什么“
  • 文字:权重调节器和语义补全器,表达“要什么”

设计原则:明确检索意图、锚定视觉焦点、属性调节校准

上下文学习

  • Zero-shot(零样本提示):只告诉模型任务和要求,不给示例
  • Few-shot(少样本提示):除任务要求外,再提供少量“输入—输出”示例,让模型模仿其中的规则、标签含义和输出格式

GPT-3 论文系统比较了 zero-shot、one-shot 和 few-shot 三种设置

Language Models are Few-Shot Learners

一般来说先零样本测试,输出不稳定或错误较多时,再添加代表性示例,构成小样本提示,但是不能无限增加示例

Responses API

最新版不再推荐使用 system

1
2
3
developer = 规则和业务逻辑
user = 用户问题、动态数据、检索上下文
assistant = 模型历史回答或示例答案

在 Responses API 中有几类数据

1
2
3
developer / user / assistant -> 普通对话消息 Message
function_call -> 模型生成的工具调用项
function_call_output -> 程序返回给模型的工具执行结果

标准流程是:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
用户输入

第一次调用 LLM
├─ 判断是否需要工具
├─ 选择工具
└─ 抽取工具参数

程序执行工具

返回 function_call_output

第二次调用 LLM

生成最终回答

函数定义:放在 tools

1
2
3
4
5
6
7
8
tools = [
{
"type": "function",
"name": "get_weather",
"description": "查询天气",
"parameters": {...},
}
]

函数调用:放在 response.output,模型返回样式:

1
2
3
4
5
6
{
"type": "function_call",
"call_id": "call_abc123",
"name": "get_weather",
"arguments": '{"city":"北京"}',
}

模型只是生成了调用请求,并没有执行 Python 函数,程序必须读取 namearguments,然后自己调用对应函数

函数执行结果:放在 function_call_output

1
2
3
4
5
{
"type": "function_call_output",
"call_id": "call_abc123",
"output": '{"temperature":26}',
}

完整结构可以理解为:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
input_items = [
# 普通消息
{
"role": "developer",
"content": "系统规则",
},
{
"role": "user",
"content": "用户问题",
},

# 模型上一次产生的工具调用
{
"type": "function_call",
"call_id": "call_123",
"name": "get_weather",
"arguments": '{"city":"北京"}',
},

# 你的程序提供的工具结果
{
"type": "function_call_output",
"call_id": "call_123",
"output": '{"temperature":26}',
},
]

最后,模型才会生成普通的 assistant 消息

1
2
3
4
{
"role": "assistant",
"content": "北京当前气温为 26℃。",
}

定义区分:

  • Function Calling 是模型与外部工具交互的接口机制
  • RAG 是利用外部知识增强生成结果的信息处理架构
  • Agent 是由模型负责决策和控制任务执行过程的系统架构

在 RAG 和 Agent 中都可以用到 Function Calling

区分:

  • RAG 的典型流程:用户问题 → 检索文档 → 拼接上下文 → 生成回答,不一定利用Function Calling
  • Agent 的典型流程:目标 → 决策 → 调工具 → 观察结果 → 再决策,根据目标和中间结果,动态选择步骤和工具

Tokenizer

主要有以下几类切分方法

  • 词级切分(Word-level)
  • 字符级切分(Character-level)
  • 子词切分(Subword-level)
  • 字节级切分(Byte-level)

目前大多数大语言模型采用子词切分,部分模型使用以字节为基础的子词算法,例如 byte-level BPE;

词级切分

将文本按空格和标点符号分割成单词

英文中可以近似按照空格和标点切分,但中文、日文等语言通常需要额外的分词算法

  • 优点在于简单直接;
  • 缺点很明显,词表通常非常大,单词形态变化会占用不同词表项,遇到词表外单词时只能映射为 [UNK]
1
2
I love machine learning
["I", "love", "machine", "learning"]

现在的大语言模型很少使用纯词级 Tokenizer

字符级切分

将每个字符作为一个 Token,能非常精细地处理文本,适用于处理拼写错误、未登录词或新词

适用于对拼写、字符结构、转写和噪声文本比较敏感的任务,例如拼写纠错、字符转写、语音文本转换以及部分多语言任务

  • 优点在于词表会非常小,几乎不会出现未知的token,对拼写错误、罕见单词比较鲁棒;
  • 缺点是 token 序列变得非常长,增加了模型的计算复杂度和训练时间;
1
2
hello
["h", "e", "l", "l", "o"]

子词切分

子词方法介于词级和字符级之间

  • 高频词或高频片段保留为一个 Token
  • 低频词拆成多个较小的 Token

常见的子词分词方法包括 BPE、WordPiece 和 Unigram

BPE WordPiece Unigram
初始状态 字符或字节等基础符号 字符等基础符号 较大的候选子词词表
训练方向 自底向上合并 自底向上构建 自顶向下裁剪
训练方法 合并最高频的相邻符号对 选择使训练数据似然提升较大的子词单元 逐步删除对整体似然贡献较小的子词
切分方式 按学习到的合并规则切分 贪心最长匹配优先 动态规划寻找概率较大的整体切分
主要特点 简单、高效、切分稳定 切分阶段简单,常配合预分词 概率模型更完整,可进行子词采样
典型应用 GPT、RoBERTa BERT SentencePiece、T5

分布式训练框架

预训练一个 LLM 所需要的算力资源极其庞大,需要分布式训练

分布式训练的两条基本思路:

  • 数据并行:模型复制,数据切分
  • 模型并行:数据可以相同或分成 micro-batch,模型本身被切分

数据并行

主要就是Pytorch 的 DDP

每张 GPU:

  1. 保存一份完整模型;
  2. 读取不同的数据;
  3. 独立进行前向和反向传播;
  4. 通过 AllReduce 聚合梯度;
  5. 获得相同梯度后,各自更新模型。
1
2
3
4
GPU 0: batch 0 → gradient G0 ─┐
GPU 1: batch 1 → gradient G1 ─┤
GPU 2: batch 2 → gradient G2 ─┼→ AllReduce → 平均梯度(g1,g2,g3,g4) → 各自更新
GPU 3: batch 3 → gradient G3 ─┘
$$ g=\frac{1}{P}\sum_{p=1}^{P}G_p $$

前提是:完整模型及其训练状态能放入单卡

模型并行

当数据并行分片仍然不够,或者单个层本身太大,就需要模型并行

张量并行 TP:将一个矩阵运算拆到多张 GPU

比如:

  • MLP层:将权重矩阵 W 按列切分为 [W1,W2] ,分别放置于不同GPU
  • Attention层:将多头注意力切分到不同设备

流水线并行 PP:按照模型层切分

1
2
3
4
5
6
7
输入


GPU 0 GPU 1 GPU 2 GPU 3
Layer 0~7 ---> Layer 8~15 ---> Layer 16~23 ---> Layer 24~31
│ │
└────────────────────── 反向传播 <────────────────────┘

PP 的主要问题是 pipeline bubble,流水线气泡

1
2
3
4
5
6
时间      0    1    2    3    4    5    6    7
---------------------------------------
GPU 0 F -- -- -- -- -- -- B
GPU 1 -- F -- -- -- -- B --
GPU 2 -- -- F -- -- B -- --
GPU 3 -- -- -- F B -- -- --
1
2
3
F  = 整个 batch 的前向传播
B = 整个 batch 的反向传播
-- = GPU 空闲

大部分时间只有一张 GPU 在工作,设备利用率很低

把一个 batch 拆成多个 micro-batch,使不同流水线阶段同时处理不同 micro-batch

1
2
3
4
5
6
时间      0    1    2    3    4    5    6    7    8    9   10   11   12   13
---------------------------------------------------------------------
GPU 0 F0 F1 F2 F3 -- -- -- -- -- -- B3 B2 B1 B0
GPU 1 -- F0 F1 F2 F3 -- -- -- -- B3 B2 B1 B0 --
GPU 2 -- -- F0 F1 F2 F3 -- -- B3 B2 B1 B0 -- --
GPU 3 -- -- -- F0 F1 F2 F3 B3 B2 B1 B0 -- -- --

增加 micro-batch 数量可以摊薄填充和排空气泡,但也会导致激活显存、调度复杂度和优化器更新延迟增加

序列/上下文并行

对于长上下文模型,激活和注意力矩阵可能成为显存瓶颈,因此可以沿序列维度切分

1
2
3
4
5
6
7
8
9
完整序列:
[token 0 ... token 32767]

GPU 0:
[token 0 ... token 8191]

GPU 1:
[token 8192 ... token 16383]
...

但注意力需要 token 之间相互交互,因此会引入 AllGatherAllToAll 或其他通信

  • AllGather:每张 GPU 把自己的数据分享给所有 GPU,最终所有 GPU 都拥有完整数据
  • AllToAll:每张 GPU 把不同的数据块分别发送给不同 GPU,最终数据在 GPU 之间重新分布

专家并行 EP

MoE的专家通常不在一张 GPU 上,Token 根据路由结果被发送到对应 GPU

1
2
3
4
GPU 0: Experts 0, 1
GPU 1: Experts 2, 3
GPU 2: Experts 4, 5
GPU 3: Experts 6, 7

token 需要从原 GPU 去往专家 GPU,专家计算完成后又要返回原 GPU,所以 MoE 的 FLOPs 可能较低,但端到端速度不一定按理论比例提升

零冗余优化器(ZeRO)

ZeRO 在数据并行基础上,通过“分片保存模型状态”降低单张 GPU 显存占用

普通 DDP 中,每张 GPU 都保存相同的三类东西

1
2
3
P:模型参数 Parameters
G:梯度 Gradients
O:优化器状态 Optimizer States
1
2
3
4
5
6
7
8
9
GPU 0:
参数 p0 p1 p2 p3 p4 p5 p6 p7
梯度 g0 g1 g2 g3 g4 g5 g6 g7
优化器 O0 O1 O2 O3 O4 O5 O6 O7

GPU 1:
参数 p0 p1 p2 p3 p4 p5 p6 p7
梯度 g0 g1 g2 g3 g4 g5 g6 g7
优化器 O0 O1 O2 O3 O4 O5 O6 O7

ZeRO 将数据并行中的冗余状态逐步分片,分为三种不同的配置

ZeRO-1

1
2
3
4
GPU 0:[完整参数] [完整梯度] [O0 O1]
GPU 1:[完整参数] [完整梯度] [O2 O3]
GPU 2:[完整参数] [完整梯度] [O4 O5]
GPU 3:[完整参数] [完整梯度] [O6 O7]

每张 GPU 只保存四分之一的优化器状态,参数和梯度依然是完整复制的

ZeRO-2

1
2
3
4
GPU 0:[完整参数] [g0 g1] [O0 O1]
GPU 1:[完整参数] [g2 g3] [O2 O3]
GPU 2:[完整参数] [g4 g5] [O4 O5]
GPU 3:[完整参数] [g6 g7] [O6 O7]

这里的g已经是平均后的了

使用 ReduceScatter聚合完成后直接拆开

ZeRO-3

把参数也分片,每张 GPU 平时只保存一部分模型状态

1
2
3
4
GPU 0:[p0 p1] [g0 g1] [O0 O1]
GPU 1:[p2 p3] [g2 g3] [O2 O3]
GPU 2:[p4 p5] [g4 g5] [O4 O5]
GPU 3:[p6 p7] [g6 g7] [O6 O7]

但是计算前需要 AllGather

大模型组件

现代大模型通常会在原始 Transformer 上采用一组工程化改进

组件 作用
RMSNorm 简化归一化计算并稳定激活尺度
RoPE 旋转 Q/K,使注意力分数包含相对位置关系
SwiGLU 使用门控结构增强 FFN 表达能力
MoE 通过稀疏激活扩大模型参数容量

RMSNorm

$$ \operatorname{RMSNorm}(x)=\gamma\odot\frac{x}{\sqrt{\frac{1}{d}\sum_{i=1}^{d}x_i^2+\epsilon}} $$

RMSNorm 可以看作 LayerNorm 的简化版本

相比 LayerNorm,RMSNorm:

  • 减少了均值计算和相关运算;
  • 保留输入的整体方向信息;
  • 对输入的正比例缩放具有不变性;
  • 在许多模型中能以更低计算成本取得与 LayerNorm 相近的效果

RMSNorm 已被 Llama、Mistral、Qwen、Gemma 等许多现代 Decoder-only LLM 广泛采用

旋转位置编码 RoPE

位置编码的目的是让注意力计算感知 token 的顺序和位置关系

旋转位置编码不把位置向量直接加到 token 表示上,而是根据 token 的位置,对 Query 和 Key 的二维子空间进行旋转

假设 Query 中有两个维度,RoPE 根据位置 $m$,将它旋转 $m\theta$

$$ R_{m,\theta}=\begin{bmatrix}\cos(m\theta)&-\sin(m\theta)\\\sin(m\theta)&\cos(m\theta)\end{bmatrix} $$
旋转后的 Query 为 $\widetilde q_m=R_{m,\theta}q_m$

同理,位置为 $n$ 的 Key $\widetilde k_n=R_{n,\theta}k_n$

位置不同,向量朝向不同

因为二维旋转矩阵是一个正交矩阵,RoPE 改变向量的方向和相位,但不改变其欧氏范数

利用重要性质:

$$ R_m^\top R_n=R_{n-m} $$
$$ \tilde{q}_m^{\top}\tilde{k}_n=(R_m q_m)^{\top}(R_n k_n)=q_m^{\top}R_m^{\top}R_n k_n $$

RoPE 之后的注意力内积为:

$$ \color{red}\boxed{\widetilde q_m^\top\widetilde k_n=q_m^\top R_{n-m}k_n} $$
这就是 RoPE 最关键的性质,这种旋转使注意力内积自然包含相对位置 $n-m$

https://arxiv.org/abs/2104.09864

实际的 Query 和 Key 不是二维向量,而是 $d_h$ 维,RoPE将相邻维度两两分组,每一对维度都视为一个二维平面,然后分别旋转

$$ (q_0,q_1),\quad(q_2,q_3),\quad\dots,\quad(q_{d_h-2},q_{d_h-1}) $$

省略了位置下标,其实应该是 $(q_{m,2i},q_{m,2i+1})$

对于第 $i$ 个二维分组,常见频率为

$$ \theta_i=10000^{-\frac{2i}{d_h}},\qquad i=0,1,\dots,\frac{d_h}{2}-1 $$
那么第 $i$ 组变换为
$$ \begin{bmatrix} \widetilde q_{m,2i}\\ \widetilde q_{m,2i+1} \end{bmatrix}= \begin{bmatrix} \cos(m\theta_i)&-\sin(m\theta_i)\\ \sin(m\theta_i)&\cos(m\theta_i) \end{bmatrix} \begin{bmatrix} q_{m,2i}\\ q_{m,2i+1} \end{bmatrix} $$
比较位置时每个维度组产生不同的相位差,模型通过所有维度组的相位组合,判断两个 token 的相对距离

不同频率的作用类似于传统正弦位置编码:

  • 较高角频率的维度组($i$ 小):相邻位置产生较明显的相位变化,提供较细的位置分辨率;
  • 较低角频率的维度组($i$ 大):相位变化较慢,在较长距离内仍能保持不同的变化尺度。

可以类比秒针,分针,时针

优点 局限
注意力分数自然包含相对位置关系 超出训练长度时外推性能下降
不引入额外可学习位置参数 不同位置可能产生相似的组合相位
只需旋转 Q、K,计算简单 长序列下会多次旋转,容易产生相位混叠
保持旋转部分的欧氏范数 更擅长相对位置,不直接强调“第几个位置”
兼容标准点积注意力,可直接用于 MHA、GQA、MQA 等 与 MLA 的低秩矩阵吸收存在结构冲突

有一些主流的训练方法:

  • NTK-aware scaling:不同频率采用不同程度的缩放
  • Dynamic NTK:只在长度超过原窗口时动态缩放
  • LongRoPE:搜索每个维度的非均匀缩放系数,再逐阶段长上下文训练

NTK-aware scaling

假设目标上下文长度超过训练上下文长度

  • 原始训练上下文长度为 $L$
  • 目标上下文长度为 $L'$
  • 上下文扩展倍数为 $s$

最直接的扩展方法是位置插值 PI,把位置 $m$ 压缩为 $m/s$

等价于把所有频率统一缩小,虽然把长度压缩回训练范围,但是所有频率都受到相同程度的压缩

NTK-aware Scaling 的核心思想是:不直接统一缩放位置,而是增大 RoPE 的基数,使高频维度少缩放、低频维度多缩放

将基数 $b$ 修改为

$$ b'=b\cdot s^{\frac{d_h}{d_h-2}} $$
于是
$$ \theta_i'=(b')^{-\frac{2i}{d_h}}=\theta_i\cdot s^{-\frac{2i}{d_h-2}} $$
因此不同维度的缩放程度不同
维度组 原始频率 缩放程度 主要效果
i=0 附近 高频 几乎不缩放 保留局部位置分辨率
中间维度 中频 部分缩放 在局部信息和长距离建模之间折中
i 较大 低频 接近缩放 (s) 倍 适应更长的相对距离

Dynamic NTK

NTK 需要固定缩放比例,而且会影响短长度的性能

Dynamic NTK 的思路是:根据当前实际序列长度动态决定缩放倍数;没有超过原始窗口时不缩放,超过后再逐渐增加缩放,$s$ 动态变化

$$ s(l)=\frac{l}{L} $$
但这个策略会影响 KV cache

LongRoPE

LongRoPE 的核心思想是:不再通过一条固定公式确定所有维度的缩放比例,而是为每个 RoPE 二维分组分别搜索缩放系数

对于第 i 个二维分组,LongRoPE 使用独立的缩放系数 $\lambda_i$

$$ \theta_i'=\frac{\theta_i}{\lambda_i} $$
位置 m 的旋转角为
$$ \alpha_{m,i}=\frac{m\theta_i}{\lambda_i} $$
这里的
$$ \lambda_0,\lambda_1,\dots,\lambda_{\frac{d_h}{2}-1} $$
不再由简单公式生成,需要通过搜索,使模型在长文本上的语言模型损失或困惑度尽可能低

主要利用了两种非均匀性

  • 不同 RoPE 维度需要不同缩放
  • 不同 token 位置需要不同缩放

总结三者

1
2
3
4
5
6
7
8
9
NTK-aware 
└── 缩放多少由维度 i 决定

Dynamic NTK
└── 缩放多少由维度 i 和当前序列长度 l 共同决定

LongRoPE
└── 缩放多少不再使用固定规律,而是为每个维度搜索
同时考虑序列开头 token 与其他位置的差异

DroPE

目前更加关注局部层使用 RoPE,全局层使用 NoPE,为了解决 RoPE 会约束超长距离注意力的问题

NoPE 删除的是显式位置编码,但因果掩码本身仍然是一种依赖位置的结构

因此 decoder-only Transformer 即使没有显式位置编码,也能从因果结构中形成一定的绝对位置和相对顺序信息

2501.00073

NoPE 虽然经常比未经扩展的 RoPE 更容易外推,但它仍然存在有效上下文长度上限,所以现在基本都是混合结构

SwiGLU

传统 Transformer FFN 通常写成:

$$ \operatorname{FFN}(x)=W_2\phi(W_1x+b_1)+b_2 $$
普通 FFN 只有一条升维分支

SwiGLU加入门控,门控 FFN 则使用两条分支

1
2
3
4
5
           +--> ---------内容分支 ----------+
| |
输入 x ----+ * --> 降维 --> 输出
| |
+--> 门控分支 --> 激活函数 -------+

一条分支产生内容,一条分支决定这些内容应该保留、放大还是抑制

标准 GLU 可以写成:

$$ \operatorname{GLU}(x)=(W_v x+b_v)\odot\sigma(W_g x+b_g) $$
SwiGLU 将门控分支中的激活函数换成 SiLU
$$ \operatorname{SwiGLU}(x)=\operatorname{SiLU}(W_g x)\odot(W_u x) $$

混合专家结构 MoE

在普通稠密 Transformer 中,每个 Token 都经过同一个 FFN

导致的问题:

  • 所有 Token 使用相同参数
  • 增大 FFN 参数量时,每个 Token 的计算量也同步增加
  • 总参数量与单 Token 计算量紧密绑定

MoE 的目标是增加大量参数容量,但每个 Token 只激活其中一小部分,通过训练一个 Router,让输入选择少量专家

实现总参数量大幅增长,而单样本计算量不会按同样比例增长

1
2
3
4
5
6
7
8
9
10
11
12
13
14
输入 Token x

Router
|
+--------> Expert 1
+--------> Expert 2
+--------> Expert 3
|
...
|
+--------> Expert N
Router 只选择 Top-K 个专家

专家输出加权求和

一个 Token 会经过:

  1. 始终启用的 Shared Expert,共享专家;(非必须)
  2. 经过 Router 选出的 Routed Expert,路由专家

Router 通常是一个很小的线性层

$$ z=W_r x\qquad W_r\in\mathbb{R}^{E\times d} $$
${z}$ 中包含每个专家的路由分数,再根据具体模型使用 Softmax、Sigmoid 或其他归一化方式选出 Top-k

在 MoE Router 中 Softmax gating 和 Sigmoid gating都有,Softmax更常见一些

第 $i$ 个专家可以写成

$$ E_i(x) = W_{\mathrm{down},i} \left[ \operatorname{SiLU}(W_{\mathrm{gate},i}x) \odot (W_{\mathrm{up},i}x) \right] $$
得到输出
$$ y(x)=\sum_{i\in \operatorname{TopK}(g(x))}p_i(x)E_i(x) $$
Top-1 优点 Top-1 局限
每个 token 只计算一个专家 路由错误缺少其他专家补偿
Dispatch 与聚合较简单 路由决策更离散
通信量通常较低 单专家负载压力更大
激活计算量较低 容易出现专家过载或 token 丢弃

专家分工是训练中自动形成的,并不一定具有稳定、清晰的人类语义标签

主要训练问题:路由不均衡,热门专家过载,冷门专家得不到充分训练;

部分现代 MoE 不只使用路由专家,还增加一个或多个共享专家,减少路由专家间的知识冗余

1
2
3
                    +--> Router --> Routed FFN Experts --+
x --> RMSNorm ------+ +--> 相加
+--> Shared FFN Experts -------------+

忽略残差连接时,MoE 分支输出为:

$$ y(x)=\sum_{j=1}^{N_s}E^{\mathrm{shared}}_j(x)+\sum_{i\in\operatorname{TopK}(g(x))}p_i(x)E^{\mathrm{routed}}_i(x) $$

DeepSeekMoE 将这个设计称为 Shared Expert Isolation,共享专家隔离

注意力

方法 历史信息组织方式 缓存特点 核心目标
MHA 每个 Query 头独立使用 K/V 最大 保留完整多头表达能力
MQA 所有 Query 头共享一组 K/V 最小 最大程度降低 KV Cache
GQA 每组 Query 头共享一组 K/V 中等 折中表达能力与缓存
MLA 每个 token 保存压缩 KV 潜变量 显著小于 MHA 低秩压缩并保留多头表达
线性注意力 全部历史压缩为固定状态 不随序列长度增长 同时降低计算量和缓存量

标准注意力 MHA

标准多头注意力 MHA 对第 t 个 token 计算将其拆成 $n_h$ 个注意力头

$$ q_t=[q_{t,1};\ldots;q_{t,n_h}],\qquad k_t=[k_{t,1};\ldots;k_{t,n_h}],\qquad v_t=[v_{t,1};\ldots;v_{t,n_h}] $$
每个 head 执行
$$ o_{t,i}=\sum_{j=1}^{t}\operatorname{softmax}_j\left(\frac{q_{t,i}^{\top}k_{j,i}}{\sqrt{d_h}}\right)v_{j,i} $$
自回归生成时,每生成一个新 token,都需要让当前 query 与全部历史 key 计算相似度,并读取相应的 value

因此所有 token 的 (K,V) 都必须保留在 KV Cache 中

如果每个 token 有 ($n_h$) 个 head,每个 head 维度为 ($d_h$),上下文长度为 ($T$)、层数为 ($L$) 时,需要缓存

$$ O(2LTn_hd_h) $$
长上下文、大 batch 推理时,KV Cache 会占用大量显存和显存带宽

发展的两条路线:

模型结构路线 历史信息保存方式 代表方法 主要目标
固定状态压缩 全部历史压入固定状态 线性注意力、GLA 同时降低序列计算量和缓存量
逐 token KV 压缩 仍为每个 token 保存独立信息 MQA、GQA、MLA 保留 Softmax 检索并减少 KV Cache

FlashAttention 不属于上述两条模型结构路线,而是与它们正交的计算实现优化,主要降低标准 Softmax Attention 的中间张量存储和显存 IO

FlashAttention

不是一种新的注意力机制,而是 Softmax Attention 的一种高效、精确实现

GPU的存储空间

存储位置 特点
HBM,全局显存 容量大,但访问相对慢
SRAM,共享内存、寄存器 容量小,但访问非常快

普通 Attention 的执行过程大致是

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
从 HBM 读取 Q、K

计算 S = QKᵀ

把完整 S 写入 HBM

从 HBM 重新读取 S

计算 Softmax,得到 P

把完整 P 写入 HBM

从 HBM 读取 P 和 V

计算 O = PV

虽然矩阵乘法本身可以利用 GPU Tensor Core 高效执行,但完整的矩阵需要多次在 HBM 和计算单元之间移动,因此 Attention 往往受到 内存带宽 限制

FlashAttention 的基本思想就是:不要生成并保存完整的 $S$ 和 $P$,在快速的片上存储中分块完成计算

对于一个 Query 块 $Q_i$,依次读取各个 $K_j,V_j$ 块

$$ S_{ij}=\frac{Q_iK_j^\top}{\sqrt d} $$
不会把 $S_{ij}$ 写回显存,而是在 SRAM 中立即进行:
  1. mask
  2. Softmax 局部更新
  3. 与 $V_j$ 相乘
  4. 更新输出块 $O_i$

然后丢弃局部 $S_{ij}$

普通 Softmax 的分母需要已知所有的结果,而且为了数值稳定通常减去最大值

如果每次只能看到一个 Key 块,就不能提前知道全局最大值

FlashAttention 使用在线算法,在读取每个块时动态更新最大值和归一化因子

在线 Softmax

保存两个统计量:

  • 当前最大值 $m$
  • 当前指数和 $\ell$
$$ \ell=\sum_j\exp(s_j-m) $$

现在读取一个新的分块,其局部最大值为:

$$ m_{\text{block}}=\max_j s_j $$
新的全局最大值为:
$$ m_{\text{new}}=\max(m,m_{\text{block}}) $$
由于基准最大值发生变化,之前累积的指数和需要重新缩放
$$ \ell_{\text{new}} = \exp(m-m_{\text{new}})\ell + \sum_{j\in\text{block}} \exp(s_j-m_{\text{new}}) $$
这样就不需要预先知道整行的最大值

同样,输出向量也同时在线更新

$$ u=\sum_j\exp(s_j-m)v_j $$
更新最大值以后
$$ u_{\text{new}} = \exp(m-m_{\text{new}})u + \sum_{j\in\text{block}} \exp(s_j-m_{\text{new}})v_j $$
全部 Key 块处理完成后
$$ o=\frac{u}{\ell} $$
所以对于每个 Query 行,只需要维护:
  • 最大值 $m$
  • 归一化分母 $\ell$
  • 输出累积向量 $u$

但是反向传播时需要计算局部的 $S_{ij}$,虽然增加了一部分重复计算,但是大幅减少了显存读写和中间张量保存,空间复杂度转为 $O(Nd)$

完整注意力矩阵从未同时存在

线性注意力

RNN 视角

普通 RNN 按时间递推隐藏状态:

$$ h_t=F(h_{t-1},x_t) $$
输出为
$$ y_t=G(h_t,x_t) $$
RNN 的核心思想是将截止到时刻的历史信息压缩进固定大小的隐藏状态 $h_t$

传统 RNN 的主要限制:

  • $h_t$ 依赖 $h_{t-1}$,时间维度存在严格的串行依赖
  • 固定大小的隐藏状态需要压缩全部历史,容易形成信息瓶颈
  • 当前输入通常不能像注意力那样,对历史信息进行显式的内容寻址

线性注意力同样可以写成递归状态模型,但它有两个重要区别

  • 隐藏状态通常是一个矩阵,而不是一个向量;
  • 当前 query 可以对矩阵状态进行内容相关的读取

因此训练时可以采用并行或分块并行形式

核函数

标准注意力可以抽象成:

$$ y_t=\frac{\sum_{i=1}^{t}\operatorname{sim}(q_t,k_i)v_i}{\sum_{i=1}^{t}\operatorname{sim}(q_t,k_i)} $$
线性注意力要求相似度函数可以分解为特征映射的内积:
$$ \operatorname{sim}(q,k)=\phi(q)^\top\phi(k),\qquad\phi:\mathbb R^{d_k}\rightarrow\mathbb R^m $$
核心就是将 query-key 相似度表示为有限维特征空间中的内积,从而利用矩阵乘法结合律重新安排计算顺序,这样的核函数是一个有限秩核

记忆混合主要来自有限维状态压缩:多个历史 key 经特征映射后被累加到同一个矩阵状态中

若不同 key 的核特征方向相似,它们写入的 value 会落入相近的状态方向,后续 query 难以将其精确分离

特征映射可以保持维度不变 $\color{red}\phi:\mathbb R^{d_k}\rightarrow\mathbb R^{d_k}$,甚至可以采用恒等映射

常用

$$ \phi(x)=\operatorname{ELU}(x)+1= \begin{cases} x+1,&x>0,\\ e^x,&x\leq 0. \end{cases} $$

如果核函数直接取为 $\kappa(q,k)=\exp\left(\frac{q^\top k}{\sqrt d}\right)$,代入后正好得到 Softmax 注意力

$$ \displaystyle y_t=\frac{\sum_i\kappa(q_t,k_i)v_i}{\sum_i\kappa(q_t,k_i)} $$
但指数点积核包含所有阶数,对应的 $\phi$ 是无限维的,因为要用到泰勒展开,所以只能近似无法等价
$$ \boxed{\exp\left(\frac{q^\top k}{\sqrt d}\right)=\phi_\infty(q)^\top\phi_\infty(k)}\qquad \boxed{\exp\left(\frac{q^\top k}{\sqrt d}\right)\approx\phi_m(q)^\top\phi_m(k)} $$

代数重排

简化符号,定义:

$$ q'_t=\phi(q_t), k'_i=\phi(k_i) \in\mathbb R^m\qquad v_i\in\mathbb R^{d_v} $$
于是
$$ y_t=\frac{\sum_{i=1}^{t}(q'_t{}^\top k'_i)v_i}{\sum_{i=1}^{t}q'_t{}^\top k'_i} $$
根据矩阵乘法,分子可以改写为
$$ \sum_{i=1}^{t}(q'_t{}^\top k'_i)v_i^\top=q'_t{}^\top\left(\sum_{i=1}^{t}k'_iv_i^\top\right) $$
分母同理
$$ \sum_{i=1}^{t}q'_t{}^\top k'_i=q'_t{}^\top\left(\sum_{i=1}^{t}k'_i\right) $$
定义矩阵状态:
$$ S_t=\sum_{i=1}^{t}k'_iv_i^\top\in\mathbb R^{m\times d_v} $$
定义归一化状态:
$$ z_t=\sum_{i=1}^{t}k'_i\in\mathbb R^m $$
最终得到
$$ \boxed{y_t^\top=\frac{q_t^{\prime\top}S_t}{q_t^{\prime\top}z_t+\varepsilon}} $$
这一重排将原本的计算方式改成先将所有历史 key-value 压缩到固定大小的状态 $S_t$ 中,再由当前 query 读取状态

递归更新

状态更新:

$$ S_t=S_{t-1}+k'_tv_t^\top\\ z_t=z_{t-1}+k'_t $$
因此,可以把线性注意力的隐藏状态定义为:
$$ h_t^{\mathrm{LA}}=(S_t,z_t) $$
其状态更新函数为:
$$ h_t^{\mathrm{LA}}=F(h_{t-1}^{\mathrm{LA}},x_t)=\left(S_{t-1}+k'_tv_t^\top,z_{t-1}+k'_t\right) $$
从这个角度看,线性注意力本质上是一个具有矩阵隐藏状态的 RNN

关键问题

所有 token 合并成一个固定大小状态,每一步都会增加新的 key-value 外积

最大的问题是旧记忆只能被持续累积,不能被选择性删除、衰减或者覆盖

标准因果 Softmax 注意力为:

$$ Y=\operatorname{softmax}\left(\frac{QK^\top}{\sqrt d}+M\right)V $$

$M$ 为因果掩码

Softmax 注意力保存每个历史 token 的独立 key 和 value,线性注意力则保存压缩状态,历史 token 已经不再作为独立条目存在,而是被叠加进固定大小状态

同时 Softmax 对分数使用指数变换,分数之间并不大的差异,可以被指数函数显著放大,很容易形成接近 one-hot 的注意力分布,从大量历史 token 中选出少数位置;并且由于行归一化还引入了显式竞争

因此 Softmax 注意力更适合:

  • 从长上下文中查找某个实体;
  • 精确复制前文中的数字或字符串;
  • 匹配括号和代码变量;
  • 进行指针式读取;
  • 根据不同 query 选择完全不同的历史位置;
  • 形成接近离散路由的检索行为
维度 Softmax 注意力 基础线性注意力
历史存储 保存全部历史 (K,V) 压缩为固定状态 ($S_t,z_t$)
状态大小 随序列长度增长 与序列长度无关
查询方式 重新查询每个历史 key 查询压缩后的矩阵状态
检索能力 精确、动态、尖锐 易发生混合和干扰
归一化 Softmax 全局竞争 核归一化
训练复杂度 通常为 $O(L^2)$ 固定特征维度下为 $O(L)$
推理缓存 KV Cache 随 $L$ 线性增长 固定大小递归状态
主要优势 精确检索与复杂路由 长流式序列和低缓存推理
主要问题 长序列计算与显存成本高 固定容量、缺少遗忘、记忆碰撞

门控线性注意力 GLA

核函数选用恒等映射

基础线性注意力最明显的问题是:

$$ S_t=S_{t-1}+k_tv_t^\top $$
旧状态只能累积,无法根据当前输入选择性遗忘

一种简单改进是加入固定衰减(类似强化学习里的回报折扣因子):

$$ S_t=\gamma S_{t-1}+k_tv_t^\top,\qquad 0<\gamma<1 $$
最一般的门可以是同样大小的矩阵
$$ S_t=G_t\odot S_{t-1}+k_tv_t^\top $$
GLA根据当前输入计算遗忘门
$$ \alpha_t=f_\alpha(x_t),\qquad\alpha_t\in(0,1)^{d_k} $$
状态更新为
$$ S_t=\operatorname{Diag}(\alpha_t)S_{t-1}+k_tv_t^\top $$
GLA 仍然是一个具有矩阵隐藏状态的 RNN,只是其状态转移矩阵由当前输入动态控制

第 $i$ 个 token 写入的内容,到达第 $t$ 步时,必须经过后续所有门,只要某个中间位置将某个通道的门设得很小,该通道中更早的信息就会被显著削弱

分组查询注意力 GQA

GQA 不改变注意力本身公式,但 Query head 的数量和 Key/Value head 的数量不再必须相同

多个 Query heads 被分成若干组,每一组共享一个 Key head 和一个 Value head,主要目的是为了减少 KV Cache

MHA MQA GQA(4Q共享为例)
Q1 → K1, V1
Q2 → K2, V2
Q3 → K3, V3
Q4 → K4, V4
Q5 → K5, V5
Q6 → K6, V6
Q7 → K7, V7
Q8 → K8, V8
Q1 ─┐
Q2 ─┤
Q3 ─┤
Q4 ─┤
Q5 ─┤──> 唯一的 K1, V1
Q6 ─┤
Q7 ─┤
Q8 ─┘
Q1 ─┐
Q2 ─┤
Q3 ─┤──> K1, V1
Q4 ─┘

Q5 ─┐
Q6 ─┤
Q7 ─┤──> K2, V2
Q8 ─┘
  • MQA 的 KV Cache 最小,但共享程度太高,可能限制不同注意力头表达不同类型的信息
  • GQA 使用了“分组共享”的思想,但其与 ResNeXt 的具体计算对象和设计目标不同

每个 KV head 对应的 Query head 数量为

$$ G = \frac{H_Q}{H_{KV}} $$
多个 Query 头可以提出不同的问题,但必须使用同一套索引方式和信息表示

很多情况下:

  • 多个 KV 头的投影空间相似;
  • 不同头可能关注相近的位置;
  • Query 头的差异比 KV 头的差异更重要;

所以共享KV不会带来太大的性能下降,反而因为KV cache的减少能带来更多收益

多头潜在注意力 MLA

MLA 不直接缓存完整的 KV Cache,把当前 token 的隐藏表示 $h_t$ 压缩成一个较低维的潜变量 $c_t^{KV}$

$$ h_t\in\mathbb R^d,\qquad c_t^{KV}=W^{DKV}h_t\in\mathbb R^{d_c},\qquad d_c\ll n_hd_h $$

这里的 $h_t$ 就是当前注意力层在位置接收到的输入表示,因为$x_t^{(2)} = h_t^{(1)}$

通过不同的上投影矩阵生成 Content Key 和 Content Value

$$ k_t^C=W^{UK}c_t^{KV},\qquad v_t^C=W^{UV}c_t^{KV} $$
对应的完整映射为
$$ k_t^C=W^{UK}W^{DKV}h_t,\qquad v_t^C=W^{UV}W^{DKV}h_t $$
可以将其理解为对原来的 ($W_K,W_V$) 做低秩分解
$$ W_K= W^{UK}W^{DKV},\qquad W_V= W^{UV}W^{DKV} $$

MLA 对每个历史 token 都保存一个独立的潜变量

$$ \mathcal C_t=\left\{c_1^{KV},c_2^{KV},\ldots,c_t^{KV}\right\} $$
因此 MLA 的缓存仍然随序列长度线性增长,但每个 token 的缓存维度显著小于 MHA

MLA 仍然保留每个 token 的独立地址,因此可以执行 token 级 Softmax 检索

Query 低秩压缩

KV 压缩是 MLA 降低 KV Cache 的核心。Query 并不必须压缩,因为 Query 只在当前计算步骤使用,不需要作为历史缓存长期保存

DeepSeek-V2 进一步对 Query 投影进行低秩分解,主要用于降低训练阶段的激活内存

$$ c_t^Q=W^{DQ}h_t $$
从 Query 潜变量 $c_t^Q$ 生成 Content Query:
$$ q_t^C=W^{UQ}c_t^Q $$
投影完成后切分为多个头
$$ q_t^C= \left[ q_{t,1}^C;\ldots;q_{t,n_h}^C \right] $$
如果不使用 Query 低秩压缩,也可以直接由 $h_t$ 生成全部 Content Query:
$$ q_t^C=W^{QC}h_t $$

解耦 RoPE

如果直接对 Content Key 应用标准 RoPE,则位置相关的旋转矩阵位于 Key 上投影之后,注意力分数包含

$$ q^\top R_j W^{UK} c_j^{KV} $$
$R_j$ 随位置产生变化,所以无法让 Query 直接吸收 $W^{UK}$

因此,MLA 将 Query 和 Key 拆分为内容部分与位置部分

位置 Query 由一个完整投影矩阵生成

$$ q_t^R= \operatorname{RoPE}_t \left( W^{QR}c_t^Q \right) $$
每个注意力头有独立的位置 Query
$$ q_t^R= \left[ q_{t,1}^R;\ldots;q_{t,n_h}^R \right] $$
位置 Key 则直接由隐藏状态生成
$$ k_t^R= \operatorname{RoPE}_t \left( W^{KR}h_t \right) $$
所有注意力头共享同一个位置 Key,因此每个历史 token 只需要缓存一份

第 $i$ 个注意力头实际使用的 Query 和 Key 为:

$$ q_{t,i}= \left[ q_{t,i}^C; q_{t,i}^R \right], \qquad k_{j,i}= \left[ k_{j,i}^C; k_j^R \right] $$
“解耦”是指:将内容表征和 RoPE 位置表征解耦,而不是将 Query 与 Key 解耦

拼接后的内积可以直接拆开:

$$ \color{red}\boxed{q_{t,i}^{\top}k_{j,i} = \left(q_{t,i}^C\right)^\top k_{j,i}^C + \left(q_{t,i}^R\right)^\top k_j^R} $$
两项分别承担内容匹配与带相对位置的匹配

相当于用额外缓存的代价,保留矩阵吸收优化

Query 优化

为了简化 KV 的投影恢复,将投影直接融入到 Q 中,减少计算

为了推导第 $i$ 个头,将 $W^{UK}$ 按输出维度分成若干行块

$$ W^{UK} = \begin{bmatrix} W_1^{UK}\\ W_2^{UK}\\ \vdots\\ W_{n_h}^{UK} \end{bmatrix} $$
当前 Query 与历史 Content Key 的内积为:
$$ \widetilde q_{t,i}^C = \left(W_i^{UK}\right)^\top q_{t,i}^C $$
利用矩阵乘法结合律:
$$ \left(q_{t,i}^C\right)^\top W_i^{UK}c_j^{KV} = \left( \left(W_i^{UK}\right)^\top q_{t,i}^C \right)^\top c_j^{KV} $$
定义吸收 Key 上投影后的 Content Query:
$$ \widetilde q_{t,i}^C = \left(W_i^{UK}\right)^\top q_{t,i}^C $$
于是内容分数可以直接写为:
$$ \left(q_{t,i}^C\right)^\top k_{j,i}^C = \left(\widetilde q_{t,i}^C\right)^\top c_j^{KV} $$
因此,优化推理时不需要为每个历史 token 恢复完整的 $k_{j,i}^C$,而是让每个头的吸收后 Query 直接查询缓存的 $c_j^{KV}$

不能简单使用:

$$ \left(W^{UK}\right)^\top q_t^C $$
因为这会将所有 Query 头的结果相加并混合,矩阵吸收在实现中可以批量完成,但必须保留独立的 Head 维度

大模型训练

目前通用文本大模型的主流预训练范式可以概括为

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
大规模原始语料

清洗、去重、质量筛选、数据配比

Tokenizer 编码与样本 Packing

Decoder-only Transformer

自回归下一 Token 预测

AdamW + 学习率调度 + 混合精度 + 分布式训练

持续预训练 / 长上下文扩展

Base Model

SFT、DPO、RL 等后训练

Encoder-only 模型仍广泛用于文本分类、Embedding 和重排序等任务

Encoder-Decoder 模型仍适合翻译和输入输出转换任务,但在通用对话大模型领域,Decoder-only 已成为主要路线

模型预训练

预训练阶段并不是直接教模型“回答问题”,而是让模型从海量文本中学习:

  • 语言的语法与语义规律;
  • 词语、事实和概念之间的关联;
  • 代码、数学和自然语言的结构;
  • 根据上文预测后文的能力;
  • 上下文学习和模式泛化能力

训练完成后得到的是 Base Model,基础模型,擅长续写,但不能稳定理解用户指令,也没有经过完整的偏好和安全对齐

大模型通常不会全部使用 FP32 训练,因为显存和计算成本过高

精度 特点
FP32 数值稳定,但显存与计算成本高
FP16 范围较小,通常需要 Loss Scaling
BF16 指数范围接近 FP32,训练稳定性较好
FP8 显存和计算成本更低,但实现和校准更复杂

当前较普遍的是 BF16 混合精度训练

  • 矩阵乘法:BF16
  • 部分归一化与关键操作:FP32
  • 优化器状态:FP32 或经过验证的低精度格式

预训练后

训练类型 数据形式 目标 主要目标
持续预训练 CPT 领域原始文本 下一 Token 预测 学习领域语言和知识
监督微调 SFT 指令—回答数据 模仿目标回答 学习如何回答、遵循格式
偏好优化 DPO 优选回答—劣选回答 提高优质回答相对概率 学习回答偏好
强化学习 RL 奖励模型或可验证奖励 最大化期望奖励 进一步优化行为或推理

DPO 和 RL 属于对齐 / 偏好优化

典型流程通常是:

1
2
3
预训练模型
→ SFT:先学会遵循指令和生成合理回答
→ DPO:再学习“两个合理回答中更偏好哪一个”

SFT

典型样本为:

1
2
3
4
5
6
7
8
9
10
System:
你是一个严谨的数学助手。

User:
求解方程 x² - 5x + 6 = 0。

Assistant:
将其因式分解:
(x - 2)(x - 3) = 0
因此 x = 2 或 x = 3。

模型会先通过 Chat Template 将消息转换成特殊 Token

1
2
3
4
5
6
7
8
9
<system>
你是一个严谨的数学助手。
</system>
<user>
求解方程……
</user>
<assistant>
将其因式分解……
</assistant>

SFT 的本质是行为克隆,对数据集的质量有要求

单条数据评估视角:规范性、完整性、及时性、安全性、有效性、唯一性

数据集评估视角:多样性(领域知识)、平衡性(领域数据集数量)、一致性、正交性(领域能力互不影响)、必要性(数据价值)

在数据筛选时需要用到

  • 困惑度:形式上和交叉熵损失一致,困惑度与交叉熵损失是指数关系

    $$ P_i(x)=\exp\left(-\frac{1}{T}\sum_{t=1}^{T}\log p(x_t\mid x_{<t})\right) $$
  • 学习百分比:判断一条 SFT 数据的学习难度

    $$ \operatorname{LP}(i)=\frac{P_{i-1}-P_i}{P_0-P_n} $$
    分子表示第 $i$ 个epoch内,困惑度的下降,一般选择 $LP(1)$
  • 数据难度可转移性:不同规模模型对数据难易程度的相对排序具有较强一致性,所以可以用小模型来筛选,降低成本

    Smaller Language Models are capable of selecting Instruction-Tuning Training Data for Larger Language Models - ACL Anthology

RL

不再要求模型逐 Token 模仿标准答案,而是让模型生成完整回答,根据回答获得奖励,再提高高奖励回答的生成概率

对于语言模型,可以将生成过程建模为一个序列决策过程:

  • 状态 State:输入指令以及当前已经生成的 Token;
  • 动作 Action:选择下一个 Token;
  • 策略 Policy:语言模型 $\pi_\theta$;
  • 轨迹 Trajectory:模型生成的完整回答;
  • 奖励 Reward:奖励模型、规则验证器或环境返回的分数。

SFT 优化的是给定标准答案的似然:

$$ \mathcal{L}_{\mathrm{SFT}} = -\sum_{t=1}^{T} \log \pi_\theta \left( y_t\mid x,y_{<t} \right) $$
RL 则优化模型生成回答的期望奖励:
$$ J(\theta) = \mathbb{E}_{x\sim\mathcal{D},\, y\sim\pi_\theta(\cdot\mid x)} \left[ R(x,y) \right] $$

RLHF

RLHF 全称为 Reinforcement Learning from Human Feedback,基于人类反馈的强化学习

经典 RLHF 通常包含三个阶段:

  1. 使用指令—回答数据进行 SFT,得到初始策略模型;
  2. 收集人类对多个候选回答的偏好排序,训练奖励模型;
  3. 使用 PPO 等强化学习算法优化策略模型,使其获得更高奖励。

奖励模型 Reward Model 接收问题和回答,并输出一个标量奖励,通常使用成对偏好数据训练奖励模型

$$ P(y_w\succ y_l\mid x) = \sigma \left( r_\phi(x,y_w)-r_\phi(x,y_l) \right) $$
对应的奖励模型损失为:
$$ \mathcal{L}_{\mathrm{RM}} = -\log \sigma \left( r_\phi(x,y_w)-r_\phi(x,y_l) \right) $$
该目标并不要求奖励模型预测绝对正确的分数,而是要求:
$$ r_\phi(x,y_w)>r_\phi(x,y_l) $$
奖励模型通常由预训练语言模型加一个标量输出头构成,其训练质量主要取决于:
  • 偏好标注是否一致;
  • Chosen 与 Rejected 是否存在有效区分;
  • 数据是否覆盖实际部署分布;
  • 标注是否受到长度、格式和表达风格等表面特征干扰

PPO

PPO 全称为 Proximal Policy Optimization,近端策略优化,是一种限制策略更新幅度的策略梯度算法

核心思想:

  • 提高优势为正的动作概率;
  • 降低优势为负的动作概率;
  • 限制新策略相对于旧策略的变化幅度,避免单次更新过大。

PPO 训练通常涉及以下组成部分

组成部分 作用
Actor / Policy 根据当前状态选择动作
Critic / Value Function 估计当前状态的期望累计回报
Old Policy 生成当前训练数据时使用的旧策略
Environment 执行动作并返回奖励和下一状态

在时刻 t,智能体处于状态 $s_t$,根据策略 $\pi_\theta$ 选择动作:

$$ a_t \sim \pi_\theta(\cdot \mid s_t) $$
环境执行动作后,返回即时奖励 $r_t$ 和下一状态 $s_{t+1}$

一次完整交互可以表示为轨迹:

$$ \tau=(s_0,a_0,r_0,s_1,a_1,r_1,\ldots,s_T) $$
PPO 的目标是调整策略参数 $\theta$,使策略获得的期望累计奖励最大

累计回报

从时刻 $t$ 开始的折扣累计回报定义为:

$$ G_t=\sum_{k=0}^{T-t-1}\gamma^k r_{t+k} $$
状态价值函数表示从状态 $s_t$ 出发,按照当前策略继续行动时能够获得的期望累计回报
$$ V^\pi(s_t)=\mathbb{E}_\pi[G_t\mid s_t] $$
动作价值函数表示在状态 $s_t$ 下执行动作 $a_t$ 之后按照当前策略继续行动时能够获得的期望累计回报:
$$ Q^\pi(s_t,a_t)=\mathbb{E}_\pi[G_t\mid s_t,a_t] $$
**优势函数**
$$ A^\pi(s_t,a_t)=Q^\pi(s_t,a_t)-V^\pi(s_t) $$
优势函数衡量某个动作相对于当前状态下平均动作水平的好坏:
  • $A^\pi>0$:该动作比预期更好,应提高其概率;
  • $A^\pi<0$:该动作比预期更差,应降低其概率;

实际训练中无法直接获得真实的 $Q^\pi$ 和 $A^\pi$,因此需要使用采样数据进行估计

一种简单的优势估计为

$$ \hat A_t=G_t-V_\psi(s_t) $$
其中,$V_\psi$ 是由 Critic 参数化的价值函数

新旧策略概率比率

对于当前策略 $\pi_\theta$ 和生成数据时使用的旧策略 $\pi_{\theta_{\mathrm{old}}}$,定义概率比率:

$$ \rho_t(\theta) = \frac{ \pi_\theta(a_t\mid s_t) }{ \pi_{\theta_{\mathrm{old}}}(a_t\mid s_t) } $$
**PPO 裁剪目标**

如果直接根据优势更新策略,策略可能在一次训练中变化过大,导致训练不稳定

PPO 使用裁剪目标限制概率比率变化带来的收益

$$ J_{\mathrm{PPO}}(\theta) = \mathbb{E}_t \left[ \min \left( \rho_t(\theta)\hat A_t,\, \operatorname{clip} \left( \rho_t(\theta),1-\epsilon,1+\epsilon \right)\hat A_t \right) \right] $$
其中,$\epsilon$ 用于控制策略更新幅度

裁剪并不保证概率比率一定处于 $[1-\epsilon,1+\epsilon]$,而是限制超出该范围后对优化目标产生的有利影响

因为 PyTorch、TensorFlow 等框架通常执行的是梯度下降,所以将裁剪目标取负,得到 Actor 的策略损失

$$ \mathcal{L}_{\mathrm{PPO}} = -\mathbb{E}_t \left[ \min \left( \rho_t(\theta)\hat{A}_t,\, \operatorname{clip} \left( \rho_t(\theta), 1-\epsilon, 1+\epsilon \right) \hat{A}_t \right) \right] $$
以下两项属于常见的可选训练稳定化技术,不是 PPO 裁剪目标的必要组成部分

价值函数裁剪

PPO 的裁剪目标主要限制 Actor 的策略更新,但 Critic 的价值估计同样可能在一次训练中变化过大

可以将当前价值预测裁剪为:

$$ V_{\mathrm{clip}}(s_t) = V_{\psi_{\mathrm{old}}}(s_t) + \operatorname{clip} \left( V_\psi(s_t)-V_{\psi_{\mathrm{old}}}(s_t), -\epsilon_v, \epsilon_v \right) $$
$$ \mathcal L_{\mathrm{value}} = \frac{1}{2} \mathbb{E}_t \left[ \max \left( \left(V_\psi(s_t)-\hat G_t\right)^2, \left(V_{\mathrm{clip}}(s_t)-\hat G_t\right)^2 \right) \right] $$

价值裁剪的作用是限制 Critic 相对于旧价值预测的变化幅度,防止价值函数在一次更新中发生剧烈变化

KL Early Stopping

即使使用裁剪目标,当前策略仍然可能逐渐偏离生成数据时使用的 Old Policy

因此,可以在每轮更新后估计当前策略与 Old Policy 之间的 KL 散度

$$ \widehat D_{\mathrm{KL}} = \mathbb{E}_t \left[ \log \pi_{\theta_{\mathrm{old}}}(a_t\mid s_t) - \log \pi_\theta(a_t\mid s_t) \right] $$
根据概率比率的定义可以简写为
$$ \widehat D_{\mathrm{KL}} = \mathbb{E}_t \left[ -\log \rho_t(\theta) \right] $$
如果估计的 KL 散度超过预设阈值,则提前停止当前批次剩余的训练轮次

KL Early Stopping 是 PPO 裁剪之外的辅助保护机制

PPO 训练流程

  1. 将当前策略参数保存为 Old Policy;
  2. 使用 Old Policy 与环境交互,采集状态、动作和奖励;
  3. 根据奖励和 Critic 计算累计回报与优势;
  4. 计算当前策略与 Old Policy 的概率比率;
  5. 使用裁剪目标更新 Actor;
  6. 使用回报目标更新 Critic;
  7. 用更新后的策略重新采集数据并重复训练。

GRPO

对同一个问题,GRPO 从旧策略中采样一组回答分别计算奖励,然后使用组内奖励的均值和标准差构造相对优势,奖励高于组内平均值的回答得到正优势,奖励低于组内平均值的回答得到负优势

GRPO 随后使用类似 PPO 的裁剪目标更新策略,但不再训练独立的 Critic

GRPO 与 PPO 的主要区别为:

对比项 典型 PPO GRPO
优势基准 Critic 预测的状态价值 同一问题下的组内平均奖励
是否需要 Critic 通常需要 不需要
数据组织 按轨迹计算 Token 级优势 同一问题采样一组回答并比较
显存开销 需要维护 Value Model 去掉 Value Model,通常更低
奖励形式 标量奖励 标量奖励,并在组内归一化
主要风险 Critic 估计偏差 组内奖励缺乏差异

GRPO 需要组内存在有效差异,如果同组回答奖励全部相同,则归一化优势接近零,当前组基本无法提供有效策略梯度信号