Zhangzhe's Blog

URL

• paper: https://arxiv.org/pdf/2501.12948
• code: https://github.com/deepseek-ai/DeepSeek-R1

TL;DR

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

Algorithm

总体流程

Benchmark

DeepSeek-R1-Zero

• DeepSeek-R1-Zero 是用 DeepSeek v3 base (pre-train) 直接通过 RL 训练得到的，不经过 SFT
• 用了以下的 Template 让模型生成 CoT 推理结果：
  
• 奖励建模分成两个部分：
  • 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 包含 CoT 的 SFT 数据

Distill model

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

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 成功的关键，甚至可以说如果 DeepSeek 开源了 80K 条合成数据，那么 CoT 领域的研究将会有一个新的起点

Reference

• https://zhuanlan.zhihu.com/p/22103896675

URL

• paper: https://arxiv.org/pdf/2412.19437
• code: https://github.com/deepseek-ai/DeepSeek-V3

TL;DR

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

Algorithm

总体流程

• Pre-Training
• Context Extension
• Post-Training
  

Multi-token predicition (MTP)

• 上图给出了一次预测 next token 和 $next^2$ token 和 $next^3$ token 的 MTP 模型结构。
• 具体来说，MTP 是在标准 Decoder only Transformer 的基础上，额外引入了 k-1 个 token 预测头，每个 token 预测头都是一个 Norm + Concat + Projection + Transformer block 结构。
• 假设图上从左到右的 token head 分别记作 $head_0$、$head_1$、$head_2$，则 $head_0$ 预测的是 next token，$head_1$ 预测的是 $next^2$ token，$head_2$ 预测的是 $next^3$ token。
• MTP 过程必须要保持因果性：
  • $head_0$ 输入 $token_0$，主干网络输出 $h_0$，output layer 将 $h_0$ 映射为 next token 预测，记作 $token_1'$。
  • $head_1$ 输入 $h_0$ 和 $token_1$，提取特征记作 $h_1$，output layer 将 $h_1$ 映射为 $next^2$ token 预测，记作 $token_2'$。
  • $head_2$ 输入 $h_1$ 和 $token_2$，提取特征记作 $h_2$，output layer 将 $h_2$ 映射为 $next^3$ token 预测，记作 $token_3'$。
• embedding layer 和 output layer 在 MTP 中是共享的。
• MTP 通常只用在训练阶段，即在推理阶段只使用 $head_0$ 的输出。
• 也可以将 inference MTP 和 speculative decoding 结合，提高模型生成效率。

Auxiliary-loss-free strategy for load balancing

训练阶段负载均衡策略

• 训练阶段负载均衡策略并不是本文原创，而是来自于 Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts

$$g'_{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}$$

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

推理阶段

• 推理阶段无需额外的负载均衡策略，直接使用 $s_{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 之间的负载均衡。
• 具体来说，需要关注：
  • 离散选择频率：每一个 Expert 在 Sequence 中的激活次数
  • 连续激活强度：每一个 Expert 在 Sequence 中平均专家分数
• 公式表述：

$$L_{Bal}=\alpha\sum_{i=1}^{N_r}f_iP_i$$

$$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))$$

$$s_{i,t}'=\frac{s_{i,t}}{\sum_{j=1}^{N_r}s_{j,t}}$$

$$P_i=\frac{1}{T}\sum_{t=1}^{T}s_{i,t}'$$

• 其中：
  • $L_{Bal}$ 是 Complementary Sequence-Wise Auxiliary Loss
  • $\alpha$ 是一个超参数，是一个很小的权重
  • $f_i$ 是 Expert $i$ 的离散选择频率
  • $P_i$ 是 Expert $i$ 的连续激活强度
  • $T$ 是 Sequence 的长度
  • $N_r$ 是 Expert 的数量
  • $K_r$ 是每个 token 激活的 Expert 数量
  • $\mathbb I(\cdot)$ 是指示函数
  • $s_{i,t}$ 是 $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 算法，减少了空泡，基本上实现了计算和通信的高度重叠。
  
