Zhangzhe's Blog

The projection of my life.

0%

URL

TL;DR

  • FlashMLA 是针对 DeepSeek 提出的 MLA (Multi-head Latent Attention) 模块,在 Nvidia Hopper 架构 GPU 上的加速算法,尤其对边长序列推理做了优化
  • FlashMLAMLA 的关系大概可以类比到 FlashAttentionAttention 的关系

关键特性和能力

  1. 仅对 Hopper (sm_90) 架构做优化:充分挖掘了硬件 (sm) 的计算能力和 Memory Hierachy 的 存储 / IO 能力
  2. 支持可变长度序列:和现实世界中推理场景更贴合
  3. 分页 KV cache:使用 block size = 64 的分页存储(这里的 64 的含义是:一个 block 存储某一层某个头的连续 64token 对应的 kv cache
  4. 高性能带宽和算力:在 H800 SXM5 设备上,在内存带宽受限配置环境下,内存带宽可达 3000 GB/s,在算力受限配置下,算力可达 580 TFLOPS
  5. 支持多种精度:支持 BF16FP16

技术实现

  1. 基于 Nvidia cutlass 库 + cuda 实现主要计算
  2. 用干净的 python API 包装,方便集成到 PyTorch-base 框架中
  3. 用元数据管理的方式支持变长序列

Thoughts

  • 做大模型不懂 cuda 是不行了,cutlass 要开始学起来了…
  • FlashMLA 只在 sm_90 上实现了,其他显卡编译不通,DeepSeek 只打高端局

URL

TL;DR

  • 本文提出了一种新的稀疏注意力机制,称为 Native Sparse Attention,该机制在硬件上对齐,并且可以直接训练,而无需额外的稀疏化技术
  • 目前没有开源代码,只能通过论文中的公式和描述来实现

Algorithm

总体流程

nsa.png

  • 本质上是用不同的 pattern 组合来替代 full attention,以减少计算量

代码实现

  • DeepSeek-R1 根据论文中的公式和描述实现
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
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
import torch
import torch.nn as nn
import torch.nn.functional as F
class NSAModule(nn.Module):
def __init__(
self,
d_model,
n_heads,
compress_block=32,
select_block=64,
select_topk=16,
window_size=512,
dropout=0.1,
):
super().__init__()
self.d_model = d_model
self.n_heads = n_heads
self.head_dim = d_model // n_heads
# 参数设置
self.compress_block = compress_block
self.select_block = select_block
self.select_topk = select_topk
self.window_size = window_size
# 压缩用MLP
self.compress_mlp = nn.Sequential(
nn.Linear(self.head_dim * compress_block, 256),
nn.GELU(),
nn.Linear(256, self.head_dim),
)
# 门控机制
self.gate_mlp = nn.Sequential(nn.Linear(d_model, 3 * n_heads), nn.Sigmoid())
# 投影层
self.q_proj = nn.Linear(d_model, d_model)
self.k_proj = nn.Linear(d_model, d_model)
self.v_proj = nn.Linear(d_model, d_model)
self.out_proj = nn.Linear(d_model, d_model)
self.dropout = nn.Dropout(dropout)
def _compress_tokens(self, k, v):
"""压缩KV序列到块级别"""
# 调整输入维度 (batch, seq_len, n_heads, head_dim)
b, t, nh, hd = k.shape # 修改为四维解包
block_size = self.compress_block
num_blocks = (t + block_size - 1) // block_size
pad_len = num_blocks * block_size - t
# 填充并分块
k = F.pad(k, (0, 0, 0, 0, 0, pad_len)) # 添加头部维度的填充
v = F.pad(v, (0, 0, 0, 0, 0, pad_len))
# 调整维度: [batch, num_blocks, block_size, n_heads, head_dim]
k_blocks = k.view(b, num_blocks, block_size, nh, hd)
v_blocks = v.view(b, num_blocks, block_size, nh, hd)
# 压缩处理 (保持头部分离)
k_compressed = self.compress_mlp(
k_blocks.permute(0, 1, 3, 2, 4).flatten(3)
) # [b, num_blocks, nh, hd]
v_compressed = self.compress_mlp(v_blocks.permute(0, 1, 3, 2, 4).flatten(3))
return k_compressed, v_compressed
def _select_blocks(self, q, k_compressed, v_compressed):
"""基于注意力分数选择关键块"""
# 计算压缩注意力分数
scores = torch.einsum("bthd,bkhd->bthk", q, k_compressed) / (
self.head_dim ** 0.5
)
probs = F.softmax(scores, dim=-1)
# 选择topk块
topk_scores, topk_indices = torch.topk(
probs.mean(dim=2), self.select_topk, dim=-1
)
# 收集选中的块
k_selected = torch.gather(
k_compressed,
1,
topk_indices.unsqueeze(-1).expand(-1, -1, -1, self.head_dim),
)
v_selected = torch.gather(
v_compressed,
1,
topk_indices.unsqueeze(-1).expand(-1, -1, -1, self.head_dim),
)
return k_selected, v_selected
def forward(self, x, attn_mask=None):
b, t, d = x.shape
# 投影QKV并保持四维结构
q = self.q_proj(x).view(b, t, self.n_heads, self.head_dim)
k = self.k_proj(x).view(b, t, self.n_heads, self.head_dim)
v = self.v_proj(x).view(b, t, self.n_heads, self.head_dim)
# 压缩路径
k_compressed, v_compressed = self._compress_tokens(k, v)
# 选择路径
k_selected, v_selected = self._select_blocks(q, k_compressed, v_compressed)
# 滑动窗口
k_window = k[:, max(0, t - self.window_size) :]
v_window = v[:, max(0, t - self.window_size) :]
# 门控权重
gate = self.gate_mlp(x).view(b, t, 3, self.n_heads).permute(2, 0, 3, 1, 2)
# 三路注意力计算
attn_outputs = []
for branch_k, branch_v in [
(k_compressed, v_compressed),
(k_selected, v_selected),
(k_window, v_window),
]:
scores = torch.einsum("bthd,bkhd->bthk", q, branch_k) / (
self.head_dim ** 0.5
)
if attn_mask is not None:
scores = scores.masked_fill(attn_mask == 0, -1e9)
probs = F.softmax(scores, dim=-1)
probs = self.dropout(probs)
output = torch.einsum("bthk,bkhd->bthd", probs, branch_v)
attn_outputs.append(output)
# 门控融合
weighted = sum(g * o for g, o in zip(gate, attn_outputs))
output = weighted.contiguous().view(b, t, d)
return self.out_proj(output)
def test_nsa_attention_shapes():
# 测试参数
batch_size = 2
seq_len = 128
d_model = 256
n_heads = 8
nsa_attn = NSAModule(d_model, n_heads)
x = torch.randn(batch_size, seq_len, d_model)
print("输入形状:", x.shape)
# 打印中间形状
q = nsa_attn.q_proj(x).view(batch_size, seq_len, n_heads, -1)
k = nsa_attn.k_proj(x).view(batch_size, seq_len, n_heads, -1)
v = nsa_attn.v_proj(x).view(batch_size, seq_len, n_heads, -1)
print("\nQ/K/V形状:", q.shape) # 应该输出 [2, 128, 8, 32]
k_comp, v_comp = nsa_attn._compress_tokens(k, v)
print("压缩后KV形状:", k_comp.shape) # [2, 4, 8, 32]
if __name__ == "__main__":
test_nsa_attention_shapes()

Thoughts

  • 速度快并不惊讶,但效果竟然比 Full Attention 还好,如果实验比较 solid,那么 NSA 确实比较有前途
  • 由于没有开源代码,所以只能通过论文中的公式和描述来实现,这对于一些实验复现和工程应用来说是一个挑战

URL

TL;DR

  • 本文提出了三个模型和其对应的训练方法,目的是提高大模型的 CoT 推理能力,三个模型分别是:
    • DeepSeek-R1-Zero:由 DeepSeek V3 base 直接通过 RL 训练得到
    • DeepSeek-R1:由 DeepSeek V3 base 在合成的高质量 CoT 数据集上 SFT + RL 训练得到,效果优于 DeepSeek-R1-Zero
    • Distill model:用相同的合成高质量 CoT 数据训练的开源模型(参数量小且非 MoE
  • 其中训练方法是本文讲述的重点,DeepSeek R1 中的 R 表示 Reasoning(推导)

Algorithm

总体流程

deepseek_r1.jpg

Benchmark

deepseek.png

DeepSeek-R1-Zero

  • DeepSeek-R1-Zero 是用 DeepSeek v3 base (pre-train) 直接通过 RL 训练得到的,不经过 SFT
  • 用了以下的 Template 让模型生成 CoT 推理结果:
    deepseekr1zero.png
  • 奖励建模分成两个部分:
    • Accuracy rewards: 用于评估模型生成的 CoT 推理结果的准确性
    • Format rewards: 用于评估模型生成的 CoT 推理结果的格式是否符合上述 prompt 要求的 CoT 的格式
  • RL 过程不使用 NN reward model,而是使用了 Rule base reward
  • RL 过程中使用了 GRPO 算法
  • DeepSeek-R1-Zero 仅仅用上述简单的 CoT 性能提升流程,即可大幅提高 benchmark 上的指标(相比于 DeepSeek v3),但仍然存在一些问题,例如可读性差和语言混合

DeepSeek-R1

  • DeepSeek-R1 用了更复杂的数据处理流程和训练流程,主要分成四个阶段(流程图上的四个 stage):
    • 冷启动:用数千个来自于 DeepSeek-R1-Zero 输出且经过人工处理的 CoT 数据,对 DeepSeek V3 base 进行 SFT
    • 面向推导的强化学习:用编码、数学、科学和逻辑推理等推理密集型任务数据,对冷启动微调后的模型进行 RL 训练,同时引入了语言一致性和可读性的奖励
    • 拒绝采样和监督微调:用上一个 stage 得到的模型生成更广泛领域的推导 CoT 数据 600K 条,用 DeepSeek V3 模型推理部分非推导数据(例如:写作、翻译等)得到潜在的思维链数据 200K 条,用这 800K 条数据对 DeepSeek V3 base 进行 SFT
    • 所有场景的强化学习:用了和 DeepSeek V3 中使用的 RL 相同的 pipeline,数据和 DeepSeek-R1-Zero 中的 RL 数据相同,对上一个 stage 得到的模型进行 RL 训练
  • 综上所述,DeepSeek-R1 最重要的是 800K 包含 CoTSFT 数据

Distill model

  • 由于 DeepSeek R1 使用了 DeepSeek V3 相同的模型结构,因此包含 671B 参数,每个 token 激活 34B 参数,因此这个模型非常大,不适合部署
  • 为了解决这个问题,本文提出了 Distill model,用开源的较小参数量的稠密模型(非 MoE 模型)在 800K 包含 CoTSFT 数据 进行 SFT 微调
  • 主要蒸馏了 Qwenllama 模型,且完全不做 RL,即可大幅提高原模型在 Benchmark 上的指标
    distill.png

Thoughts

  • 需要重点关注本论文提到的 CoT 数据合成方式,这是 DeepSeek R1 的核心
  • 能从本文中可以看出 数据重要性远大于模型Post-train 只包含一次 SFT 和 一次 RL,比如论文中 stage 1 模型会大胆抛弃 DeepSeek R1 Zero 模型,而是退回到了 DeepSeek V3 base 模型,只用了 Zero 生成数据;stage 3 模型会大胆抛弃之前 stage 的模型,而是退回到了 DeepSeek V3 base 模型,只用了 stage 2 模型生成数据
  • 合成数据是最重要的,合成高质量的数据是 DeepSeek R1 的核心

Reference

URL

TL;DR

  • 本文提出了一种新的混合专家语言模型 DeepSeek-V3,共计 671B 总参数,对于每个 token 仅仅激活 37B 个参数,可以说又强又快。
  • 主要的改进点包括:
    • Multi-head Latent Attention (MLA)MLADeepSeek 系列的核心模块(Deepseek v2 就有了),通过 MLPhidden state 映射到 latent space,降低计算和存储复杂度。可参考 Deepseek v2 MLA 详解
    • Deepseek MOE:一种比 Gshard 更优的 MoE 架构,可参考 Deepseek MoE 详解
    • Multi-token predicition (MTP):一种新的 token 预测方式,与传统的 predict next token 自回归预训练方式不同,MTP 通过一次预测多个 token(例如一次预测 next tokennext2next^2 tokennext3next^3 token),仅在训练阶段使用 MTP,推理阶段只使用 next token 的预测结果,即可大幅提高模型在 benchmark 上的综合表现。
    • auxiliary-loss-free strategy for load balancing:一种新的 MoE 间负载均衡策略,无需额外的辅助损失,简化了训练过程,提高了训练效率。
  • 性能与稳定性:
    • 性能:DeepSeek-V3 的表现超过了其他开源模型,其性能已经可以与一些领先的闭源模型相媲美。
    • 稳定性:尽管性能卓越,模型全程训练仅花费了 278.8H800 GPU 小时,而且在整个训练过程中表现出极高的稳定性,没有出现不可恢复的损失暴增或需要回滚的情况。

Algorithm

总体流程

  • Pre-Training
  • Context Extension
  • Post-Training
    deepseekv3_table_1.png

Multi-token predicition (MTP)

deepseekv3.png

  • 上图给出了一次预测 next tokennext2next^2 tokennext3next^3 tokenMTP 模型结构。
  • 具体来说,MTP 是在标准 Decoder only Transformer 的基础上,额外引入了 k-1token 预测头,每个 token 预测头都是一个 Norm + Concat + Projection + Transformer block 结构。
  • 假设图上从左到右的 token head 分别记作 head0head_0head1head_1head2head_2,则 head0head_0 预测的是 next tokenhead1head_1 预测的是 next2next^2 tokenhead2head_2 预测的是 next3next^3 token
  • MTP 过程必须要保持因果性:
    • head0head_0 输入 token0token_0,主干网络输出 h0h_0output layerh0h_0 映射为 next token 预测,记作 token1token_1'
    • head1head_1 输入 h0h_0token1token_1,提取特征记作 h1h_1output layerh1h_1 映射为 next2next^2 token 预测,记作 token2token_2'
    • head2head_2 输入 h1h_1token2token_2,提取特征记作 h2h_2output layerh2h_2 映射为 next3next^3 token 预测,记作 token3token_3'
  • embedding layeroutput layerMTP 中是共享的。
  • MTP 通常只用在训练阶段,即在推理阶段只使用 head0head_0 的输出。
  • 也可以将 inference MTPspeculative decoding 结合,提高模型生成效率。

Auxiliary-loss-free strategy for load balancing

训练阶段负载均衡策略

gi,t={si,t,si,t+biTopk({si,t+bj1jNr},Kr)0,otherwiseg'_{i,t}=\begin{cases}s_{i,t},&s_{i,t}+b_i\in Topk(\{s_{i,t}+b_j|1\le j\le N_r\},K_r)\\0,& otherwise\end{cases}

  • 其中:
    • si,t=Sigmoid(utTei)s_{i,t}=Sigmoid(u_t^Te_i)
    • bib_ibias 参数
    • NrN_rExpert 的数量
    • KrK_r 是每个 token 保留的 Expert 数量
    • bjb_j 是根据每一个 Expert 被激活的次数动态调整的
      • 如果一个 Expert 被激活的次数过多(大于专家激活平均值),则 bj=bjγb_j=b_j-\gamma
      • 如果一个 Expert 被激活的次数过少(小于专家激活平均值),则 bj=bj+γb_j=b_j+\gamma
      • 其中 γ\gamma 是一个超参数,表示 bias update speed
  • 关于 bias 的详细更新策略,可参考下图
    auxiliary_loss_free.png
    • u 和本文中的 γ\gamma 等价

推理阶段

  • 推理阶段无需额外的负载均衡策略,直接使用 si,ts_{i,t} 即可。

Complementary Sequence-Wise Auxiliary Loss

  • 上面提到的 Auxiliary-loss-free strategy for load balancing 损失只能保证 Expert 的负载均衡,但可能出现一个 Sequence 中每一个 token 都激活了相同的 Expert,这种情况下,Expert 之间的负载仍然不均衡。因此 Complementary Sequence-Wise Auxiliary Loss 通过引入 Auxiliary Loss 来保证 Sequence 内部 token 之间的负载均衡。
  • 具体来说,需要关注:
    • 离散选择频率:每一个 ExpertSequence 中的激活次数
    • 连续激活强度:每一个 ExpertSequence 中平均专家分数
  • 公式表述:

LBal=αi=1NrfiPiL_{Bal}=\alpha\sum_{i=1}^{N_r}f_iP_i

fi=NrKrTt=1TI(si,tTopk({sj,t1jNr},Kr))f_i=\frac{N_r}{K_rT}\sum_{t=1}^{T}\mathbb I(s_{i,t}\in Topk(\{s_{j,t}|1\le j\le N_r\},K_r))

si,t=si,tj=1Nrsj,ts_{i,t}'=\frac{s_{i,t}}{\sum_{j=1}^{N_r}s_{j,t}}

Pi=1Tt=1Tsi,tP_i=\frac{1}{T}\sum_{t=1}^{T}s_{i,t}'

  • 其中:
    • LBalL_{Bal}Complementary Sequence-Wise Auxiliary Loss
    • α\alpha 是一个超参数,是一个很小的权重
    • fif_iExpert ii 的离散选择频率
    • PiP_iExpert ii 的连续激活强度
    • TTSequence 的长度
    • NrN_rExpert 的数量
    • KrK_r 是每个 token 激活的 Expert 数量
    • I()\mathbb I(\cdot) 是指示函数
    • si,ts_{i,t}Sigmoid(utTei)Sigmoid(u_t^Te_i),表示专家的分数
  • 道理也很简单,就是希望在一个 Sequence 中,每一个 Expert 被激活的次数和激活的强度都是均衡的。

训练框架设计

  • 模型训练基于自主研发的 HAI-LLM 框架,这是一个经过优化的高效轻量级训练系统,目前没有开源。
  • DeepSeek-V3 的并行策略包含三个层面:
    • 16 路流水线并行(Pipeline Parallelism, PP)
    • 8 个节点的 64 路专家并行(Expert Parallelism, EP)
    • ZeRO-1 数据并行(Data Parallelism, DP)
  • 开发了 DualPipe 流水线并行算法,相比于现有的 PP 算法,减少了空泡,基本上实现了计算和通信的高度重叠。
    dualpipe.png
  • 优化了跨节点全对全通信内核,充分利用节点内 GPU 间的 NVLink 带宽和节点间 InfiniBand 带宽,同时减少了通信所需的流式多处理器(SMs)资源占用
  • 通过精细的内存管理优化,使得模型训练无需依赖开销较大的张量并行(Tensor Parallelism, TP)技术

FP8

  • DeepSeek-V3 开发了 FP8 (E4M3) 混合精度训练框架,首次在超大规模模型上验证了 FP8 训练的可行性和效果。
  • 通过算法、框架和硬件的综合优化,突破了跨节点 MoE 训练中的通信瓶颈,实现了计算与通信的高度重叠。这种优化大幅提升了训练效率,降低了训练成本,同时支持了更大规模模型的训练而无需额外开销。
    fp8.png

Thoughts

  • 实际上 DeepSeek V3 技术报告中还有很多细节,尤其是很多关于工程优化的部分,此处不再一一列举
  • 工程能力,尤其是开发一套分布式训练框架且根据模型特性对其深度优化的能力,是一个大模型成功的不可或缺因素
  • 如果想要在大模型领域有所建树,不仅需要有强大的算法研究能力,还需要有强大的工程能力,比如 CUDA 编程能力等

References

URL

TL;DR

  • 这篇论文首次提出了 Denoising Diffusion Probabilistic Models (DDPM) 模型,该模型是一种去噪扩散模型,可以用于生成图像。
  • DDPM 有详细的热力学理论基础(Langevin 朗之万动力学),对热力学扩散过程进行了建模。
  • 本质上,DDPM 是一个自回归生成模型,通过自监督方式,对图像进行添加噪声和逐步去噪的过程,来学习图像的分布。

Algorithm

DDPM 总体结构

ddpm_1.png

  • 上图中 x0x_0 是来自图像数据集的真实图像,从右到左被称为前向过程,即对图像逐步添加随机高斯噪声,直到得到噪声图像 xTx_T
  • 从左到右被称为逆向过程,即模型预测噪声,代入公式逐步去噪

前向过程

xt=1βtxt1+βtϵt1 ,  ϵt1N(0,I)x_t=\sqrt {1-\beta_t} \cdot x_{t-1}+\sqrt \beta_t\cdot \epsilon_{t-1}\ ,\ \ \epsilon_{t-1}\sim N(0, I)

  • 其中:
    • ϵt\epsilon_t 是随机高斯噪声,即 torch.randn
    • βt\beta_t 是超参数,控制噪声的强度,随着 tt 的增大,βt\beta_t1e-4 逐渐增大到 2e-2
    • 总时间步数 T=1000T=1000

前向过程的累积表示

xt=αˉtx0+1αˉtϵx_t=\sqrt {\bar\alpha_t}\cdot x_0+\sqrt{1-\bar\alpha_t}\cdot \epsilon

  • 具体推理过程:
    • x1=1β1x0+β1ϵ0x_1 = \sqrt{1-\beta_1}\cdot x_0+\sqrt{\beta_1}\cdot \epsilon_0
    • x2=1β2x1+β2ϵ1x_2 = \sqrt{1-\beta_2}\cdot x_1+\sqrt{\beta_2}\cdot \epsilon_1
    • x1x_1 代入 x2x_2 的公式中,得到 x2x_2 的累积表示: x2=1β2(1β1x0+β1ϵ0)+β2ϵ1x_2 = \sqrt{1-\beta_2}\cdot (\sqrt{1-\beta_1}\cdot x_0+\sqrt{\beta_1}\cdot \epsilon_0)+\sqrt{\beta_2}\cdot \epsilon_1
    • x2=1β21β1x0+1β2β1ϵ0+β2ϵ1x_2 = \sqrt{1-\beta_2}\sqrt{1-\beta_1}x_0+\sqrt{1-\beta_2}\sqrt{\beta_1}\cdot \epsilon_0+\sqrt{\beta_2}\cdot \epsilon_1
    • 以此类推,得到 xtx_t 的累积表示:xt=i=1t(1βi)x0+i=1tj=i+1t(1βj)βiϵx_t=\sqrt{\prod_{i=1}^t(1-\beta_i)}\cdot x_0+\sum_{i=1}^t\sqrt{\prod_{j=i+1}^t(1-\beta_j)}\cdot \sqrt{\beta_i}\cdot \epsilon
    • αˉt=i=1t(1βi)\bar\alpha_t=\prod_{i=1}^t(1-\beta_i),得到 xt=αˉtx0++i=1tβij=i+1t(1βj)ϵx_t=\sqrt{\bar\alpha_t}\cdot x_0++\sum_{i=1}^t\sqrt{\beta_i\cdot\prod_{j=i+1}^t(1-\beta_j)}\cdot \epsilon
    • 需要将噪声项合并成一个正态分布,需要确定方差,对于每一项 βij=i+1t(1βj)\sqrt{\beta_i\cdot\prod_{j=i+1}^t(1-\beta_j)},其方差为 βij=i+1t(1βj)\beta_i\cdot\prod_{j=i+1}^t(1-\beta_j)
    • 总噪声的方差为 i=1tβij=i+1t(1βj)\sum_{i=1}^t\beta_i\cdot\prod_{j=i+1}^t(1-\beta_j)
    • 代入 αˉt\bar\alpha_t 到噪声项的方差中,得到总噪声的方差为 i=1tβiαˉtαˉi=αˉti=1tβiαˉi\sum_{i=1}^t\beta_i\cdot\frac{\bar\alpha_t}{\bar\alpha_i}=\bar\alpha_t\cdot\sum_{i=1}^t\frac{\beta_i}{\bar\alpha_i}
    • 由于 βiαˉi=1αiαˉi=1αˉi1αˉi1\frac{\beta_i}{\bar\alpha_i}=\frac{1-\alpha_i}{\bar\alpha_i}=\frac{1}{\bar\alpha_i}-\frac{1}{\bar\alpha_{i-1}},因此使用 裂项求和 公式可得:i=1tβiαˉi=1αˉt1αˉ0\sum_{i=1}^t\frac{\beta_i}{\bar\alpha_i}=\frac{1}{\bar\alpha_t}-\frac{1}{\bar\alpha_0}
    • 由于 αˉ0=1\bar\alpha_0=1,因此 i=1tβiαˉi=1αˉt1\sum_{i=1}^t\frac{\beta_i}{\bar\alpha_i}=\frac{1}{\bar\alpha_t}-1
    • 代入噪声总方差公式中,得到总噪声的方差为 αˉt(1αˉt1)=1αˉt\bar\alpha_t\cdot(\frac{1}{\bar\alpha_t}-1)=1-\bar\alpha_t
    • 因此,累加形式的前向过程可表示为:xt=αˉtx0+1αˉtϵx_t=\sqrt{\bar\alpha_t}\cdot x_0+\sqrt{1-\bar\alpha_t}\cdot \epsilon
  • 累积表示的好处:
    • 在给定 tt 时,可以直接从 x0x_0 生成 xtx_t,而不需要逐步生成
    • 通过单步表示到累积表示的数学推导过程,可以更好地理解噪声的累积方差

逆向过程

xt1=1αt(xtβt1αˉtϵθ(xt,t))x_{t-1}=\frac{1}{\sqrt{\alpha_t}}(x_t-\frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\cdot \epsilon_\theta(x_t,t))

  • 其中:
    • ϵθ(xt,t)\epsilon_\theta(x_t,t) 是模型预测的噪声,即 model(x_t, t)
    • αt=1βt\alpha_t=1-\beta_t
    • αˉt=s=1tαs\bar\alpha_t=\prod_{s=1}^t\alpha_s
  • 由上述公式可知:预测得到的噪声并不会直接叠加到噪声图像上,而是利用扩散原理,根据噪声添加相关的超参数进行去噪
  • 由于 ϵθ(xt,t)\epsilon_\theta(x_t,t) 是模型对于 xtx_t 的预测,且 xtx_t 也和模型预测有关,因此逆向过程是个数据依赖的过程,因此逆向过程的累积表示没有意义

训练过程

ddpm_2.png
DDPM 的训练过程基于扩散模型的前向过程和逆向过程。训练的目标是让模型学习如何从噪声数据中恢复出原始数据。具体步骤如下:

  • 输入数据:从数据集中随机选择一批原始图像 x0x_0
  • 随机选择扩散步数:对于每张图像,随机选择一个扩散步数 tt,范围从 11TT 随机选择
  • 生成噪声数据:通过前向过程生成噪声图像 xtx_t
  • 预测噪声:使用模型 ϵθ(xt,t)\epsilon_\theta(x_t,t) 预测添加到 x0x_0 中的噪声
  • 计算损失:通过比较预测噪声和真实噪声计算损失函数(通常是均方误差,MSE
  • 更新模型参数:通过反向传播更新模型参数

伪代码

1
2
3
4
5
6
7
for epoch in range(num_epochs):
for x0 in data_loader:
t = random.randint(1, T) # 随机选择扩散步数
xt, noise = forward_diffusion(x0, t) # 生成噪声数据
predicted_noise = model(xt, t) # 预测噪声
loss = mse_loss(predicted_noise, noise) # 计算损失
update_model_parameters(loss) # 更新模型参数

模型输入输出

  • 输入
    • 原始图像 x0x_0,通常被归一化到 [1,1][-1, 1] 之间
    • 扩散步数 tt,随机选择,会变成 time embedding 输入模型
  • 输出
    • 预测噪声 ϵθ(xt,t)\epsilon_\theta(x_t,t),模型输出是预测的噪声,用于逆向过程去噪

损失函数

  • 通常使用均方误差(MSE)作为损失函数

L(θ)=Ex0,t[ϵθ(xt,t)ϵ2]\mathcal L(\theta)=\mathbb E_{x_0,t}\left[\left\|\epsilon_\theta(x_t,t)-\epsilon\right\|^2\right]

  • 其中:
    • ϵ\epsilon 是真实噪声
    • ϵθ(xt,t)\epsilon_\theta(x_t,t) 是模型预测的噪声
    • Ex0,t\mathbb E_{x_0,t} 表示对所有的 x0x_0tt 求期望

推理过程(采样过程)

用 pre-trained diffusion model 生成图像

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
28
from diffusers import UNet2DModel
from diffusers import DDPMScheduler
from PIL import Image
import numpy as np
import tqdm
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def display_sample(sample, i):
image_processed = sample.cpu().permute(0, 2, 3, 1)
image_processed = (image_processed + 1.0) * 127.5
image_processed = image_processed.numpy().astype(np.uint8)
image_pil = Image.fromarray(image_processed[0])
image_pil.save(f"sample_{i}.png")
# load pretrain model
repo_id = "google/ddpm-church-256"
model = UNet2DModel.from_pretrained(repo_id).to(device)# build scheduler
scheduler = DDPMScheduler.from_config(repo_id)# random noise as input
noisy_sample = torch.randn(
1, model.config.in_channels, model.config.sample_size, model.config.sample_size
)
sample = noisy_sample.to(device)# step through the scheduler
for i, t in enumerate(tqdm.tqdm(scheduler.timesteps)):
# 1. predict noise residual
with torch.no_grad():
residual = model(sample, t).sample # 2. compute less noisy image and set x_t -> x_t-1
sample = scheduler.step(residual, t, sample).prev_sample # 3. optionally look at image
if (i + 1) % 50 == 0:
display_sample(sample, i + 1)
  • 通过 UNet2DModel 加载预训练模型
  • 通过 DDPMScheduler 构建调度器(主要作用是生成噪声和去噪)
  • 通过 torch.randn 生成随机噪声作为输入
  • 通过 scheduler.step 逐步生成图像
  • 最终生成一张完全去噪的新图像,这张图像并不在任何数据集中(只是风格看起来像 Church-256 数据集),是模型自己生成的(这也意味着没办法去指导模型生成特定的图像)

用 pre-trained diffusion model 对已逐步添加噪声的图像去噪

  • 理论上,可以将一张数据集中的图像逐步添加噪声,在 t=Tt=T 步之后,使用模型去噪,得到一张和原始图像非常相似的图像
  • 但实际测试中,当 tt 较大时(大于 300TT 等于 1000),得到的去噪图像和原始图像相差较大,在 t<300t<300 时,去噪图像和原始图像相似度较高
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
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
from diffusers import UNet2DModel
from diffusers import DDPMScheduler
from diffusers import DDPMPipeline
from PIL import Image
import numpy as np
import tqdm
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def preprocess_image(image):
"""Preprocess an image to model input format."""
image = np.array(image).astype(np.float32) / 127.5 - 1.0 # Normalize to [-1, 1]
image = (
torch.tensor(image).permute(2, 0, 1).unsqueeze(0).to(device)
) # Shape: (1, 3, H, W)
return image
def display_sample(sample, name):
"""Save a sample as an image."""
image_processed = sample.cpu().permute(0, 2, 3, 1)
image_processed = (image_processed + 1.0) * 127.5
image_processed = image_processed.numpy().astype(np.uint8)
image_pil = Image.fromarray(image_processed[0])
image_pil.save(f"{name}.png")
# Load pretrained pipeline to generate an image
repo_id = "google/ddpm-celebahq-256"
image_pipe = DDPMPipeline.from_pretrained(repo_id).to(device)
image = image_pipe().images[0]# Save original image
image.save("original_image.png")# Preprocess original image
original_image = preprocess_image(image)# Load pretrained model and scheduler
model = UNet2DModel.from_pretrained(repo_id).to(device)
scheduler = DDPMScheduler.from_config(repo_id)# Number of timesteps
T = scheduler.num_train_timesteps# Step 1: Gradually add noise to the image
alpha_cumprod = scheduler.alphas_cumprod[300]
noisy_image = torch.sqrt(alpha_cumprod) * original_image + torch.sqrt(
1 - alpha_cumprod
) * torch.randn_like(original_image)
display_sample(noisy_image, "noisy_image")
sample = noisy_image.clone()
for t in tqdm.tqdm(scheduler.timesteps[-300:]):
with torch.no_grad():
residual = model(sample, t).sample # Predict noise residual
sample = scheduler.step(residual, t, sample).prev_sample # Denoise step by step# Save final denoised image
display_sample(sample, "final_image")# Compare original and final image
l2_diff = torch.norm(original_image - sample).item()
print(f"L2 difference between original and final image: {l2_diff}")

300.png

上图表示 t=300 时,原图 -> 噪声图 -> 去噪图
500.png
上图表示 t=500 时,原图 -> 噪声图 -> 去噪图
1000.png
上图表示 t=1000 时,原图 -> 噪声图 -> 去噪图

Thoughts

  • DDPM 是一种非常有趣的生成模型,数学基础非常扎实
  • 但仍然存储三个主要问题:
    • 时间步 TT 太大,通常是 1000
    • 无法指定模型生成特定的图像,只能生成和数据集相似的图像
    • 作用在图像域,计算量较大
  • 针对以上三个问题,stable diffusion 诞生了

URL

TL;DR

  • BLIPSalesforce 提出的一种图文多模态对齐算法,全称是 Bootstrapping Language-Image Pre-training,是一种自举的图文多模态预训练算法
  • 相较于 CLIP 只能对齐图像和文本表征,BLIP 将任务范式进行了扩展,基本已经有了多模态大模型的雏形

Algorithm

BLIP 整体架构

blip_1.png

  • BLIP 在预训练阶段,包含三种损失函数
    • Image-Text Contrastive Loss:对齐图像和文本表征,和 CLIP 的对齐方式类似,不同的是 BLIP 引入了 Momentum Model,可以使得训练更加稳定
    • Image-Text Macthing Loss:图像和文本匹配损失,输出 TrueFalse,本质是交叉熵损失
    • Language Modeling Loss:语言建模损失,预测下一个词(大模型范式),用于文本内容生成
  • 从上图可以看出,模型实际包含四个部分:
    • Image Encoder:图像编码器,将图像编码为 Image Embedding
    • Text Encoder:文本编码器,将文本编码为 Text Embedding
    • Image-grounded Text Eecoder:图像驱动文本编码器,用文本特征 Cross-Modal Attention 去索引图像特征,输出文本图像特征是否匹配(True / False
    • Image-grounded Text Decoder:图像驱动文本解码器,用文本特征 Cross-Modal Attention 去索引图像特征,用逐个 Token 预测的方式生成文本,回答文本对应的问题
  • 从实现代码看,Text Encoder / Image-grounded Text Eecoder / Image-grounded Text Decoder 实际上共用了一个 Bert 预训练模型的一组参数(仅仅在部分位置不共享参数)
    • Text Encoder
      • 原生 Bert 模型,用于文本编码
      • 不做 Cross-Modal Attention,只做 Self-Attention
      • 最后输出层为表征层
      • 用特殊的 [CLS] 作为任务标识符
    • Image-grounded Text Eecoder
      • 魔改 Bert 模型,加入 Cross-Modal Attention
        • 这种状态处于 encoderdecoder 之间,比 encoder 多了 Cross-Modal Attention,比 decoder 少了 Causal Self-Attention
        • 用文本特征去索引图像特征
      • 最后输出层为二分类层
      • 用特殊的 [Encoder] 作为任务标识符
    • Image-grounded Text Decoder
      • 魔改 Bert 模型,将其转化成一个完全的 decoder 形态,主要包括两点:
        • 加入 Cross-Modal Attention,用文本特征去索引图像特征
        • 由于是生成任务,所以需要用 Causal Self-Attention 取代 Bert 默认的 Bidirection Self-Attention
      • 最后输出层为 vocab prob 层,用于生成文本(下一个文本在词表中的概率)
      • 用特殊的 [Decoder] 作为任务标识符
    • 三个模型仅仅在 Head 层和 Attention 层参数上有所不同,其余参数共享

Bootstrapping 策略

  • BLIP 中的 B 就是 Bootstrapping 的意思,即自举,表示此算法有一种自我增强的方式
  • 自我增强的方式主要是通过 CapFilt 方法实现的:
    • Captioner:用于为图像生成高质量的文本描述
    • Filter:用于过滤掉低质量的图像-文本对,保留高质量的对用于训练
  • CapFilt 工作流程:
    1. 生成伪标签:使用 Captioner 为无标签或弱标签图像生成文本描述,形成伪标签
    2. 过滤伪标签:使用 Filter 评估生成的伪标签的质量,保留高质量的伪标签
    3. 更新训练数据集:将过滤后的高质量伪标签与原始训练数据集合并,形成新的训练数据集
    4. 重新训练模型:使用更新后的训练数据集重新训练模型,从而改进模型的性能

Thoughts

  • 在某种程度上,BLIP 统一了几乎所有的图文多模态任务,除了无法做图像生成任务外,其他任务都可以通过 BLIP 的不同模块组合实现
  • BLIPBootstrapping 策略是一种自我增强的方式,可以使得模型在无监督的情况下不断提升性能,如今这几乎是多模态预训练算法的标配
  • 从实验的角度告诉我们,encoderdecoder 做充分的参数共享是可行的,仅仅在少量的位置上不共享参数,就可以实现完全不同任务的模型

TL;DR

  • 本文是大模型DPO的入门教程,目的是用尽可能简单的方式介绍大模型DPO的基本概念和原理,不涉及过多实现细节。

什么是 DPO

  • DPODirect Performance Optimization (直接偏好优化) 的缩写,是 RLHF 的替代方案,目的是使用不涉及强化学习的方法来优化模型性能。

DPO 的数据格式

  • RLHF 一致,DPO 的数据格式为:prompt + response 1 + response 2 + label
    • prompt:问题描述
    • response 1:模型生成的回答 1
    • response 2:模型生成的回答 2
    • label:人类标注的哪个回答更好(0 / 1

DPO 损失函数的设计

L(θ)=E(x,y+,y)D[logσ(sθ(x,y+)sθ(x,y))]L(\theta)=\mathbb E_{(x,y^+,y^-)\sim D}[\log\sigma(s_\theta(x,y^+)-s_\theta(x,y^-))]

sθ(x,y)=logpθ(yx)s_\theta(x,y)=\log p_\theta(y|x)

pθ(yx)=t=1Tpθ(yty<t,x)p_\theta(y|x)=\prod_{t=1}^T p_\theta(y_t|y_{<t},x)

sθ(x,y)=t=1Tlogpθ(yty<t,x)s_\theta(x,y)=\sum_{t=1}^T \log p_\theta(y_t|y_{<t},x)

  • L(θ)L(\theta):损失函数
  • θ\theta:模型参数
  • DD:数据集
  • xx:输入
  • y+y^+:人类标注的更好的回答
  • yy^-:人类标注的更差的回答
  • sθ(x,y)s_\theta(x,y):模型对 (x,y)(x,y) 的得分
  • pθ(yx)p_\theta(y|x):模型对 xx 生成 yy 的概率
  • σ\sigmasigmoid 函数
  • TT:序列长度
  • y<ty_{<t}yy 的前 t1t-1 个元素
  • sθ(x,y)s_\theta(x,y) 可经过修改变成大名鼎鼎的 困惑度 (Perplexity, PPL)PPL(θ)=exp(sθ(x,y)len(y))PPL(\theta)=\exp(-\frac{s_\theta(x,y)}{len(y)})

Thoughts

  • DPO 事实上就是利用正负样本的困惑度来对比学习,从而优化模型性能
  • 原理非常简单,本质就是用对比学习直接对 困惑度 进行训练,收敛速度比 RLHF
  • 但可以预见的是:DPO 的损失计算方式比 RLHF,那么必然对数据集的质量要求更高,否则会导致模型性能下降
  • 一个常用的做法是在 SFT 之后,在 RLHF 之前使用 DPO 快速将模型性能提升到一个较高水平,然后再使用 RLHF 进一步优化

TL;DR

  • 本文是大模型RLHF的入门教程,目的是用尽可能简单的方式介绍大模型RLHF的基本概念和原理,不涉及过多实现细节。

什么是 RLHF

  • RLHFReinforcement Learning with Human Feedback (人类反馈强化学习) 的缩写,目的是在强化学习任务中引入人类反馈,以提高强化学习模型的性能
  • 大多数分类体系下,SFT 属于 RLHF 中的一个步骤,但本文不讨论 SFT,假设起点模型已经做好了 pre-trainSFT
  • RLHF 的目标是:
    • helpful
    • honst
    • harmless

RLHF 相关术语

  • policy model:策略模型,即待训练的强化学习模型
  • reference model:参考模型(基准模型),通常是冻结参数的原始 policy model,防止模型在 RL 过程中出现退化(可理解为正则化项)
  • reward model:奖励模型,用于评估 policy model 生成的结果的好坏,相当于强化学习的 environment
  • value model:价值模型,通常是 policy model 的一个子模块,用于评估状态的价值(也可以理解为对环境的估计),辅助 policy model 的决策,仅仅在 RLHF 过程中有效,训练结束后会被丢弃

RLHF 的基本实现原理

  • RLHF 的实现包括两个阶段:
    • 训练奖励模型阶段
    • 强化学习阶段

1. 训练奖励模型阶段

  • 数据格式:prompt + response 1 + response 2 + label
    • prompt:问题描述
    • response 1:模型生成的回答 1
    • response 2:模型生成的回答 2
    • label:人类标注的哪个回答更好(0 / 1
  • 模型:
    • 模型结构:通常是 LLM,可以由待 RLHF 的策略模型初始化,也可以重新定义模型结构和随机初始化
    • 输入:prompt + response
    • 输出:reward,是一个标量,表示 response 相对于 prompt 的好坏(logits
    • 损失函数:L=1Ni=1Nlogσ(R(prompt,responsechosen)R(prompt,responsereject))L=-\frac{1}{N}\sum_{i=1}^N \log \sigma(R(prompt, response_{chosen}) - R(prompt, response_{reject}))
      • R(prompt,response)R(prompt, response):奖励模型就当前 promptresponse 的奖励(优秀程度)
      • σ\sigmasigmoid 函数
      • NNbatch size
      • responsechosenresponse_{chosen}:人类标注的更好的回答
      • responserejectresponse_{reject}:人类标注的更差的回答

2. 强化学习阶段

  • RLHF 的强化学习阶段通常使用 PPO 算法
  • RLHF 中使用的 PPO 与传统 PPO 有以下几点不同:
    1. 数据格式不同:
      • 传统 PPO 需要的数据格式为: (st,at,rt,st+1,done)(s_t, a_t, r_t, s_{t+1}, done)
      • RLHF 需要的数据格式为: (st,at,rt)(s_t, a_t, r_t),其中:
        • sts_tprompt + 已生成的部分 response
        • ata_tresponse 的下一个词
        • rtr_treward(s_t, a_t)
    2. reward 的计算方式不同:
      • 传统 PPO 中的 reward 由环境给出
      • RLHF 中的 rewardreward model 给出
    3. policy model 的更新方式不同:
      • 为了防止 policy modelRL 过程中退化,需要在损失函数中加入 policy modelreference modelKL 散度作为正则化项
      • reference model 通常是冻结参数的原始 policy model 或者是一个不同架构的经过 pre-train + SFT 的模型
  • 综上,RLHF 强化学习阶段的损失函数为:
    • LRLHF=LPPO+βDKL(policy,reference)L_{RLHF} = L_{PPO} + \beta \cdot D_{KL}(policy, reference)
      • LPPOL_{PPO}:传统 PPO 的损失函数
      • β\beta:正则化系数
      • KLdiv(policy,reference)KL_{div}(policy, reference)policy modelreference modelKL 散度
      • LPPO=E[min(rt(θ)A^t,clip(rt(θ),1ϵ,1+ϵ)A^t)]+c1MSE(Vt,Vttarg)+c2H[πθ(atst)]L_{PPO} = \mathbb{E}[\min(r_t(\theta) \hat{A}_t, \text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon) \hat{A}_t)] + c_1 \cdot \text{MSE}(V_t, V_t^{targ}) + c_2 \cdot H[\pi_{\theta}(a_t|s_t)]
        • At=λδt+lA_t = \lambda \delta_{t+l}:优势函数
        • δt=Reward(st,at)+V(st+1)V(st)\delta_t=Reward(s_t,a_t)+V(s_{t+1})- V(s_t)
        • rt(θ)=πθ(atst)πθold(atst)r_t(\theta)=\frac{\pi_{\theta}(a_t|s_t)}{\pi_{\theta_{old}}(a_t|s_t)}policy model 新旧策略概率比值
        • VtV_t:状态价值函数
        • VttargV_t^{targ}:目标状态价值函数,Vttarg=Reward(st,at)V_t^{targ}=Reward(s_t,a_t)
        • H[πθ(atst)]H[\pi_{\theta}(a_t|s_t)]:熵
        • c1,c2c_1, c_2:权重系数

Thoughts

  • 用对比学习的方式训练奖励模型,可以减少人工标注数据量,同时降低数据标注的主观性噪声
  • PPORLHF 中的应用,重点是搞清楚损失函数中每一项的意义,尤其是状态、动作、下一个状态、奖励的定义

TL;DR

  • 本文是大模型RAG的入门教程,不是来自某一篇 paper / survey,目的是用尽可能简单的方式介绍大模型RAG的基本概念和原理,不涉及过多实现细节。

什么是 RAG

  • RAGRetrieval-Augmented Generation (检索增强生成) 的缩写,目的是在生成任务中引入检索模块,将检索到的信息通过 prompt 模板的方式传递给生成模块,以提高生成模型的性能。

RAG 想要解决的问题

  • 通常来说,RAG 想要解决的问题可以概括为三点:
    1. 大模型知识的局限性:大模型在训练过程中只能学习到非实时的、公开的知识,对于实时的、私有的知识无法获取。
    2. 幻觉问题:大模型无法记住所有的信息细节,即使这些信息曾在训练数据中出现过。
    3. 数据安全问题:敏感私域数据用于训练大模型可能导致数据泄露,而 RAG 可以通过检索模块实现数据的隔离。

RAG 的基本实现原理

rag_1.webp

  • RAG 要分为两步:
    • 数据准备阶段:
      1. 数据提取
      2. 文本切分
      3. 向量化
      4. 数据入库
    • 应用阶段:
      1. 用户提问
      2. 数据检索(召回)
      3. 注入 prompt
      4. LLM 生成答案

1. 数据提取

  • 数据加载:包括多格式数据加载、不同数据源获取等,根据数据自身情况,将数据处理为同一个范式
  • 数据处理:包括数据过滤、压缩、格式化等
  • 元数据获取:提取数据中关键信息,例如文件名、Title、时间等

2. 文本切分

  • 文本切分主要考虑两个因素:
    1. embedding 模型的 sequence length 限制情况
    2. 语义完整性对整体的检索效果的影响
  • 一些常见的文本分割方式包括:
    1. 句分割:以 “句” 的粒度进行切分,保留一个句子的完整语义,常见切分符包括:句号、感叹号、问号、换行符等
    2. 固定长度分割:根据 embedding 模型的 sequence length 长度限制,将文本分割为固定长度(例如 256 / 512token),这种切分方式会损失很多语义信息,一般通过在头尾增加一定冗余量来缓解

3. 向量化

  • 向量化是将文本转换为向量的过程,目的是将不同长度的文本转换为固定长度的向量,有许多开源的 embedding 模型可以使用,例如:

模型名称 描述 获取地址
ChatGPT-Embedding ChatGPT-Embedding由OpenAI公司提供,以接口形式调用。 https://platform.openai.com/docs/guides/embeddings/what-are-embeddings
ERNIE-Embedding V1 ERNIE-Embedding V1由百度公司提供,依赖于文心大模型能力,以接口形式调用。 https://cloud.baidu.com/doc/WENXINWORKSHOP/s/alj562vvu
M3E M3E是一款功能强大的开源Embedding模型,包含m3e-small、m3e-base、m3e-large等多个版本,支持微调和本地部署。 https://huggingface.co/moka-ai/m3e-base
BGE BGE由北京智源人工智能研究院发布,同样是一款功能强大的开源Embedding模型,包含了支持中文和英文的多个版本,同样支持微调和本地部署。 https://huggingface.co/BAAI/bge-base-en-v1.5

Transformer base 的模型天然可以将不同长度的文本转换为固定长度的向量,因此在向量化模型本质就是一个去掉 projection 头的 LLM(当然训练阶段的目标设置也不同)

4. 数据入库

  • 数据向量化之后,作为索引存储到专用的向量检索数据库中,常用的向量检索数据库包括:
    • FaissFacebook 开源的高性能相似向量检索库,支持多种索引结构(如 IVFHNSWPQ),不支持分布式存储和检索
    • MilvusZilliz 开源的分布式高性能相似向量检索库,较 Faiss 更为复杂,支持分布式存储和检索
    • ChromaChroma AI 开源的轻量化相似向量检索库
    • ESElasticSearch 是一款全文搜索引擎,支持文本检索,不原生支持向量检索,但可以通过插件实现

5. 用户提问

graph TD
    A[User Query]
    B[Generate Similar Queries By LLM]
    C1[Vector Search Query 1]
    C2[Vector Search Query 2]
    C3[Vector Search Query 3]
    C4[Vector Search Query 4]
    C5[Vector Search Original Query]
    D[Reciprocal Rank Fusion]
    E[Re-ranked Results]
    F[Generative Output]
    A --> B
    B --> C1
    B --> C2
    B --> C3
    B --> C4
    A --> C5
    C1 --> D
    C2 --> D
    C3 --> D
    C4 --> D
    C5 --> D
    D --> E
    E --> F
  • 大模型首先会根据用户的提问生成一系列相似的查询(或拆分查询内容等),得到一系列的检索查询,这样做的目的是为了提高检索的召回率。

6. 数据检索(召回)

  • 常见的数据检索方法包括:相似性检索、全文检索等,根据检索效果,一般可以选择多种检索方式融合,提升召回率。
  • (向量)相似性检索:即计算查询向量与所有存储向量的相似性得分,返回得分高的记录。常见的相似性计算方法包括:余弦相似性、欧氏距离、曼哈顿距离等。
  • 全文检索:全文检索是一种比较经典的检索方式,在数据存入时,通过 关键词 构建倒排索引;在检索时,通过关键词进行全文检索,找到对应的记录。

7. 注入 prompt

  • prompt 是一种模板化的文本,用于将检索到的信息传递给生成模块,prompt 的构建需要考虑以下几点:
    1. prompt 的内容应该尽可能简洁,不应该包含过多冗余信息
    2. prompt 的内容应该尽可能完整,包含检索到的信息的核心内容
    3. prompt 的内容应该尽可能通用,适用于多种检索结果
  • 以下是一个 prompt 的示例:
1
2
3
4
5
6
【任务描述】
假如你是一个专业的客服机器人,请参考【背景知识】,回答【问题】
【背景知识】
{content} // 数据检索得到的相关文本
【问题】
石头扫地机器人P10的续航时间是多久?

一些高级 RAG 技巧

1. 检索

1.1 向量存储索引

  • 最简答的索引方式是向量存储索引,即一段文本对应一个向量,通过向量相似性检索得到相似文本,如下图所示:
    rag_2.webp

1.2 分层索引

  • 分层索引是对向量存储索引的一种优化,通过对文本经过不同等级的抽象,得到不同层次的向量,从而提高检索效率,如下图所示:
    rag_3.webp

1.3 假设性问题和假设性回答

  • 假设性问题是让 LLM 提前给每一段知识库中的文本生成一个假设性问题,当用户提问时,可以直接通过假设性问题进行检索,从而提高检索效率。
  • 假设性回答也被称为 Hypothetical Document Embeddings(假设性文本嵌入),是用 LLM 给用户的 query 生成一个假设性回答,用假设性回答去检索,从而提高检索效率。

1.4 内容增强索引

  • 内容增强索引原理非常简单,即在检索时,将检索到的文本内容相关的其他文本(例如这段文本前后内容等)内容也一并返回,从而提高检索召回率。

1.5 融合索引

  • 将向量索引和关键词索引融合,提高召回

2. 重排和过滤

  • 重排和过滤是在检索结果返回后,对检索结果进行进一步处理,提高生成效果的一种方法。
  • 过滤/重排的常用方法包括:
    • 通过元数据过滤/重排:例如只保留某个时间段的数据
    • 通过相似性分数过滤/重排
    • 使用其他 LLM 进行过滤/重排

3. 查询转换

  • 查询转换是指在用户提问时,将用户的提问转换为更适合检索的形式,提高检索效率。
  • 例如,用户提问为:“在 Github 上,LangchainLlamaIndex 这两个框架哪个更受欢迎?”,则可以转换为:
    • “Langchain 在 Github 上有多少星?”
    • “Llamaindex 在 Github 上有多少星?”

4. 结合 RAG 的聊天引擎

  • 总体上,结合 RAG 的聊天引擎可以分成两种范式,如下图所示:
    rag_4.webp

5. 智能体(Agent)

  • 例如多文档智能体,就是给每个文档初始化一个智能体,该智能体能进行文档摘要制作和传统问答机制的操作,还有一个顶层智能体,负责将查询分配到各个文档智能体,并综合形成最终的答案。总体流程如下图所示:
    rag_5.webp

6. 响应合成

  • 这是任何 RAG 管道的最后一步——根据检索的所有上下文和初始用户查询生成答案。响应合成的主要方法有:
    • 通过将检索到的上下文逐块发送到 LLM 来优化答案
    • 概括检索到的上下文,以适应提示
    • 根据不同的上下文块生成多个答案,然后将它们连接或概括起来

Thoughts

  • 看起来 RAG 是一种自由度很高的范式,可以根据具体的任务需求,设计出不同的检索和生成策略,但是也正是因为这种自由度,RAG 的实现难度也相对较高,需要综合考虑检索和生成的各种因素,才能设计出一个高效的 RAG 系统。

URL

TL;DR

  • 谷歌提出了一种新的模型分布策略 GShard,通过 条件计算自动分片 的方式,提高了模型的训练效率和推理速度,算是大模型 MOE 的开山之作。
  • 本质就是传统 MoE 替换掉 Transformer 中的 FFN

Architecture

总体架构

Gshard.png

  • GShard 将传统 LLMTransformer 层的 FFN 替换成 MOE 模块,每个 token 只激活固定数量的 expert,大幅降低了计算复杂度。
  • GShardMOE 模块分为 routingexpert 两部分
    • routing 负责选择激活的 expert
    • expert 负责计算 FFN 的输出。

详细计算流程

  • 输入:hRseq_len×dimh \in \mathbb R^{seq\_len\times dim},表示 token 序列的输入特征,seq_len 表示序列长度,dim 表示特征维度
  • 经过 self-attention + norm,这一步和传统 Transformer 模型一致,输出 hRseq_len×dimh' \in \mathbb R^{seq\_len\times dim}
  • 计算 routing:
    1. 计算每个 tokenrouting 分数,表示每一个 token 和每一个专家的相关程度,s=h×Wgs=h'\times W_g,其中 WgRdim×num_experts, sRseq_len×num_expertsW_g \in \mathbb R^{dim\times num\_experts},\ s \in \mathbb R^{seq\_len\times num\_experts}
    2. 每个 token 取固定数量的 expertkk 表示激活的 expert 数量,kk 在本文中取 22,其它非激活的 expertrouting 分数设为 -\infty
    3. routing 分数归一化变成概率,g=Softmax(s)g=Softmax(s)gg 表示 gate,表示每个 token 激活的 expert 的输出采纳程度
  • 计算 expert:
    1. 计算当前 token 已激活的 expert 的输出,hexpert=FFN(hexpert)h_{expert}=FFN(h'_{expert})
    2. expert 的输出按照 gate 加权求和,hfinal=i=1kgihexpertih_{final}=\sum_{i=1}^{k}g_i\cdot h_{expert_i}
  • 本文中 MoE Transformer layer 和标准 Transformer layer 是交替出现的

Thoughts

  • GShard 确实在不减小模型参数量的情况下,大幅减小计算复杂度,对于推理速度的提升有很大帮助,带动了一波 LLM-MoE 的发展。
  • 其中 MoE 的部分并不新颖,重点在于 MoE 放置到 Transformer Block 中确实比较有创新。