• 优化了跨节点全对全通信内核，充分利用节点内 GPU 间的 NVLink 带宽和节点间 InfiniBand 带宽，同时减少了通信所需的流式多处理器(SMs)资源占用
• 通过精细的内存管理优化，使得模型训练无需依赖开销较大的张量并行(Tensor Parallelism, TP)技术

FP8

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

Thoughts

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

References

• https://zhuanlan.zhihu.com/p/14890557782

URL

• paper: https://arxiv.org/pdf/2006.11239
• code: https://github.com/hojonathanho/diffusion

TL;DR

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

Algorithm

DDPM 总体结构

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

前向过程

$$x_t=\sqrt {1-\beta_t} \cdot x_{t-1}+\sqrt \beta_t\cdot \epsilon_{t-1}\ ,\ \ \epsilon_{t-1}\sim N(0, I)$$

• 其中：
  • $\epsilon_t$ 是随机高斯噪声，即 torch.randn
  • $\beta_t$ 是超参数，控制噪声的强度，随着 $t$ 的增大，$\beta_t$ 从 1e-4 逐渐增大到 2e-2
  • 总时间步数 $T=1000$

前向过程的累积表示

$$x_t=\sqrt {\bar\alpha_t}\cdot x_0+\sqrt{1-\bar\alpha_t}\cdot \epsilon$$

• 具体推理过程：
  • $x_1 = \sqrt{1-\beta_1}\cdot x_0+\sqrt{\beta_1}\cdot \epsilon_0$
  • $x_2 = \sqrt{1-\beta_2}\cdot x_1+\sqrt{\beta_2}\cdot \epsilon_1$
  • 将 $x_1$ 代入 $x_2$ 的公式中，得到 $x_2$ 的累积表示： $x_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$
  • $x_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$
  • 以此类推，得到 $x_t$ 的累积表示：$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$
  • 令 $\bar\alpha_t=\prod_{i=1}^t(1-\beta_i)$，得到 $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$
  • 需要将噪声项合并成一个正态分布，需要确定方差，对于每一项 $\sqrt{\beta_i\cdot\prod_{j=i+1}^t(1-\beta_j)}$，其方差为 $\beta_i\cdot\prod_{j=i+1}^t(1-\beta_j)$
  • 总噪声的方差为 $\sum_{i=1}^t\beta_i\cdot\prod_{j=i+1}^t(1-\beta_j)$
  • 代入 $\bar\alpha_t$ 到噪声项的方差中，得到总噪声的方差为 $\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}$
  • 由于 $\frac{\beta_i}{\bar\alpha_i}=\frac{1-\alpha_i}{\bar\alpha_i}=\frac{1}{\bar\alpha_i}-\frac{1}{\bar\alpha_{i-1}}$，因此使用 裂项求和 公式可得：$\sum_{i=1}^t\frac{\beta_i}{\bar\alpha_i}=\frac{1}{\bar\alpha_t}-\frac{1}{\bar\alpha_0}$
  • 由于 $\bar\alpha_0=1$，因此 $\sum_{i=1}^t\frac{\beta_i}{\bar\alpha_i}=\frac{1}{\bar\alpha_t}-1$
  • 代入噪声总方差公式中，得到总噪声的方差为 $\bar\alpha_t\cdot(\frac{1}{\bar\alpha_t}-1)=1-\bar\alpha_t$
  • 因此，累加形式的前向过程可表示为：$x_t=\sqrt{\bar\alpha_t}\cdot x_0+\sqrt{1-\bar\alpha_t}\cdot \epsilon$
• 累积表示的好处：
  • 在给定 $t$ 时，可以直接从 $x_0$ 生成 $x_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))$$

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

训练过程

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

• 输入数据：从数据集中随机选择一批原始图像 $x_0$
• 随机选择扩散步数：对于每张图像，随机选择一个扩散步数 $t$，范围从 $1$ 到 $T$ 随机选择
• 生成噪声数据：通过前向过程生成噪声图像 $x_t$
• 预测噪声：使用模型 $\epsilon_\theta(x_t,t)$ 预测添加到 $x_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)  # 更新模型参数

模型输入输出

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

损失函数

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

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

• 其中：
  • $\epsilon$ 是真实噪声
  • $\epsilon_\theta(x_t,t)$ 是模型预测的噪声
  • $\mathbb E_{x_0,t}$ 表示对所有的 $x_0$ 和 $t$ 求期望

推理过程（采样过程）

用 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
29
30
31
32
33
34

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 = (
            sample - (1 - alphas[t]) / torch.sqrt(1 - alphas_cumprod[t]) * residual
        ) / torch.sqrt(alphas[t])   # 3. optionally look at image
        # Add random noise only for t > 1
        if t > 1:
            noise = torch.randn_like(sample).to(device)
            sample += torch.sqrt(1 - alphas[t]) * noise    
    if (i + 1) % 50 == 0:
        display_sample(sample, i + 1)

• 通过 UNet2DModel 加载预训练模型
• 通过 DDPMScheduler 构建调度器（主要作用是生成噪声和去噪）
• 通过 torch.randn 生成随机噪声作为输入
• 通过 scheduler.step 逐步生成图像
• 最终生成一张完全去噪的新图像，这张图像并不在任何数据集中（只是风格看起来像 Church-256 数据集），是模型自己生成的（这也意味着没办法去指导模型生成特定的图像）

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

• 理论上，可以将一张数据集中的图像逐步添加噪声，在 $t=T$ 步之后，使用模型去噪，得到一张和原始图像非常相似的图像
• 但实际测试中，当 $t$ 较大时（大于 300，$T$ 等于 1000），得到的去噪图像和原始图像相差较大，在 $t<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
45
46
47
48
49
50

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 = (
            sample - (1 - alphas[t]) / torch.sqrt(1 - alphas_cumprod[t]) * residual
        ) / torch.sqrt(alphas[t])   # 3. optionally look at image
        # Add random noise only for t > 1
        if t > 1:
            noise = torch.randn_like(sample).to(device)
            sample += torch.sqrt(1 - alphas[t]) * noise  
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}")

上图表示 t=300 时，原图 -> 噪声图 -> 去噪图

上图表示 t=500 时，原图 -> 噪声图 -> 去噪图

上图表示 t=1000 时，原图 -> 噪声图 -> 去噪图

Thoughts

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

URL

• paper: https://arxiv.org/pdf/2201.12086
• code: https://github.com/salesforce/BLIP

TL;DR

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

Algorithm

BLIP 整体架构

• BLIP 在预训练阶段，包含三种损失函数
  • Image-Text Contrastive Loss：对齐图像和文本表征，和 CLIP 的对齐方式类似，不同的是 BLIP 引入了 Momentum Model，可以使得训练更加稳定
  • Image-Text Macthing Loss：图像和文本匹配损失，输出 True 或 False，本质是交叉熵损失
  • 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
      • 这种状态处于 encoder 和 decoder 之间，比 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 的不同模块组合实现
• BLIP 的 Bootstrapping 策略是一种自我增强的方式，可以使得模型在无监督的情况下不断提升性能，如今这几乎是多模态预训练算法的标配
• 从实验的角度告诉我们，encoder 和 decoder 做充分的参数共享是可行的，仅仅在少量的位置上不共享参数，就可以实现完全不同任务的模型

URL

• paper: https://arxiv.org/pdf/2305.18290

TL;DR

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

什么是 DPO

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

DPO 的数据格式

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

DPO 损失函数的设计

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

$$s_\theta(x,y)=\beta(\log p_\theta(y|x) - \log p_{ref}(y|x))$$

$$p_\theta(y|x)=\prod_{t=1}^T p_\theta(y_t|y_{<t},x)$$

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

• $L(\theta)$：损失函数
• $\theta$：模型参数
• $ref$：参考模型，DPO 之前的冻结模型
• $\beta$：温度系数（通常在 0.1 - 0.5 之间），平衡探索和守成
• $D$：数据集
• $x$：输入
• $y^+$：人类标注的更好的回答
• $y^-$：人类标注的更差的回答
• $s_\theta(x,y)$：模型对 $(x,y)$ 的得分
• $p_\theta(y|x)$：模型对 $x$ 生成 $y$ 的概率
• $\sigma$：sigmoid 函数
• $T$：序列长度
• $y_{<t}$：$y$ 的前 $t-1$ 个元素
• $s_\theta(x,y)$ 可经过修改变成大名鼎鼎的 困惑度 (Perplexity, PPL)：$PPL(\theta)=\exp(-\frac{s_\theta(x,y)}{len(y)})$

Thoughts

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

TL;DR

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

什么是 RLHF

• RLHF 是 Reinforcement Learning with Human Feedback (人类反馈强化学习) 的缩写，目的是在强化学习任务中引入人类反馈，以提高强化学习模型的性能
• 大多数分类体系下，SFT 属于 RLHF 中的一个步骤，但本文不讨论 SFT，假设起点模型已经做好了 pre-train 和 SFT
• 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 的流程

1. 训练奖励模型

• 奖励模型通常选比原模型尺寸小的预训练模型，将 head 换成一个标量输出的 head
• 奖励模型的输入是 prompt + response，输出是一个奖励标量，即：奖励模型作用的对象是一个序列，不是一个 token
• 训练数据格式是：(prompt, response_chosen, response_reject)
  • prompt：输入给策略模型的问题
  • response_chosen：期待的大模型输出
  • response_reject：不期待的大模型输出
• 训练损失函数：pairwise ranking loss，本质就是对比学习

$$L=-E_{(x,y_w,y_l)}\log \sigma(R(x, y_w) - R(x, y_l))$$

• 其中：
  • x：prompt
  • y_w：win response，也就是 response_chosen
  • y_l：lose response，也就是 response_reject

2. 强化学习阶段

2.1 Rollout 阶段

对每个 Prompt x：

1. 用当前 Actor $\pi_\theta$ 自回归采样一条完整回答 y=(y_1, …, y_T)，直到遇到结束符
2. 把 (x, y) 拼成一整句喂给冻结的 Reward Model，得到一个标量 $r_T$
3. 同时用 $\pi_{ref}$ 算同样这条回答的 log-prob：$\log \pi_{ref}(y|x)$

2.2 Reward Shaping 阶段

• 由于 Reward Model 只给整句打分，而 PPO 需要每一步都有 Advantage，所以要先构造 “即时奖励” 再让 GAE 去算回报：

$$r_t=\left\{ \begin{aligned} r_T-\beta\cdot KL_t && \text{if}\ t=T \\ -\beta\cdot KL_t && \text{if}\ t<T \end{aligned} \right.$$

$$KL_t=\log \pi_\theta(y_t|x,y<t) - \log \pi_{ref}(y_t|x,y<t)$$

• 其中：
  • KL 表示 KL 散度损失，用于方式模型偏离参考模型太远
  • 参考模型：经过 SFT 之后冻结的模型
  • $\beta$ 是 KL 散度系数， 要调得够大（0.1–0.5），否则 $\pi_\theta$ 会迅速 “讨好” RM，生成乱码或重复文本

Thoughts

• 用对比学习的方式训练奖励模型，可以减少人工标注数据量，同时降低数据标注的主观性噪声
• 强化学习的部分设计的挺巧妙的，不过仍然是单回合强化学习，如果从 Agentic RL 角度看，需要多回合强化学习来微调模型使得模型更像一个 Agent

TL;DR

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

什么是 RAG

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

RAG 想要解决的问题

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

RAG 的基本实现原理

• 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 / 512 个 token），这种切分方式会损失很多语义信息，一般通过在头尾增加一定冗余量来缓解

3. 向量化

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

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

4. 数据入库

• 数据向量化之后，作为索引存储到专用的向量检索数据库中，常用的向量检索数据库包括：
  • Faiss：Facebook 开源的高性能相似向量检索库，支持多种索引结构（如 IVF、HNSW、PQ），不支持分布式存储和检索
  • Milvus：Zilliz 开源的分布式高性能相似向量检索库，较 Faiss 更为复杂，支持分布式存储和检索
  • Chroma：Chroma AI 开源的轻量化相似向量检索库
  • ES：ElasticSearch 是一款全文搜索引擎，支持文本检索，不原生支持向量检索，但可以通过插件实现

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 向量存储索引

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

1.2 分层索引

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

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

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

1.4 内容增强索引

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

1.5 融合索引

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

2. 重排和过滤

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

3. 查询转换

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

4. 结合 RAG 的聊天引擎

• 总体上，结合 RAG 的聊天引擎可以分成两种范式，如下图所示：
  

5. 智能体（Agent）

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

6. 响应合成

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

Thoughts

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

URL

• paper: https://arxiv.org/pdf/2006.16668
• code: 没找到

TL;DR

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

Architecture

总体架构

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

详细计算流程

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

Thoughts

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

URL

• paper: https://arxiv.org/pdf/2401.06066
• code: https://github.com/deepseek-ai/DeepSeek-MoE

TL;DR

• 本文在 GShard 的基础上，提出了一种新的混合专家语言模型 DeepSeekMoE，通过 孤立共享专家 和 细分专家 的方式，提高了模型性能并降低了计算复杂度。
• 替换的也是 Transformer 中的 FFN 层。

Architecture

传统 MoE 模型（例如 GShard）

• 传统 MoE 如图 a 所示，核心思想是将 Transformer 的 FFN 替换为 MoE，每个 token 通过 Gate 机制选择不同的 Expert 来处理。
• 用公式表示为：
  • $h_t^l=\sum_{i=1}^{N}(g_{i,t}FFN_i(u^l_t))+u^l_t$
  • $g_{i,t}=\left\{\begin{array}{ll}{s_{i,t},} & {s_{i,t}\in TopK({s_{j,t}|1\le j\le N}, K),} \\ {0,} & {otherwise,}\end{array}\right.$
  • $s_{i,t}=Softmax({u_t^{l}}^Te_i^l)$
  • 其中：
    • l 表示第 l 层
    • N 表示 Expert 的数量
    • K 表示每个 token 保留的 Expert 数量

细粒度 MoE 模型

• 如图 b 所示，和传统 MoE 的区别是将专家切分的更小，专家数量更多，也可以理解为传统 MoE 中的 Expert 也是由多个 Sub-Expert 组成。
• 用公式表示为：
  • $h_t^l=\sum_{i=1}^{mN}(g_{i,t}FFN_i(u^l_t))+u^l_t$
  • $g_{i,t}=\left\{\begin{array}{ll}{s_{i,t},} & {s_{i,t}\in TopK({s_{j,t}|1\le j\le mN}, mK),} \\ {0,} & {otherwise,}\end{array}\right.$
  • $s_{i,t}=Softmax({u_t^{l}}^Te_i^l)$
  • 其中：
    • l 表示第 l 层
    • N 表示 Expert 的数量
    • m 表示每个 Expert 中包含的 Sub-Expert 的数量
    • K 表示每个 token 保留的 Expert 数量

细粒度 MoE + 孤立共享专家

• 如图 c 所示，在细粒度 MoE 的基础上，引入了 Isolated Shared Expert，这种专家不参与 Gate 选择，而是在所有 token 之间共享。
• 用公式表示为：
  • $h_t^l=\sum_{i=1}^{K_s}FFN_i(u^l_t)+\sum_{i=K_s+1}^{mN}(g_{i,t}FFN_i(u^l_t))+u^l_t$
  • $g_{i,t}=\left\{\begin{array}{ll}{s_{i,t},} & {s_{i,t}\in TopK({s_{j,t}|K_s+1\le j\le mN}, mK-K_s),} \\ {0,} & {otherwise,}\end{array}\right.$
  • $s_{i,t}=Softmax({u_t^{l}}^Te_i^l)$
  • 其中：
    • l 表示第 l 层
    • N 表示 Expert 的数量
    • m 表示每个 Expert 中包含的 Sub-Expert 的数量
    • K 表示每个 token 保留的 Expert 数量
    • $K_s$ 表示 Isolated Shared Expert 的数量

Thoughts

• 这篇论文名字起的有点大《迈向终极专家专业化的MoE语言模型》，但是实际上只是在 GShard 的基础上做了一些小的改进。

URL

• paper: https://arxiv.org/pdf/2405.04434
• code: https://github.com/deepseek-ai/DeepSeek-V3

TL;DR

• DeepSeek 系列是国内一家名叫深度求索的科技公司推出的一种混合专家语言模型，这家公司背后是幻方量化。
• DeepSeek-V2 的核心思想是 Multi-head Latent Attention (MLA) 和 DeepSeekMoE 两个模块，这两个模块分别替换了传统 Transformer 层的 Attention Module 和 FFN，是本博客想要探讨的重点。

Architecture

Multi-head Latent Attention (MLA)

MLA 总体计算流程

• 核心思想是将 hidden feature 通过 MLP 映射到 latent space，降低计算（attention）和存储（kv cache）的复杂度。
• 和主流的 kv cache 方案不同，MLA 只需要 cache $c_t^{KV}$ 和 $k_t^R$ 两部分。

MLA 详细计算流程

• 输入：$h_t \in \mathbb R^d$ ，表示第 t 个 token 在某个 attention layer 层的输入特征
• 计算 key / value：
  • $c_t^{KV}=W^{DKV}h_t$ ，其中 $W^{DKV} \in \mathbb R^{d_c\times d}$，DKV 表示 down-projection key value
  • $k_t^C=W^{UK}c_t^{KV}$ ，其中 $W^{UK} \in \mathbb R^{d_hn_h\times d_c}$，UK 表示 up-projection key，$d_c \ll d_hn_h$
  • $v_t^C=W^{UV}c_t^{KV}$ ，其中 $W^{UV} \in \mathbb R^{d_hn_h\times d_c}$，UV 表示 up-projection value，$d_c \ll d_hn_h$
• 计算 query：
  • $c_t^Q=W^{DQ}h_t$ ，其中 $W^{DQ} \in \mathbb R^{d_c'\times d}$，DQ 表示 down-projection query
  • $q_t^C=W^{UQ}c_t^Q$ ，其中 $W^{UQ} \in \mathbb R^{d_hn_h\times d_c'}$，UQ 表示 up-projection query，$d_c' \ll d_hn_h$
• 计算 RoPE：
  • RoPE 是一种在输入 attention layer 之前，对 query 和 key 做 position encoding 的方法
  • RoPE 和 MLA 在设计上是冲突的，因此 MLA 对 RoPE 做了一些修改，主要是：额外计算 multi-head query rotary 和 shared key rotary
  • $[q^R_{t,1},q^R_{t,2},...,q^R_{t,n_h}]=q_t^R=RoPE(W^{QR}c_t^Q)$ ，其中 $W^{QR} \in \mathbb R^{d_h^Rn_h\times d_c'}$，$QR$ 表示 query rotary
  • $k_t^R=RoPE(W^{KR}k_t^C)$ ，其中 $W^{KR} \in \mathbb R^{d_h^R\times d}$，$KR$ 表示 key rotary
  • $q_t,i=[q^C_{t,i},q^R_{t,i}]$ ，将 query 和 query rotary 拼接，得到正式的 query
  • $k_t,i=[k^C_{t,i},k^R_{t}]$ ，将 key 和 key rotary 拼接，得到正式的 key
• 计算 attention：
  • $o_{t,i}=\sum^t_{j=1}Softmax_j(\frac{q^T_{t,i}k_{j,i}}{\sqrt{d_h+d_h^R}})v^C_{j,i}$
• 计算 output：
  • $u_t=W^O[o_{t,1},o_{t,2},...,o_{t,n_h}]$ ，其中 $W^O \in \mathbb R^{d\times d_hn_h}$

MLA 优势

• 大幅降低了 kv cache 的存储空间
  

DeepSeekMoE

• 见 DeepSeekMoE 详解

Thoughts

• MLA 算是对 self-attention 比较深度的改进，兼顾了 kv cache 和 RoPE 的设计，比 MQA 和 GQA 设计看上去更用心。
• DeepSeekMoE 感觉在 GShard 上的改进并不大，主要是在 Expert 的基础上加入了 Sub-Expert 的概念，以及常开的 Isolated Shared Expert 设计。
• 这二者结合，给了 DeepSeek-V2 一个很好的性能提升，可以在相同计算复杂度下，塞下更多的参数量，提高模型表现。