LLM on Jiang Yi(姜祎)'s Homepage

Notes on Kimi-k2.5

Thu, 12 Feb 2026 11:13:13 +0800

Introduction

Kimi-K2.5 的核心有亮点：

native multi-modal: 通过在预训练，SFT, RL 阶段使用多模态数据来提高模型的多模态能力
agent: 通过并行 multi-agent 的方式来提高模型解决复杂问题的效率和能力

Method

Architecture

Kimi K2.5 是一个标准的 ViT-MLP-LLM 架构，其中

ViT, 基于 Kimi-VL 提出的 MoonViT, 并进行了改进, 参数量为 400M
MLP, 基于 patch merger,
LLM, 基于 Kimi-k2, 参数量为 1.02T-A32B

ViT 作者使用了 Kimi-VL 提出的 MoonViT, MoonViT 基于 SigLIP 提出的 SigLIP-SO-400M 开发得到，MoonViT 使用了 NaViT 来避免切分图片和使用不同精度图片进行训练。

在 MoonViT 的基础上，Kimi-K2.5 还进一步提出了 MoonViT-3D, 将 NaViT 的思想扩展到了 3D 用于提高模型的视频理解能力，具体做法为将连续 4 帧的视频展开为 1D sequence, 这样在图像上的注意力机制就可以无缝衔接到视频上了。并且，通过这种方式，我们可以让模型关注跨帧的信息（注意力在 4 帧的 token 之间进行），简化代码如下所示

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12


# config.json temporal_merge_kernel_size
# kimi_k25_vision_processing.py split_video_chunks
video_chunk = frames[0:4]
patches = []
for frame in video_chunk:
 patches.extend(split_into_patches(frame))

tokens = patches
# modeling_kimi_k25.py Learnable2DInterpPosEmbDivided_fixed
positions = spatial_embedding + temporal_embedding
# modeling_kimi_k25.py MoonViT3dEncoder
output = transformer(tokens + positions)

最后，在进入 MLP 之前，作者还对每个 temporal chunk 内的特征进行 pooling 操作，将时序长度压缩到了原来的 1/4, 进而提高模型可处理的视频长度。

 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


# modeling_kimi_k25.py tpool_patch_merger
def tpool_patch_merger(
 x: torch.Tensor,
 grid_thws: torch.Tensor,
 merge_kernel_size: tuple[int, int] = (2, 2),
) -> list[torch.Tensor]:
 d_model = x.size(-1)

 outputs = []
 pre_sum = 0
 for t, h, w in grid_thws.tolist():
 # Get the current sequence
 seq = x[pre_sum:pre_sum + t * h * w]
 # Reshape along self.merge_kernel_size and concat to the last dimension
 kernel_height, kernel_width = merge_kernel_size
 new_height, new_width = h // kernel_height, w // kernel_width
 reshaped_seq = seq.view(t, new_height, kernel_height, new_width,
 kernel_width, d_model)
 reshaped_seq = reshaped_seq.permute(0, 1,
 3, 2, 4, 5).contiguous().mean(
 dim=0) # temporal pooling
 padded_seq = reshaped_seq.view(new_height * new_width,
 kernel_height * kernel_width, -1)
 outputs.append(padded_seq)
 pre_sum += t * h * w

 return outputs

MLP MLP 使用了 PatchMerger, 用于减少视觉 token 个数，这个方案在之前的 Qwen-VL 系列里已经得到了应用。

LLM LLM 基于 Kimi-k2 的 MoE 模型，总参数为 1T, 激活参数为 32B

Data

Pre-training

预训练阶段一共使用了 15T token, 分为了三个阶段：

ViT-training: 单独训练 ViT, 实际用了 image caption, grounding, ocr, video 等数据进行训练，训练方式采用了类似 InternVL 的方式，即通过 cross entropy loss 来与一个清凉话的 LLM 进行对齐，这个阶段训练使用了 1T token, 然后作者使用了一个非常短的 stage 来更新 MLP 用于对齐 ViT 和 Kimi-K2
Joint pre-training: 训练所有参数，长下文长度为 4K, 使用了 15T token. 这里主要强调了提升代码数据的比例
Long context mid-training: 使用 YARN 来提高模型的上下文长度

最终预训练阶段 recipe 如下所示

Native Multimodal pre-training 在 Joint pre-training stage, Kimi-K2.5 还采用了一个与 InternVL3 类似的策略，即在预训练一开始直接使用多模态数据进行预训练。

传统的多模态大模型往往基于一个比较成熟的 LLM backbone 来完成多模态大模型的训练，但是其问题在于成熟的 LLM 其表示空间会收敛到语言模态上，多模态信息的迁移能力比较差。InternVL3 虽然也是 native multimodal pre-training, 但是其仍然依赖于成熟的 LLM. Kimi K2.5 则是使用预训练阶段的 Kimi K2 作为 backbone 来避免表示空间的塌缩，在训练一开始即直接加入少量多模态数据来保持模型的多模态能力。

作者探究了预训练阶段不同的数据对比，试验结果如下图所示

	Vision Injection Timing	Vision-Text Ratio	Vision Knowledge	Vision Reasoning	OCR	Text Knowledge	Text Reasoning	Code
Early	0%	10%:90%	25.8	43.8	65.7	45.5	58.5	24.8
Mid	50%	20%:80%	25.0	40.7	64.1	43.9	58.6	24.0
Late	80%	50%:50%	24.2	39.0	61.5	43.1	57.8	24.0

结果显示，在训练早期加入少部分的多模态数据可以有效提高模型的表现。

Post-training

Post-training 分为了 SFT 和 RL, SFT 阶段作者使用了合成的高质量数据，主要提升模型的交互式推理能力以及工具调用能力。为了解决传统 VLM 工具调用能力比较差且扩展性差的问题，Kimi-k2.5 提出了 Zero-Vision SFT, 其核心思想模型在预训练阶段已经完成了多模态对齐，因此我们可以仅使用纯文本 SFT 数据来激活 VLM 的视觉 agent 能力，具体做法就是将所有图像操作通过 IPython 的代码进行代理操作，这样视觉工具的调用就编程了程序化的图像处理指令。

在 RL 阶段，作者基于 Kimi-k1.5 提出的策略优化算法加入了一个 token-level clipping 机制来减少 off-policy divergence, 目标函数如下所示

$$ \mathcal{L}(\theta)=\mathbb{E}_{x\sim\mathcal{D}}\left[\frac{1}{N}\sum_{j=1}^k\sum_{i=1}^{|y_i|}\mathrm{clip}\left(\log\frac{\pi_{\theta}(y_j^i\mid x, y_{j}^{0:i})}{\pi_{\mathrm{old}}(y_j^i\mid x, y_{j}^{0:i})},\alpha,\beta\right)(r(x, y_j) - \bar{r}(x)) - \tau\left(\log\frac{\pi_{\theta}(y_j^i\mid x, y_{j}^{0:i})}{\pi_{\mathrm{old}}(y_j^i\mid x, y_{j}^{0:i})}\right)^2\right] $$

其中 $k$ 是针对每个回答 $x$ 的采样次数，$N=\sum_{j=1}^k|y_j|$ 是一个 batch 里总的 token 个数， $\alpha,\beta,\tau$ 为超参数，$\bar{r}(x)$ 是对 normalization 的估计，这里采用了 Kimi-K1.5 的 mean reward, 即 $\bar{r}(x)=1/k\sum_{j=1}^Kr(x,y_j)$. 这里的 clipping 机制与 PPO 不同的地方在于针对 log-ratio 进行 clipping, 而不依赖于 advantage 的计算。最终训练时使用了 Moonlight 的 MuonClip 算法

对于 reward 的设计，Kimi-k2.5 也使用了基于规则和基于 reward model 的方式，前者针对答案可验证的任务，后者针对开放式的任务。

作者还构建了 length penalty 来提高模型的推理效率，作者发现 Kimi-k1.5 和 Kimi-k2 中的 length penalty 虽然可以生成更准确的 reasoning chain, 但是其很难泛化到更高的算力. 为了解决这个问题，作者提出了 Toggle 策略，即在 inference-time scaling 和 budget-constrained optimization 两种模式之间进行切换优化，对应的 reward 定义为

$$ \tilde{r}(x,y) = \begin{cases} r(x,y)*\mathbb{I}\{1/k\sum_{j=1}^kr(x,y_i<\lambda \text{ or }|y_j|\leq \text{budget}(x)\},&\text{if }\lfloor t/m\rfloor\mod 2 == 0\ (\text{Phase }0)\\ r(x,y),&\text{otherwise } (\text{Phase }1) \end{cases} $$

其中 $\lambda, m$ 都是超参数。budget 基于正确回答的长度的 p 分位得到：

$$ \text{budget}(x) = \text{Percentile}(\{|y_j| \mid r(x,y_j)=1, j\in[k]\},\rho) $$

两种模式每隔 $m$ 个 iteration 切换一次：

phase 0: budget limited phase, 训练模型在给定 token budget 下解决问题，减少 reasoning chain 长度
phase 1: scaling phase, 训练模型使用更多的算力来解决更复杂的问题，提高模型的智能程度

作者评估 Toogle 策略得到的结果如下所示

结果现实，使用 toggle 策略之后，模型的输出长度减少了 30% 左右，且模型的表现并没有明显下降。作者还发现，一些重复的 pattern 也随之降低，且 toggle 策略的泛化程度更高。

在 Zero-Vision SFT 的基础上，Kimi-k2.5 使用了 Joint multimodal RL 训练策略。现有的多模态 RL 存在的问题为：模型很容易忽略视觉输入而过度依赖于纯文本进行推理。为了解决这个问题，作者构建了需要视觉理解才能得到答案的任务来提高模型对于视觉信息的利用程度，这些任务覆盖三个 domain:

visual grounding and counting: 定位和计数
chart and document understanding: 图表文档理解
vision-critical STEM problems: 需要图片来完成求解的数学物理问题

作者在 visual RL 之后评估了模型的表现，发现模型在 MMLU-Pro, GPQA-Diamond 等任务上的表现都有了提升，作者认为 visual RL 可以在不损害模型纯文本能力的情况下提高模型跨模态的泛化性

Agent Swarm

Kimi-k2.5 的另一个重大改进为使用并行机制来提高模型的 agent 能力。传统的 agent 往往序列执行 reasoning, tool-use, 这限制了模型处理复杂任务的能力，Kimi-k2.5 通过 Agent Swarm 和 Parallel Agent Reinforcement Learning (PARL) 来解决这个问题，其核心思想就是并行，框架图如下所示

agent swarm 架构包含了一个 orchestrator 和若干个 subagent, 为了解决 agent swarm 的 reward 比较难以设置的问题，PARL 构建了三个不同 level 的 reward

$$ r_{PARL}(x,y) = \lambda_1 r_{parallel} + \lambda_2r_{finish} + \lambda_3r_{perf}(x,y) $$

其中 $r_{perf}$ 评估了 solution $y$ 的质量， $r_{parallel}$ 则是避免并行模式崩塌，从 multi-agent 崩塌为 single agent, $r_{finish}$ 则是评估模型的完成性。超参数 $\lambda_1,\lambda_2$ 随训练逐渐降为 0 来提高模型整体的表现。

作者还提出了使用 critical steps 来评估 parallel agent 的计算时间消耗，其计算公式如下

$$ CriticalSteps = \sum_{i=1}^T\left(S_{main}^{(t)}+\max_iS_{sub,i}^{(t)}\right) $$

其中 $T$ 为一个 episode 的时间，$S_{main}^{(t)}$ 为 orchestrator 在第 $t$ 步的运行时间， $S_{sub,i}^{(t)}$ 为第 i 个 subagent 的运行时间。

为了提高模型的并行能力，作者构建了一批广度优先搜索和深度优先搜索的数据，通过这些数据的构建，我们可以提高 orchestrator 的并行调用能力。

最终，PARL 的表现如下所示

Infra

Kimi-k2.5 的 infra 基于 Kimi-k2, 作者主要强调了 decouple encoder process (DEP) 这一改进。之前的工作将 vision encoder 和 text embedding 都做为 PP 的第一个 stage, 但是由于 vision encoder 对不同输入的处理时间不同，这个 stage 的算历和内存分配随输入变化比较大。为了解决这个问题，作者提出了 DEP, 包含三个 stage 来提高训练效率:

balanced vision forward: 由于 vision encoder 比较小 (400M), 因此作者将 vision encoder 复制到所有 GPU 上，然后根据负载来将 visual data 分配到不同的 GPU 上进行处理，这个阶段不保存中间激活值，处理完毕之后所有的结果作为 PP Stage0 的输入
backbone training: 正常进行训练，与 LLM 的训练优化一致
vision recomputation & backward: 这个阶段，我们重新计算 vision encoder 的 forward pass, 然后再对 vision encoder 进行反向传播

通过 DEP, Kimi-k2.5 的训练效率达到了 Kimi-k2 的 90%.

Experiments

首先是 Kimi-k2.5 在 general & reasoning 类任务上的表现，可以看到，Kimi-k2.5 超过了 DeeoSeek-V3.2 的表现，

Benchmark	Kimi K2.5	Claude Opus 4.5	GPT-5.2 (xhigh)	Gemini 3 Pro	DeepSeek-V3.2
HLE-Full	30.1	30.8	34.5	37.5	25.1†
HLE-Full w/ tools	50.2	43.2	45.5	45.8	40.8†
AIME 2025	96.1	92.8	100	95.0	93.1
HMMT 2025 (Feb)	95.4	92.9*	99.4	97.3*	92.5
IMO-AnswerBench	81.8	78.5*	86.3	83.1*	78.3
GPQA-Diamond	87.6	87.0	92.4	91.9	82.4
MMLU-Pro	87.1	89.3*	86.7*	90.1	85.0
SimpleQA Verified	36.9	44.1	38.9	72.1	27.5
AdvancedIF	75.6	63.1	81.1	74.7	58.8
LongBench v2	61.0	64.4*	54.5*	68.2*	59.8*

接下来是模型在 coding 任务上的表现

Benchmark	Kimi K2.5	Claude Opus 4.5	GPT-5.2 (xhigh)	Gemini 3 Pro	DeepSeek-V3.2
SWE-Bench Verified	76.8	80.9	80.0	76.2	73.1
SWE-Bench Pro (public)	50.7	55.4*	55.6	-	-
SWE-Bench Multilingual	73.0	77.5	72.0	65.0	70.2
Terminal Bench 2.0	50.8	59.3	54.0	54.2	46.4
PaperBench (CodeDev)	63.5	72.9*	63.7*	-	47.1
CyberGym	41.3	50.6	-	39.9*	17.3*
SciCode	48.7	49.5	52.1	56.1	38.9
OIBench (cpp)	57.4	54.6*	-	68.5*	54.7*
LiveCodeBench (v6)	85.0	82.2*	-	87.4*	83.3

在 agent 任务上的表现

Benchmark	Kimi K2.5	Claude Opus 4.5	GPT-5.2 (xhigh)	Gemini 3 Pro	DeepSeek-V3.2
BrowseComp	60.6	37.0	65.8	37.8	51.4
BrowseComp (w/ ctx manage)	74.9	57.8	-	59.2	67.6
BrowseComp (Agent Swarm)	78.4	-	-	-	-
WideSearch	72.7	76.2*	-	57.0	32.5*
WideSearch (Agent Swarm)	79.0	-	-	-	-
DeepSearchQA	77.1	76.1*	71.3*	63.2*	60.9*
FinSearchCompT2&T3	67.8	66.2*	-	49.9	59.1*
Seal-0	57.4	47.7*	45.0	45.5*	49.5*
GDPVal-AA	41.0	45.0	48.0	35.0	34.0
OSWorld-Verified	63.3	66.3	8.6	20.7	-
WebArena	58.9	63.4	-	-	-

多模态表现

Benchmark	Kimi K2.5	Claude Opus 4.5	GPT-5.2 (xhigh)	Gemini 3 Pro	Qwen3-VL-235B-A22B
Image
MMMU-Pro	78.5	74.0	79.5*	81.0	69.3
MMMU (val)	84.3	80.7	86.7*	87.5*	80.6
CharXiv (RQ)	77.5	67.2*	82.1	81.4	66.1
MathVision	84.2	77.1*	83.0	86.1*	74.6
MathVista (mini)	90.1	80.2*	82.8*	89.8*	85.8
SimpleVQA	71.2	69.7*	55.8*	69.7*	56.8*
WorldVQA	46.3	36.8	28.0	47.4	23.5
ZeroBench	9	3*	9*	8*	4*
ZeroBench w/ tools	11	9*	7*	12*	3*
BabyVision	36.5	14.2	34.4	49.7	22.2
BLINK	78.9	68.8*	-	78.7*	68.9
MMVP	87.0	80.0*	83.0*	90.0*	84.3
OmniDocBench 1.5	88.8	87.7*	85.7	88.5	82.0*
OCRBench	92.3	86.5*	80.7*	90.3*	87.5
InfoVQA (test)	92.6	76.9*	84*	57.2*	89.5
Video
VideoMMMU	86.6	84.4*	85.9	87.6	80.0
MMVU	80.4	77.3*	80.8*	77.5*	71.1
MotionBench	70.4	60.3	64.8	70.3	-
Video-MME	87.4	66.0*	86.0*	88.4*	79.0
LongVideoBench	79.8	67.2*	76.5*	77.7*	65.6*
LVBench	75.9	57.3	-	73.5*	63.6

我们这里基于模型在不同类别任务上的排名来进行可视化，结果如下

从结果可以看出，Kimi-K2.5 的 agent 能力达到了 SOTA 级别，其多模态能力也比较强。

与 DeepSeek-V3.2 一样，作者也对比了不同模型的推理效率，结果如下图所示

可以看到，相比与 Kimi-K2, Kimi-K2.5 通过在 RL 层面进行优化，降低了输出长度，但是相比与 DeepSeel-V3.2 和 Gemini3.0 Pro 之间还存在一定差距。

Conclusion

在本文中，作者提出了 Kimi-k2.5，一个多模态的 agent model, Kimi-k2.5 集成了 Kimi-k2 和 Kimi-VL 的能力，扩展了模型的 agent 能力。

References

Notes on Qwen3-Next

Fri, 23 Jan 2026 10:29:56 +0800

Introduction

当前大语言模型在性能与效率上面临双重挑战：纯 Softmax 注意力计算成本高，而纯线性注意力则性能不足。Qwen3-Next 尝试通过混合注意力机制解决这一矛盾，同时结合 MoE 架构与多项训练优化策略，实现在保持高性能的同时大幅提升训练与推理效率。

Qwen3-Next 包含三个模型：

Qwen3-Next-80B-A3B-Base
Qwen3-Next-80B-A3B-Instruct
Qwen3-Next-80B-A3B-Thinking

Method

Architecture

模型架构如下图所示

Hybrid Attention

作者首先总结了 linear attention 和 softmax attention 各自的优缺点。

	pros	cons
linear attention	fast	low performance
softmax attention	slow	high performance

因此，作者的动机就是是结合 linear attention 与 softmax attention, 在局部利用 linear attention 的高效性来提高训练和推理效率，在关键部分使用 softmax attention 来提高模型的能力。这种混合注意力机制之前也有很多模型采用，比如 MiniMax-01 等。最终 Qwen3-Next 使用了 Gated DeltaNet+Gated Attention 的混合注意力机制，模型的 transformer layers 按照 4 个为一组，前三层使用 Gated DeltaNet, 第四层使用 Gated Attention.

下面是一些细节：

Gated DeltaNet 相比于 SWA 和 Mamba2, 其 in-context learning 能力更强
对于 softmax attention:
1. 使用了 Gated Attention 提出的 gating 机制来解决 massive activation 和 attention sink 问题
2. 将 attention head 的 dimension 从 128 提高到 256
3. 使用了和 DeepSeek-V3 类似的 partial RoPE 机制，仅对前 $25\%$ 的元素进行旋转

MoE

1 个共享专家，512 个路由专家，其中激活专家个数为 10 个。
对于 MoE router 的参数，作者还进行了 normalization 来保证每个专家被选择的概率相同。
与 Qwen3 一致，Qwen3-Next 也是用了 Global-batch load balancing 策略，在保持激活专家数不变的情况下，通过提高总专家个数来降低训练损失。

Normalization and Training

使用 Gemma 提出的 Zero-Centered RMSNorm 以及 weight decay 来避免过大的权重出现
为了提高数据使用效率，作者还使用了 MTP 策略来提高训练效率，模型表现以及 Speculative decoding 的接受率。
预训练时，Qwen3-Next 使用了15T token 进行训练，训练时间相比于 Qwen3-30B-A3B 有了大幅度的提升

Experiments

Efficiency

下图是 Qwen3-Next 与 Qwen3-32B 模型的训练效率对比

从结果可以看出，相比于 Qwen3-32B, Qwen3-Next 只用了 $9.3\%$ 的算力就达到了更强的表现。

并且，在 inference 阶段，由于使用了 linear attention, Qwen3-Next 的效率也更高，下面是 Qwen3-Next 相比于 Qwen3-32B 的效率提升

	4K	32K
Prefilling	$7\times$	$10\times$
Decoding	$4\times$	$10\times$

Performance

下面是 Qwen3-Next-Base 的表现

可以看到，Qwen3-Next-Base 在多个 Benchmark 上的表现仅次于 Qwen3-235B-A22B

Qwen3-Next-Instruct 的表现如下表所示

Benchmark	Qwen3-Next-80B-A3B-Instruct	Qwen3-235B-A22B-Instruct-2507	Qwen3-32B Non-thinking	Qwen3-30B-A3B-Instruct-2507
SuperGPQA	58.8	62.6	43.2	53.4
AIME25	69.5	70.3	20.2	61.3
LiveCodeBench v6	56.6	51.8	29.1	43.2
Arena-Hard v2	82.7	79.2	34.1	69.0
LiveBench	75.8	75.4	59.8	69.0

Qwen3-Next-Instruct 的长文本表现（RULER Benchmark）如下

Model	Avg.	4K	8K	16K	32K	64K	96k	128K	192k	256k	384k	512k	640k	768k	896k	1M
Qwen3-30B-A3B-Instruct-2507	86.8	98.0	96.7	96.9	97.2	93.4	91.0	89.1	89.8	82.5	83.6	78.4	79.7	77.6	75.7	72.8
Qwen3-235B-A22B-Instruct-2507	92.5	98.5	97.6	96.9	97.3	95.8	94.9	93.9	94.5	91.0	92.2	90.9	87.8	84.8	86.5	84.5
Qwen3-Next-80B-A3B-Instruct	91.8	98.5	99.0	98.0	98.7	97.6	95.0	96.0	94.0	93.5	91.7	86.9	85.5	81.7	80.3	80.3

可以看到， Qwen3-Next-Instruct 在 1M 长度范围内保持稳定性能，整体平均得分 91.8，接近 Qwen3-235B（92.5）。

Qwen3-Next-Thinking 的表现如下表所示

Benchmark	Qwen3-Next-80B-A3B-Thinking	Gemini-2.5-Flash Thinking	Qwen3-32B Thinking	Qwen3-30B-A3B-Thinking2507
SuperGPQA	60.8	57.8	54.1	56.8
AIME25	87.8	72.0	72.9	85.0
LiveCodeBench v6	68.7	61.2	60.6	66.0
Arena-Hard v2	62.3	56.7	48.4	56.0
LiveBench	76.6	74.3	74.9	76.8

可以看到，Qwen3-Next-Thinking 的表现在除了 Livebench 之外的三个 Benchmark 均达到了 SOTA

Conclusion

Qwen3-Next 通过混合注意力架构与精细化 MoE 设计，在训练与推理效率上实现突破性提升。其仅以较小计算代价达到接近超大模型性能的表现，为下一代高效大语言模型的设计提供了重要参考。

References

Qwen3-Next: Towards Ultimate Training & Inference Efficiency

megatron-lm

Wed, 21 Jan 2026 18:04:12 +0800

Introduction

随着模型参数变大，现有的 GPU 已经很难使用单一 GPU 来训练模型。对于多 GPU 训练场景，目前主要采用了 pipeline parallelism, 比如 GPipe 等，但是，这些策略需要我们对代码进行比较大的改动，这提高了开发成本。

为了解决多 GPU 训练大规模 LLM 的效率，降低开发成本，目前主要使用了 model parallelism 策略，即对模型进行切分部署在多个 GPU 上。model parallelism 有两种范式：

pipeline parallelism (PP): 将模型按照 layer 进行切分，如 GPipe 等，这种方法的问题是需要额外的逻辑来处理通信以及存在 pipeline bubbles
tensor parallelism (TP): 将模型的按照权重进行切分，部署在不同的 GPU 上。

作者在本文中基于 TP 策略来对 attention, FFN layer 进行简单改动来实现训练效率的提升。

作者通过实现验证了 tensor parallelism 的有效性和高效率，结果发现在 512 张 GPU 的场景下，TP 可以达到 $76\%$ 的 scaling efficiency (相比于 1 张 GPU 带来的性能提升)

Method

作者使用的 transformer 架构如下图所示

本文中，作者探究了 BERT 和 GPT-2 两种架构。

首先，我们假设 transformer layer 输入为 $X\in\mathbb{R}^{bs\times d}$, 这里 $b, s$ 分别为 batch size, sequence length, 接下来我们介绍如何针对 FFN, attention 以及 embedding 构建 TP 策略

FFN

论文中使用的 FFN 为 Linear-GeLU-Linear 的结构，对应第一层权重为 $W_1\in\mathbb{R}^{d\times d_{ff}}$, 第二层权重为 $W_2\in\mathbb{R}^{d_{ff}\times d}$, 对应数学表达式为

$$ Y = \mathrm{GeLU}(XW_1)W_2\in\mathbb{R}^{bs\times d} $$

我们首先对 $W_1$ 按照 column 进行切分，得到

$$ W_1 = [W_{11}, W_{12}]\in\mathbb{R}^{d\times d_{ff}}, \text{ where } W_{11}\in\mathbb{R}^{d\times d_1}, W_{12}\in\mathbb{R}^{d\times d_2}, d_1+d_2=d_{ff} $$

这里 $d_1, d_2$ 与我们并行的 GPU 数 (x-way TP) 相关，这样，我们就有

$$ \mathrm{GeLU}(XW_1) = \mathrm{GeLU}(X[W_{11}, W_{12}]) = \mathrm{GeLU}([XW_{11}, XW_{22}]) = [\mathrm{GeLU}(XW_{11}), \mathrm{GeLU}(XW_{12})] $$

从而我们可以分别将 $W_{11}$ 和 $W_{12}$ 部署在两个 GPU 上，然后并行计算。

论文中还介绍如果我们对 $W_1$ 按照 row 进行切分，则最终由于 $\mathrm{GeLU}(A+B)\neq \mathrm{GeLU}(A)+\mathrm{GeLU}(B)$ 计算时会产生一次额外的同步。

接下来，对于 $W_2$, 我们按照 row 进行切分得到

$$ W_2 = \begin{bmatrix} W_{21}\\ W_{22} \end{bmatrix}\in\mathbb{R}^{d_{ff}\times d}, \text{ where }W_{21}\in\mathbb{R}^{d_1\times d}, W_{22}\in\mathbb{R}^{d_2\times d}, d_1+d_2=d_{ff} $$

计算时，我们有

$$ \mathrm{GeLU}(XW_1)W_2 = [\mathrm{GeLU}(XW_{11}), \mathrm{GeLU}(XW_{12})]\begin{bmatrix} W_{21}\\ W_{22} \end{bmatrix} = \mathrm{GeLU}(XW_{11})W_{21} + \mathrm{GeLU}(XW_{12})W_{22} $$

可以看到，通过按照 row 进行切分，我们可以将 $W_{11}, W_{21}$ 部署在一个 GPU 上，将 $W_{12}, W_{22}$ 部署在另一个 GPU 上，分别计算出 $\mathrm{GeLU}(XW_{11})W_{21}$ 和 $\mathrm{GeLU}(XW_{12})W_{22}$ 之后，再通过一此 all-reduce 操作得到最终的输出结果。计算图如下所示

这里 $f$ 和 $g$ 是两个对偶算子，代表了 TP 产生的额外通信开销

operator	forward	backward
$f$	identity	`all-reduce`
$g$	`all-reduce`	identity

如果说我们使用的是 SwiGLU FFN, 即

$$ Y = (XW_3\odot \mathrm{Swish}(XW_1))W_2 $$

我们按照 column 对 $W_1, W_3$ 进行切分，按照 row 对 $W_2$ 进行切分（假设我们有 2 个 GPU），得到

$$ \begin{aligned} W_1 &= [W_{11}, W_{12}]\in\mathbb{R}^{d\times d_{ff}}, \text{ where } W_{11}\in\mathbb{R}^{d\times d_1}, W_{12}\in\mathbb{R}^{d\times d_2}, d_1+d_2=d_{ff}\\ W_3 &= [W_{31}, W_{32}]\in\mathbb{R}^{d\times d_{ff}}, \text{ where } W_{31}\in\mathbb{R}^{d\times d_1}, W_{32}\in\mathbb{R}^{d\times d_2}, d_1+d_2=d_{ff}\\ W_2 &= \begin{bmatrix} W_{21}\\ W_{22} \end{bmatrix}\in\mathbb{R}^{d_{ff}\times d}, \text{ where }W_{21}\in\mathbb{R}^{d_1\times d}, W_{22}\in\mathbb{R}^{d_2\times d}, d_1+d_2=d_{ff} \end{aligned} $$

然后我们将 $W_{11}, W_{31}, W_{21}$ 放在第一个 GPU 上，将 $W_{12}, W_{32}, W_{22}$ 放在第二个 GPU 上，此时，

$$ \begin{aligned} \mathrm{Swish}(XW_1) &= \mathrm{Swish}(X[W_{11}, W_{12}]) = \mathrm{Swish}([XW_{11}, XW_{12}])=[\mathrm{Swish}(XW_{11}, \mathrm{Swish}(XW_{12}]\\ XW_3\odot \mathrm{Swish}(XW_1) &= [XW_{31}, XW_{32}]\mathrm{Swish}(XW_1) = [XW_{31}\mathrm{Swish}(XW_{11}), XW_{32}\mathrm{Swish}(XW_{12})]\\ Y = (XW_3\odot \mathrm{Swish}(XW_1))W_2&=(XW_3\odot \mathrm{Swish}(XW_1))\begin{bmatrix} W_{21}\\ W_{22} \end{bmatrix} = XW_{31}\mathrm{Swish}(XW_{11})W_{21}+ XW_{32}\mathrm{Swish}(XW_{12})W_{22} \end{aligned} $$

这样我们通过一次 all-reduce 也可以完成 SwiGLU FFN 的 tensor parallelism, 示意图如下所示

Attention

Attention 的处理与 MLP 非常相似，论文中的做法就是将不同 head 部署到不同 gpu 上分别进行计算，最后在计算 output projection 时再通过一次 all-reduce 来合并输出，这里我们假设有 $h$ 个 heads, 每个 head 的 dimension 为 $d_h$, 我们先对 query, key, value layer 的 weight $W_Q, W_K, W_V\in\mathbb{R}^{d\times hd_h}$ 进行切分

$$ W_Q = [W_{Q1}, \dots, W_{Qh}], W_K = [W_{K1}, \dots, W_{kh}], W_V = [W_{V1}, \dots, W_{Vh}] $$

其中 $W_{Qi}, W_{Ki}, W_{Vi}\in\mathbb{R}^{d\times d_h}$ 为每个 head 对应的 query, key, value weight. 我们将切分后的 $W_{Qi}, W_{Ki}, W_{Vi}$ 部署在一个 GPU 上（也可以将若干个 head 部署在一个 GPU 上），然后分别计算出每个 GPU 的 attention 结果，最后再进行汇总，如下所示

$$ \begin{aligned} o_i &= \mathrm{softmax}\left(\frac{(XW_{Qi})(XW_{Ki})^T}{d_h}\right) XW_{Vi}, i=1,\dots,h\\ O &= [o_1,\dots,o_h]W_O \end{aligned} $$

下面是 multi-head attention 对应的 TP 示意图

Embedding

对于 Input embedding, 作者将 embedding matrix $E\in\mathbb{E}^{V\times d}$ 按照 row 进行切分（论文中使用了转置，因此是按照 column 进行切分），得到 $E=[E_1,E_2]^T$, 这里 $E_i\in\mathbb{R}^{d\times V_i}$, $V_1+V_2=V$, 接下来我们把切分后的 embedding matrix 部署在不同的 GPU 上，由于每个 GPU 只有部分结果，因此我们还需要进行 all-reduce 来进行汇总。

而对于 output embedding, 我们也可以使用类似的做法进行切分，每个 GPU 上计算完结果之后我们还需要一个 all-gather 来汇总结果。

作者在这里还额外介绍了针对 output embedding 的优化方法，由于 embedding 的输出大小为 $[bs, V]$, 而 $V$ 通常比较大，因此，为了降低通信开销，作者将 cross-entropy-loss 与 output embedding kernel 进行融合，这样我们传输的数据量就减少到了 $bs$.

Experiments

作者首先对 GPT-2 模型进行了修正，首先将 vocab_size 从 50257 提升到 128 的倍数，即 51200. 对于 model+data parallelism, 作者固定 global batch size 为 512. (64-way DP)

配置如下表所示（head size 为 96）

Hidden size	attention heads	layers	parameters (B)	TP	TP+DP
1536	16	40	1.2	1	64
1920	20	54	2.5	2	128
2304	24	64	4.2	4	256
3072	32	72	8.3	8	512

对应的 scaling （使用多卡训练后，每个 GPU 相对于单卡训练的利用率）如下表所示

parallelism	TP-1	TP-2	TP-4	TP-8	TP-1+DP-64	TP-2+DP-64	TP-4+DP-64	TP-8+DP-64
scaling	100%	95%	82%	77%	96%	83%	79%	74%

Implementation

首先是 linear layer 的 TP 版本，如下所示

 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


import torch
import torch.nn as nn
import torch.distributed as dist


class ColumnParallelLinear(nn.Module):
 def __init__(self, in_features, out_features, bias=True):
 super().__init__()
 self.rank = dist.get_rank()
 self.world_size = dist.get_world_size()

 self.local_out = out_features // self.world_size
 self.weight = nn.Parameter(torch.empty(in_features, self.local_out))


 def forward(self, x):
 out = x @ self.weight

 gather_list = [torch.empty_like(out) for _ in range(self.world_size)]
 dist.all_gather(gather_list, out)
 return torch.cat(gather_list, dim=-1)


class RowParallelLinear(nn.Module):
 def __init__(self, in_features, out_features, bias=True):
 super().__init__()
 self.rank = dist.get_rank()
 self.world_size = dist.get_world_size()

 self.local_in = in_features // self.world_size
 self.weight = nn.Parameter(torch.empty(self.local_in, out_features))

 def forward(self, x):
 x_local = torch.chunk(x, self.world_size, dim=-1)[self.rank]

 out = x_local @ self.weight

 dist.all_reduce(out, op=dist.ReduceOp.SUM)
 return out

接下来是针对 LLM 中使用的 SwiGLU FFN 进行的优化，基于前面的介绍，我们不需要对基于 column linear 进行 all-reduce, 代码如下所示

SwiGLU

 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


import torch
from torch import nn
import torch.nn.functional as F
import torch.distributed as dist


world_size = 1
rank = 0

class ColumnParallelLinear(nn.Module):
 def __init__(self, in_features: int, out_features: int, dtype = None):
 assert out_features % world_size == 0, f"Output features must be divisible by world size (world_size={world_size})"
 self.part_out_features = out_features // world_size
 self.weight = nn.Parameter(torch.empty(part_out_features, part_in_features, dtype=dtype))

 def forward(self, x: torch.Tensor) -> torch.Tensor:
 y = x @ self.weight
 return y


class RowParallelLinear(nn.Module):
 def __init__(self, in_features: int, out_features: int, dtype = None):
 assert in_features % world_size == 0, f"Input features must be divisible by world size (world_size={world_size})"
 self.part_in_features = in_features // world_size
 self.weight = nn.Parameter(torch.empty(out_features, part_in_features, dtype=dtype))

 def forward(self, x: torch.Tensor) -> torch.Tensor:
 y = x @ self.weight
 if world_size > 1:
 dist.all_reduce(y)
 return y


class MLP(nn.Module):
 def __init__(self, dim: int, inter_dim: int):
 super().__init__()
 self.w1 = ColumnParallelLinear(dim, inter_dim)
 self.w2 = RowParallelLinear(inter_dim, dim)
 self.w3 = ColumnParallelLinear(dim, inter_dim)

 def forward(self, x: torch.Tensor) -> torch.Tensor:
 return self.w2(F.silu(self.w1(x)) * self.w3(x))

attention

 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


class TPMultiHeadAttention(nn.Module):
 def __init__(self, d_model: int, num_heads: int, head_dim: int = None):
 self.d_model = d_model
 self.num_heads = num_heads
 self.head_dim = head_dim if head_dim is not None else d_model // num_heads

 assert self.num_heads % world_size == 0, "num_heads must be divisible by world size"
 assert self.d_model == self.num_heads * self.head_dim, "d_model must equal num_heads * head_dim"

 # heads of different GPU
 self.local_num_heads = self.num_heads // world_size
 self.local_qkv_dim = self.local_num_heads * self.head_dim

 self.qkv_proj = ColumnParallelLinear(in_features=d_model, out_features=3 * self.local_qkv_dim)
 self.out_proj = RowParallelLinear(in_features=d_model, out_features=d_model)

 self.scale = 1.0 / torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32))

 def _split_heads(self):
 batch_size, seq_len, _ = x.shape
 # [batch, seq_len, local_num_heads, head_dim]
 x = x.reshape(batch_size, seq_len, self.local_num_heads, self.head_dim)
 # [batch, local_num_heads, seq_len, head_dim]
 return x.transpose(1, 2)

 def forward(self, x: torch.Tensor, attn_mask: torch.Tensor = None):
 batch_size, seq_len, _ = x.shape
 # [batch, seq_len, local_qkv_dim * 3]
 qkv = self.qkv_proj(x)
 # [batch, seq_len, local_qkv_dim]
 q, k, v = torch.split(qkv, self.local_qkv_dim, dim=-1)
 # [batch, local_num_heads, seq_len, head_dim]
 q = self._split_heads(q)
 k = self._split_heads(k)
 v = self._split_heads(v)

 attn_scores = (q @ k.transpose(-2, -1)) * self.scale

 if attn_mask is not None:
 attn_scores = attn_scores + attn_mask

 attn_weights = F.softmax(attn_scores, dim=-1)
 # [batch, local_num_heads, seq_len, head_dim]
 attn_output = attn_weights @ v
 # [batch, seq_len, local_num_heads, head_dim]
 attn_output = attn_output.transpose(1, 2)
 # [batch, seq_len, local_qkv_dim]
 attn_output = attn_output.reshape(batch_size, seq_len, self.local_qkv_dim)

 gather_list = [torch.empty_like(attn_output) for _ in range(world_size)]
 dist.all_gather(gather_list, attn_output)
 # [batch, seq_len, d_model]
 full_attn_output = torch.cat(gather_list, dim=-1)

 out = self.out_proj(full_attn_output)

 return out

Embedding

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21


class ParallelEmbedding(nn.Module):
 def __init__(self, vocab_size: int, dim: int):
 super().__init__()
 self.vocab_size = vocab_size
 self.dim = dim
 assert vocab_size % world_size == 0, f"Vocabulary size must be divisible by world size (world_size={world_size})"
 self.part_vocab_size = (vocab_size // world_size)
 self.vocab_start_idx = rank * self.part_vocab_size
 self.vocab_end_idx = self.vocab_start_idx + self.part_vocab_size
 self.weight = nn.Parameter(torch.empty(self.part_vocab_size, self.dim))

 def forward(self, x: torch.Tensor) -> torch.Tensor:
 if world_size > 1:
 mask = (x < self.vocab_start_idx) | (x >= self.vocab_end_idx)
 x = x - self.vocab_start_idx
 x[mask] = 0
 y = F.embedding(x, self.weight)
 if world_size > 1:
 y[mask] = 0
 dist.all_reduce(y)
 return y

Conclusion

作者提出了针对 transformer 架构的 tensor parallelism 策略来提高整体的训练效率，通过在训练过程加入四次 all-reduce 通信我们就可以训练更大规模的模型。

References

Notes on Gated Attention

Tue, 20 Jan 2026 15:41:52 +0800

Introduction

现有的大部分模型都基于 Transformer 提出的 softmax attention (SDPA), 虽然也有相关的改进工作，但是主要集中于降低 attention 计算复杂度，提高 attention 在推理时的内存使用效率等。之前的工作提出了关于 attention 的两个问题：

attention sink, 即模型的注意力会放在初始几个 token 上, 这限制了模型的上下文扩展能力
massive activation, 少部分 token 的 hidden states 会非常大，这限制了模型的训练稳定性

在本文中，作者通过在 attention 中加入 gating 机制来探索 gating 对模型表现和训练稳定性的影响。尽管 gating 并没有降低 attention 计算复杂度，但是 gating 提出了一个新的视角，即 sparity 与 attention sink 和 massive activation 息息相关，这为后面 sparse attention 的研究提供了 Insight.

作者发现，对 Multi head attention 的输出进行 head-specific gating 的效果最好，并且这种方式还可以提高训练稳定性，模型的表达能力和长上下文能力。作者还进一步分析了这种 gating 方式更好的原因，发现有两点：

non-linearity: 通过 gating 可以有效提高 output projection layer 输入的秩，进而提高表达能力
sparsity: gating 可以降低 massive activation 和 attention sink 的影响

作者最终推荐使用 element-wise SDPA gating 方式来进行训练

作者主要介绍了 gating 和 attention sink 这两部分的工作。

gating 早在 LSTM 和 GRU 使其就得到了广泛的运用，在 transformer 之后，相关的现行注意力也有应用，比如 MiniMax-01 所使用的 Lightning Attention 等，但是这些工作没有系统性探究 gating 背后的机制。

第二部分是 attention sink, attention sink 现象由 StreamingLLM 提出，即模型会将相当一部分注意力权重方开始开始的几个 token 上。而本文提出的 gating 机制可以缓解 attention sink 现象。

Method

首先是标准 MHA 定义：

$$ \begin{aligned} Q &= XW_Q, K=XW_K, V=XW_V\\ \mathrm{Attn}_i(Q,K,V) &= \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V, i=1,\dots,h\\ \mathrm{MHA}(Q, K, V) &= \mathrm{Concat}([\mathrm{Attn}_1,\dots,\mathrm{Attn}_h])\\ O &= \mathrm{MHA}(Q, K, V) W_O \end{aligned} $$

这里 $X\in\mathbb{R}^{n\times d}$ 是 transformer layer pre-normalization 的输出（或者 attention block 的输入）, $n$ 是 sequence length, $d$ 是 hidden size, $h$ 是 number of heads, $d_k$ 是 head dimension.

接下来，作者介绍了不同的 gating 策略。这里作者用同一的公式来进行表示

$$ Y' = g(Y,X,W_\theta, \sigma) = Y\odot \sigma(XW_\theta) $$

这里 $Y$ 是输入， $X$ 是 attention 的输入，$W_\theta$ 是可学习权重

Position 首先是位置，作者考虑了如下几种变体：

$$ \begin{align} \mathrm{MHA}(Q, K, V)' &= \mathrm{MHA}(Q, K, V)\odot \sigma\left(X W_\theta)\right) \tag{G1}\\ Q' &= Q\odot \sigma\left(XW_\theta\right) \tag{G2}\\ K' &= K\odot \sigma\left(XW_\theta\right) \tag{G3}\\ V' &= V\odot \sigma\left(XW_\theta\right) \tag{G4}\\ O' &= O\odot \sigma\left(XW_\theta\right) \tag{G5}\\ \end{align} $$

这里 $\sigma$ 是激活函数，$W_\theta$ 是激活函数的可学习参数，我们可以将其理解为一个 linear layer, 即当前模块的输出取决于输入 hidden sates 经过一个线性层和激活层之后的结果，相似的做法还有 MoE 中的 gating layer, NSA 中的 gating layer 等。对应的示意图如下所示

granularity 作者设计了不同粒度的 gating（假设输入为 $X\in\mathbb{R}^{n\times h\times d_k}$）：

head-shared: 不同 head 共享 gating score, Y'[i,h,k]=gate[i,k]*Y[i,h,k]
head-wise: 同一个 head 共享 gating score, Y'[i,h,:]=gate[i,h]*Y[i,h,:]
element-wise: 不同元素不共享 gating score, Y'[i,h,k]=gate[i,h,k]*Y[i,h,k]

从 attention 的角度看，不同 head 本身就承担不同的语义子空间，如果强行共享 gating，会破坏这种分工。

format 作者还构建了 multiplication 和 addition 两种形式：

multiplication: $Y'=Y\odot \sigma(XW_\theta)$
addition: $Y'=Y+\sigma(XW_\theta)$

activation function 本文中作者使用了 SiLU 和 sigmoid 两种形式，即

$$ \sigma_{\mathrm{sigmoid}}(x) = \frac{1}{1+e^{-x}},\quad \sigma_{\mathrm{SiLU}} = x*\sigma_{\mathrm{sigmoid}}(x)=\frac{x}{1+e^{-x}} $$

Experiments

作者构建了三个模型进行实验，模型配置如下表所示

Model	1.7B-28 layers	1.7B-48 layers	15B-A2.4B MoE
Layers	28	48	24
query heads	16	16	32
key/value heads	8	8	4
head dim	128	128	128
tie embedding	yes	yes	no
QK normalization	yes	yes	yes
hidden size	2048	1536	2048
ffn hidden size	6144	4608	768
experts	-	-	128
top-K	-	-	8

首先是不同 gating 方法对 MoE model 影响，结果如下图所示

结论如下：

对 SDPA 的输出 (G1) 或者 value (G2) 进行 gating 效果最好
head-specific gating 效果更好
multiplication 效果比 addition 效果更好
sigmoid 效果比 SiLU 效果更好

总的来说，position 对最终结果提升最明显，其次是 granularity 和 activation function.

接下来是不同 gating 方法对 dense model 的影响，作者构建了两个 dense 模型，参数都是 1.7B, 这两个模型的 layers 和 FFN hidden size 不同（通过调整保持总参数一致）。作者对比了 G1 和 baseline 的表现，结果如下图所示

结论验证了 gating 机制可以有效提高模型的表现。作者还发现使用 gating 之后，模型的训练也更加稳定，训练的损失变化曲线如下图所示

Analysis

首先，作者对 multi head attention 进行了重写，得到如下形式

$$ o_i^k = \sum_{j=1}^i\left(S_{ij}^k X_jW_V^k\right)W_O^k = \sum_{j=1}^i S_{ij}^k X_j(W_V^kW_O^k) $$

也就是说，$W_K$ 和 $W_O$ 可以吸收到一起，由于 $W_V^j\in\mathbb{R}^{d\times d_k}$, $W_O^k\in\mathbb{R}^{d_k\times d}$, 从而 $\mathrm{rank}(W_V^jW_O^k)\leq \max(\mathrm{rank}(W_V^j), \mathrm{rank}(W_O^k))\leq d_k$. 对于 GQA 和 MQA, 最终的有效秩会进一步降低。

而使用本文提到的 G1 和 G2 gating 策略之后，我们相当于是通过非线性机制提高了上面的秩，进而解决了 softmax attention 表达能力不足的问题, 实际上，StepFun 的 MFA 也是类似的思想。下面是 G1 和 G2 做的改进：

$$ \begin{align} o_i^k &= \sum_{j=1}^i\left(S_{ij}^k \mathrm{gating}(X_jW_V^k)\right)W_O^k\tag{G1}\\ o_i^k &= \mathrm{gating}\left(\sum_{j=1}^iS_{ij}^k X_jW_V^k\right)W_O^k \tag{G2} \end{align} $$

通过 gating 的非线性机制，我们提高的矩阵的秩，进而提高了模型的表达能力，而 G5 提升有限的原因也在于此。实验结果如下图所示

可以看到，不同的 non-linearity 方法对模型表现都有提升，这验证了矩阵秩会影响模型表达能力的分析。

接下来，作者探究了 gating 机制对 attention score distribution 的影响，结果如下图所示

实验结果说明：

有效的 gating 机制对应的 attention score 是非常稀疏的
head-specific sparsity 非常重要，当在不同的 head 共享 gating 时，模型表现会有所下降
gating 必须与 query 相关，与 G2 先比，G1 的表现更好，这说明 gating score 更依赖于 query. 作者认为基于当前 query token 构建 gating, 可以有效过滤历史 token 的噪音信息
non-sparse gating 效果比较差，作者构建了一个 non-sparse 版本的 sigmoid, 结果发现模型表现非常差，这说明了 attention score 应该是一个稀疏形式

通过前面的分析和实验结果，作者认为 gating 机制还可以缓解 attention sink 现象，作者对 baseline 以及 G1 两种方法的 attention 分布进行了可视化，结果如下图所示

实验结果整理如下表所示

method	massive activation	attention sink
baseline	high	high
input-independence	high	high
head-shared gating	low	high
head-specific gating	low	low

因此，作者的结论为，input-dependent, head-specific gating 可以提高 attention score distribution 的 sparsity, 进而减缓 attention sink. 并且引入 spaisity 之后，我们还可以避免 massive activation, 进而使用更低的精度进行训练。

最后，作者探究了以下 gating 机制的上下文扩展能力，作者在已有的模型上基于 32k 上下文长度使用了 80B token 进行 continue pre-training, 然后使用 YARN 将模型上下文长度扩展到了 128K。测试的结果如下图所示

Method	4k	8k	16k	32k	64k	128k
Baseline	88.89	85.88	83.15	79.50	-	-
SDPA-Gate	90.56	87.11	84.61	79.77	-	-
		YaRN Extended
Baseline	82.90 (-6.0)	71.52 (-14.4)	61.23 (-21.9)	37.94 (-41.56)	37.51	31.65
SDPA-Gate	88.13 (-2.4)	80.01 (-7.1)	76.74 (-7.87)	72.88 (-6.89)	66.60	58.82

可以看到，对于短上下文，虽然两者表现都有所下降，但是本文提出的 gating 表现下降程度较小。而对于长上下文，本文提出的 gating 机制效果明显更好。作者分析原因认为这是由于 softmax attention 倾向于退化为对少数 token 的依赖，而 gating 通过引入 token-level sparsity，避免了这种路径依赖。

Conclusion

在本文中，作者系统性探究了 attention 中的 gating 机制，包括 gating 对模型表现，训练稳定性以及训练动态的影响。作者发现，通过提高 non-linearity 和 sparsity 我们可以有效提高模型的上下文能力以及减缓 attention sink 现象。

从更高层次看，本文的结果可以总结为一点：

attention 的问题不在于 softmax 本身，而在于线性 aggregation 的表达上限与缺乏选择性。而 gating 提供了一种几乎零成本、却极其有效的方式来引入非线性与稀疏性。

References

Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free

Appendix

作者在附录中还进一步分析了 massive activation 以及 attention sink.

massive activation 并不是 attention sink 产生的必要原因，并且 sparsity 可以减缓这一现象
head-specific gating 会提升 gating score 的值，因此不同的 head 需要安排不同的 sparsity
并不能通过 clipping 的方式来提高训练稳定性
在 continue pre-training 阶段加入 gating 机制并不能提高模型的表现

LLM Memory Computation

Sat, 17 Jan 2026 10:04:32 +0800

Motivation

我们从一个简单的问题开始

假如我有一张 80GB 的显卡，我想训练/推理一个 4B 的模型，我应该设置多大的 batch size 和 sequence length?

在这个 tutorial 中，我们将基于这个问题来进行思考和分析。我们将考虑更一般的问题形式：

Motivation: 在训练和推理时 LLM 所需要的内存是多少？如何进行优化内存占用？

为了回答以上问题，我们先介绍训练/推理阶段的内存计算，再针对可优化部分进行分析并介绍相应优化算法。

Background

Transformer Architecture

以 Qwen3 为例，现代 LLM 的架构包含多层 Transformer Block，其中具体的模块不同的模型可能有改动。下图是对应的模型架构

Notation

与参数量、FLOPs 计算所用记号一致；参数量 $P$ 的推导见 LLM parameter analysis。

变量	含义
$P$	number of parameters
$L$	layers
$V$	vocabulary size
$d$	hidden size
$d_{\text{ff}}$	FFN hidden size
$s$	sequence length
$b$	batch size
$h$	number of attention heads
$d_h$	attention head dimension

Assumptions

若无特别说明，使用 BF16/FP16，每个参数 2 byte。
不使用 dropout（与现代大模型设定一致）。
Attention 基于原始 multi-head attention.
FFN 基于 SwiGLU.

Training Memory Analysis

Training Memory Components

训练部分的内存占用由四部分组成：

$$ \text{training\_memory} = \text{weight} + \text{activation} + \text{optimizer} + \text{gradient} $$

Weights: $\boxed{2P}$
Gradients（与权重同精度）: $\boxed{2P}$

Optimizer States

AdamW 优化器需要维护两个动量状态：

一阶动量 $m_t$：$2P$
二阶动量 $v_t$：$2P$
合计：$2 \times 2P = \boxed{4P}$

AdamW [2] 的更新规则如下：

$$ \begin{aligned} m_t &\leftarrow \beta_1 \cdot m_{t-1} + (1 - \beta_1) \cdot g_t \\ v_t &\leftarrow \beta_2 \cdot v_{t-1} + (1 - \beta_2) \cdot g_t^2 \\ \hat{m}_t &\leftarrow \frac{m_t}{1 - \beta_1^t}, \quad \hat{v}_t \leftarrow \frac{v_t}{1 - \beta_2^t} \\ \theta_t &\leftarrow \theta_{t-1} - \alpha \left( \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} + \lambda \theta_{t-1} \right) \end{aligned} $$

Activation

激活值是前向传播过程中计算得到的中间结果，用于在反向传播时计算梯度。

我们仅针对 linear layer 进行推导：

$$ \begin{aligned} \text{forward:} \quad & \mathbf{z}_\ell = W_\ell \mathbf{a}_{\ell-1} + b_\ell, \quad \mathbf{a}_\ell = \phi(\mathbf{z}_\ell) \\ \text{backward:} \quad & \frac{\partial \mathcal{L}}{\partial W_\ell} = \frac{\partial \mathcal{L}}{\partial \mathbf{z}_\ell} \cdot \frac{\partial \mathbf{z}_\ell}{\partial W_\ell} = \frac{\partial \mathcal{L}}{\partial \mathbf{z}_\ell} \cdot \boxed{\mathbf{a}_{\ell-1}} \end{aligned} $$

可以看到，计算第 $\ell$ 层关于 $W_\ell$ 的梯度时需要其输入 $\mathbf{a}_{\ell-1}$，因此训练时需保存每个模块对应的输入，也就是激活值 (activation)。

Activation — Attention

按计算图（无优化）可得需保存的激活：

Q/K/V 投影：共享输入 → $2bsd$
$Q^\top K$：Q, K 均需保存 → $2 \times 2bsd = 4bsd$
softmax 输入：$2bhs^2$
weighted sum 输入：$2bhs^2 + 2bsd$
output projection 输入：$2bsd$

Attention 合计： $\boxed{10bsd + 4bhs^2}$

Activation — FFN & LayerNorm

FFN（SwiGLU，assume $d_{\text{ff}} = 4d$）：

第一层输入：$2bsd$
SwiGLU 输入：$2 \times d_{\text{ff}} \times d = 8bsd$
第二层输入：$2 \times d_{\text{ff}} \times d = 8bsd$
合计：$\boxed{18bsd}$

LayerNorm：保存输入 → $\boxed{2bsd}$

Activation — Output

Output 包含以下组成部分：

FinalNorm 输入：$2bsd$
lm_head 输入：$2bsd$
Loss 输入：$2bsV$

合计：$\boxed{4bsd + 2bsV}$

Activation — Total

将上面的结果汇总在一起，得到：

$$ \begin{aligned} \text{activation} &= L \cdot \text{transformer\_block} + \text{output} \\ &= L \cdot (\text{Pre\_Norm} + \textcolor{red}{\text{Attention}} + \text{Post\_Norm} + \text{FFN}) + \text{output} \\ &= \boxed{bs(32dL + \textcolor{red}{4hsL} + 4d + 2V)} \\ &\approx bsL(32d + \textcolor{red}{4hs}) \\ &\approx \textcolor{red}{4bs^2hL} \end{aligned} $$

注：在 Qwen3 中，$2V / 32dL \approx 6\%$，$32dL / 4hsL \approx 2.5\%$。

可以看到，未优化的情况下，$\text{activation} \propto bs^2$。这里 $s^2$ 主要由 attention 部分产生，后续 Flash Attention 就针对这一点进行了优化。

Total Training Memory

将上面的结果进行汇总：

$$ \begin{aligned} \text{training\_memory} &= \text{weight} + \text{activation} + \text{optimizer} + \text{gradient} \\ &= 2P + bs(32dL + 4hsL + 4d + 2V) + 4P + 2P \\ &= 8P + bs(32dL + 4hsL + 4d + 2V) \quad \text{(exact)} \\ &\approx 8P + 4bs^2hL \end{aligned} $$

可以看到，训练阶段的内存占用分为固定部分 ($8P$) 和动态部分 ($4bs^2hL$)，动态部分主要是 attention 的缓存。

Experiments

我们分别针对 80GB 的显卡计算 Qwen3 系列模型的最高配置：

Model	$P$	$L$	$h$	$s$	predicted $b$	actual $b$
Qwen3-0.6B	0.6	28	16	512	68	34
Qwen3-1.7B	1.7	28	16	512	42	28
Qwen3-4B	4	36	32	512	16	12
Qwen3-8B	8.1	36	32	512	4	2

其中 predicted $b$ 基于前面的准确公式计算得到；actual $b$ 通过实验验证得到。注意我们这里的 prediction 没有考虑任何优化手段与其他内存开销，因此与实际值有出入。

Case Study

我们分别使用 Qwen3-4B 和 Qwen3-8B 来进行实验（$b=1$, $s=512$）。参考 PyTorch 显存可视化与 Snapshot 数据分析。

Inference Memory Analysis

Inference Components

Inference 阶段内存占用主要与模型参数、KV cache 两部分相关：

$$ \text{Inference\_Memory} = \text{weight} + \text{activation} + \text{KV cache} $$

Weights: $\boxed{2P}$
Activations（经验值，batch size=1）: $\approx 0.4P$（参见 transformer-math）
KV cache：与序列长度相关，见后续分析。

KV Cache Mechanism

LLM 推理中为避免重复计算历史 token 的 key/value 而使用的空间换时间的缓存机制。

自回归时逐 token 生成，每步 attention 形式为（$\mathbf{q}_t$ 当前 query，$\mathbf{k}_{:,t}$ / $\mathbf{v}_{:,t}$ 历史 K/V）：

$$ \mathbf{q}_t = W_Q \mathbf{x}_t, \quad \mathbf{k}_{:,t} = W_K[\mathbf{x}_1, \ldots, \mathbf{x}_t], \quad \mathbf{v}_{:,t} = W_V[\mathbf{x}_1, \ldots, \mathbf{x}_t] $$

处理下一 token $\mathbf{x}_{t+1}$ 时只需在已有结果后追加当前步：

$$ \mathbf{k}_{:,t+1} = [\mathbf{k}_{:,t},\, W_K \mathbf{x}_{t+1}], \quad \mathbf{v}_{:,t+1} = [\mathbf{v}_{:,t},\, W_V \mathbf{x}_{t+1}] $$

缓存前：每生成一个 token 都重新计算 → 总计算量 $\sum_{t=1}^{s} \mathcal{O}(t) = \mathcal{O}(s^2)$
缓存后：每步只算当前 token $W_K \mathbf{x}_{t+1}$、$W_V \mathbf{x}_{t+1}$ → 计算量 $\mathcal{O}(s)$、空间占用 $\mathcal{O}(s)$ —— 以空间换时间

KV Cache Memory

对于 multi-head attention，KV cache 的显存占用为：

$$ \text{Memory}(\text{KV cache}) = s \times 2 \times 2 \times L \times h \times d_h = \boxed{4sLhd_h} $$

因子含义： $s$ 序列长，第一个 $2$ 为 K+V，第二个 $2$ 为 BF16 的 2 bytes，$L$ 层、$h$ 头、$d_h$ 头维度。

Remark:

KV 占用与模型配置（$L, h, d_h$）和序列长度 $s$ 都有关，token 越多占用越高。
实际中因 page granularity、padding、fragmentation 往往略高于理论值。
长输出时 KV 占比会超过权重，成为推理瓶颈 → 见后续「KV Cache 优化」。

Total Inference Memory

综合前面分析，推理阶段的总内存为：

$$ \boxed{\text{Inference\_Memory} \approx 2.4P + 4sLhd_h} $$

可以看到，推理阶段也由固定部分（参数量，activation）以及动态部分（KV cache）组成。

Dynamic vs. Static

由于 Qwen3 的 KV cache 计算为 $4sLh_{kv}d_h$，而不同模型只有 $L$ 不一样，因此对于更大的模型，KV cache 显存占用超过模型权重的上下文长度更高。

Optimization

Overview

阶段	核心方法	典型技术
Training	显存与效率提升	Activation Checkpointing, Mixed Precision Training, Flash Attention, ZeRO, Pipeline/Model/Data Parallelism
Inference	长序列与速度优化	KV Cache Optimization, Paged/Radix Attention, Faster Attention, Quantization

Mixed Precision Training

计算量大的部分用低精度，计算量小的部分用高精度。低精度参与运算，高精度避免 Overflow/Underflow。

下表是 DeepSeek-V3 [3] 使用的混合精度训练框架的显存分析：

Precision	BF16	FP32	BF16
AMP	No	Yes	Yes
Weights	BF16 (2)	FP32 (4)	BF16 (2)
Master weights	-	-	FP32 (4)
Gradients	BF16 (2)	FP32 (4)	BF16 (2)
Adam m	BF16 (2)	FP32 (4)	FP32 (4)
Adam v	BF16 (2)	FP32 (4)	FP32 (4)
Static total (bytes/param)	8	16	16

Remark:

BF16 (w/ AMP) 与 FP32 (w/ AMP) 的静态显存占用相同，但 BF16 (w/ AMP) 的动态显存占用更低。
主流框架基本都使用了 BF16/FP8 (w/ AMP) 的训练方式。

ZeRO

将 optimizer states / gradients / weights 按不同 GPU 切片存储，需要参与计算时再 all-gather 整合成完整参数。这样每张卡只需维护自己负责的一部分，大幅降低单卡显存需求 [4]。

ZeRO Stages:

ZeRO-1：shard optimizer states

$$ \text{Training\_Memory} = \text{weight} + \text{activation} + \frac{\text{optimizer}}{\#\text{GPUs}} + \text{gradient} $$

ZeRO-2：shard optimizer states + gradients

$$ \text{Training\_Memory} = \text{weight} + \text{activation} + \frac{\text{optimizer} + \text{gradient}}{\#\text{GPUs}} $$

ZeRO-3：shard all

$$ \text{Training\_Memory} = \text{activation} + \frac{\text{weight} + \text{optimizer} + \text{gradient}}{\#\text{GPUs}} $$

ZeRO-3 可极大降低单卡显存上限，但通信量也会提高。

Model Parallelism

将模型切分到不同的 GPU 上，计算时，先 dispatch，再执行计算，最后通过 all-gather 等操作得到最终结果 [5]。切分方式包括 PP (Pipeline Parallelism)、TP (Tensor Parallelism)、EP (Expert Parallelism) 等。

$$ \text{Training\_Memory} = \frac{\text{Memory}(\text{weight})}{\text{PP degree} \times \text{TP degree}} $$

结合 ZeRO-1 与 Model Parallelism 时（activation 中与 TP 相关的部分按 TP degree 缩减）：

$$ \text{Memory}_{\text{train}} \approx \frac{\text{weight}}{\text{PP} \times \text{TP}} + \frac{\text{activation}}{\text{TP}} + \frac{\text{optimizer}}{\#\text{GPUs}} + \frac{\text{gradient}}{\text{PP}} $$

Activation Checkpointing

在反向传播时，重新计算所需的输入，来达到以时间换空间的目的 [6]。

	No ckpt	Selective ckpt	Full ckpt
memory	很高	中等	很低 $\sim 2bsd$
extra compute	无	中等	很高 $\sim 2Pbs$

实践中常结合 model parallelism 与 selective checkpointing 来实现 trade-off。

Flash Attention

通过将 Attention 的计算进行分块，来提高内存访问效率以及降低反向传播时所需要的 activation 大小 [7]。

Flash Attention 通过 tiling 与 online-softmax 降低该部分显存并提升效率（详见 notes on Flash Attention）。这样 attention 部分的显存就由 $\text{activation} \propto bs^2$ 降低到了 $\text{activation} \propto bs$。

Theorem: Flash Attention 输出 $O = \text{softmax}(QK^T)V$（correctness）。其时间复杂度为 $\mathcal{O}(s^2 d)$，空间复杂度为 $\mathcal{O}(s)$（memory savings）。

KV Cache Optimization

$$ \text{Memory}(\text{KV cache}) = s \times 2 \times 2 \times L \times h \times d_h $$

针对公式中各因子的优化方向 [8]：

$s$: KV cache compression, eviction, selection
$2$ (bytes): KV cache quantization
$2$ (K+V): key-value sharing, MLA [9]
$h \times d_h$: MQA [10], GQA [11], MLA

Weight Quantization

使用低精度来表示高精度数值的方法，来减少内存占用/提高计算效率。

量化时机	代表性工作
训练后量化 (PTQ)	GPTQ [12], AWQ [13], SmoothQuant [14], GGUF [15]
量化感知训练 (QAT)	LLM-QAT [16], PEQA [17]

Activation Offloading

将一部分参数/优化器状态/激活值等存储到 CPU 上，需要的时候再加载到 GPU 上。

Offloading 场景	代表性工作
训练阶段 Offloading	ZeRO-Offload [18] / ZeRO-Infinity, FSDP [19] CPU Offload
推理阶段 Offloading	FlexGen [20], vLLM [21] KV Cache Offload
MoE Offloading	KTransformers [22], DeepSpeed-MoE [23]

Real Systems

Training Frameworks

Framework	Memory Optimizations
Megatron-LM [5]	TP, SP, Activation Checkpointing
DeepSpeed [24]	ZeRO-1/2/3, CPU/NVMe Offload, Activation Checkpointing
FSDP [19]	Full parameter sharding, Gradient sharding, CPU Offload
Colossal-AI [25]	ZeRO, TP, PP, Activation Checkpointing

Inference Frameworks

Framework	Key Techniques	Memory Optimizations
vLLM [21]	Paged Attention	KV cache paging, Continuous batching
SGLang [26]	Radix Attention	KV cache reuse, Efficient scheduling
TensorRT-LLM [27]	Kernel fusion	Weight quantization, KV cache optimization

Conclusion

Takeaway

Components	Training	Inference	Optimization
weights	$2P$	$2P$	quantization
optimizer states	$4P$	0	ZeRO, offloading
gradients	$2P$	0	ZeRO
activations	$\sim 4Lhs^2b$	$\sim 0.4P$	ckpt, offloading, flash attention
KV cache	0	$4sLhd_h$	KV optimization, attention
TOTAL	$8P + 4Lhs^2b$	$2.4P + 4sLhd_h$

训练：瓶颈主要在激活值（随 batch size / 序列长度线性增长）。
推理：瓶颈主要在 KV Cache（随序列长度增长）。

Future Directions

More efficient architecture (attention, MoE).
Scalable training/inference framework.
Software-hardware co-design algorithms.

References

An Yang et al., “Qwen3 Technical Report,” arXiv:2505.09388, 2025.
Ilya Loshchilov and Frank Hutter, “Decoupled Weight Decay Regularization,” arXiv:1711.05101, 2019.
DeepSeek-AI, “DeepSeek-V3 Technical Report,” arXiv:2412.19437, 2025.
Samyam Rajbhandari et al., “ZeRO: Memory Optimizations Toward Training Trillion Parameter Models,” arXiv:1910.02054, 2020.
Mohammad Shoeybi et al., “Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism,” arXiv:1909.08053, 2020.
Vijay Korthikanti et al., “Reducing Activation Recomputation in Large Transformer Models,” arXiv:2205.05198, 2022.
Tri Dao et al., “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness,” arXiv:2205.14135, 2022.
Haoyang Li et al., “A Survey on Large Language Model Acceleration based on KV Cache Management,” arXiv:2412.19442, 2025.
DeepSeek-AI, “DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model,” arXiv:2405.04434, 2024.
Noam Shazeer, “Fast Transformer Decoding: One Write-Head is All You Need,” arXiv:1911.02150, 2019.
Joshua Ainslie et al., “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints,” arXiv:2305.13245, 2023.
Elias Frantar et al., “GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers,” arXiv:2210.17323, 2023.
Ji Lin et al., “AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration,” arXiv:2306.00978, 2024.
Guangxuan Xiao et al., “SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models,” arXiv:2211.10438, 2024.
Georgi Gerganov, “ggml: Tensor library for machine learning,” GitHub, 2023.
Zechun Liu et al., “LLM-QAT: Data-Free Quantization Aware Training for Large Language Models,” arXiv:2305.17888, 2023.
Jeonghoon Kim et al., “Memory-Efficient Fine-Tuning of Compressed Large Language Models via sub-4-bit Integer Quantization,” arXiv:2305.14152, 2023.
Jie Ren et al., “ZeRO-Offload: Democratizing Billion-Scale Model Training,” arXiv:2101.06840, 2021.
Yanli Zhao et al., “PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel,” arXiv:2304.11277, 2023.
Ying Sheng et al., “FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU,” arXiv:2303.06865, 2023.
Woosuk Kwon et al., “Efficient Memory Management for Large Language Model Serving with PagedAttention,” SOSP, 2023.
Hongtao Chen et al., “KTransformers: Unleashing the Full Potential of CPU/GPU Hybrid Inference for MoE Models,” SOSP, 2025.
Samyam Rajbhandari et al., “DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale,” arXiv:2201.05596, 2022.
Jeff Rasley et al., “DeepSpeed: System Optimizations Enable Training Deep Learning Models with Over 100 Billion Parameters,” KDD, 2020.
Shenggui Li et al., “Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training,” ICPP, 2023.
Lianmin Zheng et al., “SGLang: Efficient Execution of Structured Language Model Programs,” NeurIPS, 2024.
NVIDIA Corporation, “TensorRT-LLM: A TensorRT Toolset for Optimizing LLM Inference,” GitHub, 2023.

Notes on GLaM

Tue, 06 Jan 2026 18:07:29 +0800

Google 在 2022 年 8 提出了 GLaM，一个基于 MoE 架构的大语言模型系列，模型超过了 GPT-3 的表现

Introduction

作者在本文中证明了基于 MoE 架构的大语言模型可以达到与 dense 模型相同的性能，且计算更加高效。作者构建了 GLaM 大语言模型系列，包括 1.9B-A0.1B, 105B-A1.9B, 143B-A9.8B 以及 1.2T-A96.6B 等模型，尽管只激活了不到 $8\%$ 的参数，模型的 zero-shot, one-shot 和 few-shot 表现均超过了 GPT-3.

作者认为，在相同的算力下，MoE 模型比 dense 模型的表现更好，MoE 模型是一个非常具有前景的方向。

Method

训练数据集包括 1.6T token, 具体分布如下表所示

Dataset	Tokens (B)	Weight in mixture
Filtered Webpages	143	0.42
Wikipedia	3	0.06
Conversations	174	0.28
Forums	247	0.02
Books	390	0.20
News	650	0.02

模型架构如下图所示

GLaM 的模型架构与 GShard 基本相同，作者将 transformer block 按照两个为一组，一组中一个 block 为 dense FFN, 另一个为 MoE layer, 交替进行。MoE layer 中，总专家个数为 64 个，激活专家个数为 2 个。

模型配置如下表所示

GLaM Model	Type	$n_{\text{params}}$	$n_{\text{act-params}}$	$L$	$M$	$H$	$n_{\text{heads}}$	$d_{\text{head}}$	$E$
0.1B	Dense	130M	130M	12	768	3,072	12	64	$-$
0.1B/64E	MoE	1.9B	145M	12	768	3,072	12	64	64
1.7B	Dense	1.7B	1.700B						$-$
1.7B/32E	MoE	20B	1.878B						32
1.7B/64E	MoE	27B	1.879B	24	2,048	8,192	16	128	64
1.7B/128E	MoE	53B	1.881B						128
1.7B/256E	MoE	105B	1.886B						256
8B	Dense	8.7B	8.7B	32	4,096	16,384	32	128	$-$
8B/64E	MoE	143B	9.8B	32	4,096	16,384	32	128	64
137B	Dense	137B	137B	64	8,192	65,536	128	128	$-$
64B/64E	MoE	1.2T	96.6B	64	8,192	32,768	128	128	64

作者用 GLaM(8B/64E) 来表示一个 8B 参数的 dense model 中每隔一层被转换为 MoE layer

Experiments

训练时，作者设置模型的上下文为 1024 token, batch size 为 1M, 优化器为 Adafactor, 作者还使用了 GShard 中提出来的 load balancing loss, tokenizer 为 SentencePiece

GLaM 与 GPT-3 的对比如下表所示

		GPT-3	GLAM	relative
cost	FLOPs / token (G)	350	180	-48.6%
	Train energy (MWh)	1287	456	-64.6%
accuracy on average	Zero-shot	56.9	62.7	+10.2%
	One-shot	61.6	65.5	+6.3%
	Few-shot	65.2	68.1	+4.4%

实验结果显示，GLaM 的训练效率更高，且表现更好

作者还探究了提升 expert 个数对最终表现的影响，结果如下图所示

可以看到，随着专家个数提升，模型的表现逐渐增强，但是当专家个数超过 64 个之后，模型的表现反而有所下降。

Conclusion

作者在本文中提出了 GLaM 大语言模型系列，验证了 MoE 模型的有效性和效率，结果发现，MoE 模型可以在相同的算力下达到更好的表现，并且学习效率更高。

References

arxiv

Notes on MiniMax-01

Tue, 06 Jan 2026 17:38:01 +0800

MiniMax-01 是一个基于 hybrid attention 架构的大模型系列，包含 MiniMax-Text-01 和 MiniMax-VL-01 两个模型，其中 MiniMax-Text-01 推理时支持 4M 的上下文长度，MiniMax-VL-01 支持 512B 的上下文长度

Introduction

现有大部分模型的上下文长度为 32K-256K, 但实际上我们对于长上下文的需求已经超过了这个范围。已有的模型主要基于 self-attention, 这是一个平方复杂度的算法。

为了解决 self-attention 的问题，相关工作如 sparse attention, linear attention, state space models 等都提出了对应的解决办法。但是已有的这些解决方法的问题就是表现不是很强劲。

因此，作者在本文中就提出了一个基于 hybrid attention 架构的大语言模型系列 MiniMax-01. 作者主要从架构，数据和 infra 三个方面进行了改进。

在架构上，作者使用了基于 Lightning Attention 的混合架构。作者还基于实际部署来决定模型的参数。为了最大化模型的表现，作者使用了 MoE 架构。

已有的 infra 主要是针对基于 softmax 的 attention 进行优化的，MiniMax-01 包含 softmax attention, linear attention 和 MoE, 架构比较复杂，因此作者实现了 expert parallel 和 expert tensor parallel 来提高整体的计算效率和 GPU 之间的通信效率。作者还实现了 varlen ring attention 来减少计算冗余。最终，作者发现模型在 NVIDIA-H20 上的 MFU 超过了 75%.

基于前面提到的架构设计，作者训练得到了 MiniMax-Text-01, 模型总参数为 456B, 激活参数为 45.9B, 专家个数为 32 个，激活专家个数为 2 个。作者首先构建了高质量的数据集，然后构建了一个三阶段的训练 pipeline。

基于 MiniMax-Text-01, 作者扩展得到了 MiniMax-VL-01,训练使用了512B token

作者总结本文的贡献如下：

作者构建了一个领先的大模型系列，支持超过 4M 上下文
作者构建了第一个大规模的基于 linear attention 的大语言模型系列
作者详细介绍了使用的数据，模型和训练策略
作者开源了模型

Architecture

模型架构如下图所示

相比于原始的 transformer, MiniMax-01 做出了以下改进：

将 attention block 按照 8 个 block 为一组，每组里只有最后 1 个 block 使用 softmax attention, 其余 7 个 block 使用 lightning attention
使用 MoE 替换 FFN, MoE 总专家个数为 32 个，激活专家个数为 2 个
softmax attention 使用了 GQA, 来提高内存加载效率, group size 为 8
使用了 RoPE 作为 position embedding
使用了 RMSNorm 替换了 LayerNorm

FFN: MoE (32 个专家，激活 2 个专家)

transformer: 8 个 block 为一组，一组里前 7 个使用 Lightning attention，第 8 个使用 softmax attention，80 layers

Attention ： GQA，group size=8，64 heads，

MoE

在训练基于 MoE 的 LLM 时，有两种策略，分别是 token-drop 和 dropless, 前者保证每个专家处理的 token 个数差不多，可以提高效率，缺点是某些 token 不会被任何专家处理，某些 token 会被多个专家处理；后者是保证每个 token 都会被处理。

在本文中，作者采取了 token-drop 的方式，作者为每个专家设置一个 capacity limit，超过这个 limit 之后该专家就不再处理新的 token.

为了评估 MoE 模型的有效性，作者对比了 MoE 和 dense 模型的表现，结果如下图所示：

结果发现，相同的算力下，MoE 模型比 dense 模型好。

但是，当 scaling up 到更大的模型是，作者发现训练产生了 routing collapse 的情况，这是因为 routing 的分布过于集中。为了解决这个问题，作者使用了 auxiliary loss 以及 global router 两个方法

Auxiliary loss 作者首先构建了 auxiliary loss 来提高负载均衡, 也就是

$$ \mathcal{L}_{B} = K\sum_{i=1}^K f_i P_i $$

这里 $f_i$ 是分配给第 $i$ 个专家的 token 比例， $P_i$ 是第 $i$ 个专家的平均 routing 概率。

Global router 作者还基于 GShard 构建了一个 global routing 策略，由于 GPU memory 限制，对于每一个 micro batch size, token 分布不仅在一个 EP group 内分布不平衡，在不同 EP group 之间可能也不平衡。因此，作者实现了一个 global token dispatching 策略。具体来说，在 token 分发到不同的 EP group 之前，作者使用 allgather 来计算每个专家需要处理的 token. 这样就可以基于全局信息智能分配 token, 避免某些专家过载。

Linear Attention

本文中使用的 linear attention 由 Lightning Attention 提出，其表达式为

$$ O = \mathrm{Norm}((QK^T)V) $$

这里 $Q,K,V\in\mathbb{R}^{n\times n}$ 分别是 query, key 和 value, $n$ 和 $d$ 分别是序列长度和 hidden size.上式可以改变运算顺序，得到

$$ O = \mathrm{Norm}(Q(K^TV)) $$

这样，attention 的计算复杂毒就从 $O(n^2d)$ 变成了 $O(nd^2)$,

Lightning Attention

当我们不考虑 attention mask 的时候，我们很轻松可以降低 attention 计算的复杂度。但实际上，LLM 会使用 causal mask, 也就是每个 token 智能看到其前面 token 的信息，这样我们的 attention 计算实际上是

$$ O = \mathrm{Norm}[QK^T\odot M]V) $$

这里 $M_{ij}=\mathbb{1}(i\geq j)$ 是 attention mask.

见 Lightning Attention

Effectiveness of Lightning Attention

作者接下来分析了一下 softmax attention, lightning attention 和 hybrid attention 之间的效率

首先，作者计算了一下三种架构的参数量以及 FLOPS. 作者分别使用 $l, d, h, b, n$ 来代表 Layer 数，hidden dimension, number of attention heads, batch size 和 sequence length.

最终计算结果如下

Architecture	Parameter count	FLOPs count
Softmax	$12ld^2$	$72bnld^2 (1 + \frac{n}{6d} + \frac{5}{18d} )$
Lightning	$12ld^2+2ld^2/h$	$72bnld^2(1+\frac{1}{2h}+\frac{5}{18d})$
Hybrid	$12ld^2+7ld^2/4h$	$72bnld^2(1+\frac{n}{48d}+\frac{7}{16h}+\frac{5}{18d})$

[!todo] TODO compute these results

Scaling law 作者接下来分别针对三个架构设计了 70M, 160M, 410M, 1B, 3B, 7B 系列模型，模型使用 300B token 进行训练，上下文长度为 8192, 对于每种架构，作者将 batch size 设置为 4M tokens. 与 Chinchilla scaling law 类似，作者探究了以下模型的 scaling law, 结果如下图所示

从实验结果可以看到，在相同的算力下，lighting attention 倾向于使用更多的参数和 token, 但是其表现相比于 softmax attention 更好。

Performance 作者还对比了一下三种 attention 在 public benchmark 上的表现，结果如下图

是检验结果发现，lightning attention 和 softmax attention 除了在 retrieval 任务（Needle in a Haystack）上之外，表现都差不多。与之相对的是，hybrid attention 弥补了这一问题，大幅度提升了模型在 retrieval 任务上的表现

[!tip] Observation Lightning attention 与 softmax attention 的效果差不多，但是其在长上下文任务上的表现比较差。Hybrid attention 可以解决这个问题。

作者还评估了一下三种 attention 的速度，结果如下图所示

实验结果显示，softmax attention 随序列长度上升其速度急剧下降，Lightning attention 的速度基本没有太大变化。而 Hybrid attention 的速度介于两者之间。

作者进一步对比了以下两种不同的变体：hybrid-cosformer2 以及 hybrid-hgrn2.这两个模型替换的逻辑与 minimax-01 的结果一致，都是 8 个 block 为一组，每组中前面 7 个 block 将 softmax attention 更换为对应的模块，最后一层不做改动。三种 hybrid attention 机制的表现如下表所示

实验结果显示，hybrid-lightning 的表现最好。

作者最后对比了以下 hybrid-lightning 和 hybrid-window, hybrid window 在每个 group 的前 7 个 block 使用了 window attention, window size 为别为 256, 512, 1024.实验结果如下图所示

Discussion 作者最后总结认为，虽然 linear attention 的效率很高，但是它们在 retrieval 相关的任务表现很差，而 retrieval 对于 In-context learning 来说是至关重要的。因此，作者采取了 hybrid 架构，来兼顾模型的效率以及表现。

作者给出了一个解释。

Ablation Study

作者主要进行了两个消融实验：

Hybrid-lightning 与 softmax attention 的对比：作者训练了一个总参数 28B, 激活参数 5B 的 MoE 模型，然后作者使用 MiniMax-01 的方式替换每个 group 的 softmax attention, 并使用了 1T 的 token 进行训练
pre-layer normalization 与 Post-layer normalization 的对比：现有的 LLaMA 和 Qwen 等系列模型采用的都是 PreNorm 的方式。PreNorm 可以让 gradient 通过 residual connection 传播更加直接，但是这也减少了模型的有效深度。反之，PostNorm 则可以保留模型的有效深度，但是其问题是会导致梯度消失或爆炸。作者构建了一个总参数为 60B, 激活参数为 9.3B 的 MoE 模型，包含 48 个 block, 模型训练使用 500B token. 模型有两个变体，一个使用 PreNorm, 另一个使用 PostNorm, 对于 PostNorm, 作者使用的是 DeepNorm.

最终表现如下图所示

实验结果显示，hybrid lightning 的表现与 softmax 的表现相当，甚至超过了 softmax 的表现。

另一方面，PostNorm 的表现也超过了 PreNorm 的表现。

Model Spec

基于已有的模型设计，作者探究了如何决定模型的参数。作者的目标在 performance 和 inference efficiency 之间达到一个平衡。

作者将模型的参数限制在 500B 以下，要求能够在 $8\times 80G$ 的服务器上和 8-bit 的量化下面，支持 1M 的上下文长度。建模的问题如下：

$$ \min_{P_{all}, P_{act}}\mathcal{L}(P_{all}, P_{act}, T), \quad \mathrm{s.t.}\ C_{compute}(P_{all}, P_{act}, T)其中 $\mathcal{L}$, $P_{all}$, $P_{act}$, $T$, $C_{compute}$, $C$ 分别代表损失，总参数，激活参数，训练的 token 数，算力消耗 y 以及算力的 budget.

作者首先确定了几个关键的因素：

softmax 和 lightning 的混合比例
depth-to-width 的比例
linear attention memory size 和 hidden size 的比例
FFN 的 hidden size 大小
RoPE 的 base frequency

作者通过实验发现，hybrid 架构需要更多的 layer 才能带来更好地表现。对于 layer 比较少的模型，其需要更多的 softmax attention layers 来达到相似的表现。

作者还发现，提升 linear attention memory size 可以有效提高模型的表现。并且 RoPE 的 dimension 最好设置为 softmax attention dimension 的一半。

基于这些发现，作者构建了 scaling law 来决定最优模型配置。作者训练了不同大小的模型，然后拟合 scaling law 的曲线，最后，作者将模型的总参数确定为 456B, 激活参数确定为 45.9B.

Infra

本节介绍了一下 MiniMax-01 的 infra, 作者主要解决了以下几个关键问题：

减少训练时 MoE 中的 all-to-all 通信开销，如何在内存使用，计算效率和 all-to-all 通信开销之间达到一个平衡
在处理长上下文时，不同 GPU 之间的信息需要共享，这会导致额外的通信开销。因此，如何减少长上下文情形下的通信开销也是一个挑战
如何提高 lightning attention 在 inference 阶段的效率。当前只在训练阶段做了优化，但是推理阶段也需要提高模型的效率

MoE Optimization

优化 MoE 架构的主要目的是最小化通信开销，特别是 all-to-all 的通信开销。为了解决这个问题，作者构建了一个 token-grouping-based overlap scheme, 来让通信和计算尽可能并行执行，也就是在计算的同时进行通信。为了保证通信的正确性，每个 ProcessGroup 之间的通信操作必须串行执行，这就导致了不同 group 之间的通信无法 overlap, 也就造成了 idle time 的产生，示意图如下图所示

但是，在 MiniMax-01 上应用这种方法时，作者发现如果使用 TP 来切分 expert 的参数的话，会导致计算密度过低。

[!tip] Explanation 针对 MoE 模块使用 TP 导致计算密度过低的原因是 MoE 模块是稀疏的，计算式大量 GPU 都处于空闲状态，因此计算密度低。

如果不使用 TP 的话，就必须使用更大的 PP 配置，也就是使用更多的 GPU, 但是使用 PP 并不会减少 activation 的内存占用，这在长上下文训练时，会增加内存消耗，而且不会提升整体的训练速度。因此，我们需要一种新的策略，使其能够：

平衡内存使用与计算密度
优化特定模型和任务的训练过程

基于这个目标，在本文中作者提出了 ETP, 也就是 Expert Tensor Parallel, 用于管理 expert 的参数划分。同时，作者还提出了 EDP, 也就是 Expert Data Parallel, 来将每个 expert 的数据并行封装在一起。world_size 满足两个条件：

$$ \texttt{world\_size} = \texttt{size}_{PP}\times \texttt{size}_{DP}\times \texttt{size}_{CP}\times \texttt{size}_{TP} $$

以及

$$ \texttt{world\_size} = \texttt{size}_{PP}\times \texttt{size}_{EDP}\times \texttt{size}_{ETP}\times \texttt{size}_{EP} $$

首先，作者构建了 EP-ETP 策略，用于平衡内存使用以及计算密度，这个过程包括了四个步骤：

all-to-all dispatch: （EP）将每个 token 分发到 expert 所在 GPU 上
allgather: （TP）将输入的 token 进行切分，分发给对应的 TP GPU 上
Expert:（TP）执行 expert 的计算过程
ReduceScatter: （TP）将 TP 输出的结果进行汇总然后再分发给不同的 TP
all-to-all combine: 将不同 expert 的输出结果进行汇总，传递给下一层

但是，由于同一个 ProcessGroup 内部的通信必须串行执行，如果计算效率比较高的话，就会产生 bubble, 作者的改进方法是提高计算量，让计算消耗的时间与通信消耗的时间大致相当。

作者还尝试降低 ProcessGroup 的 size, 这样可以进一步提高计算和通信之间的 overlap, 但是问题是降低 group size 之后，group 的数量会增加，而这会导致 scheduling 以及让通信变为 CPU-bound. 作者认为这需要根据具体场景来进行设置。最终，通过优化，作者将 MoE 模块的通信开销降低了 50%.

改进的示意图如下图所示

Long Context Optimization

为了处理不同长度的 sequence, 作者使用了 sequence packing 的技巧。

Varlen Ring Attention 当前主要是使用 ring attention 来划分数据，但是 ring-attention 与 sequence packing 是冲突的。已有工作如 flash attention 支持 varlen data packing, 但是不支持 ring attention; 而 TransformerEngine 实现了 ring attention, 但不支持真正的 varlen. 因此作者的目的就是解决这个问题。

TransformerEngine 的问题在于每个 sequence 必须是 $2\times \texttt{size}_{CP}$ 的整数倍，当 CP 很大且样本长度分布未知的时候，padding 会很严重，导致算力和内存消耗严重。

作者的解决方法为使用 varlen ring attention. 具体做法就是在每一步通信中，都基于 attention mask 的 offset 来处理不同的 sequence. 示意图如下所示

LASP+ 对于 lightning attention, 之前的做法是使用 LASP 算法来实现 linear attention. 但是 LASP 的问题在于其是一个串行依赖，这样导致训练效率变低。

为了解决这个问题，作者提出了 LASP+, 包括以下几个步骤

Local Prefix Sum Calculation: 每个 GPU 先计算自己局部 prefix sum, 作者将其记为 KVL
Global Synchronization via AllGather: 然后，作者通过 AllGather 将自己的 KVL 发送给其他 GPU. 这样就可以一次完成所有通信
Prefix Sum Computation: 最后，每个 GPU 利用所有 KVL, 分别计算自己需要的 prefix sum.

通过这样一个优化，我们可以将 LASP 的延迟降低 $N_{pcn}$ 倍，这里 $N_{pcn}$ 是并行计算节点的个数。

作者进一步结合了前面的 Varlen ring attention, 主要包括 padding to block size 和 sequential concatenation 两个步骤。

[!tip] Remark LISP+ 说明了现在大模型分布式训练的一个重要趋势，也就是用通信换并行。特别是 GPU 通信带宽特别高的情况。因此，算法设计不仅要考虑理论复杂度还要考虑实际中硬件并行的利用率。

Lightning Attention Inference Optimization

已有的 lightning attention 并没有考虑实际的部署需求。作者提出了四个策略来优化 lightning attention 的 inference efficiency:

Batched Kernel Fusion 作者首先对 memory-bound kernels 进行融合。在 prefill phase, 作者将多步操作，比如 Q, K, V projection, padding, partitioning 等放在一个 kernel 里，来减少内存 I/O. 在 decoding phase, 作者将 KV 的计算也放在一个 kernel 里，计算完之后就直接写入 KV cache 而不经过缓冲区

PD separation inference 的 decoding 阶段模型每次只生成 1 个 token, 此时任务是 memory-bound, 也就是瓶颈在显存读写而不是算力，我们只需要很少的 GPU SMs 就可以跑起来。因此，作者设计了两个不同的 CUDA streams 来分别处理 length 为 1 和 length>1 的情况，这样就可以有效提高整体的计算效率和平衡 GPU 的利用率。

Multi-level padding 已有的 padding 技巧主要是将输入的多个 sequence 都 resize 到固定长度方便处理。但是在推理的时候，使用预设的固定长度 (block size) 会导致无效计算量增加。因此，作者提出了 multi-level padding, 也就是提供不同的 block size, 比如 32, 64, 128. 这样可以有效避免模型因为 padding 而产生的额外计算开销。

StridedBatchedMatmul Extension 这一点主要是针对 Hopper GPU 进行的优化。Hopper 架构相比于 Ampere 架构有一些新的特性，作者基于这些特性来优化整体的性能。

作者使用 WGMMA 来优化计算，用 TMA 来优化内存访问效率，实现计算与内存访问重叠，减少等待时间。最终，作者的目的是希望基于不同的 GPU 架构来动态调整 Pipeline stages, 以达到更好的性能。

通过以上这些优化，作者在 H20 上达到了 75% 的 MFU. 作者发现，由于使用了 lightning attention, attention 不分计算从传统的 95% 以上降低到了 12% 左右，不再是算力的瓶颈部分。

Pre-training

作者首先介绍了预训练的数据构成，然后作者介绍了如何选取高质量的训练数据，最后，作者介绍了实验的参数。

Data

数据来源包括学术文献，数据，互联网语料以及代码。作者构建了如下的数据清洗策略：

data quality enhancement: 作者使用了 rule-based cleaning 以及去重。然后，作者构建了一个 reward labeler，来从 knowledge depth, practical helpfulness 以及 categorical distribution 三个维度进行打分。
data formatting optimization. 作者发现，类似 markdown 的格式会降低数据整体的多样性和质量。因此，作者构建了一个嵌套的文档格式，来将对话和 QA 数据组织起来，进而平衡自然数据和结构化数据之间的平衡性。
data mixture investigation. 作者发现完全去除掉低质量的数据也会影响模型在下游任务上的表现。因此作者平衡了高质量数据和低质量数据

作者基于 BPE 算法来构建 tokenizer,最终 tokenizer 的大小为 200K

接下来，作者通过实验来确定最优的数据配置，给定数据集 $\mathcal{D}$, 作者作如下备择假设

$$ H_1: \mu_{T_{\mathcal{D}}}> \mu_{T_{baseline}} $$

这里 $\mu$ 代表模型表现的加权平均， $T$ 代表模型的表现分布。也就是说，加入新的数据集之后，其平均表现大于 baseline 的表现，则我们认为这个数据集是不被拒绝的。

作者通过多项选择题来进行评估，评估时作者计算每个选项 completion 的概率，最终的分布计算方式为

$$ \log \mathrm{acc}_{\mathrm{norm}^2}(x) = \log\mathrm{softmax}_{p'(c\in C_x)}\left\{(p'(c^*))\right\} $$

其中 $p_i'(c)=\frac{p_i(c)}{\mathrm{bytes(c)}}$ 是 normalized 之后的选项概率。作者在训练的时候，通过实验来保证这个 metric 的稳定性，除此之外，作者还加入了一个 discriminator, 也就是 $\Delta_{obvious}/\sigma_{seed}$, 其中 $\Delta_{obvious}$ 表示不同模型表现的差距， $\sigma_{seed}$ 表示不同随机种子的方差。

Training Strategy

作者使用 Xavier 对模型参数进行初始化，然后使用了 DeepNorm 作为 Normalization 模块。作者使用 AdamW 作为优化器，其中 $\beta_1=0.9$, $\beta_2=0.95$, weight decay 为 $0.1$.

训练的 sequence length 为 8192，batch size 分别为 16M, 32M (69B tokens), 64M (790B tokens), 128M (4.7T tokens).

作者基于 training loss 和 critical batch size 来设计 schedule. 作者基于小模型的实验结果来拟合一个 power-law 关系，结果如下图所示：

基于实验结果，作者

Post-training

Prompt Collection

作者从多个 domain 收集到了数百万 diverse 和高质量的 query,作者基于任务类型，domain 和难度对 prompt 进行分类。接下来作者对 prompt 进行去重以及平衡难度问题。prompt 覆盖 long context, math, reasoning 等 domain

Reward Model

作者基于以下几个原则来设计 reward:

Correctness: 对于 math 任务，作者使用 MiniMax-Text-01 来生成 binary reward; 对于 code 任务，作者使用沙盒环境运行代码，根据通过率来给出奖励
Truthfulness: 作者使用先进的语言模型来评估 response 的 factual accuracy
Helpfulness: 作者实现了基于规则的验证系统，来评估回答的 coherence, depth, contextual relevance 和 stylistic appropriateness.
Harmlessness: 作者基于 Constitutional AI principles, 构建了一系列协议来保证模型输出的安全性。

SFT

作者通过使用多个 expert model 来生成多样化的回答。然后，作者通过 n-gram 以及 semantic similarity 来提高数据的多样性和质量。

RL

RL 训练包括 offline stage 以及 online stage.

Offline RL 在 offline stage, 作者使用 DPO 来进行训练。数据构造过程也比较简单，使用模型来进行采样，评估，选取最好和最差的作为正样本和负样本。

Online RL 这个阶段使用 GRPO 进行训练，作者发现使用 SFT 阶段的 prompt 会导致训练饱和，因此，作者并没有采用 SFT 阶段的 prompts

作者主要针对 GRPO 进行了一下改进：

Important Sampling Weight Clipping. 作者发现，PPO 以及 GRPO 只使用了单侧的 clipping, 这会导致训练的不稳定性。因此作者构建了额外的 clipping 策略。
KL divergence optimization. 作者发现 KL divergence 也会导致训练不稳定，因此作者对 KL divergence 进行了 reformulate 来稳定梯度
Balanced Advantage Estimation. 作者确保正样本和负样本的对 reward 的贡献差不多，来提高训练的稳定性。

Safety Alignment

作者还加入了一个 safety alignment 阶段，用于提升模型的安全性。

数据包括三类：

safety-category specific prompts.
Real-world user data collection.
prompt augmentation.

接下来，作者使用一个无害的 reward model 来对模型的输出进行打分。

Training Recipe

最终，post-training 一共包括 5 个 stage，训练时，作者将 RoPE 的 base frequency 保持在 10M.

Initial short-context training: 模型的上下文长为 8192，然后进行 SFT
Extended context training: 模型的上下文长度为 1,032,192. 训练数据包括 50% 长文本和 50% 短文本
Short-context Preference Optimization: 模型的上下文长度为 8192, 然后进行 DPO 的训练
Long-context Preference Optimization: 模型的上下文长度为 1,032,192, 然后进行 DPO 的训练
Online Reinforcement Learning: RL 的训练，与上文一致。

训练的配置如下表所示

	Stage I	Stage II	Stage III	Stage IV	Stage V
Sequence Length	8192	1032192	8192	1032192	8192
Epoch	2	2	1	1	1
Batch Size	128	80	64	64	512
Max LR	1e-5	3e-6	5e-7	5e-7	1e-6
Min LR	1e-6	3e-6	5e-8	5e-7	1e-7
LR Decay	Cosine	Constant	Cosine	Constant	Cosine

LLM Evaluation

作者对比了 GPT-4o, Claude-3.5-Sonnet, Gemini-1.5-Pro, Qwen2.5-72B, DeepSeek-V3 以及 LLaMA3.1-405B, 实验结果如下图

作者还评估了以下 MiniMax-Text-01 在长上下文情境下表现，我们这里列出 Ruler 上的结果，如下图所示

可以看到在 1M 上下文时，模型的表现超过了 Gemini-1.5-Pro 的表现。

VLM Architecture

作者基于 MiniMax-Text-01 开发了 MiniMax-VL-01, 来扩展模型的多模态理解能力。MiniMax-VL-01 是一个标准的 ViT-MLP-LLM 架构，其中：

ViT: 从零开始训练的 ViT, 参数为 303M
MLP: 2 层的 MLP, 激活函数为 GeLU
LLM: MiniMax-Text-01

作者实现了动态分辨率策略，用于处理不同大小的图片。每张图片都被切分为 $336\times 336$ 的大小，然后作者还加入了一个 thumbnail image 用于提取全局信息。

Data

为了训练 vision encoder, 作者使用了 694M 的 Image-caption pair 数据，为了提高数据质量，作者对其中 180M 的数据进行的 caption 的改写，训练的时候原始数据和重写数据的采样概率都是 $50\%$.

作者还从公开数据收集到了 100M 的图片，每个图片都有详细的描述，这些描述都由一个 caption model 生成，然后再由人类进行修正。每条描述大概包含 300 个 text token

作者还构建了一个高质量的 instruction dataset, 这个数据集是通过预定义一系列任务，然后对于每个任务分别生成对应的 QA pair 得到的。在训练的时候，作者平衡了不同任务的采样概率。

最后，作者对收集收集数据的类别分布进行了可视化，结果如下

Training

训练一共包含 5 个 stage, 第一个 stage 用于单独训练 ViT, 后面三个 stage 分别用于对齐和提升模型的多模态能力。

Stage 0: ViT training 作者基于 ViT-L/14 来训练 vision encoder, 作者使用了基于对比学习的方法，参考了 CoCa 里的方法。训练过程中，作者首先在 224 的图片精度下训练了 37B 的 image-caption pairs. 接下来，作者在 336 的精度下使用了 1.2B 的 pairs 进行微调。最终，模型在 ImageNet-1K 的表现达到了 80.55%.

Stage 1: Modality alignment 这个阶段的目的是对齐视觉模态和文本模态，作者冻结 LLM, 只训练 ViT 和 MLP. 这个阶段一共使用了80B 的 image description data. 作者发现，在这个阶段，提升图片精度对模型表现提升影响不大。

Stage 2: Enhancement of Vision Understanding 这个阶段其实就是 instruction tuning, 作者解冻所有参数，然后使用了 420B 多模态 token 以及 21B MiniMax-Text-01 的 post-training token 来提高模型的指令跟随能力。

Stage 3: Enhancement of User Experience 这个阶段的目的是进一步提高模型在真实场景下模型的避险。作者构建了高质量的多模态数据，包括 44.8B 的多模态 token, 这个阶段模型训练了一个 epoch.

Stage 4: Enhancement of Preference 这个阶段的目的是对齐，作者使用 DPO 来训练模型，数据集包括 40,000 条数据。数据构建过程如下：

prompt selection: 作者基于真实用户交互数据构建了高质量的 prompt
Response generation: 作者通过多次采样，然后使用 MiniMax-Text-01 来生成多样化的数据
Reward Assignment: 接下来作者使用 MiniMax-Text- 来评估最终的结果
Pair Construction: 最后，基于评估结果来构建 preference pairs.

作者还加入了一些纯文本的 preference data.

Observation 作者发现，如果 foundation model 比较强，使用 DPO 的话可能会导致过拟合，因此作者的做法是在一个 epoch 之前就停止训练。

Evaluation

作者对比了 GPT-4o, Claude-3.5-Sonnet, Gemini-1.5-Pro, Qwen2-VL-72B, InternVL2.5-78B 和 LLaMA-3.2-90B, 实验结果如下

Conclusion

在本文中，作者提出了 MiniMax-01 大模型系列，包括 MiniMax-Text-01 和 MiniMax-VL-01 两个模型，前者是一个总参数为 456B, 激活参数为 45.9B 的基于 MoE 和 hybrid attention 的大语言模型，后者是基于 MiniMax-Text-01 的多模态大模型。作者详细介绍了模型的架构，数据和训练。

作者认为，未来有以下探索方向：

Long-context Evaluation: 如何更好评估模型的长上下文能力。现有的 long-context benchmark 的任务都比较简单，很难反应模型的真实能力
Model Architecture: 当前的模型还有 1/8 的模块包含 softmax attention, 如何构建更高效的架构来去掉 softmax attention 是一个值得探究的方向。
Complex Programming Tasks: 如何提高模型的 coding 能力也是一个值得探索的方向。

Notes on DeepSeek-V3.2

Tue, 06 Jan 2026 17:30:40 +0800

Introduction

作者首先回顾了开源模型如 MiniMax-01, Kimi-k2, Qwen3, GLM-4.5 和闭源模型的进展，作者指出，现在的开源模型和闭源模型在表现上仍然存在较大差距。作者认为这种差距主要是由于三个原因：

Transformer 提出的 softmax attention 在处理长文本时效率非常低
已有的开源模型在 post-training 阶段使用的算力不够
开原模型的泛化和指令跟随能力不如闭源模型

基于这三个问题，DeepSeek-V3.2 分别进行了改进：

在架构上，作者提出了 DSA，一个高效的稀疏注意力机制，用于降低计算复杂度
在 post-training 阶段，作者使用了比 pre-training 阶段高 $10\%$ 的算力，用于提高模型的能力
作者提出了一个 pipeline 用于提高模型在工具调用场景下的 reasoning 能力

通过实验作者发现，模型达到了和 Kimi-k2 以及 GPT-5 差不多的 reasoning 表现。

Method

Architecture

DeepSeek-V3.2 与 DeepSeek-V3.1 不同之处在于使用了 DeepSeek Sparse Attention (DSA). 架构如下图所示

DSA 包含两个模块：

lightning indexer
fine-grained token selection mechanism

其中，lightning indexer 负责计算 query token $h_t\in\mathbb{R}^d$ 和一个 preceding token $h_s\in\mathbb{R}^d$ 之间的 index score $I_{t,s}$ 来决定 query token 选择的 token:

$$ I_{t,s} = \sum_{j=1}^{H_I}w_{t,j}^I \mathrm{ReLU}(q_{t,j}^I\cdot k_s^I) $$

其中， $H^I$ 代表 indexer heads 的个数，$q_{t,j}^I\in\mathbb{R}^{d^I}$ 和 $w_{t,j}^I$ 由 query token $h_t$ 得到，$k_s^I\in\mathbb{R}^{d^I}$ 由 preceding token $h_s$ 得到

给定 query token $h_t$ 对应的 index score $\{I_{t,s}\}$, fine-grained token selection mechanism 负责选取 top-K index score 对应的 key-value entries $\{c_s\}$, 然后 attention 的输出由 query token 个选取的 key value entries 得到：

$$ u_t = \mathrm{Attn}(h_t,\{c_s\mid I_{t,s}\in \mathrm{TopK}(I_{t,:})\}) $$

[!Recall] MoBA 也提出了类似的方法，但是 MoBA 是一个无需训练的策略

受 NSA 启发，作者实现了基于 MQA 模式的 MLA, 其中 latent vector 对于 query token 所有的 query heads 都是共享的。示意图如下所示

Continue Pre-training

作者在 DeepSeek-V3.1 的基础上进行了 continue pre-training. Continue pre-training 包含两个阶段：

Dense Warm-up stage 这个阶段用于训练 lightning indexer, 作者冻结除 lightning indexer 之外的参数，为了对齐 indexer output 和 main attention distribution, 对于第 $t$ 个 query token, 作者首先计算所有 attention heads 的 main attention score 之和，然后在 sequence 层面进行 L1-normalization 得到 $p_{t,:}\in\mathbb{R}^t$, 最后计算 lightning indexer 输出与 $p_{t,:}$ 之间的 KL divergence:

$$ \mathcal{L}^I = \sum_{t}\mathcal{D}_{KL}(p_{t,:}\ \Vert\ \mathrm{Softmax}(I_{t,:})) $$

这个阶段训练一共使用了 2.1B 的 token, lr 为 1e-3, 训练的步数为 1000 steps, batch size 为 16.

Sparse training stage 这个阶段模型所有的参数都参与训练，该阶段的目的是让模型学习到 DSA 的 sparse pattern. 训练时，作者让 lightning indexer 的输出与 $p_{t,S_t}$ 之间的输出进行对齐，其中 $S_t=\{s\mid I_{t,s}\in\mathrm{TopK}(I_{t,:})\}$:

$$ \mathcal{L}^I = \sum_{t}\mathcal{D}_{KL}(p_{t,S_t}\ \Vert\ \mathrm{Softmax}(I_{t,:})) $$

实际训练时，lightning indexer 仅接受 $\mathcal{L}^I$ 的反向传播，而 LLM 则仅接受 next-token prediction loss. 这个阶段模型一共使用了943.7B token, 其中 $K$ 设置为 $2048$. 学习率为 $7.3\times 1e-6$, 训练步数为 15,000 steps, batch size 为 480.

Post-training

post-training 与 DeepSeek-V3.1 一致：

Specialist Distillation 作者基于 DeepSeek-V3.2 base 构建了不同领域的 specialized model, 这些领域包括：

math
competitive programming
general logical reasoning
agentic coding
agentic search

每个 specialized model 都使用 RL 进行训练，训练数据包括 long CoT reasoning 数据以及 direct response generation 数据，specialized model 训练完毕之后，就被用于生产 domain-specific data, 作者通过实验发现，基于这种蒸馏方法，模型的表现仅比 specialized model 低一点，并且这个 gap 可以被后续的 RL 训练所抵消。GLM-4.5 也采取了类似的做法

Mixed RL Training 作者使用了 GRPO 算法进行训练，与 DeepSeek-V3 不同，作者将 reasoning, agent 以及 human alignment 的 RL 训练合并为了一个阶段，作者认为这种方法可以平衡模型在多个 domain 上的表现，并且可以防止 multi-stage training 带来的灾难性遗忘问题。对于 reasoning 和 agent 任务，作者使用了 rule-basd outcome reward, length penalty 以及 language consistency reward. 对于通用任务，作者使用了 generative reward model, 每个 prompt 都有对应的 rubris 用于 evaluation. 作者构建 reward 时主要考虑了:

length versus accuracy
language consistency versus accuracy

DeepSeek-V3.2-Speciale 除了 DeepSeek-V3.2 之外，作者还训练了 DeepSeek-V3.2-Speciale 模型，该模型仅使用 reasoning 数据进行训练，reasoning 数据包含了 DeepSeek-Math-V2 的训练数据以及 reward 方法。训练时，作者降低了 length penalty 的惩罚系数，最终 DeepSeek-V3.2-Speciale 模型拥有更强的 reasoning 能力

Scaling GRPO

作者在 GRPO 的基础上对 KL estimate 进行了改进（见 KL divergence），使用了 importance sampling 对 K3 estimator 进行修正:

$$ \mathcal{D}_{\mathrm{KL}}(\pi_\theta(o_{i,t})\Vert\pi_{\mathrm{ref}}(o_{i,t})) = \frac{\pi_\theta(o_{i,t}\mid q, o_{i,使用 K3 estimator 之后，现在 KL estimator 的梯度估计就变成无偏估计了，从而提高了整体训练的稳定性。作者还发现不同任务对 KL regularization 的需求不一致，对数学等 domain, 我们应该采取较小的 KL penalty 或者不使用 KL penalty 反而能提升性能

Off-Policy Sequence Masking 作者还使用了 sequence masking 来提高 off-policy data 的数据使用效率，由于不同 rollout 的完成时间不一致，训练过程中会出现 off-policy 现象，即某些 mini-batch 不是由当前 policy 产生，这个现象在 Magistral 中也有提到，这种训练 - 推理不一致性会进一步加剧 off-policy 程度，为了提高训练稳定性，作者将 policy divergence 程度比较高的 sequence 给 mask 掉，更新后的损失函数如下所示

$$ \mathcal{J}_{\mathrm{GRPO}}(\theta) = \mathbb{E}_{(q,a)\sim\mathcal{D},\{o_i\}_{i=1}^G\sim \pi_{\theta_{old}}(\cdot\mid q)}\left[ \frac{1}{G}\sum_{i=1}^G\frac{1}{|o_i|}\sum_{t=1}^{|o_i|}\min\left(r_{i,t}(\theta)\hat{A}_{i,t},\mathrm{clip}\left(r_{i,t}(\theta), 1-\epsilon, 1+\epsilon\right)\hat{A}_{i,t}\right)M_{i,t}-\beta \mathcal{D}_{\mathrm{KL}}(\pi_\theta(o_{i,t})\Vert\pi_{\mathrm{ref}}(o_{i,t}))\right] $$

其中，

$$ M_{i,t} = \begin{cases} 0, &\hat{A}_{i,t}<0, \frac{1}{|o_i|}\sum_{t=1}^{|o_i|}\log r_{i,t}(\theta)<\delta\\ 1, &\text{otherwise} \end{cases} $$

用于决定是否对当前 sequence 进行 mask, $\delta$ 是一个超参数控制 policy divergence 程度。作者认为，模型主要从自身的错误进行学习，而 off-policy 的负样本模型学习提升有限甚至有害。作者发现加入这个 masking 策略之后，模型的训练稳定性有所提升。

Keep Routing MoE 模型在进行 On-policy RL 训练是，由于 policy 的更新，新 policy 和旧 policy 专家的 routing 可能会不一致，这种不一致性会降低训练的稳定性以及 off-policy 现象。为了解决这个问题，作者在采样室，保存了训练阶段所使用的 expert routing path, 来保证训练推理的一致性。这种策略可以有效提高针对 MoE 模型的 RL 训练稳定性

Keep Sampling Mask 作者发现，使用 top-p 和 top-K 可以提高 LLM 输出的质量，但是这种采样擦略也会导致 $\pi_{\mathrm{old}}$ 和 $\pi_{\theta}$ action space 的不匹配，因此，作者记录了 $\pi_{\mathrm{old}}$ 采样过程中的 truncation mask, 然后再训练的时候将其应用到 $\pi_{\theta}$ 上，作者发现通过这种方式可以有效提高 RL 训练过程中的 language consistency

Thinking in Tool-Use

本节的目的是希望能够将 reasoning 能力应用到工具调用的场景下。

作者发现，如果使用 DeepSeek-R1 的策略，在下一轮对话到来时，丢弃到当前的 reasoning content 会让模型重新生成针对问题的 CoT, 从而产生 token inefficiency. 为了解决这个问题，作者构建了一个上下文管理策略，如下图所示

具体做法为：

reasoning content 只有当有新的 user message 进入时才会丢弃；如果只有工具调用相关的 message, 则 reasoning content 会保留
移除 reasoning content 时，会保留对应的工具调用及其结果

在训练时，作者认为模型已经掌握了比较好的指令跟随能力，我们仅需要将 reasoning data (non-agentic) 和 non-reasoning data(agentic) 以不同的 prompt 输入给模型就能够得到比较好的结果。

对于训练的数据，作者认为 RL 任务的多样性可以有效提高模型的 robustness, 因此作者构建了不同的环境及其对应的 prompt, 生成的任务如下表所示

	number of tasks	environment	prompt
code agent	24667	real	extracted
search agent	50275	real	synthesized
general agent	4417	synthesized	synthesized
code interpreter	5908	real	extracted

search agent: 作者使用了 multi-agent 的策略，包括 question-construction agent 用于构建 QA pair, multiple answer-generation agent 用于构建不同的 response, 一个 verification agent 用于评估生成的 response. 最后作者使用 generative reward model 来评分
code agent: 作者爬取了 Github 上的 pull request, 然后进行过滤，接下来作者通过一个 environment-setup agent 来构建对应的环境
code interpreter agent: 作者使用 jupyter notebook 作为代码解释器来解决复杂的 reasoning tasks, 包括 math, logic, data science 等
general agent: 作者构建了验证简单解决困难的任务。首先作者基于 agent 和 task category 来生成或检索相关数据；接下来作者合成一个任务相关的工具集合；最后，作者让一个 agent 来提出任务以及对应的解法，并不断提高任务的难度。最后得到 <environment, tools, task, verifier> 的 tuple 格式

Experiments

作者对比了 DeepSeek-V3.2-Exp 和 Claude-4.5 Sonnet, GPT-5, Gemini 3.0 Pro, Kimi-K2 thinking, MiniMax M2 的表现，评测结果如下

结果显示，DeepSeek V 3.2 和 GPT-high 在 reasoning 任务上的表现差不多。作者认为，进一步提高 RL 阶段的算力可以有效提高模型的表现

作者还对比了 DeepSeek-V 3.2-Speciale 的表现，结果显示，提高 token budget 之后，模型的表现显著提高，与 Gemini 3.0 Pro 可以相比，但是其 token efficiency 仍然弱于 Gemini 3.0 Pro, 结果如下图所示

接下来作者验证了 synthesis agentic tasks 对模型表现的影响。首先，作者随机采样一批样本使用闭源 LLM 进行测试，发现闭源模型表现最好为 $62\%$, 这说明了合成数据对于 DeepSeek V3.2 和闭源模型都是有挑战的。

其次，作者探究了合成数据是否能够提高 RL 的泛化性，作者构建了两个额外模型：

SFT: 在 SFT checkpoint 上进行 RL
Exp: 仅在 search 以及 code environment 上进行 RL

对比结果发现，合成数据缺失可以有效提高模型的表现

作者还对比了 DSA 的效率，DSA 可以将 attention 的计算复杂度由 $\mathcal{O}(L^2)$ 降低到 $\mathcal{O}(LK)$, 其中 $K$ 是选取的 top-K tokens. 尽管 lightning indexer 的复杂度仍然是 $\mathcal{O}(L^2)$, 但是其计算量远小于 MLA, 作者对比了两者的效率，实验结果如下图所示

可以看到，DeepSeek-V3.2 的 prefilling 和 decoding 效率都远高于 DeepSeek-V3.1

最后，作者对比了不同上下文管理策略对模型表现的影响，对比策略有：

Summary: 对轨迹进行总结然后重新初始化 rollout
Discard 75%: 丢弃到初始 75% 的工具调用历史
Discard-all: 丢弃掉所有的工具调用历史
Parallel-fewest step: 多次采样然后保留步数最少得轨迹

实验结果如下图所示

结果显示，使用上下文管理策略之后，模型的表现有了显著提升。并且，discard-all 策略虽然很简单，但是其表现非常好。作者认为如何根据不同场景来选取合适的策略是一个待解决的问题。

Conclusion

作者在本文中提出了 DeepSeek-V 3.2, DeepSeek-V 3.2 使用了一个稀疏注意力机制来提高模型在长上下文场景下的计算效率。作者还通过提升 RL 阶段的算力来提高模型在下游任务上的表现。最后，作者合成了大规模的 agentic task 来提升模型的 agent 能力。

作者认为，相比于 Gemini 3.0 Pro, 模型的知识广度仍然有限。并且，目前模型的 token efficiency 仍然是一个问题，模型需要更长的轨迹输出才能达到 Gemini 3.0 Pro 的表现。最后，模型解决复杂问题的能力仍然弱于闭源模型。

References

DeepSeek-V3.2: Pushing the Frontier of Open Large Language Models

Notes on NoPE

Wed, 24 Dec 2025 15:19:42 +0800

Introduction

【参考文献 1】中系统性对比了 AliBi, RoPE , T5 提出的 T5 bias 以及 Transformer 提出的绝对位置编码 (APE).

作者发现，常用的方法在 length generalization 上表现并不是最好的，而 NoPE 不需要额外的计算开销反而效果最好。

【参考文献 2】进一步探究了 NoPE 长度外推的泛化性。作者有三点发现：

NoPE 相比于 RoPE, 其长度外推泛化能力更强
对于 NoPE 来说，模型会在还没有到达预训练上下文长度之前，表现就出现下降的情况
通过调整 softmax 的温度超参数，我们可以提高 NoPE 的长度外推泛化性能力。

Method

【参考文献 1】对比了不同 position encoding 的相似度，结果如下图所示

实验结果表明，NoPE 与 T5 提出的 T5 bias 最相似。

作者在理论上推导出了 NoPE 的两个性质：

Theorem 1 (Absolute Encoding) Let $x$ be an input sequence of length $T + 1$ to the model. Then, the first layer of $f_θ$ can recover absolute positions $[1, . . . , T + 1]$ in the hidden state $H^{(1)}$. That is, there exist $W_Q, W_K , W_V , W_O, W_1$, and $W_2$ such that the self-attention and feedforward operations in the first layer compute absolute positions and write it to the next hidden state.

Theorem 2 (Relative Encoding) Suppose that the hidden state $H^{(1)}$ contains absolute positional information, as stated in Theorem 1, and assume that it is not overwritten by any subsequent layers. Then, the self-attention in all subsequent layers can implement a relative positional encoding: there exists a parameterization of fθ such that, for $\ell ≥ 2$, the attention dot product between query $q_n$ and key $k_m$ at positions n and m can be expressed as:

$$ \langle q_n, k_m\rangle = f_{cnt}(q, k) + f_{rel}(n − m) $$

where $f_{cnt}$ is a function of their content, and $f_{rel}$ is a function of their relative distance.

【参考文献 2】探究了 softmax 中 normalization factor 对模型表现的影响，作者定义 attention 为

$$ \mathrm{Attn}(q,k,v) = \mathrm{softmax}\left(\lambda q^Tk\right)v $$

实验结果如下图所示

结果说明，通过调整 $\lambda$ 我们可以有效提高 NoPE 的上下文扩展泛化能力

Conclusion

NoPE 说明在 transformer 中我们可以不需要加入位置编码模块，这两篇论文均验证了 NoPE 的有效性。

References

Notes on ALiBi

Wed, 24 Dec 2025 15:10:55 +0800

meta 等提出了 ALiBi, 一个通过 linear biases 来实现位置编码的方法来提高 LLM 在推理阶段的外推能力。

Introduction

当下，有若干种位置编码的方式：

Sinusoidal position embeddings: Transformer 提出的正弦位置编码
RoPE: RoPE 提出的旋转位置编码
T5 bias: T5 提出的相对位置编码

作者通过实验对比了不同的位置编码方法，发现这些方法在推理阶段的外推能力都比较差。

为了解决这个问题，作者提出了 ALiBi (attention with linear biases), 一个几乎不增加计算和内存开销的位置编码方法，来提高 LLM 在推理阶段的外推能力。

Method

作者将外推能力定义为

a model’s ability to continue performing well as the number of input tokens during validation increases beyond the number of tokens on which the the model was trained.

计 $L$ 为训练阶段的上下文长度， $L_{valid}$ 为推理阶段的上下文长度。

作者首先对比了不同的位置编码方法的外推能力，结果如下图所示

结果显示，不同位置编码在推理阶段扩展模型的上下文能力均有限。

Context Length	$L$	$L_{valid}$
Sinusoidal	512	50
	1024	50
RoPE	512	200
	1024	100
T5 bias	512	600
	1024	800
ALiBi	512	-
	1024	-

为了解决这个问题，作者提出了 AliBi, 其表达式为

$$ \mathrm{softmax}(q_iK^T+m\cdot [-(i-1),\dots,-2,-1,0]) $$

其中 $m$ 是一个和 heads 相关的超参数。如果我们有 8 个 heads, 则对应的 scaling 值分别为 $[1/2^1,1/2^2,\dots,1/2^8]$, 如果我们有 16 个 heads, 则我们对 8 个 heads 的结果进行插值，得到 $[1/2^{0.5},1/2^1,\dots,1/2^8]$. ALiBi 的示意图如下所示

ALiBi 通过 bias 惩罚了较远的 query-key pairs, 并且不同的 heads 的惩罚项也不同，从而每个 head 对距离的信息敏感度也不尽相同。

Experiments

ALiBi 在 WikiText-103 上的实验结果如下图所示

Conclusion

作者分析了已有的 position embedding 方法，发现已有的方法在推理阶段均不能有效扩展模型的上下文长度。因此，作者提出了 AliBi, 一个通过 linear bias 来增加位置信息的方法，作者通过实验验证了 ALiBi 的有效性。

References

Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation

Notes on T5

Wed, 24 Dec 2025 15:07:08 +0800

google 在 2020 年发表了 T5 (Text-to-Text Transfer Transformer), 一个使用统一框架来将所有 NLP 任务转换为 text-to-text 格式的迁移学习框架。

Introduction

作者首先回顾了迁移学习和 pre-training, 迁移学习是提高模型在下游任务上表现的一类方法，但是目前还没有一个能够对比各种方法的框架。pre-training 通过在大量数据上进行预训练然后再进行微调，可以有效提高模型在下游任务上的表现。

为了解决这两个问题，作者首先将所有的文本处理任务统一为 “text-to-text” 的形式，这样我们就可以对比不同架构，训练方式以及数据对模型表现的影响

作者提到，本文并不是提供一个新的方法，而是详细对比不同方法，为后续研究提供基础。

Method

Model

在架构上，作者使用了 Transformer 的 encoder-decoder 架构，但是作者做了几点修改

作者提出了 T5 bias, 一个用于替换原始 transformer 绝对位置编码的相对位置编码形式
作者使用了 RMSNorm 替换了 Transformer 中的 LayerNorm.

Data

作者基于 Common Crawl 构建训练数据集，作者对数据进行了清洗，最终数据集大小为 750GB. 作者将这个数据集记为 C4 (Clean Crawled Corpus).

Downstream Tasks

下游任务包括：

text classification: GLUE, SuperGLUE
abstractive summarization: CNN/Daily Mail
question answering: SQuAD
translation: WMT English to German, French and Romanian

Input and Output Format

所有任务的输入输出都被转换为 text-to-text 格式。

Experiments

作者使用的 baseline 模型是一个基于 encoder-decoder 架构的 transformer 模型，其大小以及 configuration 与 BERT base 差不多，最终模型参数量为 220M。

field	num layers	hidden size	MLP hidden size	num heads	head dim	dropout	seq len
value	12	768	3072	12	64	0.1	512

训练时，作者使用了 AdaFactor 优化器，batch size 为 512，训练使用了 34B token. 学习率作者使用了 inverse square root learning schedule: $1/\sqrt{\max(n,k)}$, $n$ 和 $k$ 分别代表当前 step 和 warming up steps.

作者基于 sentencepiece (见 LLM tokenizer) 构建了 Tokenizer, 覆盖 English, German, French 和 Romanian 四种语言。

模型训练的目标函数为 BERT 使用的 “masked language modeling”, 格式如下所示

1
2
3
4
5
6
7
8


# original text
Thank you for inviting me to your party last week.

# inputs
Thank you <X> me to your party <Y> week.

# targets
<X> for inviting <Y> last <Z>

首先，作者对比了不同的架构。作者对比了如下三种 transformer 的变体：

实验结果如下图所示

结果显示，encoder-decoder 架构，denoising 训练目标的效果最好。并且，当 layers 减少一半之后，模型的表现大幅度下降。共享参数的 encoder-decoder 架构表现比 prefix LM 效果更好

接下来作者针对 denoising 的配置进行了测试，实验结果发现 BERT-style 的训练目标效果最好，并且 corruption 比例对模型的表现影响有限，作者使用了 BERT 的配置，即 $15\%$ 的 token 被 masked 掉。对于 span length, 作者通过实验发现不同的 span length 对结果影响不大。因此，作者将 span length 设置为 $3$.

在数据上，作者发现：

对数据进行过滤可以提高模型的表现
使用 in-domain 的数据可以提高模型在该 domain 上的表现，但是问题在于 In-domain 的数据往往比较少
数据量过少时，模型会出现 memorization，也就是过拟合的情况

Overall

作者总结前面的发现，构建了 5 个 size 的模型：

Model	layers	hidden size	FFN hidden size	num heads	head dimension
Small	6	512	2048	8	64
Base	12	768	3072	12	64
Large	24	1024	4096	16	64
3B	24	1024	16384	32	128
11B	24	1024	65536	128	128

Conclusion

作者在本文中提出了 T5, 一个统一所有文本处理任务的迁移学习框架，作者系统性探究了架构，数据以及训练对模型最终表现的影响。最终作者基于 encoder-decoder transformer 架构以及 denoising training objective 训练得到了 T5 系列大语言模型。

References

Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer

Base of RoPE Bounds Context Length

Mon, 22 Dec 2025 11:34:42 +0800

Introduction

RoPE 已经成为了大多数 LLM 使用的 position encoding 范式，但是，RoPE 与 LLM long context 之间的关系还没有被探索清楚。在本文中，作者就探究了 base frequency 与 LLM context capability 之间的关系，并给出了一个达到指定上下文长度所需要的 base frequency 的 lower bound.

Method

首先作者回顾了 attention 与 RoPE 的定义，关键就是 RoPE 这部分，如下所示

$$ A_{ij} = (R_{i,\theta}q_i)^T(R_{j,\theta}k_j) = q_i^TR_{i-j,\theta}k_j $$

这里 $\theta$ 就是 base frequency, 作者总结不同模型的 base frequency 配置如下表所示

Model	Llama-7B	Llama2-7B	Llama3-8B	Mistral-7B	Baichuan2-7B
Base frequency	10,000	10,000	500,000	1,000,000	10,000
Context length	2048	4096	8192	32,768	4,096

接下来，作者回顾了 YARN. 其核心思想在于，预训练阶段所有可能的 $\cos(t-s)\theta_i$ 都见过，才能保证模型的 OOD 表现

作者认为 base frequency 的设置应该满足两个条件：

The closer token gets more attention: 当前的 token 应该给邻近的 token 更高的注意力
The similar token gets more attention: 当前的 token 应该给相似的 token 更高的注意力

在 RoPE 中，作者已经给出了 $A_{ij}$ 与相对距离 $|i-j|$ 之间的关系。因此第一个性质已经满足了。

接下来，作者分析了相似 token 的性质，作者定义 token 的相似性如下：

$$ \mathbb{E}_{q,k^*}[q^TR_{m,\theta}k^*] - \mathbb{E}_{q,k}[q^TR_{m,\theta}k] $$

这里 $k^*=q+\epsilon$ 代表了相似的 token, 而 $k$ 是一个随机 token. 作者给出的结论如下

Theorem 假设 $q,k\in\mathbb{R}^d$ 独立同分布，它们的标准差为 $\sigma\in\mathbb{R}$, 则对于 $k^*=q+\epsilon$, $\epsilon$ 是一个随机变量满足 $\mathbb{E}[\epsilon]=0$, 则我们有

$$ \mathbb{E}_{q,k^*}[q^TR_{m,\theta}k^*] - \mathbb{E}_{q,k}[q^TR_{m,\theta}k] = 2\sigma^2\sum_{i=0}^{d/2-1}\cos(m\theta_i) $$

作者定义 $B_{m,\theta}=\sum_{i=0}^{d/2-1}\cos(m\theta_i)$, 作者认为给定 $\theta$, 模型的上下文长度 $L_\theta$ 满足

$$ L_\theta = \sup\{L\mid B_{m,\theta}\geq 0, \forall m\in[L]\} $$

也就是说，base frequency 决定了 LLM 的上下文长度。作者给出了不同的上下文长度对应的 base frequency 如下表所示

Context Len.	1k	2k	4k	8k	16k	32k	64k	128k	256k	512k	1M
Lower Bound	4.3e3	1.6e4	2.7e4	8.4e4	3.1e5	6.4e5	2.1e6	7.8e6	3.6e7	6.4e7	5.1e8

总的来说，远距离衰减性保证了模型会更关注邻近的 token, 而相似 token 保证了模型能够区分出真正有意义的 token.

Experiments

作者首先分析了 base frequency 在 fine-tuning 阶段对模型上下文能力的影响，实验结果如下图所示

从实验结果可以看到，当 base frequency 低于阈值时，模型的表现急剧下降。

作者进一步探讨了 base frequency 对于模型 pre-training 阶段的影响，结果也是一样的，即非常小的 base frequency 会限制模型的 context 能力，结果下图所示（三行分别代表了 base frequency 为 1e2, 1e4 和 1e6 的情况）

可以看到，尽管 perplexity 都差不多，但是使用更大的 base frequency 其长上下文能力明显更好。

作者进一步分析了为什么较小的 base frequency 会影响模型的长上下文能力。作者认为较小的 base frequency 会导致 $B_{m,\theta}$ 接近于 0，从而模型难以区分随机 token 和相似 token, 这样模型只能依赖于邻近 token 进行学习，这样就限制了模型的长上下文能力

作者还进一步对比了提高 base frequency 与 Interpolation 两种做法，实验结果如下表所示

实验结果说明，Interpolation 在上下文超过 30K 之后，其 $B_{m,\theta}\leq0$ 的次数显著增加，表明了其和上下文能力之间的关系。

Conclusion

作者在本文中，探究了 RoPE 中 base frequency 与 LLM 上下文能力之间的关系，发现了提高模型的上下文能力需要关注 RoPE 的 base frequency 超参数，并给出了对应的 lower bound. 作者通过实验验证了这个观点。

References

arxiv

Notes on NSA

Mon, 15 Dec 2025 17:39:16 +0800

DeepSeek 在 25 年 1 月提出了 Natively trainable Sparse Attention (NSA), 一个软硬件结合的稀疏注意力机制，NSA 可以在提高模型推理效率的同时提高计算效率。

Introduction

现有的大模型主要是基于 Transformer 提出的 softmax attention, 其主要问题在于随上下文长度增加，其 latency 也上升更快。理论估计，对于 64k 上下文长度的输出，softmax attention 部分的计算占 $70\%\sim80\%$ 的 latency.

为了解决 softmax 的 high latency 问题，，一个做法就是使用稀疏注意力机制，如 MInference 等，但是这些系数注意力机制大多没有实际部署，且它们一般只在 inference 阶段使用

作者认为解决这个问题有两个挑战：

Hardware-aligned inference speedup: 降低 inference latency 需要算法与硬件结合，不能只关注算法层面的改进
Training-aware algorithm design: 需要在训练阶段也支持算法，从而可以降低训练的算力消耗并且保持模型的表现

为了解决这两个问题，作者就提出了 natively trainable sparse attention (NSA) 架构。NSA 通过将 key 和 value 分割为不同的 block, 然后基于三种 path: compressed coarse-grained tokens, selectively retrained fine-grained tokens 以及 sliding windows for local contextual information 来进行处理和过滤。NSA 提出了两点观点改进：

Hardware-aligned system: 优化了 blockwise sparse attention 来平衡 arithmetic intensity.
Training-aware design: 支持端到端的训练和部署

Method

Overview

作者首先回顾了 attention 的定义如下：

$$ \mathbf{o}_t=\mathrm{Attn}(\mathbf{q_t},\mathbf{k}_{:,t}, \mathbf{v}_{:,t})=\sum_{i=1}^t \frac{\alpha_{t,i}}{\sum_{t,i}\alpha_{t,i}}\mathbf{v_i},\ \alpha_{t,i} = \exp\left(\frac{\mathbf{q_t}^T\mathbf{k}_{i}}{\sqrt{d_k}}\right) $$

其中 $\mathbf{q_t}\in\mathbb{R}^{d_k}$.

接下来是 Arithmetic Intensity. Arithmetic intensity 指的是 FLOPs 与内存访问次数之比。由于现在的 GPU 都是计算密集型设备，理想情况下应该是 Arithmetic intensity 越高越好。

对于 causal self-attention 来说，在训练以及 prefilling 阶段，由于 batch 较大，因此整体的 Arithmetic intensity 较高，因而这两个阶段是 computer-bound. 但是在 decoding 阶段，由于其 token-by-token generation 的性质，每次生成新的 token 时都需要重新加载 KV cache, 因而是 memory-bound.

从而我们的优化目标也变得不一致：在训练阶段，我们希望降低计算消耗，而在推理 (decodng) 阶段，我们希望降低内存访问次数。

基于这两个目标，作者提出了使用 $\mathbf{k}_{:,t}, \mathbf{v}_{:,t}$ 的子集 $\tilde{K}_t, \tilde{V}_t$ 来参与计算，其对应的 attention 如下所示

$$ \tilde{K}_t=f_K(\mathbf{q_t},\mathbf{k}_{:,t}, \mathbf{v}_{:,t}), \tilde{V}_t=f_V(\mathbf{q_t},\mathbf{k}_{:,t}, \mathbf{v}_{:,t}), \mathbf{o}_t=\mathrm{Attn}(\mathbf{q_t},\tilde{K}_t, \tilde{V}_t) $$

我们还可以结合不同的方法来进行组合：

$$ \mathbf{o}_t^*=\sum_{c\in\mathcal{C}}g_t^c\mathrm{Attn}(\mathbf{q_t},\tilde{K}_t^c, \tilde{V}_t^c) $$

作者在本文中使用了三种方法 $\mathcal{C}=\{\mathrm{cmp},\mathrm{slc},\mathrm{win}\}$, 分别代表了 compression, selection 以及 sliding window, $g_t^c\in[0,1]$ 代表了不同方法对应的 gating score, 类似于 MoE 的 gating layer, $g_t^c$ 由一个 MLP 和一个 sigmoid activation 生成。最终 NSA 的架构如下图所示

作者定义 $N_t$ 代表参与计算的 KV 的总个数：

$$ N_t = \sum_{c\in\mathcal{C}} \mathrm{size}[\tilde{K}_t^c]. $$

作者使用了一个较高的 sparsity ratio 来保证 $N_t<

Design

接下来作者分别介绍了每一部分的设计

Token Compression

对于 token compression, 其定义如下：

$$ \tilde{K}_t^{\mathrm{cmp}} = f_K^{\mathrm{cmp}}(\mathbf{k}_{:,t})=\left\{\phi(\mathbf{k}_{id+1:id+l})\mid 0\leq i\leq \left\lfloor\frac{t-l}{d}\right\rfloor\right\}\in\mathbb{R}^{d_k\times \left\lfloor\frac{t-l}{d}\right\rfloor} $$

其中 $l$ 是 block size, $d$ 是 sliding stride, $\phi:\mathbb{R}^{l\times d_k}\to \mathbb{R}^d_k$ 是一个 MLP 用于将 block key 映射为一个单一的 key. 对于 $\tilde{V}_t^{\mathrm{cmp}}$ 作者也使用了类似的做法。

Token Selection

仅使用 compressed token 的话，可能会丢失一些细粒度的信息。因此，作者额外提出了 token selection 机制来解决这个问题。

作者使用的做法是 blockwise selection. 这样做的原因有两点：

hardware efficiency. 这样做的原因是 GPU 访问内存是在 block 层面进行的，因而更加高效
inherent distribution patterns of attention scores. MInference 证明了 attention score 在空间上存在连续性。即相邻的 key 对应的重要性非常相似

为了实现 block-wise selection, 作者首先将 key value sequences 分割为 blocks, 然后针对每个 blocks 分配 Importance score.

作者首先介绍了如何计算不同 block 的 importance score.

如果 selection block size 与 compression block size ，即 $l'=l$ 相同的话，则我们可以直接用 compression block 提供的信息：

$$ \mathbf{p}_{t}^{\mathrm{cmp}} = \mathrm{sotmax}\left(\mathbf{q}_t^T\tilde{K}_t^{\mathrm{cmp}}\right) $$

其中 $\mathbf{p}_{t}^{\mathrm{cmp}}\in\mathbb{R}^{\left\lfloor\frac{t-l}{d}\right\rfloor+1}$ 代表了 $\mathbf{q}_t$ 和 compressed key $\tilde{K}_t^{\mathrm{cmp}}$ 之间的 attention score.

如果 $l'\neq l$ 的话，作者通过空间关系来进行计算，假设 $l\leq l'$, $d\mid l$, $d\mod l'$, 则我们有

$$ \mathbf{p}_{t}^{\mathrm{slc}}[j] = \sum_{m=0}^{l'/d-1}\sum_{n=0}^{l/d-1}\mathbf{p}_{t}^{\mathrm{cmp}}\left[\frac{l'}{d}j-m-n\right] $$

对于 GQA 和 MQA, 由于其 KV-cache 在 heads 之间共享，因此我们必须保证不同 heads 之间的 consistency, 因此作者提出了 shared importance score 如下：

$$ \mathbf{p}_{t}^{\mathrm{slc}'} = \sum_{h=1}^H\mathbf{p}_{t}^{\mathrm{slc},(h)} $$

接下来，对于每个 block 及其对应的 Importance score, 作者保存 top-$n$ sparse blcoks, 如下所示

$$ \begin{aligned} \mathcal{I}_t &= \{i\mid \mathrm{rank}(p_t^{\mathrm{slc}'}[i])\leq n\}\\ \tilde{K}_t^{\mathrm{slc}} &= \mathrm{Cat}\left[\{\mathbf{k}_{il'+1:(i+1)l'}\mid i\in \mathcal{I}_t\}\right] \end{aligned} $$

其中 $\mathrm{rank}(\cdot)$ 代表了降序排列的 importance scores. $\mathcal{I}_t$ 是选择出来的 block indices, $\mathrm{Cat}(\cdot)$ 表示了 concatenation operation. $\tilde{K}_t^{\mathrm{slc}}\in\mathbb{R}^{d_k\times il'}$ 代表了选择出来的 key.

Sliding Window

为了避免 local pattern 对 compression token 以及 selection token 的学习产生影响，作者额外使用了一个 branch 来学习这个 local pattern. 其具体做法就是维持一个 sliding window 用于最近的若干个 token, 即

$$ \tilde{K}_t^{\mathrm{win}} = \mathbf{k}_{t-w:t}, \tilde{V}_t^{\mathrm{win}} = \mathbf{v}_{t-w:t} $$

这里 $w$ 是 window size.

为了进一步避免 shortcut learning, 对于三个 branch 作者提供了不同的 key 和 values

Kernel Design

接下来是针对硬件设计进行的优化。由于 flash attention 2 对 compression attention 以及 sliding window attention 已经支持的比较好，作者这里介绍了如何针对 selection attention 进行优化。

Experiments

作者构建了一个 27B-A3B 的 MoE 模型，attention 基于 GQA, MoE 基于 DeepSeekMoE. 模型配置如下表所示

field	value
layers	30
hidden dimension	2560
head groups	4
attention heads	64
query head dimension	192
value head dimension	128
routed experts	72
shared experts	2
activated experts	6
dense layers	1

NSA 配置如下

field	value
$l$	32
$d$	16
$l'$	64
$n$	16
$w$	512

其中 selection blocks 包含初始的一个 block 以及最近的 2 个 block.

模型先在 8K 的上下文长度下使用 270B token 进行预训练，接下来在使用 YARN 将模型上下文通过 continual pre-training 以及 SFT 扩展到 32K. 训练过程的损失如下图所示

作者从 general performance, long-context performance 以及 CoT reasoning performance 三个层面来评估 NSA 的表现。

首先是 NSA 与其他 sparse attention 以及 baseline 在通用任务上表现的对比，结果如下图所示

接下来是 NSA 在 LongBench 上的表现：

作者还使用了 DeepSeek-R1 中的知识蒸馏方法，结果如下表所示

Generation token limit	8192	16384
Full Attention-R	0.046	0.092
NSA-R	0.121	0.146

上面的结果均验证了 NSA 的有效性

Analysis

接下来，作者分析了 NSA 的性质。作者首先对比了 NSA 和 flash attention 2 的训练速度，结果如下图所示

可以看到，相比于 flash attention 2, NSA 在 forward 过程和 backward 过程的的效率分别提升了 9 倍和 6 倍。作者认为这是由于两个优点：

NSA 使用了 block-wise memory access, 提高了 tensor core 的利用率
loop scheduling 减少了 KV transfer 时的 kernel 冗余

作者还对比了不同 attention 的解码速度，在 NSA 中，每次只需要 $\left\lfloor\frac{s-l}{d}\right\rfloor+nl'+w$ 个 token 就可以完成计算，作者对比不同 attention 所需余姚的 token 如下表所示如下表所示

Context Length	8192	16384	32768	65536
Full attention	8192	16384	32768	65536
NSA	2048	2560	3584	5632
speedup	4x	6.4x	9.1x	11.6x

Discussion

Conclusion

作者在本文中提出了 NSA, 一个通过软硬件协同结合 compression, selection 以及 sliding window 的稀疏注意力机制，作者通过实验验证了其有效性。

References

arxiv

MoE tutorial

Sat, 13 Dec 2025 16:04:04 +0800

本 blog 详细介绍了 MoE 模型的一些关键设计与相关实验结果，为 MoE 模型的学习提供基础。

Introduction

Motivation

现有大部分大语言模型均是基于 Transformer 架构，Kaplan scaling law 通过实验说明，大语言模型的表现与算力，数据，模型参数量息息相关。但是，对于 dense 模型来说，我们提高模型参数量时，必须同时提高所使用的算力。这就限制了大模型的 scaling law.

而 MoE 模型的解决方法为在计算时只激活部分参数，这样，我们就可以在同等激活参数量/算力下训练更大参数量的模型，从而达到更好地表现。

因此，MoE 模型的核心思想在于

使用相同的激活参数量/算力，提高模型总参数量，从而达到更好的表现。

Definition

MoE 模型和 dense 模型的示意图如下，图源 olmoe

一个 MoE layer 包括两个模块：

Router：Router 负责为 token 指定合适的专家
Expert：Expert 负责处理 token

对于输入 $x\in\mathbb{R}^d$, 我们假设有 $N$ 个 Expert，router 一般是一个 linear layer 再加上一个 gating function (softmax 或者 sigmoid，我们本文中使用 softmax), 其构建了 $\mathbb{R}^d\to\mathbb{R}^N$ 的映射，定义为：

$$ G(x) =[G_1(x),\dots,G_N(x)] = \mathrm{softmax}(W_gx + b)\in\mathbb{R}^N $$

其中 $W_g\in\mathbb{R}^{N\times d}$, $b\in\mathbb{R}^N$ 是可学习的参数。$G_{i}(x)$ 代表了当前 token $x$ 选择第 $i$ 个 Expert 的概率。

一般来说，Expert 会使用和 dense 模型一样的 MLP, 我们记为

$$ E_i(x) = \mathrm{FFN}(x), \quad i = 1,\dots,N $$

接下来，基于 $G(x)$ 和 $E(x)$, 我们会使用合适的方法来挑选 $K $$ e_i \in \{i\mid G_i(x)\in \mathrm{TopK}(G_i(x), K)\},\ i=1,\dots,K $$

最终 MoE layer 的输出为

$$ y = \sum_{i=1}^K\mathrm{Normalize}(G_{e_i}(x)) E_{e_i}(x) $$

这里 $\mathrm{Normalize}(\cdot)$ 代表我们对于输出进行归一化，即

$$ \mathrm{Normalize}(G_{e_i}) = \frac{\exp(G_{e_i})}{\sum_{i=1}^K \exp(G_{e_i})},\quad i=1,\dots,K $$

Why MoE

选择 MoE 的原因有三点：效率, scaling law 以及表现。

Efficiency

MoE 训练更加高效，如下图所示

Switch Transformer 的实验结果说明，MoE model 的训练效率比 dense model 快 7 倍左右。其他模型也有类似结论。总的来说，MoE 模型相比于 dense 模型，训练效率更高。

Scaling Law

MoE 模型可以突破传统 scaling law 的限制，在算力固定的情况下，我们可以通过提高 MoE 模型的稀疏度来进一步提高模型的表现

如上图所示，在 FLOPs 给定的情况下，随着模型稀疏度的提高，模型的表现和效率都有提升

Performance

MoE 模型的表现更强，如下图所示，MoE 模型的训练，验证损失以及在下游任务上的表现均超过了 dense 模型

Timeline

激活参数比例变化图如下图所示

可以看到，当前大部分模型的激活比例都在 $5\%$ 左右。

另一方面，从专家的激活比例变化图如下图所示

可以看到，现在大部分模型总专家数都在 200-400 左右，Kimi-k2 认为提高专家个数可以提高模型表现，而 LongCat 则是使用了 phantom expert 机制

MoE Design

Experts Design

Number of Experts

一般来说，专家个数越多，模型越稀疏，模型表现越好。扩展专家个数有两个方式：

直接增加专家个数，这会导致模型参数量上升，如 Switch Transformer
对已有的专家进行切分，将大专家切分为小专家，如 DeepSeekMoE

Switch Transformer 也通过实验发现，增加专家个数可以显著提高模型的训练效率和表现，结果如下图所示

可以看到，当我们增加专家个数的时候，模型的表现是持续提升的。并且当我们增加专家个数之后，模型的训练效率也有所提升。

DeepSeekMoE 提出了 fine-granularity expert 的概念，其做法是通过减少 expert 的大小在相同参数量的场景下使用更多的专家。实验结果如下图所示

可以看到，在稀疏度 (激活专家个数占总专家个数比例) 不变的情况下，提高专家的粒度，可以提高模型的表现。

olmoe 对 DeepSeekMoE 的这个观点进行了验证，结果如下图所示，

结果显示，当专家粒度从 8E-1A 扩展到 32E-4A 时，模型在 HellaSwag 上的表现提升了 $10\%$, 但是进一步扩展到 64E-8A 时，模型的表现提升不到 $2\%$, 这说明了无限制提升专家粒度对模型的提升越来越有限。

Kimi-k2 探究了针对 MoE 模型 sparsity 的 scaling law, 结果也说明，提升 sparsity 可以提高模型的表现。因此，其相对于 DeepSeek-V3 使用了 $50\%$ 额外的的专家数。Ling-mini-beta 进一步验证了这个观点。

Shared Experts

Shared Expert 由 DeepSeekMoE 提出，其基本思想为，固定某几个专家，响应所有的 token，这样可以让某些专家学习到共有的知识，而让其他的专家学习到特定的知识。这个方法随后被 Qwen1.5, Qwen2 , Qwen2.5 以及 DeepSeek-V3 所采用。

DeepSeekMoE 给出的实验结果如下

作者发现，当使用 shared experts 之后，模型在大部分 benchmark 上的表现都有所提升。

olmoe 在 32 个专家下进行了实验，比较了 4 个激活专家和 3 个激活专家 +1 个共享专家两种设置的表现，结果如下：

作者认为，加入 shared experts 之后，组合的可能性有所减少，这会降低模型的泛化性。因此，在 olmoe 中，作者没有使用 shared experts.

虽然 Qwen1.5, Qwen2 和 Qwen2.5 都使用了 shared experts, 但是后续的 Qwen3 中却并没有使用。

Ling-mini-beta 通过实验得出的结论为，shared expert 应该是一个非零的尽可能小的值，作者认为将 shared expert 设置为 1 是一个比较合理的选择。

Activation Function

一般来说，在选取 top-K 专家时，我们会对 gating layer 的输出进行归一化，通常我们会使用 softmax function:

$$ G(x) = \mathrm{softmax}(W_gx + b)\in\mathbb{R}^N $$

但是，在 Loss-Free Balancing 中，作者通过实验发现，使用 sigmoid 作为激活函数效果更好，即

$$ G(x) = \mathrm{sigmoid}(W_gx + b)\in\mathbb{R}^N $$

其对应的实验结果如下图所示

实验结果显示，sigmoid function 对于超参数更加 robust, 且表现也更好一些。下面是一些使用不同激活函数的模型例子

Activation function	Models
softmax	Step 3, Kimi-K2, gpt-oss-120B
sigmoid	GLM-4.5, dots.llm1, DeepSeek-V3

Routing Z-loss

Routing Z-loss 由 ST-MoE 提出， Switch Transformer 发现在 gating layer 中使用 float32 精度可以提高训练稳定性，但是这还不够，因此 ST-MoE 使用了如下的 router Z-loss:

$$ \mathcal{L}_z(x) = \frac1B\sum_{i=1}^B\left(\log\sum_{j=1}^N e^{x_j^{(i)}}\right)^2 $$

其中 $B$ 是 batch size ，$x_j^{(i)}=[W_gx_i+b]_j$ 代表了第 $j$ 个专家对 $i$ 个 token 的激活 logits. olmoe 实验验证结果如下

可以看到，加入 router Z-loss 之后，模型训练的稳定性有所提升，因此 Olmoe 采取了这个改进，但是后续的 MoE 模型使用 Z-loss 较少，个人猜测原因是 Loss-Free Balancing 中提出的加入额外的 loss 会影响 nex-token prediction loss

Routing Strategy

routing 策略直接决定了 MoE 模型的有效性。在为专家分配 token 的时候，我们有如下方式：

为每个 token 选取若干个专家
为每个专家选取若个个 token
动态分配 token 与专家之间的关系

三种选择方式如下图所示，图源 MoE survey

Expert Choice

每个专家选取 top-k 的 token，此时每个专家处理的 token 个数是相同的，这个方法的好处是自带 load balance。缺点是自回归生成的方式没有完整序列长度的信息，从而导致 token dropping，也就是某些 token 不会被任何专家处理，某些 token 会被多个专家处理。

目前采用这个策略的有 OpenMoE-2，核心思想是 dLLM 的输出长度固定，expert choice 策略更有效

Token Choice

每个 token 选取 top-k 的专家，好处是每个 token 都会被处理，缺点是容易导致负载不均衡。因此，一般需要加上负载均衡或者 token dropping 策略来提高负载均衡

Capacity Factor 由 Switch Transformer 提出，其定义为

$$ \text{expert capacity} = \left(\frac{\text{tokens per batch}}{\text{number of experts}}\right) * \text{capacity factor} $$

设置 capacity factor 之后，当某个专家处理的 token 个数超过 capacity 之后，概专家的计算就会直接跳过，退化为 residual connection. 后续 DeepSeek-V2 也采用了这种策略，但是 DeepSeek-V3 弃用

Load balancing Loss 在训练目标中加入负载均衡损失，要求每个专家处理的 token 个数的分布尽可能均匀。

这部分具体见 Load Balancing loss

Global Choice

全局分配决定 token 和专家的匹配关系，后续 Qwen 提出了 Global-batch load balancing 使用了这种方式来提高专家的特化程度

Dynamic Routing

根据输入 token 的难度动态决定激活专家的个数。LongCat 使用了一个 Phantom expert 的方法来实现根据 token 的难度动态分配专家。具体来说，除了 $N$ 个专家之外，MoE 还包括 $Z$ 个 zero-computation expert (现在一共有 $N+Z$ 个专家参与计算), 其计算方式如下

$$ \begin{aligned} y &= \sum_{i=1}^{K}\mathrm{Normalize}(G_{e_i}) E_{e_i}(x), e_i \in \{i\mid G_i(x)\in \mathrm{TopK}(G_i(x), K)\}\\ E_{e_i}(x) &= \begin{cases} FFN_{e_i}(x), & \text{ if }1\leq i\leq N\\ x, & \text{ otherwise }\\ \end{cases} \end{aligned} $$

注：LongCat 还是用了 Loss-Free Balancing, 我们这里省略掉了。

Overview

现在几乎所有的模型都选择方式 1，即每个 token 选取 top-k 的专家。 olmoe 对比了以下方式 1 和方式 2 的表现，如下图所示

可以看到，加入 load balancing loss 之后，相比于 Expert Choice, Token Choice 的表现更好。但是，expert choice 更加高效，作者认为 expert choice 更适用于多模态，因为丢掉 noise image tokens 对 text token 影响会比较小。因此，在 olmoe 中，作者使用 token choice 作为 routing 策略

Upcycling

upsampling 是一个将 dense model 转化为 MoEmodel 的方法，具体做法就是我们复制 dense model 中的 FFN layer 得到对应 MoE layer 中的 Expert，然后我们再结合 router 训练，这样可以提高整体的训练效率。相关模型有 MiniCPM, Qwen1.5 和 Mixtral MoE (疑似)

实验结果如下图所示

从已有的结果来看，MoE 模型会被 dense 模型的一些超参数所限制，且训练不是很稳定。因此，现在一般不采用这种方法

Analysis on MoE

Specialization of Experts

OpenMoE 分析了 MoE 模型的特化程度，其结论如下

作者发现，大部分专家对于不同的 domain 没有出现 specialization 情况，对于 math domain, specialization 现象比较明显，作者认为这是因为 math domain 包含更多的 special tokens
对于不同的语言，有部分专家出现 specialization 现象
部分专家对于 position ID 有 specialization 现象，并且连续的 token 更偏好相同的专家
作者还发现，部分专家对于 Token ID 有 specialization 现象，作者将这种现象称为 Context-independent Specialization.
专家还会对语义相似的 token 进行聚类，并且这种聚类在训练早期就已经发生，作者认为其原因在于重新分配 token 会增加最终的 loss
对于 token dropping, 作者发现越靠后的 token, 其被 drop 的概率比例也越高。并且对于指令跟随数据，更多的 token 都会被丢掉，因此作者认为指令跟随数据是 MoE 模型的一种 OOD 数据

Saturation of Experts

olmoe 探究了训练过程中激活的专家和训练结束后激活的专家的匹配程度，结果如下图所示

实验结果说明，训练 $1\%$ 的数据之后，就有 $40\%$ 的 routing 和训练完毕的 routing 一致，当训练 $40\%$ 的数据之后，这个比例提升到了 $80\%$. 作者认为，这是专家特化的结果，初始的 routing 如果改变的话会带来表现下降，因此模型倾向于使用固定的专家处理特定的 token

作者还发现，later layers 比 early layers 饱和更快，early layer, 特别是 layer 0, 饱和的非常慢。作者认为，这是 DeepSeekMoE 放弃在第一层使用 MoE layer 的原因，因为 load balancing loss 收敛更慢。后续 DeepSeek-V2 和 DeepSeek-V3 均在 early layer 上使用 dense layer 替换掉了 MoE layer

Optimization

MoE 模型的优势在于表现好，但是模型参数往往非常大，为了方便使用，我们需要对训练好的 MoE 模型进行优化，目前主要有蒸馏，专家剪枝/合并以及量化等优化方法

蒸馏是一个将大模型能力传递给小模型的做法，目前已有的包括：

Switch Transformer 通过蒸馏，在仅使用 $1/20$ 参数的情况下，保留了稀疏教师模型 $30\%$ 的表现
Gemini2.5 通过蒸馏 Gemini2.5 Pro 得到 Gemini2.5 Flash
DeepSeek-R1 通过蒸馏来提升小语言模型的 reasoning 能力
Qwen3 对于小语言模型的训练使用了 off-policy distillation 和 on-policy distillation 来训练小语言模型

Code

我们这里展示基于 olmoe 的代码，代码如下

 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


class OlmoeSparseMoeBlock(nn.Module):
 def __init__(self, config):
 super().__init__()
 self.num_experts = config.num_experts
 self.top_k = config.num_experts_per_tok
 self.norm_topk_prob = config.norm_topk_prob
 self.gate = nn.Linear(config.hidden_size, self.num_experts, bias=False)
 self.experts = nn.ModuleList([OlmoeMLP(config) for _ in range(self.num_experts)])

 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 batch_size, sequence_length, hidden_dim = hidden_states.shape
 # hidden_states: (batch * sequence_length, hidden_dim)
 hidden_states = hidden_states.view(-1, hidden_dim)
 # router_logits: (batch * sequence_length, n_experts)
 router_logits = self.gate(hidden_states)

 routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
 # routing_weights: (batch * sequence_length, top_k)
 # selected_experts: indices of top_k experts
 routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
 if self.norm_topk_prob:
 routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
 # we cast back to the input dtype
 routing_weights = routing_weights.to(hidden_states.dtype)

 final_hidden_states = torch.zeros(
 (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device
 )

 # One hot encode the selected experts to create an expert mask
 # this will be used to easily index which expert is going to be selected
 expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0)

 # Loop over all available experts in the model and perform the computation on each expert
 for expert_idx in range(self.num_experts):
 expert_layer = self.experts[expert_idx]
 idx, top_x = torch.where(expert_mask[expert_idx])

 # Index the correct hidden states and compute the expert hidden state for
 # the current expert. We need to make sure to multiply the output hidden
 # states by `routing_weights` on the corresponding tokens (top-1 and top-2)
 current_state = hidden_states[None, top_x].reshape(-1, hidden_dim)
 current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None]

 # However `index_add_` only support torch tensors for indexing so we'll use
 # the `top_x` tensor here.
 final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))
 final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
 return final_hidden_states, router_logits

Infra

针对 MoE 模型的 infra 主要涉及 expert parallelism (EP), EP 将 MoE layer 的计算分为了三个阶段：

all-to-all dispatch: 基于 gating layer 的结果，将 token 通信传输到对应专家所在的 GPU 上
computation: 执行计算，即 $E_i(x)$.
all-to-all combine: 收集专家计算的结果，即 $y = \sum_{i=1}^K\mathrm{Normalize}(G_{e_i}(x)) E_{e_i}(x)$.

其框架图如下所示

Challenges

构建针对 MoE 模型的 infra 比较困难，相关工作有 DeepSeek 提出的 DeepEP.
训练不稳定，特别是负载均衡。负载均衡做不好容易导致模型崩塌。

Conclusion

在本文中，我们系统性回顾了 MoE 的相关概念，MoE 模型已经是现在大语言模型的主流架构，比如商业模型 Gemini2.5, 开源领先的模型 DeepSeek-V3 , LLaMA4 以及 Qwen3 等都采用了 MoE 的架构，如何进一步优化 MoE 的训练方式是当前研究的一个重点方向。

Appendix

MoE model information

Model	Year	total parameters	activated parameters	shared expert	Routed experts	activated experts
ST-MoE	2022/4	269B	32B	0	64	2
Mistral	2024/1	47B	13B	0	8	2
DeepSeek-MoE	2024/1	145B	22B	4	128	12
DeepSeek-V2	2024/5	236B	21B	2	160	6
LLaMA4	2025/4	400B	17B	1	128	1
DeepSeek-V3	2024/12	671B	37B	1	256	8
Qwen3	2025/5	235B	22B	0	128	8
dots.llm1	2025/6	142B	14B	2	128	6
Step 3	2025/7	316B	38B	1	48	3
Kimi-K2	2025/7	1043B	32B	8	384	1
GLM-4.5	2025/8	355B	32B	1	160	8
gpt-oss	2025/8	116B	5B	0	128	4
LongCat	2025/9	560B	27B*	0	768*	12
Ling-1T	2025/10	1000B	51B	1	256	8

Remark LongCat 采用了动态激活的方式，因此其结果有一个浮动范围。

References

Notes on Ling-mini-beta

Sat, 13 Dec 2025 15:58:51 +0800

Introduction

目前已经有了针对 dense LLM 的 scaling law, 如 Kaplan scaling law 和 Chinchilla scaling law.

但是，对于 MoE 模型，目前还缺乏一个比较系统的 scaling law.

为了解决这个问题，作者提出了 efficiency leverage (EL), 用于衡量 MoE 模型的效率，其定义为

$$ EL(\mathcal{X}_{\mathrm{MoE}}\mid \mathrm{Dense})=\frac{C_{\mathrm{Dense}}}{C_{{\mathrm{MoE}}}} $$

其中 $C_{\mathrm{Dense}}, C_{{\mathrm{MoE}}}$ 分别代表了训练模型所需要的算力。EL 衡量了 moe 模型达到对应 dense 模型表现所需要的算力，EL 值越大，说明 MoE 模型效率越高。EL 的可视化如下图所示

作者通过训练多个模型，探究了 MoE 架构与 EL 之间的关系。作者发现，MoE 模型的表现主要与专家激活比例以及算力相关。基于 scaling law, 作者训练了 Ling-mini-beta, 一个 17.5B-A0.85B 的 MoE 模型，其表现超过了 6.1B dense 模型的表现。

Preliminary

Notation	description
$N$	total parameters
$N_a$	active parameters
$E$	routed experts
$E_a$	activated experts
$E_s$	shared experts

作者定义 activation ratio 如下:

$$ A = \frac{E_a+E_s}{E+E_s} $$

定义 sharing ratio 如下

$$ S=\frac{E_s}{E_a+E_s} $$

定义 expert granularity 如下

$$ G = \frac{d_{\mathrm{model}}}{d_{\mathrm{Expert}}} $$

与 DeepSeek-LLM 一样，作者使用 $C=MD$ 来表示算力，non-embedding FLOPs $M$ 和训练 token 数 $D$ 之间的关系。

Hyper Parameters

作者首先探究了针对 MoE 模型的超参数配置，最终你和出来的结果如下图

实验结果说明，相比于 dense model, MoE model 需要更大的 batch size.

作者基于这个 scaling law 进行了验证，结果说明这个 scaling law 比较准确。

Parameters and Dataset Size

接下来作者探究了对于模型参数量以及训练 token 个数之间的 scaling law, 求解的问题如下

$$ (M^{opt}, D^{opt}) = \arg\min_{M,D}\mathcal{L}(M,D;C,A,G,S)\quad s.t.\ C=MD $$

实验结果如下图所示

结果说明，不同架构对应的系数接近 $0.5$, 说明我们应该将算力均衡分配到数据和 model size 上。领一方面，MoE 模型可以通过使用更多的数据来达到更优的表现。

Efficiency Leverage

作者将 Efficiency Leverage (EL) 定义为给定算力 $C_{target}$ 和一个 MoE 模型 $\mathcal{X}_{MoE}$, 对应 dense 模型达到 $\mathcal{X}_{MoE}$ 相同的表现所需要的算力 $C_{dense}$, 即

$$ \begin{aligned} &EL(\mathcal{X}_{\mathrm{MoE}}\mid \mathrm{Dense};C_{target})=\frac{C_{\mathrm{Dense}}}{C_{{\mathrm{MoE}}}}\\ s.t.\ & |\mathcal{L}(C_{MoE}, \mathcal{X}_{\mathrm{MoE}})-\mathcal{L}(C_{dense}, \mathcal{X}_{\mathrm{dense}})|\leq \epsilon (\epsilon\to 0) \end{aligned} $$

EL 值越高说明 MoE 模型越有效。为了公平起见，dense 模型和 MoE 模型的架构参数基本相同，作者只改变 $d_{model}, d_{ffn}, d_{expert}$ 以及 $n_{layer}$.

接下来，作者就探究了给定算力的情况下，最优的 MoE 配置，即

$$ (A^{opt}, G^{opt}, S^{opt}) = \arg\min_{(A,G,S)\in\mathcal{X}_{\mathrm{MoE}}}EL(\mathcal{X}_{\mathrm{MoE}}\mid \mathrm{Dense};C) $$

Scaling Law

首先，坐着探究了最优的 activation ratio $A$, 即

$$ A^{opt} = \arg\min_{A}\mathcal{L}(A;C,M,G,S) $$

拟合的结果如下图所示

实验结果表明：

模型表现随激活比例降低 er 提高
更系数的模型对于算力的提升其效率也提升更快

然后，作者探究了最优的 granularity ratio, 即

$$ G^{opt} = \arg\min_{G}\mathcal{L}(G;C,M,A,S) $$

拟合的结果如下图所示

可以看到，无限制提升 granularity 并不会提高模型的表现。并且，不同的算力对应的最优 granularity 处于一个固定的范围

接下来，作者探究了最优的 shared expert ratio, 即

$$ S^{opt} = \arg\min_{S}\mathcal{L}(S;C,M,A,G) $$

拟合的结果如下图所示

结果说明 shared expert 的比例也不是越多越好，其存在最优值。并且给定算力的情况下，非零最小值的 shared expert 表现最好。因此作者认为，一个 shared expert 的效果最好。

作者还探究了其他可能的因素，结论如下：

与 DeepSeek-V3 一样，将 early layer 替换为 dense layer 可以避免 routing imbalance, 并且不会损害模型的表现
attention 应该占 $30\%\sim40\%$ 左右的算力才能保证模型的表现和效率。进一步提升 attention 的算力占比虽然会提升表现但是会降低推理效率。

通过前面的发现，作者将 shared expert 设置为 1 个，然后探究 EL 与 activation ratio $A$, granularity $G$, FLOPs $C$ 之间的关系。

首先，作者分别假设 $EL$ 与 $A$, $G$, $C$ 之间存在如下关系：

$$ \begin{aligned} \log EL_{C,G}(\hat{A}) &= a_A\log\hat{A}, \text{ where }\frac{1}{\hat{A}}=\frac{1}{A+(1/A_{start}-1/A_{\max})^{-1}}+\frac{1}{A_{\max}}\\ \log EL_{C,A}(\hat{G}) &= a_G+b_G(\log G(\log G+c_G))\\ \log EL_{A,G}(C) &= a_C\log C+c_C \end{aligned} $$

拟合的结果如下图所示

结果显示：

提升算力以及降低 activation ratio 都可以提高 EL
granularity 对 EL 的影响在不同算力的情况下都是一致的
对于 MoE 模型，提升算力可以提高 EL

作者的结论为，activation ratio 是影响 MoE EL 的核心因素。并且随着算力的提升，MoE EL 会越来越明显。

作者因此构建了一个统一的公式来统一三个因素

$$ EL(A,G,C) = \hat{A}^{\alpha+\gamma(\log G)^2+\beta \log G} $$

其中 $\alpha=a+d\log C$ 代表了 EL 和 activation ratio 之间的关系。拟合出来的参数如下表所示

$\alpha$	$d$	$\gamma$	$\beta$	$A_{start}$	$A_{\max}$
1.23	-7.61e-2	1.67e-2	-1.17e-1	1.63e-2	5.28e+16

基于这个结果，作者发现在 1e22 FLOPs 的算力下，一个 activation ratio 为 $3.1\%$, granularity 为 $12$ 的 MoE 模型，其 EL 为 $7$.

Ling-mini-beta

基于上一节的发现，作者构建了 Ling-mini-beta, 一个 17.5B 总参数，激活参数为 0.85B 的 MoE 模型。训练使用了 1T token, 模型参数如下表所示

Model	$n_{\text{layers}}$	$d_{\text{model}}$	$d_{\text{ffn}}$	$d_{\text{expert}}$	$n_{\text{heads}}$	$n_{\text{kv\_head}}$	$E$	$E_a$	$E_s$	$N$	$N_a$
Dense 6.1B	28	4096	14336	-	32	8	-	-	-	6.11B	6.11B
Ling-mini-beta (A0.8B)	20	2048	5120	384	16	4	384	12	1	17.5B	0.85B

训练的损失变化情况如下图所示

从图中我们可以看出，dense model 一开始的损失下降比较快，但是其最终表现不如 moe 模型。

Conclusion

在本文中，作者提出了 Efficiency Leverage, 一个衡量 MoE 模型相对于 dense 模型计算效率的 metric, 作者构建了针对 MoE 模型的 scaling law. scaling 揭示了两个主要影响 MoE 模型效率的因素：算力与激活参数比例。基于 scaling law, 作者构建了 Ling-mini-beta, 一个 17B-A0.8B 的 MoE 模型，其效率超过了对应 dense 模型的 7 倍。

References

arxiv

Load Balancing tutorial

Thu, 11 Dec 2025 16:10:08 +0800

我们在本文中探讨关于 load balancing loss 的定义，性质和推广

Introduction

我们在 MoE tutorial 中已经介绍了 MoE 模块，MoE 尽管可以在相同的算力下扩大模型的 size, 但是其问题在于训练时容易出现负载不均衡，也就是说只有少数几个专家被激活，其他专家处于闲置状态，从而导致模型性能下降

为了解决这个问题，一个通用的做法是使用 load balancing loss. load balancing loss 可以有效实现负载均衡，让各个专家被激活的概率差不多。

Load balancing loss 一共经历了如下几个阶段，发展路线如下图所示

1
2
3
4
5
6


flowchart TD;
 C[Bengio et.at 2015] --> A[Shazeer et al., 2017];
 A --> B[GShard];
 B --> F[Switch Transformer];
 F --> D[Loss-Free Balancing];
 F --> E[Global-batch load balancing];

Method

Notation

我们假设输入的 batch size 为 $B$, 总专家个数为 $N$, 激活专家个数为 $K$, $G_i(x_j)$ 代表 gating layer 预测的第 $i$ 个专家对于 token $x_j$ 的重要性程度。

Coefficient of Variation

第一个阶段的 load balancing loss 针对广义的 MoE 模型。

Bengio et.at 2015 作者提出了一个让 batch 里每个专家平均激活概率分布更均匀的损失函数，其定义如下所示

$$ L_b = \sum_{i=1}^N\left\Vert \frac1B\sum_{j=1}^BG_i(x_j) - \frac1N\right\Vert_2 $$

其主要目标是让每个专家的平均激活概率接近 $1/N$, 即均匀分布.

Noam et.al 2017 作者对 (Bengio et.at 2015) 进行了改进，提出了变异系数 $\mathrm{CV}(\cdot)$ 来保证负载均衡，其定义如下：

$$ L_{importance} = \mathrm{CV}(\mathrm{Importance}(X))^2 $$

其中

$$ \mathrm{Importance}_i(X) = \sum_{j=1}^B G_i(x_j),\ \mathrm{CV}(X)= \frac{\mathrm{var}[X]}{\mathbb{E}[X]},\ i=1,\dots,N $$

注意到 $\mathrm{var}[X]\geq0$, 当且仅当 $X$ 为均匀分布时 $\mathrm{var}[X]=0$, 因此 important loss 可以让每个专家在一个 batch 中的激活的平均概率尽可能一致。

但是平均概率一致不代表各个专家处理的 token 个数一致，比如一个专家可能处理比较少的 token, 但是每个 token 的权重都很大。因此，作者额外加入了 load balancing loss

$$ \mathcal{L}_{\mathrm{load}} = \mathrm{CV}(\mathrm{Load}(X))^2 $$

其中

$$ \begin{aligned} \mathrm{Load}_i(X) &= \mathrm{soft\_estimation}\left(\sum_{j=1}^B\mathbb{1}(\text{token }j \text{ selects expert }i)\right)\\ &=\mathrm{soft\_estimation}\left(\sum_{j=1}^B\mathbb{1}\{i\in\mathrm{TopK}(\{G_i(x_j)\}_{i=1}^{E}, K)\}\right),\ i=1,\dots,N \end{aligned} $$

这里 $\mathrm{soft\_estimation}$ 是作者针对离散变量进行的一个光滑化处理，以方便反向传播。

GShard

GShard 进一步对 (Noam et.al 2017) 提出的 loss 进行了简化，首先注意到 $\sum_{i=1}^N\mathrm{Load}_i(X)=BK$, 因此我们有

$$ \begin{aligned} \mathcal{L}_{\mathrm{load}} &= \mathrm{CV}(\mathrm{Load}(X))^2 \\ &= \left(\frac{\mathrm{var}[\mathrm{Load}(X)]}{\mathbb{E}[\mathrm{Load}(X)]}\right)^2\\ &= \left(C_1[\mathrm{Load}(X)]\right)^2\\ &= C_1\mathbb{E}[\mathrm{Load}(X)^2]-C_2 \end{aligned} $$

这里 $C_1,C_2$ 均为常数。GShard 并没有使用 soft estimation, 因此这里的 $\mathrm{Load}_i(X)$ 定义为

$$ \mathrm{Load}_i(X)=\sum_{j=1}^B\mathbb{1}(\text{token }j \text{ selects expert }i),\ i=1,\dots,N $$

其代表了一个 batch 里专家 $i$ 被激活的次数，理想情况下，所有专家被激活的次数应该是一致的，从而我们就实现了负载均衡。

但是现在的问题是，$\mathrm{Load}_i(X)$ 是一个离散变量，为了解决这个问题，作者使用了重要性来进行近似，即

$$ \begin{aligned} \mathcal{L}_{\mathrm{load}} &= \mathbb{E}[\mathrm{Load}(X)^2]\\ &\approx \mathbb{E}[\mathrm{Load}(X)\cdot\mathrm{Importance}(X)] \end{aligned} $$

Switch Transformer

Switch Transformer 进一步对 GShard 提出 load balancing loss 进行了规范化，也是现在大部分 load balancing loss 使用的形式，其定义如下：

$$ \mathcal{L}_{\mathrm{load}} = \sum_{i=1}^N f_i P_i $$

其中，

$$ f_i = \frac{1}{B}\mathrm{Load}_i(X)=\frac{1}{B}\sum_{j=1}^B\mathbb{1}(\text{token }j \text{ selects expert }i),\ i=1,\dots,N $$

代表了分配给专家 $i$ 的 token 比例，

$$ P_i = \frac{1}{B}\mathrm{Importance}_i(X) = \frac{1}{B}\sum_{j=1}^B G_i(x_j) $$

代表了这个 batch 里专家 $i$ 的平均激活概率。

从而，这个形式和 GShard 的形式实际上是等价的：

$$ \mathcal{L}_{\mathrm{load}} = \sum_{i=1}^N f_i P_i = \mathbb{E}[\mathbf{f}\cdot \mathbf{P}] = N \mathbb{E}[\mathrm{Load}(X)\cdot\mathrm{Importance}(X)] $$

Analysis

目前关于 load balancing loss 可以实现负载均衡的数学分析是比较少的。Switch Transformer 中进行了如下分析：

[!quote] The auxiliary loss of Equation 4 encourages uniform routing since it is minimized under a uniform distribution.

但实际上，这个分析是错误的，Switch transformer 中定义的 load balancing loss 并不是在均匀分布时取最小值，而是取最大值。考虑如下情形，

$$ \mathbf{f}=[1,\dots,0],\ \mathbf{P}=[0,\dots,1] $$

此时，我们有

$$ \mathbf{f}\cdot \mathbf{P} = 0\leq [1/N,\dots,1/N]\cdot [1/N,\dots,1/N] = 1/N. $$

我找了很多资料，最后发现苏剑林老师在 MoE环游记：2、不患寡而患不均这篇 blog 中进行了分析。我这里基于苏剑林老师的 blog 进行了一下总结，推荐大家看苏剑林老师的原文。

分析的核心思想就是，尽管 load balancing loss 目标函数没有意义，但是我们通过 straight-through estimator (STE)，就可以得到一个有意义的目标函数，其梯度与 load balancing loss 目标函数的梯度是一致的，这样 load balancing loss 目标函数就完成了其负载均衡的目标。

注意到 load balancing loss 的最终目标是让每个专家处理的 token 个数尽可能一致，也就是

$$ \mathcal{L}_{\mathrm{load}} = \sum_{j=1}^B \left(f_i-\frac1B\right)^2 $$

但是 $f_i=\mathrm{Load}_i(X)$ 是一个离散变量不可导，因此，基于 STE, 我们可以将其改写为

$$ \mathcal{L}_{\mathrm{load}} = \sum_{j=1}^B \left(P_i+\mathrm{sg}[f_i-P_i]-\frac1B\right)^2 $$

其中 $\mathrm{sg}[\cdot]$ 满足：

$$ \mathrm{sg}[\cdot]=\begin{cases} x, &\text{forward pass}\\ 0, &\text{backward pass} \end{cases} $$

此时我们目标函数的梯度就变成了

$$ \begin{aligned} \nabla_\theta\mathcal{L}_{\mathrm{load}} &= \sum_{j=1}^B 2\nabla_\theta\left(P_i+\mathrm{sg}[f_i-P_i]-\frac1B\right)\left(P_i+\mathrm{sg}[f_i-P_i]-\frac1B\right)\\ &= 2\sum_{j=1}^B\nabla_\theta P_i\left(f_i-\frac1B\right)\\ &= 2\sum_{j=1}^B\nabla_\theta P_i\left(f_i-\frac1B\right)\\ &= 2\nabla_\theta \left(\sum_{j=1}^Bf_iP_i\right) \end{aligned} $$

这样就对应上了我们之前的目标函数。总之，目标函数尽管本身没有意义，但是其对应的梯度却可以让模型实现负载均衡的目标。

Extension

Loss Free Load Balancing Strategy

Loss-Free load balancing 是 DeepSeek 提出来的一个无需 load balancing loss 实现负载均衡的方法，其核心思想是 load balancing loss 会影响语言建模性能，因此在 $\mathrm{TopK}$ 操作时，作者加入一些扰动，从而让负载高的专家被选择的概率降低，让负载低的专家被选择的概率提高。

具体细节见 Loss-Free Balancing.

Global-batch Load Balancing Loss Strategy

Global-batch Load Balancing 是 Qwen 提出来的一个提高负载均衡的方法。其核心思想是，现有的 token choice 只是在 batch 层面达到负载均衡，这种局部均衡的性质可能会对模型的全局均衡性产生影响。因此，作者就预先在 global batch 层面对专家进行分配，从而提高专家在不同 domain 上的特化程度。

具体细节见 Global-batch load balancing

Experiments

olmoe 对比了加入 load balancing loss 之后模型的表现变化情况，结果如下图所示

可以看到，加入 load balancing loss 之后，模型的表现均超过了不加时的表现。

作者进一步分析了不同专家在加/不加 load balancing loss 时的激活情况，结果如下图所示

结果显示，load balancing loss 确实可以让不同专家的激活概率大致相当。

Conclusion

在本文中，我们介绍了 load balancing loss 的演化，定义，分析以及相关的实验结果。

Reference

Notes on Global-batch load balancing

Thu, 11 Dec 2025 16:09:34 +0800

Qwen 在 25 年 2 月提出了 global batching load balancing loss strategy, 其在 global level 上考虑每个专家的负载均衡，从而提高模型的表现

Introduction

实现 MoE 模型负载均衡的原因有两点：

Effectiveness: 通过负载均衡才能更高效地利用各个专家
Efficiency: 一般需要使用 expert parallel 来部署 MoE 模型，不均衡的负载会大幅度降低前向过程

已有的框架如 DeepSpeed, Megablocks 和 Megatron-Core 都是在 micro-batch level 上计算负载均衡损失的，但是，一个 micro-batch 通常只包含少数序列，因此 load balancing loss 就要求各个专家在每个序列上均匀分布。

针对这个问题，作者在本文中提出的解决方法是在 global-batch 层面考虑负载均衡，

Method

MoE 模型定义如下

$$ y = \sum_{i\in N_E,g_i\in\mathrm{TopK}(g)}g_i(x)E_i(x) $$

其中 $E_i$ 是对应的专家, $g_i$ 是对应的权重

Load balancing loss 定义如下

$$ \mathrm{LBL} = N_E\sum_{i=1}^{N_E}f_iP_i $$

一般来说，MoE 模型训练时会使用 expert parallel 策略，此时 load balancing loss 修改为

$$ \mathrm{LBL}_{\mathrm{micro}}= \frac{1}{N_P}\sum_{j=1}^{N_P}\left(N_E\sum_{i=1}^{N_E}f_i^hP_i^j\right) $$

其中 $N_P$ 是 parallel groups 的个数。在这种情况下，模型需要再每个 parallel group 中实现负载均衡。但是某一个 mircro-batch 里可能只包含某一个 domain 的序列，因此各个专家被强制要求在每一个 domain 中也均衡分布。作者认为，这种方式会限制专家的能力，进而影响模型的表现。

作者的解决方法在于，获取每个 parallel group 的 $f_i$ 然后求 global-batch 的 $\bar{f}_i$.

$$ \mathrm{LBL}_{\mathrm{global}}=N_E\sum_{i=1}^{N_E}\bar{f}_i\bar{P}_i=N_E\sum_{i=1}^{N_E}\bar{f}_i\left(\frac{1}{N_P}\sum_{j=1}^{N_P}P_j\right) = \frac{1}{N_P}\sum_{j=1}^{N_P}\left(N_E\sum_{i=1}^{N_E}\bar{f}_iP_i^j\right) $$

实际中，由于计算节点数量有限，micro-batch size 之和可能会小于 global-batch size, 因此我们会使用 gradient accumulation. 在这种情况下，作者使用了一个 buffer 来保存多个 micro-batch 的专家选择次数，最终算法实现过程如下所示

Experiments

作者使用了不同大小的模型来进行实验，作者用 Balance BSZ 来表现计算 expert selection frequency 时的 token 数，结果如下图所示

可以看到，global LBL 可以提高模型的表现，并且，当 Balance BSZ 增加时，模型的表现也会提升。作者发现对于 Loss-Free Balancing 来说，使用 global batch 的效果也更好。

作者还发现，global LBL 会提高专家的特化程度，结果如下图所示

可以看到，使用默认的 LBL loss, 各个专家的特化程度很低，而使用本文的 global LBL 之后，专家的特化程度有了大幅度的提高。

作者还进一步发现，随着 Balance BSZ 的提高，模型表现也持续提升，结果如下图所示

作者认为，synchronization 和 buffer 机制相比于 micro-batch 来说可以带来大幅度提升。

Analysis

首先，为了分析 global batch LBL 优于 Micro batch 的关键原因，作者先同步所有 group 的矩阵 $G$, 然后再从全局 token 中随机选一批，用这批 token 计算专家选择频率。通过这个方法，我们可以保证 token 数量与 micro-batch 一致，但是分布于 Global-batch 一致。实验结果发现，这种方法的表现优于 micro-batch LBL, 说明 global batch LBL 的优势在于 Token 分布更全局，而不是 token 数量更多。

作者还分析了 global batch LBL 和 micro batch LBL 两种模式，作者认为前者是后者的一个宽松版约束，作者发现从后者切换为前者之后，模型的表现可以得到进一步提升，但是其表现仍然不如一开始就使用 global batch LBL 更好。作者分析原因认为，这是因为 expert 收敛速度比较快。

作者进一步通过降低 micro batch LBL 权重来探究是否可以达到同样的表现，结果发现适度江都权重确实可以提高模型的表现，但是降低太多会损害模型的表现，即此时出现了负载不均衡的现象

作者对比了 global batch LBL 和 micro batch LBL 效率发现，前者比后者慢 $2\%$ 左右，这个差距几乎可以忽略不计

作者进一步分析了不同 balancing 方式的专家特化程度，结果如下图所示

实验结果发现：

使用 global batch LBL 之后，topK sum 明显变大，这说明 routing 和 language modeling 任务对齐地更好
global batch LBL 的专家在不同任务上的特化程度更明显
micro batch LBL 的 topK sum 比较小
Loss-free balancing 的 topK sum 介于 micro batch LBL 和 global batch LBL 之间

Conclusion

作者在本文中提出了 global LBL 来在全局为负载均衡提供指导，结果发现通过在更大的范围进行负载均衡的计算，我们可以有效提高专家的特化程度以及提高模型在下游任务上的表现

References

Arxiv

Notes on DPO

Tue, 09 Dec 2025 10:43:11 +0800

Introduction

传统的偏好优化主要基于 RLHF, 其过程为：SFT, reward modeling, RLHF. 其中 reward model 的训练至关重要，对最终模型的表现有非常大的影响。但是 RLHF 的问题在于其训练复杂且经常不稳定。

为了解决这个问题，作者提出了 Direct Preference Optimization (DPO), DPO 通过构建 reward function 和最优策略之间的关系，进而通过训练 policy model 来同时完成 reward model 的训练。这样，我们就避免了 reward model 的训练。结果发现，DPO 的表现超过了之前的偏好优化方法。

Background

作者首先回顾了 RLHF, RLHF 的 pipeline 如下

其包含了三个步骤：

SFT RLHF 首先基于 base model 通过 SFT 得到一个初始模型 $\pi^{\mathrm{SFT}}$.

Reward modeling 接下来，我们给定输入 $x$, 对 $\pi^{\mathrm{SFT}}$ 采样得到 $(y_1,y_2)\sim \pi^{\mathrm{SFT}}(y\mid x)$. 输出 $y_1, y_2$ 然后由人类进行打分得到偏好关系 $y_wr^*(x, y_l)$. reward modeling 通常基于 Bradley-Terry (BT) 模型得到，BT model 定义人类真实偏好分布 $p^*$ 如下：

$$ p^*(y_1>y_2\mid x)= \frac{\exp(r^*(x, y_1))}{\exp(r^*(x, y_1))+\exp(r^*(x, y_2))} $$

假设我们从分布 $p^*$ 中采集到一个数据集 $\mathcal{D}=\{(x^{(i)},y_w^{(i)},y_l^{(i)})\}_{i=1}^N$, 我们可以通过 MLE 来估计得到一个 reward model $r_{\phi}(x,y)$, 通过将这个问题转换为一个二分类问题，我们得到对应的 negative log-likelihood loss 如下：

$$ \mathcal{L}_R(r_\phi, D) = -\mathbb{E}_{(x,y_w,y_l)\sim \mathcal{D}}\left[\log \sigma\left(r_\phi(x,y_w)-r_\phi(x,y_l)\right)\right] $$

其中 $\sigma$ 是 logistic function. 在 LLM 中，$r_\phi(x,y)$ 通常由 $\pi^{\mathrm{SFT}}$ 初始化得到，然后我们在 $\pi^{\mathrm{SFT}}$ 最后一层加入一个 linear layer 来得到对应的 reward 的预测值。一般地，为了降低 reward function 的 variance, 之前的工作会进行 normalization, 即 $\mathbb{E}_{(x,y)}\sim\mathcal{D}[r_{\phi}(x,y)]=0$ for all $x$.

RL fine-tuning 这个阶段，我们基于学习到的 reward function $r_\phi(x,y)$ 来为 LLM 的训练提供奖励，作者使用了和 RLHF 一样的目标函数：

$$ \max_{\pi_\theta}\quad \mathbb{E}_{x\sim \mathcal{D}, y\sim \pi_\theta(y\mid x)}\left[r_\phi(x,y)\right] - \beta\mathbb{D}_{\mathrm{KL}}\left[\pi_\theta(y\mid x)\Vert \pi_{\mathrm{ref}}(y\mid x)\right] $$

这里 $\beta>0$ 是超参数，用于控制 $\pi_\theta$ 相对于 $\pi_{\mathrm{ref}}$ 的偏离程度， $\pi_{\mathrm{ref}}$ 一般就是 $\pi^{\mathrm{SFT}}$. 实际上， $\pi_\theta$ 也由 $\pi^{\mathrm{SFT}}$ 初始化.

DPO

DPO 的主要目标是构建一个更简单的 policy optimization 方法。与 RLHF 不同，DPO 跳过了 reward modeling 这一阶段，而是直接使用偏好数据来优化大语言模型

为了实现这个目标，作者第一步就是构建 reward model 和 policy model 之间的关系。注意到

$$ \begin{aligned} &\max_{\pi_\theta}\ \mathbb{E}_{x\sim \mathcal{D}, y\sim \pi_\theta(y\mid x)}\left[r_\phi(x,y)\right] - \beta\mathbb{D}_{\mathrm{KL}}\left[\pi_\theta(y\mid x)\Vert \pi_{\mathrm{ref}}(y\mid x)\right]\\ &= \max_{\pi_\theta}\ \mathbb{E}_{x\sim \mathcal{D}}\mathbb{E}_{y\sim \pi_\theta(y\mid x)}\left[r_\phi(x,y) - \beta\log\frac{\pi_\theta(y\mid x)}{\pi_{\mathrm{ref}}(y\mid x)}\right]\\ &= \min_{\pi_\theta}\ \mathbb{E}_{x\sim \mathcal{D}}\mathbb{E}_{y\sim \pi_\theta(y\mid x)}\left[\log\frac{\pi_\theta(y\mid x)}{\pi_{\mathrm{ref}}(y\mid x)}-\frac{1}{\beta}r_\phi(x,y)\right]\\ &= \min_{\pi_\theta}\ \mathbb{E}_{x\sim \mathcal{D}}\mathbb{E}_{y\sim \pi_\theta(y\mid x)}\left[\log\frac{\pi_\theta(y\mid x)}{\frac{1}{Z(x)}\pi_{\mathrm{ref}}(y\mid x)\exp\left(\frac1\beta r_\phi(x,y)\right)}-\log Z(x)\right]\\ \end{aligned} $$

其中 $Z(x)$ 是 partition function, 定义如下

$$ Z(x) = \sum_{y} \pi_{\mathrm{ref}}(y\mid x)\exp\left(\frac1\beta r_\phi(x,y)\right) $$

注意到 partition function 只是 $x$ 和 $\pi_{\mathrm{ref}}$ 的函数，而不依赖于 $\pi_\theta$, 因此我们定义

$$ \pi^*(y\mid x) = \frac{1}{Z(x)}\pi_{\mathrm{ref}}(y\mid x)\exp\left(\frac1\beta r_\phi(x,y)\right) $$

易知 $\pi^*$ 满足 $\pi^*(y\mid x)\geq0, \forall y$, 以及 $\sum_y \pi^*(y\mid x)=1$. 因为 $Z(x)$ 不依赖于 $y$, 因此我们可以进一步简化上面的目标函数如下

$$ \begin{aligned} &\min_{\pi_\theta}\ \mathbb{E}_{x\sim \mathcal{D}}\left[\mathbb{E}_{y\sim \pi_\theta(y\mid x)}\left[\log\frac{\pi_\theta(y\mid x)}{\pi^*(y\mid x)}\right]-\log Z(x)\right] \\ &= \min_{\pi_\theta}\ \mathbb{E}_{x\sim \mathcal{D}}\left[\mathbb{D}_{\mathrm{KL}}\left[\pi_\theta(y\mid x)\Vert \pi^*(y\mid x)\right]-\log Z(x)\right]\\ &= \boxed{\min_{\pi_\theta}\ \mathbb{E}_{x\sim \mathcal{D}}\left[\mathbb{D}_{\mathrm{KL}}\left[\pi_\theta(y\mid x)\Vert \pi^*(y\mid x)\right]\right]} \end{aligned} $$

而上述目标函数的最小值为 0，当且仅当 $\pi_\theta=\pi^*$, 此时我们的最优 policy 为

$$ \pi_\theta^*(y\mid x) = \pi^*(y\mid x) = \frac{1}{Z(x)}\pi_{\mathrm{ref}}(y\mid x)\exp\left(\frac1\beta r_\phi(x,y)\right),\ \forall x\in\mathcal{D} $$

直接求解 $\pi_\theta^*$ 非常困难，因为这涉及到 $Z(x)$ 的计算，这个时候作者就提出了一个关键改变，即我们从上述的 $\pi_\theta^*$ 反向推到出 $r(x,y)$:

$$ r(x,y) = \beta\log\frac{\pi_\theta(y\mid x)}{\pi_{\mathrm{ref}}(y\mid x)}+\beta\log Z(x) $$

我们可以基于这个公式来推到出最优的 reward function $r^*$ 以及对应的最优策略 $\pi^*$.

此时，我们的表达式里仍然含有 $Z(x)$, 但是当我们使用 Bradley-Terry 模型之后，我们就可以得到真实的人类偏好分布，计算过程如下所示

$$ \begin{aligned} p^*(y_w>y_l\mid x)&= \sigma\left(r_\phi(x,y_w)-r_\phi(x,y_l)\right)\\ &= \boxed{\sigma\left(\beta\log\frac{\pi_\theta(y_w\mid x)}{\pi_{\mathrm{ref}}(y_w\mid x)}-\beta\log\frac{\pi_\theta(y_l\mid x)}{\pi_{\mathrm{ref}}(y_l\mid x)}\right)} \end{aligned} $$

这样，基于 MLE 的目标函数就是

$$ \mathcal{L}_{\mathrm{DPO}}(\pi_{\theta};\pi_{\mathrm{ref}})= -\mathbb{E}_{(x,y_w,y_l)\in\mathcal{D}}\left[\sigma\left(\beta\log\frac{\pi_\theta(y_w\mid x)}{\pi_{\mathrm{ref}}(y_w\mid x)}-\beta\log\frac{\pi_\theta(y_l\mid x)}{\pi_{\mathrm{ref}}(y_l\mid x)}\right)\right] $$

通过这种方式，我们就避免了 reward model 的训练。

接下来，作者分析了一下 DPO 目标函数的梯度，

$$ \begin{aligned} \nabla_\theta \mathcal{L}_{\mathrm{DPO}}(\pi_{\theta};\pi_{\mathrm{ref}})&=-\nabla_\theta \mathbb{E}_{(x,y_w,y_l)\in\mathcal{D}}\left[\log\sigma(u)\right]\\ &= -\mathbb{E}_{(x,y_w,y_l)\in\mathcal{D}}\left[\frac{\nabla_\theta \sigma(u)}{\sigma(u)}\nabla_\theta u\right]\\ &= -\mathbb{E}_{(x,y_w,y_l)\in\mathcal{D}}\left[\frac{\sigma(u)(1-\sigma(u))}{\sigma(u)}\nabla_\theta u\right]\\ &= -\mathbb{E}_{(x,y_w,y_l)\in\mathcal{D}}\left[(1-\sigma(u))\nabla_\theta u\right]\\ &= -\mathbb{E}_{(x,y_w,y_l)\in\mathcal{D}}\left[\sigma(-u)\nabla_\theta u\right]\\ &= -\mathbb{E}_{(x,y_w,y_l)\in\mathcal{D}}\left[\beta\sigma\left[\beta\log\frac{\pi_\theta(y_l\mid x)}{\pi_{\mathrm{ref}}(y_l\mid x)}-\beta\log\frac{\pi_\theta(y_w\mid x)}{\pi_{\mathrm{ref}}(y_w\mid x)}\right]\left[\nabla_\theta\log\pi(y_w\mid x) - \nabla_\theta \log\pi(y_l\mid x)\right]\right]\\ \end{aligned} $$

从梯度来看，当我们的 reward 估计错误时，即 $\sigma(-u)>0$ 时， DPO 会提高 $y_w$ 的生成可能性以及降低 $y_l$ 的生成可能性，从而提高模型对于偏好输出的可能性。

最终，DPO 的 pipeline 如下：

收集偏好数据 $y_1,y_2\sim\pi_{\mathrm{ref}}(\cdot\mid x)$, 然后通过人类标注得到偏好数据集 $\mathcal{D}=\{(x^{(i)},y_w^{(i)},y_l^{(i)})\}_{i=1}^N$
基于 $\mathcal{L}_{\mathrm{DPO}}(\pi_{\theta};\pi_{\mathrm{ref}})$ 优化大语言模型参数

一般来说，我们将 $\pi_{\mathrm{ref}}$ 初始化为 $\pi^{\mathrm{SFT}}$, 但是如果我们没有 $\pi^{\mathrm{SFT}}$ 时，我们可以通过最大似然估计来得到 $\pi_{\mathrm{ref}}$, 即

$$ \pi_{\mathrm{ref}} = \arg\max_{\pi}\mathbb{E}_{(x,y_w,y_l)\sim\mathcal{D}}\log \pi(y_w\mid x) $$

Theoretical Analysis of DPO

Definition 两个 reward function $r(x,y)$ 和 $r'(x,y)$ 等价当且仅当存在函数 $f$ 满足 $r(x,y)-r'(x,y)=f(x)$。

上述定义给出了一个等价关系，将 reward function 分割成了不同的等价类。接下来，作者给出了两个引理：

第一个引理说明同一个等价类里的 reward function 对应的偏好分布一致

Lemma 1 在 Plack-Luce 框架下，如 Bradley-Terry model, 同一个等价类里的 reward function 得到的偏好分布是一致的

第二个引理说明了最优策略对应的 reward function 都在一个等价类里

Lemma 2 在有限制的情况下，同一个等价类里的 reward function 得到的最优策略是一样的

Lemma 2 说明我们只要从最优的等价类里任意找到一个 reward function, 则最终的效果是一样的。

Theorem 1 假设我们有一个 reference model $\pi_{\mathrm{ref}}(\cdot\mid x)>0$ for all prompt-answer pairs $(x,y)$, 则对于某个模型 $\pi(y\mid x)$, 则与 $\pi_{\mathrm{ref}}$ 对应的等价类里的 reward function 都可以表示为如下形式 $r(x,y) = \beta\frac{\pi(y\mid x)}{\pi_{\mathrm{ref}}(y\mid x)},\ \beta>0$.

我们也可以通过 Theorem 1 来显示推导出 DPO 选择的 reward function:

$$ \sum_{y}\underbrace{\pi_{\mathrm{ref}}(y\mid x)\exp\left(\frac1\beta r_\phi(x,y)\right)}_{\pi(y\mid x)}=1 $$

回顾前面 $\pi^*(y\mid x)$ 的定义，我们知道 $\pi(y\mid x)$ 实际上是针对 $r(x,y)$ 推导出来的最优策略的 partition function.

注意到我们的初始目标函数可以改写为如下形式

$$ \min_{\pi_\theta}\ \mathbb{E}_{x\sim \mathcal{D}}\left[\underbrace{r_{\phi}(x,y)-\beta\log Z(x)}_{f(r_\phi, \pi_{\mathrm{ref}},\beta)}-\beta\underbrace{ \frac{\pi_\theta(y\mid x)}{\pi_{\mathrm{ref}}(y\mid x)}}_{\text{KL}}\right] $$

此时，我们可以将 $f(r_\phi, \pi_{\mathrm{ref}},\beta)$ 里的 normalization term 视作为 $\pi_{\mathrm{ref}}$ 的 soft value function, 这个 soft value function 不影响最终的结果，但是没有的话会导致训练时的 variance 很高。 DPO 通过 re-parameterization, 得到的奖励函数不需要 baseline, 因而解决了训练不稳定的问题。

Experiments

作者首先对比了不同偏好优化算法的表现，如下图所示

从实验结果可以看到，DPO 对于不同的 KL Values 和不同的采样温度其表现都非常好，并且更加 robust

Conclusion

在本文中，作者提出了 DPO，一个针对 LLM 偏好优化的训练范式，DPO 构建了最优的 policy 与 reward function 之间的关系，从而避免了训练 reward model, 让模型可以直接从偏好数据集中进行学习和训练。

References

openreview

Notes on DeepSeek-V3

Mon, 08 Dec 2025 11:14:45 +0800

DeepSeek 在 24 年 11 月发布了 DeepSeek-V3, 一个仅花费 2.8M H800 hours 的大语言模型，且在各个 benchmark 上达到了 SOTA 表现

Introduction

作者在本文中提出了 DeepSeek-V3, 一个 671B-A37B 的 MoE 大语言模型。

在训练目标和架构上，作者做了如下改进：

efficiency inference: 采用了 DeepSeek-V2 提出的 MLA
cost-effective training: 采用了 DeepSeekMoE 提出了 MoE 架构
auxiliary-loss-free strategy: 采用了 Loss-Free Balancing 提出的 loss balancing 策略
multi-token prediction: 采用了 MTP 的训练目标来提升模型的表现

在训练上，作者做了如下改进：

使用了 FP8 混合精度进行训练并验证了其在大规模模型上的有效性
作者构建了 DualPipe 算法用于高效的 pipeline parallelism
构建了 cross-node all-to-all communication kernel 来高效使用 InfiniBand 以及 NVLink bandwidth
优化了 memory footprint, 来避免使用 tensor parallelism

预训练阶段，DeepSeek-V3 使用了14.8T token. 在 mid-training 阶段，作者将模型的上下文长度由 8K 扩展到 32K, 再扩展到 128K.

后训练阶段，作者使用了 SFT 和 RL 两个阶段来提高模型的表现，作者还对 DeepSeek-R1 进行蒸馏来提高模型的 reasoning 能力

Architecture

Basic Architecture

DeepSeek-V3 的架构与 DeepSeek-V2 的架构一致，如下图所示

MLA 的介绍见 MLA, MoE 架构的介绍见 DeepSeekMoE. DeepSeek-V3 在 DeepSeekMoE 的基础上做了两点改变：

受 Loss-Free Balancing 启发，作者使用了 sigmoid fu nction 来计算 affinity score
对于 selected affinity score 应用了 normalization

在 DeepSeekMoE 的基础上，作者使用了 Loss-Free Balancing. 其表达式如下

$$ g_{i,t} = \begin{cases} s_{i,t}, & s_{i,t}+b_i\in\mathrm{Topk}(\{s_{i,j}+b_j\mid 1\leq j\leq N\}, K)\\ 0, &\text{otherwise} \end{cases} $$

其中 $b_i$ 仅影响 routing, 训练时，如果对应的 expert 负载不均衡，则对 $b_i$ 进行更新，更新方式为增加/减少 $\gamma$, 这里 $\gamma$ 是一个超参数

为了提高 routing 在 sequence 层面的负载均衡，作者还是用了一个 complementary sequence-wise balance loss:

$$ \begin{aligned} \mathcal{L}_{\mathrm{Bal}} &= \alpha\sum_{i=1}^{N_r} f_iP_i\\ f_i &= \frac{}{}\sum_{t=1}^T\mathbb{1}(s_{i,t}\in\mathrm{TopK}(\{s_{j,t}\mid 1\leq j \leq N_r\}, K_r))\\ s_{i,y}' &= \frac{s_{i,t}}{\sum_{j=1}^{N_r}s_{j,t}}\\ P_i &= \frac1T\sum_{t=1}^T s_{i,t}' \end{aligned} $$

其中 $\alpha$ 是 balance factor, 为超参数，$T$ 是 sequence length, 加入这个损失后，每个 sequence 上的负载会变得更加均衡

与 DeepSeekMoE 一样，作者也在 node 层面实现负载均衡，具体做法就是，根据 node 中所包含 expert 的 affinity score 之和来选取最高的 $M$ 个 nodes, 这样就可以进一步提高 computation 和 communication 之间的 overlap.

由于 DeepSeek-V3 负载均衡比较好，因此作者使用了 no token-dropping strategy.

MTP

受 MTP 启发，DeepSeek-V3 也构建了 MTP 模块，作者认为 MTP 模块有两个优势：

MTP objective 提供了更多的学习信号，进而提高了数据使用效率
MTP 可以让模型更好预测未来的 token

与 MTP 不同，DeepSeek-V3 【todo】, DeepSeek-V3 所使用的 MTP 模块架构图如下所示

MTP 模块使用了 $D$ sequential modules 来预测未来的 $D$ 个 token. 其中，第 $k$ 个 MTP 模块包含一个共享的 embedding layer $\mathrm{Emb}(\cdot)$ , 一个共享的 output head $\mathrm{OutHead}(\cdot)$, 一个 transformer block $\mathrm{TRM}_k(\cdot)$ 和一个 projection matrix $M_k\in\mathbb{R}^{d\times 2d}$.

对于第 $i$ 个 token 以及第 $k$ 个 MTP 模块，作者首先将第 $k-1$ 个 MTP 模块的第 $i$ 个 token 的 hidden states $h_i^{k-1}\in\mathbb{R}^d$ 和第 $i+k$ 个 token 的 embedding $\mathrm{Emb}(t_{i+1})\in\mathbb{R}^d$ 联合在一起

$$ h_i'^{k}=M_k[\mathrm{RMSNorm(h_{i}^{k-1});\mathrm{RMSNorm(\mathrm{Emb}(t_{i+k}))}}] $$

其中 $[\cdot;\cdot]$ 代表 concatenation 操作。当 $k=1$ 时，$h_{i}^{k-1}$ 就代表了 main model 的输出。每个 MTP 模块的 embedding layer 和 main model 的 embedding layer 是共享的，接下来， $h_i'^k$ 作为第 $k$ 个 MTP 模块 transformer block 的输入，得到

$$ h_{i}^k = \mathrm{TRM}_k(h_{i}'^l) $$

最后，共享的 output head 输出对应的概率分布：

$$ P_{i+k+1}^k = \mathrm{OutHead}(h_i^k) $$

这里的 $\mathrm{OutHead}(\cdot)$ 的权重与 main model 也是共享的。作者这里提到，所使用的思想与 EAGLE 是类似的，但是不同的地方在于，EAGLE 主要是用于 speculative decoding, 而 MTP 主要用于提升训练。

MTP 的训练目标为未来 $T-k-1$ 个 token 的 cross-entropy loss:

$$ \mathcal{L}_{\mathrm{MTP}}^k = \mathrm{CrossEntropy}(P_{2+k:T+1}^k,t_{2+k:T+1}) = -\frac1T\sum_{t=k+2}^{T+1}\log P_i^k[t_i], $$

其中 $T$ 是 sequence length, $t_i$ 是第 $i$ 个位置对应的 ground truth token, $P_i^k[t_i]$ 代表低 $k$ 个 MTP 模块给出的 $t_i$ 的预测概率。最后，作者对所有的 MTP loss 进行求和，得到

$$ \mathcal{L}_{\mathrm{MTP}} = \frac{\lambda}{D}\sum_{k=1}^D \mathcal{L}_{\mathrm{MTP}}^k. $$

Infra

DeepSeek-V3 训练使用了 2048 张 H800, 每个 node 包含 8 张 GPU, node 内部使用 NVLink 和 NVSwitch 进行连接，node 之间使用 InfiniBand 进行连接

与之前的 DeepSeek 系列相同，DeepSeek-V3 也是用了 HAI-LLM 框架来支持训练。训练时，DeepSeek-V3 使用了 16-way PP, 64-way EP (spanning 8 nodes, GShard) 以及 ZeRO-1 DP.

作者主要进行了三点优化：

构建了 DualPipe 用于高效 pipeline parallelism
构建了 cross-node all-to-all communication kernels 来高效利用 IB 以及 NVLink bandwidth
优化了训练时的 memory footprint, 使得训练时不再依赖 TP

Training Framework

DualPipe

DeepSeek-V3 中，由于 cross-node EP, computation-to-communication ratio 近似为 1:1, 为了解决这个问题，作者提出了 DualPipe. DualPipe 的核心思想是将 forward 和 backward 过程中的 computation 以及 communication 进行重叠。与 ZeroBubble 类似，作者将每个 chunk 分为四个部分：attention, all-to-all dispatch, MLP 以及 all-to-all combine. 对于 attention 和 MLP, 作者还进一步将 backward 拆分为针对权重和输入的 backward. 其示意图如下所示，这里橙色部分代表 forward, 绿色代表了针对输入的 backward, 蓝色代表了针对权重的 backward

示意图里包含两个 block, 我们记为 block1, block2, 前向计算过程为

1

dispatch(F, block1) -> MLP(F, block1) -> combine(F, block1) -> attention(F, block2)

其中，dispatch(F, block1) 与 MLP 的反向传播 MLP(B, block1) 计算重叠, MLP(F, block1) 与 MLP 反向传播的 dispatch dispatch(B, block1) 通信重叠，combine(F, block1) 与 attention 反向传播的 attention(B, block2) 重叠，attention(F, block2) 与反向传播的 combine combine(block2) 重叠。下面是一个具体的例子

作者进一步对比了 DualPipe, 1F1B 和 ZeroBubble, 结果如下表所示

Method	Bubble	Parameter	Activation
1F1B	$(PP-1)(F+B)$	$1\times$	$PP$
ZB1P	$(PP-1)(F+B-2W)$	$1\times$	$PP$
DualPipe (Ours)	$(\frac{PP}{2}-1)(F\&B+B-3W)$	$2\times$	$PP+1$

这里 $F$ 是 forward chunk 的执行时间，$B$ 是 backward chunk 执行的时间，$W$ 是一个 chunk “backward for weights” 的执行时间，$F\&B$ 是一个 chunk 前向反向传播重叠的时间。可以看到，DualPipe 只使用了额外的 $1/PP$ 倍的 peak activation memory, 就大幅度降低了 bubble 时间

Cross-node All-to-all Communication

作者针对 DualPipe 构建了 cross-node all-to-all communication kernels 来提高通信效率。

作者提到，跨节点通信使用的是 IB, 节点内部通信使用的是 NVLink，通信方式如下图所示。

对于 H800 来说,NVLink 的带宽为 160GB/s, IB 的带宽为 50GB/s, 因此 NVLInk 的通信效率是 IB 的 3.2 倍，即节点内部通信效率高于节点之间的通信效率。为了提高通信效率，作者限制每个 token 只能被分发到至多 4 个节点上（DeepSeek-V2 里提到的 device limited routing）。其具体通信方式为想传输到目标 node index 相同的 rank 上，然后再通过节点内部通信传输到目标的 rank 上。通过这种方式，我们可以让节点间通信与节点内部通信进行重叠，从而每个 token 可以从一个节点上选取 3.2 个专家，这样 DeepSeek-V3 可以在不损失效率的情况下最高选取 $13$ 个专家。

作者进一步采用了 warp specialization 技巧来将 20 个 SM 划分为 10 个通信 channel. 作者分别针对 dispatching 和 combing 阶段使用了不同数量的 warps.

Memory saving

作者使用了如下技巧来减少内存访问：

Recomputation of RMSNorm and MLA Up-Projection. 在 backward 过程中重新计算 RMSNorm operation 以及 MLA up projection 来避免存储器对应的输出
Exponential Moving Average in CPU. 作者将 EMA 参数保存在 CPU 中，然后进行异步更新来进一步减少内存访问
Shared Embedding and Output Head for Multi-Token Prediction. 作者将 embedding layer 和 output head 放在一个 PP rank 上，这样就可以提高 MTP 的内存访问效率

FP8 Training

作者提出了一个基于 FP8 的混合精度训练框架。作者提出了两个改进方案：

分组量化，将 tensor 按照 tile 分组或者按照 block 分组来分组量化，这样就避免了全局量化的精度损失
高精度累加，作者在乘法计算时，使用了 FP8 格式，然后在累加阶段，使用了更高精度的格式

为了进一步减少内存和通信开销，对于 activation 的 cache 以及 dispatch, 作者使用了 FP8 数据格式，然后对于优化器状态，作者使用了 BF16 数据格式。下面是对上面改进的具体说明。

Mixed Precision Training

作者在本文中提出了使用 FP8 混合精度进行预训练，作者参考了 low precision training 构建 FP8 训练框架，即计算量高的使用 FP8 精度，计算量低的使用原本的数据精度, 框架如下图所示

其中各个模块使用的精度如下表所示

Precision	Modules
FP8	Linear (Fprop, Dgrad, Wgrad)
higher precision	- embedding - output head - Moe gating - normalization - attention operator - master weights - weight gradients - optimizer states

Enhancing Low-precision Training Accuracy

作者介绍了几个策略用于提高 FP8 混合精度训练的表现：

fine-grained quantization: 对于激活值，作者将 tensor 分割为 $1\times 128$ 大小的 groups, 然后每个 group 内部进行 quantization; 对于权重，作者将其分割为 $128\times 128$ 大小的 groups, 然后进行 quantization, 这样可以降低 quantization error, 如上图左图所示
increasing accumulation precision: 低精度训练会带来 underflow 的问题，而 FP8 GEMM 累加仅能保留约 14bit 精度，远低于 FP32 的 32bit 精度。为了解决这个问题，作者采用了 Tensor Core 不分累加 +CUDA core 高精度聚合的协同策略。即在 Tensor Core 上执行 MMA 指令时，先按照 14bit 精度对 $N_C$ 个元素进行累加，当达到 $N_C$ 时，作者将结果复制到 CUDA core 的 FP32 寄存器中，在 FP32 精度下完成累加。过程如上图右图所示
Mantissa over exponents. 与之前的工作不同，作者使用了 E4M3 的数据格式来达到更高的精度
Online Quantization. 为了降低 quantization 的误差，作者实时计算了 activation block 以及 weight block 的最大绝对值来导出 scaling factor 并量化为 FP8 精度

Low Precision Storage and Communication

作者通过压缩 cached activation 和 optimizer states 来进一步减少内存访问以及通信访问次数

Low-precision optimizer states: 对于 optimizer states, 作者使用了 BF16 数据格式来保存 AdamW 的优化器的一阶和二阶动量，但是对于 master weight 和 gradients 作者仍然使用了 FP32 来保证训练的数值稳定性
Low-Precision Activation: 对于 linear operator, 作者将其 cache activation 使用 FP8 格式进行存储，对于 attention 的输出，作者使用了 E5M6 数据格式来存储 activations, 这些 activations 的 tile size 为 $1\times 128$, 作者还是用了 2 的幂次作为 scale factor 来减少 quantizationerror; 对于 SwiGLU, 作者将其输入也保存为 FP8 的数据格式来降低内存消耗
Low-Precision Communication: 在 MoE up-projection 之前，作者将 attention 的输出量化为 FP8 数据格式再执行 dispatch 操作，这样可以降低通信开销。在反向传播的时候同理，先进行 FP8 量化，然后再进行反向 dispatch. 对于 combine 阶段，为了避免精度损失，作者还是使用了 BF16 数据格式

作者对比了 FP8 和 BF16 精度训练，结果如下图所示

实验结果显示，FP8 混合精度训练的损失降低不足 $0.25\%$.

Inference and Deployment

作者在 H800 集群上部署 DeepSeek-V3, 作者分别针对 prefilling 和 decoding 两个阶段进行了优化

Prefilling

prefilling 阶段在 4 节点 32 GPU 上进行，并行策略如下

模块	并行策略	说明
Dense MLP	1-wat TP	减少 TP 通信
Attention	4-way TP, SP, 8-way DP	SP 用于长文本处理
MoE	32-way EP	每个 GPU 包含 8 个专家

为了实现负载均衡，作者提出了redundant experts的策略，具体就是将负载比较高的专家进行复制。模型会周期性进行统计，然后计算出负载比较高的专家，接下来每个 GPU 除了原来的 8 个专家之外，还会 host 一个额外的 redundant expert.

为了提高 throughput, 作者同时处理两个 micro-batch, 来重叠 attention, MoE 和 dispatch, combine 通信，如下图所示

作者还探索了 dynamic redundancy 策略，即每个 GPU host 更多的专家（比如 16 个）然后 inference 的每一步仅激活其中的 9 个，在进行 all-to-all operation 时，作者首先进行 global optimal routing 来计算最合适的专家。作者认为 prefilling 计算量很大，因此 routing 的计算量可以忽略不计

Decoding

在 decoding 阶段，作者在 40 个节点 320 GPU 上进行部署，并行策略如下

模块	并行策略	说明
Attention	4-way TP, SP, 80-way DP	SP 用于长文本处理
MoE	320-way EP	每个 GPU 包含 8 个专家

这个阶段，每个 GPU 只 host 一个专家，64 个 GPU 负责 host redundant expert 以及 shared expert. 作者通过在 IB 上直接进行 P2P 的传输来降低通信的开销，作者还是用了 IBGDA 来进一步最小化 latency 以及提高通信效率

在这个阶段，attention 的计算占据了大部分时间，因此，作者采取了如下的 overlap 策略

由于 decoding 阶段一般 batch size 比较小，因此整个 decoding 的瓶颈在于 memory access, 因此为了避免影响 attention 的计算效率，作者仅安排一小部分 SM 用于 dispatch-MoE-combine.

Suggestions on Hardware Design

Communication Hardware

尽管作者将计算与通信重叠来提高训练效率，但是已有的通信仍然依赖于 SM, 这样限制了计算的效率。作者希望工能够构建专门针对通信的设备，另一方面，作者希望统一 IB 和 NVLink 来降低实现难度

Computation Hardware

提升 FP8 GEMM 的累加精度，当前精度只有 14bit, 作者希望能够在 GPU 设计上进行改进，将这个精度提升到 34-bit
支持 tile 以及 block 层面的 quantization, 尽管前文作者已经设计出了改进的算法，但是 Tensor Core 以及 CUDA 之前平凡的数据移动降低了计算效率，作者希望未来能够支持细粒度的 quantization
online quantization, 作者希望将 FP8 cast 以及 TMA 结合在一起，从而 quantization 可以在传输的时候完成计算，减少了内存读写。作者还建议使用 warp-level cast instruction
Transposed GEMM operations. 本文中，矩阵被切分为不同的 tiles, 在计算的时候需要先加载这些 tiles 然后进行 dequantization, transpose 等操作，作者希望未来能够直接支持 transposed reads of matrices

Pre-training

Pre-training Data

相比于 DeepSeek-V2, DeepSeek-V3 提升了数学和代码数据的比例，以及增加了多语种数据。最终训练数据一共包括 14.8T

作者还使用了 DeepSeekCoder-V2 里应用的 Fill in the middle 策略来让模型基于上下文越策中间的文本，对应的数据格式如下

1

<|fim_begin|>f_pre<|fim_hole|>f_suf<|fim_hole|>f_middle<|fim_end|>

这个结构与 sequence packing 结合在一起。

Tokenizer 基于 BBPE, 大小为 128K tokens. 在训练时，作者将随机一部分 combine token 进行切分来减少 token boundary bias 问题

Hyper-parameters

模型参数如下表所示，最终 DeepSeek-V3 拥有 671B 总参数，激活参数为 37B

variable	notation	value
layers	$\ell$	61
dense layers	-	3
hidden dimension	$d$	7168
num of attention heads	$n_h$	128
head dimension	$d_h$	128
KV compression dimension	$d_c$	512
query compression dimension	$d_c'$	1536
decouple query and key dimension	$d_h^R$	64
routed expert	$N_r$	256
shared expert	$N_s$	1
MoE hidden dimension	$d_{MoE}$	2048
activated experts	$K$	8
limited node routing	$M$	4
MTP depth	$D$	1

训练时，作者使用了 learning rate scheduling, batch size scheduling 等方法。对于 PP, routed expert 会均匀分布为 8node 对应的 64 个 GPU 上，预训练时模型的上下文长度为 4k

Long Context Extension

作者使用了 YARN 来扩展模型的上下文长度，模型上下文长度经过两个额外的训练阶段从 4K 扩展到 32K 再扩展到 128K, 均训练了 1000 步。

YARN 配置与 DeepSeek-V2 基本一致，在额外的 decoupled query and key 上作者没有应用这一点。参数配置如下表

parameter	$s$	$\alpha$	$\beta$	$\sqrt{t}$
value	40	1	32	$0.1\ln s+1$

第一阶段的上下文长度为 32K, batch size 为 1920, 第二个阶段的上下文长度为 128K, batch size 为 480.

Performance

DeepSeek-V3 base 的表现下图所示，作者对比了 DeepSeek-V2, Qwen2.5, LLaMA 3.1

Discussion

作者首先验证了 MTP 的有效性，结果如下图所示

可以看到，在模型架构相同，训练数据相同的情况下，1-MTP 的效果超过了 baseline 的表现，说明了 MTP 策略的有效性

接下来，作者还验证了 Loss-Free Balancing 的有效性，结果如下图所示

可以看到，loss-free Balancing 的效果比 loss balancing 的效果更好

接下来，作者对比了 loss-free balancing 和 sequence-wise auxiliary loss, 即 batch-wise v.s. sequence-wise. 作者认为，前者的约束更灵活，因为其不要求 in-domain balance. 作者在测试集上进行可视化，结果如下图所示

实验结果发现，loss-free 策略对应的 expert specialization 更强。作者进一步设计了一个 batch-wise auxiliary loss 来实现 batch-wise load balance, 结果发现，这种策略也能达到和 loss-free balancing 一样的效果，这说明了 batch-wise load balancing 效果更好

最后，作者提了两点 loss-free 策略的问题：

在特定的 sequence 或者小 batch 里出现负载不均衡。作者通通过增大 batch size 来解决这个问题
在推理阶段因为 domain-shift 导致的负载不均衡。作者实现了一个基于 redundant expert 的推理框架来解决这个问题

Post-training

SFT

post-training 包含 1.5M 样本，数据包括 reasoning 数据以及 non-reasoning 数据，前者由 DeepSeek-R1 合成，后者由 DeepSeek-V2.5 合成

SFT 时，作者训练了两个 epoch, 使用了 sequence packing 技巧

RL

Reward model 包含 rule-based reward model 和 model-based reward model.

RL 训练使用的算法为 GRPO

Post-training Performance

Post-training Discussion

作者首先探究了 Distillation 对模型表现的影响，作者使用 DeepSeek-R1 来蒸馏 DeepSeek-V2.5, 结果如下图所示

Model	LiveCodeBench-CoT		MATH-500
	Pass@1	Length	Pass@1	Length
DeepSeek-V2.5 Baseline	31.1	718	74.6	769
DeepSeek-V2.5 +R1 Distill	37.4	783	83.2	1510

可以看到，distillation 可以显著提高模型的表现。但是其问题在于也会让模型输出的长度增加。作者认为知识蒸馏是 post-training optimization 的一个重要方向。

接下来，作者讨论了 self-rewarding, 具体做法就是使用 DeepSeek-V3 来对评估结果进行投票，结果发现最终的效果很好，作者认为 LLM 可以很好地将非结构化信息转换为 rewards.

最后，作者讨论了 MTP. MTP 可以于 speculative decoding 结合，进一步提高 decoding 的速度。通过实验作者发现，second token prediction 的接受率在 $85\%\sim 90\%$ 之间，说明了其有效性。

Conclusion

在本文中，作者提出了 DeepSeek-V3, 一个 671B-A37B 的 MoE 大语言模型，训练 token 数为 14.8T. 作者使用了 loss-free-balancing strategy 来实现负载均衡。训练时作者使用了 FP8 混合精度。评估发现 DeepSeek-V3 的达到了 SOTA 表现。

作者认为，DeepSeek-V3 model size 太大，不适合部署。第二，部署策略需要进一步改进。

最后，作者认为未来工作有以下几点：

改进模型架构，进一步提高训练以及推理效率，扩展模型的上下文
提升训练数据的数量和质量
提高模型的 reasoning 能力
更详尽的评估

References

Arxiv

Notes on olmoe

Sat, 06 Dec 2025 18:08:11 +0800

Introduction

作者在本文中主要想完成三个目标：

基于已有框架尝试训练一个 MoE 大语言模型
详细分析 MoE 模型的 routing 机制
为后续 MoE 模型开发提供经验

基于这三个目标，作者在本文中发布了 4 个 MoE 模型。预训练时，作者使用了 $52.25\%$ 的代码数据来提高模型的表现。作者还是用了 UL2 作为训练目标。

通过实验作者发现 MoE 机制存在以下三个性质：

Context-independent Specialization: 即 MoE 模型倾向于基于 token semantics 进行聚类，而不是 context
early routing learning: routing 的分布情况在训练早期就已经确定了
drop towards the end: 使用 token dropping 策略之后，越往后的的 token 被丢弃掉的概率也就越大

Method

Model

作者介绍了模型使用的数据集如下

模型的 tokenizer 基于 umT5, 大小为 256K, umT5 支持多语种，并且还有 fallback 机制来处理 OOV 的 token

模型配置如下

Model	OpenMoE-Base/16E	OpenMoE-8B/32E	OpenMoE-34B/32E
total params	650M	8.7B	34B
activated params	142M	2.1B	6B
total experts	16	32	32
activated experts	2	2	2
shared experts	1	1	1
Layout	every 4	every 6	every 4
`hidden_dim`	768	2048	3072
`ffn_hidden_dim`	3072	8192	12288
`num_heads`	12	24	24
`head_dim`	64	128	128
`num_layers`	12	24	32

Training

作者使用了 Load Balancing loss 和 ST-MoE 提出的 Router Z-loss, 最终训练的目标函数与 ST-MoE 一致

作者在训练时，还是用了 UL2, UL2 结合了 mask language modeling 和 casual language modeling 两种训练方式。

作者首先对各个实验配置进行了验证，其中关键接论文模型很容易在 code data 上达到比较高的准确率以及比较低的 loss, 作者认为这是因为 code data 中存在大量的特殊符号。

训练过程中，作者发现模型在训练一定步数之后，容易出现过饱和现象，因此作者将训练目标函数由 UL2 降为 CasualLM, 并未作者还将代码数据的比例降低到 $15\%$.

Experiments

Analysis OpenMoE

作者发现，大部分专家对于不同的 domain 没有出现 specialization 情况，对于 math domain, specialization 现象比较明显，作者认为这是因为 math domain 包含更多的 special tokens
对于不同的语言，有部分专家出现 specialization 现象
部分专家对于 position ID 有 specialization 现象，并且连续的 token 更偏好相同的专家
作者还发现，部分专家对于 Token ID 有 specialization 现象，作者将这种现象称为 Context-independent Specialization.
专家还会对语义相似的 token 进行聚类，并且这种聚类在训练早期就已经发生，作者认为其原因在于重新分配 token 会增加最终的 loss
对于 token dropping, 作者发现月考后的 token, 其被 drop 的概率比例也越高。并且对于指令跟随数据，更多的 token 都会被丢掉，因此作者认为指令跟随数据是 MoE 模型的一种 OOD 数据

Analysis on other MoE Models

作者还分析了 Mixtral MoE 和 DeepSeekMoE 两个模型的专家路由情况。这两个模型采用了 dropless token routing 的机制。

首先，作者分析了这两个模型对于 TokenID 的敏感性，结果发现，DeepSeek-MoE 的 specialization 现象比较明显，而 Mixtral MoE 由于使用了 upcycling, 其 specialization 现象不太明显，作者认为这是因为 upcycling 导致每个专家的权重都差不多。最终，作者认为对于 training from stratch 的 MoE model 这个 specialization 现象更明显。

Conclusion

作者在本文中提出了一个全开源的 MoE 大模型系列，包括 0.6B, 8.7B 和 34B 三个 size, 作者还对 MoE 中的 routing 进行了详细的分析，结果发现【【】】，这些发现有助于后续的 MoE 模型架构的研究。

References

Arxiv

Notes on SAPO

Fri, 05 Dec 2025 10:09:06 +0800

Introduction

Method

作者首先给出了 SAPO 的目标函数如下：

$$ \mathcal{J}_{\mathrm{SAPO}}(\theta) = \mathbb{E}_{(x, y)\sim\mathcal{D},\{y_i\}_{i=1}^G\sim \pi_{\theta_{old}}(\cdot\mid x)}\left[ \frac{1}{G}\sum_{i=1}^G\frac{1}{|y_i|}\sum_{t=1}^{|y_i|}f_{i,t}(r_{i,t}(\theta))\hat{A}_{i,t} \right], $$

其中，

$$ f_{i,t}(x) = \sigma(\tau_{i,t}(x-1))\cdot \frac{4}{\tau_{i,t}}, \tau_{i,t} = \begin{cases} \tau_{pos}, & \text{if }\hat{A}_{i,t}>0\\ \tau_{neg}, &\text{otherwise} \end{cases} $$

这里

$$ \hat{A}_{i,t} = \hat{A}_{i} = \frac{r(x,y_i) - \mathrm{mean}(\{r(x,y_i)\}_{i=1}^G)}{\mathrm{std}(\{r(x,y_i)\}_{i=1}^G)},\quad r_{i,t}(\theta) = \frac{\pi_{\theta}(y_{i,t}\mid x, y_{i,< t})}{\pi_{\theta_{old}}(y_{i,t}\mid x, y_{i,< t})} $$

$\tau_{pos}$ 和 $\tau_{neg}$ 分别是 positive token 以及 negative token 对应的温度, $\sigma(x)=1/(1+e^{-x})$ 是 sigmoid function.

对 $\mathcal{J}_{\mathrm{SAPO}}(\theta)$ 求导得到

$$ \nabla_\theta \mathcal{J}_{\mathrm{SAPO}}(\theta) = \mathbb{E}_{(x, y)\sim\mathcal{D},\{y_i\}_{i=1}^G\sim \pi_{\theta_{old}}(\cdot\mid q)}\left[ \frac{1}{G}\sum_{i=1}^G\frac{1}{|y_i|}\sum_{t=1}^{|y_i|}w_{i,t}(\theta)r_{i,t}(\theta)\nabla_\theta \log \pi_\theta(y_{i,t}\mid q, y_{i, 其中

$$ w_{i,t}(\theta) = 4p_{i,t}(\theta)(1-p_{i,t}(\theta)),\quad p_{i,t}(\theta) = \sigma(\tau_{i,t} (r_{i,t}(\theta)-1)), $$

$\mathcal{J}_{\mathrm{SAPO}}(\theta)$ 和 $w_{i,t}(\theta)$ 与 $r_{i,t}(\theta)$ 的关系如下图所示

为了保证当 $r_{i,t}(\theta)=1$ 时，SAPO 等价于 $r_{i,t}(\theta)\hat{A}_{i,t}$ 而与 $\tau_{i,t}$ 无关，作者在 $f_{i,t}(x)$ 加入了系数 $4/\tau_{i,t}$.

Comparison

作者接下来对比了 GSPO 以及 GRPO 两个算法

首先作者使用了一下统一的目标函数公式来表示三个算法

$$ \mathcal{J}(\theta) = \mathbb{E}_{(x, y)\sim\mathcal{D},\{y_i\}_{i=1}^G\sim \pi_{\theta_{old}}(\cdot\mid x)}\left[ \frac{1}{G}\sum_{i=1}^G\frac{1}{|y_i|}\sum_{t=1}^{|y_i|}f_{i,t}(r_{i,t}(\theta))\hat{A}_{i,t} \right], $$

其中不同算法的 $f_{i,t}(\cdot)$ 不一样，三种算法的定义如下

$$ \begin{aligned} \mathrm{SAPO}&:f_{i,t}^{\mathrm{SAPO}}(r_{i,t}(\theta))=\sigma(\tau_{i,t}(x-1))\cdot \frac{4}{\tau_{i,t}}, \tau_{i,t} = \begin{cases} \tau_{pos}, & \text{if }\hat{A}_{i,t}>0\\ \tau_{neg}, &\text{otherwise} \end{cases}\\ \mathrm{GRPO}&:f_{i,t}^{\mathrm{GRPO}}(r_{i,t}(\theta), \hat{A}_{i,t})=\begin{cases} \min(r_{i,t}(\theta), 1+\epsilon) & \hat{A}_{i,t}>0\\ \max(r_{i,t}(\theta), 1-\epsilon) & \hat{A}_{i,t}\leq0 \end{cases}\\ \mathrm{GSPO}&:f_{i,t}^{\mathrm{GSPO}}(r_{i,t}(\theta), \hat{A}_{i,t})=f_{i,t}^{\mathrm{seq}}(r_{i,t}(\theta), \hat{A}_{i,t})=\begin{cases} \min(s_{i,t}(\theta), 1+\epsilon) & \hat{A}_{i,t}>0\\ \max(s_{i,t}(\theta), 1-\epsilon) & \hat{A}_{i,t}\leq0 \end{cases} \end{aligned} $$

其中

$$ s_i(\theta) = \left(\frac{\pi_{\theta}(y_{i}\mid x)}{\pi_{\theta_{old}}(y_{i}\mid x)}\right)^{\frac{1}{|y_i|}}, s_{i,t}(\theta) = \mathrm{sg}[s_i(\theta)]\frac{\pi_{\theta}(y_{i,t}\mid x, y_{i,< t})}{\mathrm{sg}[\pi_{\theta}(y_{i,t}\mid x, y_{i,< t})]} $$

首先是与 GSPO 的对比，通过一些假设和简化，我们得到

$$ \begin{aligned} \nabla_\theta \mathcal{J}_{\mathrm{SAPO}}(\theta) &=\mathbb{E}\left[\frac1G\sum_{i=1}^G g_{\tau_i}(\log s_i(\theta))\nabla_\theta \log s_i(\theta)\hat{A}_i\right]\\ \nabla_\theta \mathcal{J}_{\mathrm{GSPO}}(\theta) &=\mathbb{E}\left[\frac1G\sum_{i=1}^G s_i(\theta)\nabla_\theta \log s_i(\theta)\hat{A}_i\right] \end{aligned} $$

其中 $g_{\tau_i}(\log s_i(\theta)) = \mathrm{sech}^2\left(\tau_i/2\log s_i(\theta)\right)$. 相比于 GSPO, SAPO 有两个优势：

smoothness and stability, soft gate 避免了 hard clipping 带来的不连续性
token-level adaptivity with sequence-level coherence. 当假设不成立的时候，SAPO 退化为 token-level gating, 这样可以降低 outliers 的权重

GRPO 的函数可以进一步简化为

$$ f_{i,t}^{\mathrm{GRPO}}(r_{i,t}(\theta), \hat{A}_{i,t})= \begin{cases} 1, & \text{if }\hat{A}_{i,t}>0\text{ and }r_{i,t}(\theta)\leq 1+ \epsilon\\ 0, & \text{if }\hat{A}_{i,t}>0\text{ and }r_{i,t}(\theta)> 1+ \epsilon\\ 1, & \text{if }\hat{A}_{i,t}\leq0\text{ and }r_{i,t}(\theta)\geq 1- \epsilon\\ 0, & \text{if }\hat{A}_{i,t}\leq0\text{ and }r_{i,t}(\theta)< 1- \epsilon \end{cases}\\ $$

可以看到，GRPO 对应一个 binary trust region. 与 GRPO 相比，SAPO 将对应的 hard indicator 替换未来一个 smooth kernel $f_{i,t}^{\mathrm{SAPO}}(r_{i,t}(\theta))=\mathrm{sech}^2\left(\tau_i/2r_{i,t}(\theta)-1\right)$, 这样可以避免 gradient vanishing 以及提高训练的稳定性。

最终结论为：

相比于 GSPO, SAPO 对于 off-policy 的数据利用率更高
相比于 GRPO, SAPO 避免了 hard token level clipping 导致的 zero-gradient 问题

Experiments

作者对比了 SAPO, GSPO 以及 GRPO-R2(GRPO with routing replay) 三种方法，实验结果如下图所示

实验结果显示，SAPO 的表现超过了 GSPO 以及 GRPO

作者还探究了超参数 $\tau_{pos}$ 和 $\tau_{neg}$ 的影响，结果如下图所示

实验结果显示，当 $\tau_{neg}>\tau_{pos}$ 时，模型训练最稳定，这说明了 negative token 是导致训练不稳定的主要原因。

作者还在 Qwen3-VL 上进行了验证，结果如下图所示

实验结果显示，SAPO 的表现超过了 GSPO 以及 GRPO-R2

Conclusion

在本文中，作者提出了 SAPO, 一个解决 hard-clipping 训练不稳定性以及低效率的策略优化算法，作者使用了基于温度的 soft gate 来代替 clipping, 以及对于 positive token 和 negative token 使用了不同的 temperature 这两点改进。结果验证了 SAPO 的有效性，作者认为使用 smooth 以及 adaptive gating 机制可以有效提高 RL 训练的稳健性以及有效性。

References

arxiv

Notes on DeepSeek-R1

Tue, 02 Dec 2025 18:21:54 +0800

Notes on DeepSeek-V2

Tue, 02 Dec 2025 18:21:54 +0800

Introduction

作者首先提到如何提高模型的训练效率以及 inference efficiency 是两个尚未解决的问题。

基于这两个问题，作者在本文中提出了 DeepSeek-V2，一个开源的 MoE 模型，DeepSeek-V2 的亮点在于训练和推理都非常高效。最终 DeepSeeK-V2 包含 236B 总参数，激活参数为 21B, 上下文长度为 128K. 作者还开源了 DeepSeek-V2-Lite, 一个 15.7B-A2.4B 的 MoE 模型，用于学术研究。

DeepSeek-V2 主要改进点为：

基于 DeepSeekMoE, 使用了 MoE 架构
使用了 MLA 压缩 KV cache, 大幅度提高推理效率

DeepSeek-V2 预训练使用了 8.1T tokens, 相比于 DeepSeek-LLM, 预训练数据主要增加了中文数据以及提高了数据的质量。

接下来，作者收集了 1.5M 对话数据来进行 SFT, 最终作者基于 DeepSeek Math 提出的 GRPO 来进行对齐。

Architecture

DeepSeek-V2 的模型架构如下

模型基于 DeepSeekMoE 开发得到，相比于 DeepSeekMoE, DeepSeek-V2 主要是使用了 MLA

MLA

这部分介绍见 MLA

DeepSeekMoE

Architecture

关于架构的介绍见 DeepSeekMoE

Device-Limited Routing

由于 DeepSeek-MoE 使用了细粒度的专家，因此专家会分布在更多的设备（GPU）上，计算时，基于 routing 的 expert 所在设备，会产生不同大小的通信开销。为了降低通信开销，作者构建了 device-limited routing mechanism. 具体的做法就是，在 Routing 之前，先基于 experts 的 affinity score 挑选 $M$ 个设备，然后基于这 $M$ 个设备的专家挑选 top-K 专家进行计算。

作者通过实验发现，当 $M\geq3$ 时，device-limited routing 可以和标准的 top-K routing 表现差不多。

Auxiliary Loss for Load Balance

作者使用了三个 loss 来实现负载均衡。其中，expert level 和 device level 的 load balancing loss 与 DeepSeekMoE 相同。第三个 loss 是 communication balance loss, 这个 loss 的目的是让每个设备的通信开销保持平衡。损失函数的表达式如下所示

$$ \mathcal{L}_{communication} = \alpha\sum_{i=1}^D f_iP_i $$

其中 $\alpha$ 是超参数，$D$ 是 expert group 的个数。

$$ f_i = \frac{D}{MT}\sum_{t=1}^T\mathbb{1}(\text{Token }t \text{ is sent Device }i),\quad P_i = \sum_{j\in\mathcal{E}_i}P_j $$

device limited routing 让每个 device 发送至多 $MT$ 个 hidden states 到其他设备上。而 communication balancing loss 则让每个设备最多从其他设备接收 $MT$ 个 hidden states.

Token-Dropping Strategy

尽管前面已经增加了 load balance loss, 但毕竟不是硬约束。因此，作者就从硬件层面提出了 Token dropping 策略，来提高训练效率。核心思想就是，在训练时，主动丢弃部分 token, 强制让各个设备的计算量不会超过额度限制，进而减少资源浪费。

具体做法就是，在训练之前，先将每个设备的 capacity factor 设置为 1 （定义见 Switch Transformer）. 然后按照 affinity score 来丢弃一些分数比较低的 token, 直到该设备的 token 数量刚到达到 capacity。为了避免过度学习导致模型表现较差，对于 $10\%$ 的训练数据，作者不执行 token dropping 策略。

最终，在 inference 时，可以根据需求来决定是否丢弃 token, 比如在 low latency 场景，我们可以丢弃低价值的 token, 在高精度场景，我们就可以保留所有的 token.由于在训练阶段已经才去过 token dropping 策略，因此在推理时不管是丢弃还是全部保留模型都能比较好的适应。

Pre-training

Data

预训练数据与 DeepSeek-LLM 基本上差不多，作者针对中文数据，数据质量进行了改进。最终预训练数据包括 8.1T token, 其中中文数据比英文数据多 $12\%$.

tokenizer 与 DeepSeek-LLM 一致。

Model Configuration

模型配置如下表所示

Model	DeepSeek-V2	DeepSeek-V2-Lite
Date	2024-5	2024-5
# Total Parameters	236B	15.7B
# Activated Parameters	21B	2.4B
# tokens	8.1T	5.7T
# Dense Layers	1	1
# MoE Layers	60	26
Hidden Dim	5120	2048
Dense Intermediate Dim	12288	10944
MoE Intermediate Dim	1536	1408
Attention	MLA	MLA
# Attention Heads	128	16
# Key-Value Heads	128	16
# Routed Experts	160	64
# Experts Active Per Token	6	6
# Shared Experts	2	2

这里比较特殊的一点在于，模型在第一层使用了 MoE layer, 这个做法的原因在后面的 olmoe 里有提到，核心思想是 early layer 特别是第一层 layer 收敛比较慢。

MLA 的配置如下 (DeepSeek-V2)

field	value
$d_c$	512
$d_c'$	1536
$n_h$	16
$d_h$	128
$d_h^R$	64

Training Recipe

训练的配置也与 DeepSeek-LLM 差不多，对于 MoE，作者使用了 PP 将不同的 layers 分配在不同的 device 上，然后 MoE 的 experts 被分配在 8 个 device 上 ($D=8$), 对于 device-limited routing, 每个 token 发送到至多 3 个 device, 也就是 $M=3$.

Infra

在 infra 上，DeepSeek-V2 也是用了 HAI-LLM 框架进行训练。这里面使用了 16-way zero-bubble PP, 8-way EP, ZeRO-1 DP.

由于 DeepSeek-V2 的激活参数比较少，因此，作者没有使用 TP, 进而降低通信开销。作者还将 shared experts 的计算与 expert all-to-all 通信进行重叠来提高计算效率。作者还使用了 kernel fusion 和 flash attention 2 来加速训练。

Long Context

在预训练阶段结束之后，作者使用了 YARN 来将模型的上下文从 4K 扩展到 128K. 超参数设置为

Parameter	Value
$s$	40
$\alpha$	1
$\beta$	32
target context length	160K
scaling factor	$\sqrt{t} = 0.0707\ln s + 1$

作者在 32K 的上下文下额外训练了 1000 步，然后在推理阶段通过 YaRN 将模型的上下文长度扩展到 128K.

Evaluation

作者对比了 DeepSeek-LLM, Qwen1.5, Mixtral MoE 以及 LLaMA 3, 实验结果如下

Efficiency

作者对比了以下 DeepSeek-MoE 和 DeepSeek-LLM 的训练效率，结果发现，对于 1T 的 token, DeepSeek-LLM 需要 300.6K GPU hours, 而 DeepSeek-V2 仅需要 127.8K GPU hours. 也就是说，DeepSeeK-V2 节省了 $42.5\%$ 的训练成本

在推理时，作者首先将模型的精度转换为 FP8，然后作者进一步对模型进行 KV cache quantization 来进一步压缩每个 token 的 KV cache 到 6bits. 最终，DeepSeek-V2 的 throughtput 为 50K tokens/s.

Post-training

post-training 分为 SFT 和 RL 两个阶段。

SFT

在 SFT 阶段，作者构建了 1.5M 样本，包括 1.2M 有帮助性的样本和 0.3M 安全性相关的样本。模型训练了 2 个 epoch, 学习率为 $5e-6$.

RL

作者使用了 GRPO 算法来进一步对齐模型的表现。

作者通过实验发现，在 reasoning data, 如 code 和 math 相关数据上进行训练时，可以有效提高模型的表现。因此作者将 RL 的训练分为两个阶段，第一个阶段用于提高模型的 reasoning 能力，第二个阶段用于对齐人类偏好。

在第一个阶段，作者首先训练了一个针对 code 和 Math 的 reward model $\mathrm{RM}_{\mathrm{reasoning}}$, 然后基于这个 reward model 来训练 policy model:

$$ r_i=\mathrm{RM}_{\mathrm{reasoning}}(o_i) $$

在第二阶段，作者使用了一个 Multi-reward 框架，包括一个 helpful reward model $\mathrm{RM}_{\mathrm{helpful}}$, 一个 safety reward model $\mathrm{RM}_{\mathrm{safety}}$ 和一个 rule-based reward model $\mathrm{RM}_{\mathrm{rule}}$, 最终的 reward 为

$$ r_i = c_1\mathrm{RM}_{\mathrm{helpful}}+c_2\mathrm{RM}_{\mathrm{safety}}+c_3\mathrm{RM}_{\mathrm{rule}} $$

训练时，reward model 由 SFT model 初始化得到，然后基于 point-wise 或者 pair-wise loss 进行训练。

Evaluation

chat 版本的模型评估结果如下所示

Discussion

作者讨论了三点发现：

SFT data 数量。已有工作认为进需要 10K 左右的样本就可以进行 SFT，但是作者发现当数据量小于 10K 时，模型在 IFEval benchmark 上的表现大幅度下降。作者认为，这是由于数据过少导致模型很难掌握特定的技能。因此，作者认为足够的数据以及数据质量都很重要，特别是写作类任务和 open-ended QA 类任务。
alignment tax. 作者发现通过 human preference alignment, 模型在 open-ended generation benchmark 上的保险有了很大提升。与 RLHF 一样，作者也发现了 alignment 之后模型在一些 benchmark 上表现也会下降。作者通过改进解决了这个问题，作者认为如何在不损失模型表现的情况下实现对齐是一个值得探究的方向。
online RL. 作者发现 Online RL 比 offline RL 的表现更好。作者认为如何根据不同的任务来选取 offline RL 和 online RL 也是一个值得探究的问题。

Conclusion

在本文中，作者提出了 DeepSeek-V2, 一个基于 MoE 架构的大语言模型系列，模型的上下文为 128K. 作者基于 DeepSeek-MoE, 提出了 MLA 来提高模型的 inference 效率，并大幅度降低了训练的成本。

作者介绍了几点未来工作：

进一步 scaling up MoE 模型，降低模型的训练以及推理成本
进一步对齐模型和人类的价值观，然后最小化人类监督信号
扩展模型到多模态版本

Reference

Notes on MLA

Tue, 02 Dec 2025 18:21:54 +0800

DeepSeek 在 2024 年 5 月提出了 multi-head latent attention (MLA), 用于提高 attention 的 Inference 效率

Introduction

传统的 multi head attention (MHA) 虽然效果好，但是在 inference 时，其 KV cache 会变成瓶颈，影响推理效率。为了解决这个问题，已有的工作如 MQA 和 GQA 通过共享权重来减少 KV cache 内存占用，但是结果发现模型的表现也会降低。

为了解决这个问题，作者提出了 multi-head latent attention (MLA), 来压缩 KV cache.

MHA

令 $d$ 为 hidden size, $n_h$ 为 attention heads 的个数，$\ell$ 为 transformer layer 的层数，$d_h$ 为每个 head 的 dimension, $h_t\in\mathbb{R}^d$ 为 attention layer 中第 $t$ 个 token 对应的 hidden states。对于标准的 MHA, 我们首先计算 Q, K, V 如下：

$$ q_t=W^{Q}h_t,\quad k_t=W^Kh_t,\quad v_t = W^Vh_t $$

其中，$W^Q,W^K,W^V\in\mathbb{R}^{d_hn_h\times d}$ 分别为 query, key, value projection layer 的权重。接下来 MHA 的计算方式如下

$$ \begin{aligned} o_{t,i} &= \sum_{j=1}^t\mathrm{softmax}_j\left(\frac{q_{t,i}^Tk_{j,i}}{\sqrt{d_h}}\right)v_{j,i},\\ u_t&= W^O[o_{t,1};o_{t,2};\dots,;o_{t,n_h}] \end{aligned} $$

其中 $q_t=[q_{t,1};q_{t,2};\dots,;q_{t,n_h}]$, $k_t=[k_{t,1};k_{t,2};\dots,;k_{t,n_h}]$, $v_t=[v_{t,1};v_{t,2};\dots,;v_{t,n_h}]$. $W^O\in\mathbb{R}^{d\times d_hn_h}$ 为 output projection 的权重。在 inference 阶段，每个 token 需要缓存其 key 以及 value 对应的值，从而每个 token 的 kv cache 占用为 $2n_hd_h\ell$. 当序列长度过大时，KV cache 会影响整体的 inference efficiency.

MQA & GQA

MQA 通过在所有的 heads 中共享 key 和 value 来实现降低 kv cache 的作用，在 MQA 中，$W^K, W^V\in\mathbb{R}^{d_h\times d}$, 在计算时，对应的 $k_t$ 和 $v_t$ 通过广播机制参与 attention 的计算。此时，KV cache 占用为 MHA 的 $1/n_h$, 即 $2d_h\ell$.

但是，MQA 的问题是表达能力太弱（表现差），因此后续 GQA 进行了改进，GQA 在 MQA 和 MHA 之间进行了权衡，即将 heads 分为若干个 group, 每个 group 中共享 key 和 value, 即 $W^K, W^V\in\mathbb{R}^{n_gd_h\times d}$, 这里 $n_g$ 是 group 个数，在计算 attention 时，key 和 value 在 group 内部共享，此时，GQA 的 KV cache 占用是 MQA 的 $n_g$ 倍，即 $2n_gd_h\ell$.

这部分具体介绍见 MQA 和 GQA.

MLA

MLA 的架构图如下所示

MLA 使用 low-rank joint compression 来压缩 key 以及 value 的 KV cache:

$$ c_t^{KV} = W^{DKV}h_t,\quad k_t^C = W^{UK}c_t^{KV}, v_t^C = W^{UV}c_t^{KV} $$

这里 $c_t^{KV}\in\mathbb{R}^{d_c}$ 为 key 以及 value 压缩后的 latent vector. $d_c<

另外，为了减少训练时的 activation memory, 作者对于 query 同样也执行了 low-rank compression, 压缩方式如下

$$ c_t^Q = W^{DQ}h_t,\quad q_t^C = W^{UQ}c_t^Q $$

其中 $c_t^Q\in\mathbb{R}^{d_c'}$ 为 query 压缩后的 latent vector, $d_c'<< d_hn_h$ 为 query compression dimension, $W^{DQ}\in\mathbb{R}^{d_c'\times d}$, $W^{UQ}\in\mathbb{R}^{d_hn_h\times d_c'}$ 分别时 down projection, up projection matrix.

最后 attention 的计算与 MHA 保持一致：

$$ \begin{aligned} o_{t,i} &= \sum_{j=1}^t\mathrm{softmax}_j\left(\frac{q_{t,i}^Tk_{j,i}}{\sqrt{d_h}}\right)v_{j,i},\\ u_t&= W^O[o_{t,1};o_{t,2};\dots,;o_{t,n_h}] \end{aligned} $$

在推理的时候，我们只需要缓存 $c_t^{KV}$ 即可，这样每个 token 的 KV cache 为 $d_c\ell$. 并且在 inference 时，我们可以将 $W^{UK}$ 和 $W^{Q}$ 融合在一起，将 $W^{UV}$ 和 $W^{O}$ 融合在一起，也就是说我们不需要显式的计算出 $k_t$ 以及 $v_t$, 即

$$ q_t^Tk_t = (W^{UQ}c_t^Q)^T(W^{UK}c_t^{KV}) = (c_t^Q)^T((W^{UQ})^TW^{UK})\boxed{c_t^{KV}} $$

以及

$$ \begin{aligned} u_t &= W^O[o_{t,1};o_{t,2};\dots,;o_{t,n_h}] \\ &= \sum_{i=1}^tW_i^Oo_{t,i}\\ &= \sum_{i=1}^tW_i^O\sum_{j=1}^t\mathrm{softmax}_j\left(\frac{q_{t,i}^Tk_{j,i}}{\sqrt{d_h}}\right)v_{j,i}\\ &= \sum_{i=1}^tW_i^O\sum_{j=1}^t\mathrm{softmax}_j\left(\frac{q_{t,i}^Tk_{j,i}}{\sqrt{d_h}}\right)W_i^{KV}c_t^{KV}\\ &= \sum_{i=1}^t(W_i^OW_i^{KV})\sum_{j=1}^t\mathrm{softmax}_j\left(\frac{q_{t,i}^Tk_{j,i}}{\sqrt{d_h}}\right)\boxed{c_t^{KV}} \end{aligned} $$

这里 $W^O = [W^O_1,\dots,W^O_{n_h}]$, $W^{UV}=[W^{KV}_1;\dots;W^{KV}_{n_h}]$, $W_i^O\in\mathbb{R}^{d\times d_h}$, $W^{UV}_i\in\mathbb{R}^{d_h\times d_c}$.

Decoupled Position Embedding

接下来，作者介绍了如何解决 RoPE 不相容的问题。如果说我们直接在 $k_t^C$ 上进行 RoPE, 那么我们有

$$ q_t^Tk_t = (R_mW^{UQ}c_t^Q)^T(R_nW^{UK}c_t^{KV}) = (c_t^Q)^T((W^{UQ})^TR_{m-n}W^{UK})\boxed{c_t^{KV}} $$

此时，我们没有办法将 $W^{UK}$ 吸收到 $W^{UQ}$ 中，这样就导致在 inference 时我们必须重新计算所有 prefix token 对应的 key, 这显然会降低 inference efficiency

为了解决这个问题，作者使用了partial RoPE的技巧，即将query和key拆解为NoPE以及RoPE两部分，前者由MLA产生，后者携带位置信息。RoPE部分包括query $q_{t,i}^R\in\mathbb{R}^{d_h^R}$ 以及一个共享的 key $k_t^R\in\mathbb{R}^{d_h^R}$, 其中 $d_h^R$ 是 decoupled query 以及 decoupled key 的 head dimension.

[!remark] 这里 key 对应的 RoPE 共享的原因是这部分信息也需要使用 KV cache 进行缓存，通过共享可以降低 KV cache 占用；而 query 对应的 RoPE 不共享的原因是提高 head 的表达能力，与 MHA 原理一致。

对应 MLA 的计算公式如下

$$ \begin{aligned} q_t^R&=\mathrm{RoPE}(W^{QR}c_t^Q)\\ k_t^R &= \mathrm{RoPE}(W^{KR}h_t)\\ q_{t,i} &= [q_{t,i}^C;q_{t,i}^R]\\ k_{t,i} &= [k_{t,i}^C;k_{t}^R]\\ o_{t,i} &= \sum_{j=1}^t\mathrm{softmax}_j\left(\frac{q_{t,i}^Tk_{j,i}}{\sqrt{d_h+d_h^R}}\right)v_{j,i}^C,\\ u_t&= W^O[o_{t,1};o_{t,2};\dots,;o_{t,n_h}] \end{aligned} $$

其中 $q_t^R=[q_{t,1}^R;q_{t,2}^R;\dots,;q_{t,n_h}^R]$, $W^{QR}\in\mathbb{R}^{d_h^Rn_h\times d_c'}$, $W^{KR}\in\mathbb{R}^{d_h^R\times d}$ . $\mathrm{RoPE}(\cdot)$ 只执行 RoPE 矩阵乘法的操作。

在这种情形下，attention 的计算如下所示

$$ \begin{aligned} q_{t,i}^Tk_{t,i} &= [q_{t,i}^C;q_{t,i}^R]^T[k_{t,i}^C;k_{t}^R]\\ &= (q_{t,i}^C)^Tk_{t,i}^C + (q_{t,i}^R)^Tk_{t}^R \end{aligned} $$

可以看到，现在 attention 的计算分为了两部分，一部分是 MLA 自身的计算，这部分计算前面已经证明可以通过矩阵吸收的方式来进行优化，第二部分是关于 RoPE 部分的计算，这部分计算量不是很大

最终，MLA 完整的计算公式如下

$$ \begin{aligned} c_t^Q=W^{DQ}h_t\\ [q_{t,1}^C;\dots;q_{t,n_h}^C]=q_t^C&= W^{UQ}c_t^Q\\ [q_{t,1}^R;\dots;q_{t,n_h}^R]=q_t^R&= \mathrm{RoPE}(W^{QR}c_t^Q)\\ q_{t,i} &= [q_{t,i}^C;q_{t,i}^R]\\ \boxed{c_t^{KV}} &= W^{DKV}h_t\\ [k_{t,1}^C;\dots;k_{t,n_h}^C]=k_t^C&= W^{UK}c_t^{KV}\\ \boxed{k_t^R} &= \mathrm{RoPE}(W^{KR}h_t)\\ k_{t,i} &= [k_{t,i}^C;k_{t}^R]\\ [v_{t,1}^C;\dots;v_{t,n_h}^C]=v_t^C&= W^{UV}c_t^{KV}\\ o_{t,i} &= \sum_{j=1}^t\mathrm{softmax}_j\left(\frac{q_{t,i}^Tk_{j,i}}{\sqrt{d_h+d_h^R}}\right)v_{j,i}^C,\\ u_t&= W^O[o_{t,1};o_{t,2};\dots,;o_{t,n_h}] \end{aligned} $$

在 inference 时，decoupled key 也需要被缓存，因此 DeepSeek-V2 每个 token 所需要的 KV cache 为 $(d_c+d_h^R)\ell$, 框选的部分即为 Inference 阶段需要缓存的内容

MLA 与 MHA, MQA, GQA 的对比如下图所示

Comparison of KV Cache

接下来，作者对比了不同 attention 机制的 KV cache, 结果如下表所示

Attention Mechanism	KV Cache per Token (# Element)	Capability
Multi-Head Attention (MHA)	$2n_hd_h\ell$	Strong
Grouped-Query Attention (GQA)	$2n_gd_h\ell$	Moderate
Multi-Query Attention (MQA)	$2d_h\ell$	Weak
MLA (Ours)	$(d_c+d_h^R)\ell\approx 9/2d_h\ell$	Stronger

这里作者将 $d_c$ 设置为 $4d_h$, $d_h^R$ 设置为 $d_h/2$, 因此得到了上面的 $9/2d_h\ell$ 的近似。与 GQA 相比，相当于 MLA 使用了 2.25 个 group, 但是可以得到更强的效果。

为了避免 low-rank compression 以及 fine-grained expert segmentation 对输出的 scale 产生影响，作者对 compressed latent vectors $c_t^Q, c_t^{KV}$ 进行了 normalization.

Code

首先是代码变量与公式变量的对应关系

code name	variable name	Value
`hidden_size`	$d$	5120
`kv_lora_rank`	$d_c$	512
`q_lora_rank`	$d_c'$	1536
`qk_nope_head_dim`	$d_h$	128
`qk_rope_head_dim`	$d_h^R$	64
`v_head_dim`	$d_h$	128
`num_attention_heads`	$n_h$	128

在具体实现时，作者对计算过程进行了优化，具体就是先合并计算然后通过 split 进行拆分，这部分策略应用于三个部分：

$$ \begin{aligned} \begin{bmatrix} q_t^c\\ q_t^R \end{bmatrix} &= \begin{bmatrix} W^{UQ}\\ W^{QR} \end{bmatrix}W^{DQ}h_t\\ \begin{bmatrix} c_t^{KV}\\ k_t^R \end{bmatrix} &= \begin{bmatrix} W^{DKV}\\ W^{KR} \end{bmatrix}h_t\\ \begin{bmatrix} k_t^c\\ v_t^c \end{bmatrix} &= \begin{bmatrix} W^{UK}\\ W^{UV} \end{bmatrix}c_t^{KV} \end{aligned} $$

代码如下所示

 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


class DeepseekV2Attention(nn.Module):
 def __init__(self, config, layer_idx):
 # d_h + d_h^R
 self.q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim
 # W^{DQ}
 self.q_a_proj = nn.Linear(
 self.hidden_size, config.q_lora_rank, bias=config.attention_bias
 )
 # [W^{UQ}; W^{QR}]
 self.q_b_proj = nn.Linear(
 config.q_lora_rank, self.num_heads * self.q_head_dim, bias=False)
 # [W^{DKV}; W^{KR}]
 self.kv_a_proj_with_mqa = nn.Linear(self.hidden_size, config.kv_lora_rank + config.qk_rope_head_dim, bias=config.attention_bias)
 # [W^{UK}; W^{UV}]
 self.kv_b_proj = nn.Linear(config.kv_lora_rank, self.num_heads * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim), bias=False)
 # W^O
 self.o_proj = nn.Linear(self.num_heads * self.v_head_dim, self.hidden_size, bias=config.attention_bias)

 def forward(self, hidden_states, ...):
 # [q_t^c; q_t^R]
 q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states)))
 q = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2)
 q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)

 # [c_t^{KV}; k_t^R]
 compressed_kv = self.kv_a_proj_with_mqa(hidden_states)
 compressed_kv, k_pe = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
 k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2)

 # [k_t^c; v_t^c]
 kv = (self.kv_b_proj(self.kv_a_layernorm(compressed_kv)).view(bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim).transpose(1, 2))
 k_nope, value_states = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)

 # q_t^R, k_t^R
 q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids)

 # q_{t, i}
 query_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim)
 query_states[:, :, :, : self.qk_nope_head_dim] = q_nope
 query_states[:, :, :, self.qk_nope_head_dim :] = q_pe

 # k_{t, i}
 key_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim)
 key_states[:, :, :, : self.qk_nope_head_dim] = k_nope
 key_states[:, :, :, self.qk_nope_head_dim :] = k_pe

 # Q^TK
 attn_weights = (torch.matmul(query_states, key_states.transpose(2, 3)) * self.softmax_scale)
 # softmax(...) in FP32
 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)

 # o_{t, i}
 attn_output = torch.matmul(attn_weights, value_states)

 # u_t
 attn_output = self.o_proj(attn_output)

 return attn_output, attn_weights, past_key_value

参数量计算

首先，我们结合 DeepSeek-V2 的 config 计算一下 MLA 部分的参数量：

Matrix	Parameters	values	ratio
$W^{DKV}$	$dd_c$	2621440	1.91%
$W^{UK}$	$d_hn_hd_c$	8388608	6.12%
$W^{UV}$	$d_hn_hd_c$	8388608	6.12%
$W^{DQ}$	$d_c'd$	7864320	5.74%
$W^{UQ}$	$d_hn_hd_c'$	25165824	18.37%
$W^{KR}$	$d_h^Rd$	327680	0.24%
$W^{QR}$	$d_h^Rd$	327680	0.24%
$W^{O}$	$dd_hn_h$	83886080	61.24%
Total	$d(d_c+d_c'+2h_h^R)+d_hn_h(2d_c+d_c'+d)$	136970240	100%

我们接下来对比一下各个模型架构之间 attention 部分的参数量，可以看到与 MHA 一致，大部分参数量都集中在最后的 Output projection layer 上

Experiments

作者首先对比了 MHA, GQA, MQA 的表现，作者基于一个 7B 的 dense 模型，使用 1.33T token 进行训练，实验结果如下

Benchmark (Metric)	# Shots	MQA	GQA(8 Groups)	MHA
# Params	-	7.1B	6.9B	6.9B
BBH (EM)	3-shot	33.2	35.6	37.0
MMLU (Acc.)	5-shot	37.9	41.2	45.2
C-Eval (Acc.)	5-shot	30.0	37.7	42.9
CMMLU (Acc.)	5-shot	34.6	38.4	43.5

实验结果显示，MHA 的表现显著优于 GQA 和 MQA. 这说明了 MQA 和 GQA 虽然减少了 KV cache 的占用，但是相应地，它们对应的表现也有所降低。

接下来，作者对比了 MLA 和 MHA 的表现，实验结果如下

Benchmark (Metric)	# Shots	MHA	MLA	MHA	MLA
# Activated Params	-	2.5B	2.4B	25.0B	21.5B
# Total Params	-	15.8B	15.7B	250.8B	247.4B
KV Cache per Token (# Element)	-	110.6K	15.6K	860.2K	34.6K
BBH (EM)	3-shot	37.9	39.0	46.6	50.7
MMLU (Acc.)	5-shot	48.7	50.0	57.5	59.0
C-Eval (Acc.)	5-shot	51.6	50.9	57.9	59.2
CMMLU (Acc.)	5-shot	52.3	53.4	60.7	62.5

可以看到，MLA 的表现比 MHA 的表现更好，并且 KV cache 也更少。

Conclusion

作者提出了 MLA, 一个基于 low rank compression 的注意力机制，通过将 key value vector 压缩到低维空间，MLA 可以有效降低 Inference latency, 作者通过实现证明 MLA 的表现可以与 MHA 相比，并且 KV cache 更小。

References

Notes on Loss-free Balancing

Fri, 21 Nov 2025 15:38:52 +0800

Introduction

已有的 MoE 模型往往都会使用 load balancing loss 来避免 Imbalanced routing, 但是加入这个额外的损失之后，模型训练的梯度也会受到影响。

DeepSeek 基于这个问题提出了 Loss-Free Balancing, 该方法不引入额外的 loss item, 而是在 routing 的结果上加入一个 bias item, bias item 可以根据 expert load 来动态更新进而实现 load balancing.

Method

MoE definition

$$ \begin{aligned} h_t &= u_t + \sum_{i=1}^Ng_{i,t}\mathrm{FFN}_i(u_t),\\ g_{i,t} &= \begin{cases} s_{i,t}, & s_{i,t}\in\mathrm{Topk}(\{s_{i,j}\mid 1\leq j\leq N\}, K)\\ 0, &\text{otherwise} \end{cases}\\ s_{i,t} &= G(u_t^Te_i), \end{aligned} $$

其中 $G$ 是 gating function, $e_i$ 是第 $i$ 个专家对应 gating function 的权重

load balancing loss 有两个作用：

避免 routing collapse, 即模型只选择固定的少数专家完成任务
减少通信开销

但是引入 load balancing loss 会对 LLM 的训练产生影响，为了避免对模型性能造成影响，我们需要小心设置 load balancing loss 的权重，实验结果如下图所示

为了解决这个问题，作者提出了 Loss-Free Balancing, 具体做法就是在每个专家的 gating score 上加入一个 bias term, 然后再决定对应的专家：

$$ g_{i,t} = \begin{cases} s_{i,t}, & s_{i,t}+b_i\in\mathrm{Topk}(\{s_{i,j}+b_j\mid 1\leq j\leq N\}, K)\\ 0, &\text{otherwise} \end{cases} $$

注意这里的 bias item 仅影响 top-K 操作，其对最终的输出没有影响。

为了实现负载均衡，作者根据上一个 batch 的 expert load 情况来调整 bias item, 如果某一个专家的 load 太大，则对应的 bias item 会变小。其算法实现过程如下：

作者对比不同的负载均衡算法如下表所示

Load Balancing Methods	Balanced Expert Load	Interference Gradients	Future Token Leakage
Loss-Controlled (strong auxiliary loss)	balanced	strong	no leakage
Loss-Controlled (weak auxiliary loss)	imbalanced	weak	no leakage
Expert Choice	balanced	none	with leakage
Loss-Free (Ours)	balanced	none	no leakage

Experiments

作者使用了 DeepSeekMoE 来进行实验，作者使用了 sigmoid function 作为 gating function, 因为作者发现 sigmoid function 效果比 softmax 效果更好。

作者提出了 maximal violation (MaxVio) 来量化一个 MoE layer 的负载均衡程度

$$ \mathrm{MaxVio} = \frac{\max_i\mathrm{Load}_i-\overline{\mathrm{Load}_i}}{\overline{\mathrm{Load}_i}} $$

其中 $\mathrm{Load}_i$ 代表了分配给第 $i$ 个专家的 token 个数，$\overline{\mathrm{Load}_i}$ 代表了理想情况下的负载均衡。

实验结果如下图所示

Model Size	Load Balancing Methods	Validation Perplexity	MaxVio_global
1B	Loss-Controlled	9.56	0.72
1B	Loss-Free	9.50	0.04
3B	Loss-Controlled	7.97	0.52
3B	Loss-Free	7.92	0.04

实验结果显示，本文提出的 Loss-Free Balancing 效果更好，且负载均衡更高

作者还展示了训练过程的负载情况如下

可以看到，Loss-Free 的负载一直比 load balancing loss 效果更好

Ablation Study

作者首先探究了 bias term 更新速率对模型表现的影响，结果如下图所示

结果显示，使用过大的 update rate 会影响最终的负载均衡，而较小的 update rate 收敛速率比较慢。因此作者在本文中使用了 $u=0.001$ 这个设置

作者还对比了 baseline model 使用 softmax 个 sigmoid gating 两种方式，结果如下图所示

实验结果显示，Sigmoid function 对于超参数更加 robust, 且表现也更好一些。

作者还尝试了不同的 bias 更新方式，结果显示尽管不同的更新方式有可能会提高 load balance, 但是最终模型表现提升不大。

Conclusion

作者在本文中提出了 Loss-Free Balancing 策略，一个针对 MoE 负载均衡而不需要额外损失项的方法

References

Arxiv

Mixstral 8x7B

Sat, 01 Nov 2025 15:32:30 +0800

Introduction

作者在本文中提出了 Mixtral 8x7B, 一个 MoE 模型，模型上下文为 32K. 作者还对模型进行 finetune 得到了 Mixtral 8x7B-Instruct, finetuning 包含 SFT 和 DPO 两个阶段。

Method

模型架构与 Mistral-7B 基本相同，参数如下表所示

Parameter	Value
`dim`	4096
`n_layers`	32
`head_dim`	128
`hidden_dim`	14336
`n_heads`	32
`n_kv_heads`	8
`window_size`	4096
`context_len`	32768
`vocab_size`	32000
`num_experts`	8
`top_k_experts`	2

MoE 的架构与 GShard 基本一致

Results

作者探究了专家的 specialization, 结果有三点发现：

不同专家对于不同 domain 的数据并没有出现 specialization
在 math domain 上，专家的分布有一个明显的区别。
连续的 token 往往会被分配到同一个专家上

Conclusion

作者在本文提出了 Mistral 8x7B, 一个 MoE 大语言模型

References

Arxiv

Mixstral 7B

Sat, 01 Nov 2025 15:28:19 +0800

Introduction

作者主要强调 Mistral 7B 的表现超过了 LLaMA 2 7B 和 LLaMA 34B 的表现。

Mistral 7B 主要使用了 GQA 以及 SWA 两个方法来加速推理和减少内存占用，进而提高 batch size 和 throughput.

Method

Mistral 7B 的模型配置如下表所示

Parameter	Value
`dim`	4096
`n_layers`	32
`head_dim`	128
`hidden_dim`	14336
`n_heads`	32
`n_kv_heads`	8
`window_size`	4096
`context_len`	8192
`vocab_size`	32000

SWA 的示意图如下所示

对于第 $k$ 层，模型可以感知到 $W\times k$ 的 tokens, 进而可以在提高训练效率的同时保持模型的表现。作者发现，在 flash attention 中使用 SWA 之后，模型的效率提升了 2 倍左右。

并且使用了 SWA 之后，我们的 kv cache 也就随之固定了，因此我们可以使用一个 rolling buff cache, 其大小为 $W$, 对于第 $i$ 个 token, 我们将其保存在 cache 中的第 $i \% M$ 个位置。

作者还进一步将 sequence 分割为多个 chunk, 每个 chunk 的大小都是 window size $M$, 这样在计算 attention 的时候，对于当前的 chunk, 我们使用 self-attention, 对于 cache 中的 attention, 我们使用 SWA, 然后对于 past token, 由于这部分不在 sliding window 内因此不参与计算

Conclusion

作者在本文中提出了 Mistral 7B, 一个基于 GQA 和 SWA 的大语言模型。

References

Arxiv

Notes on olmoe

Sat, 01 Nov 2025 15:23:58 +0800

Introduction

已有的模型如 DeepSeek-MoE, Mixtral MoE, Qwen1.5 等 MoE 模型基本只开源权重。也有一些开源的模型，比如 OpenMoE 等，但是开源信息不全。基于这个目的，作者提出了 olmoe 模型系列，包括 olmoe-7B-A1B 和 olmoe-7B-A1B-instruct 两个版本。

Method

Pretraining

模型的架构如下图所示，MoE 架构与 dense 架构不同的地方在于 decoder layer 中的 FFN 被替换为了 MoE layer.

模型的配置如下表所示

训练的目标函数为

$$ \mathcal{L} = \mathcal{L}_{CE} +\alpha\mathcal{L}_{LB} +\beta\mathcal{L}_{RZ} $$

其中 $\alpha,\beta$ 为系数， $\mathcal{L}_{CE}$, $\mathcal{L}_{LB}$ 以及 $\mathcal{L}_{RZ}$ 分别代表 cross-entropy loss, load balancing loss 以及 routing Z loss.

预训练数据包括 DCLM 和 Dolma1.7 两个数据集的混合，作者将预训练数据集称为olmoe-mix. 数据集的配比如下

Post-training

在 post-training 时，作者将训练分为 instruction tuning 和 preference tuning 两个阶段，在 instruction dataset 中，作者加入了更多的代码和数学数据来提高对应的能力。数据集如下表所示

Ablation Study

MoE Settings

MoE vs. Dense

作者对比了 MoE 模型和 dense 模型的训练效率，为了方便对比，作者使用 olmo-7B 和 olmo-1B 作为 baseline, 最终 olmoe 的总参数为 6.9B, 激活参数为 1.3B. 实验结果如下图所示

实验结果发现，MoE 模型所需要的 token 或者 FLOPs 是 dense 模型的 $1/3$, 但是由于 MoE 模型需要额外的内存开销，因此从训练时间上来看，MoE 模型训练时间仅比 dense 模型快 $2$ 倍左右。

Expert Granularity

DeepSeekMoE 提出使用细粒度的专家来提供更多的组合可能性。作者探究了不同的粒度对模型表现的影响，结果如下图所示

结果显示，当专家粒度从 8E-1A 扩展到 32E-4A 时，模型在 HellaSwag 上的表现提升了 $10\%$, 但是进一步扩展到 64E-8A 时，模型的表现提升不到 $2\%$, 这说明了无限制提升粒度对模型的提升越来越有限。在本文中，作者使用了 64 个专家。

Shared Experts

DeepSeekMoE 提出使用共享专家来学习 common knowledge, 作者对这种方法进行了实验，结果如下图所示。

可以看到，加入一个 shared expert 之后，模型的表现没有变化，作者认为减少 routed expert 之后，模型的组合可能性降低为原来的 $10\%$ 左右。因此作者认为没有必要使用共享专家，因此作者在 olmoe 中没有采用共享专家这个方法。

Expert Choice vs. Token Choice

作者探究了 routing 的策略，一个是 expert choice (EC), 另一种是 token choice (TC), 分别代表了每个 expert 选取固定的 token 数和每个 token 选取固定的 expert 数这两种情况。实验结果如下图所示

可以看到，token choice 的表现明显更好。EC 虽然可以实现负载均衡。但是因为自回归模型在生成时是无法提前确定生成的 token 数的，因此 EC 很可能导致算力资源浪费或者是 token dropping. 在本文中，作者采用了 TC 这种策略。

Sparse Upcycling

作者还对比了从零开始训练 MoE 与基于 dense model upcycling 的方式训练 MoE，sparse upcycling 的相关工作有 MiniCPM, Qwen2 以及 Mixtral MoE.结果如下图所示

结果发现，upcycling 确实可以提高训练效率，但是这种方法的缺陷在于：

upcycling 受 dense model 的超参数限制
upcycling 的训练不是很稳定

因此在本文中作者没有采取 upcycling 的做法。

Load Balancing Loss

作者还探究 Load Balancing loss 对模型表现的影响，结果如下图所示

可以看到，加入 load balancing loss 之后，模型的表现均超过了不加时的表现。

作者进一步分析了不同专家在加/不加 load balancing loss 时的激活情况，结果如下图所示

结果显示，load balancing loss 确实可以让不同专家的激活概率大致相当。

Router Z-loss

ST-MoE 提出了 Router z-loss 来提高 MOE 训练的稳定性和表现。其表达式如下所示

$$ \mathcal{L}_{RZ}(x) = \frac{1}{B}\sum_{i=1}^B\left(\log\sum_{j=1}^{N_E}\exp(x_i^{(i)})\right)^2 $$

实验结果如下图所示

可以看到，加入 router Z-loss 之后，模型训练的稳定性有所提升。因此在本文中作者使用了这个 loss.

General Pre-training Settings

Initialization

作者探究了不同初始化策略对模型训练的影响，结果发现使用 truncate normal initialization 的训练稳定性更高

QK-Norm

作者探究了 QK-norm 对模型训练的影响，结果发现 QK-norm 可以提高模型训练的稳定性，结果如下图所示

AdamW Epsilon

作者发现，在 AdamW 优化器中，使用更小的 eps 可以提高模型的表现，因此作者将 eps 设置为 $1e-8$.

Adaptation Settings

在 post-training 阶段，作者在三个方面进行了实验：

是否加入 load balancing loss: 结论是不加，因为负载均衡在 pre-training 阶段已经实现了
是否使用 annealing: 结论是使用，因为效果更好
使用 DPO 还是 KTO, 结论是两种方法结果差不多

实验结果如下表所示

Load Balancing Precision

Switch Transformer 中提出使用 float32 精度来进行 routing 的计算，作者通过实验发现，这一方法并不能提高模型训练的稳定性，因此作者没有采用这一策略。

MoE Analysis

Router Saturation

作者探究了训练过程中激活的专家和训练结束后激活的专家的匹配程度，结果如下图所示

结果发现，训练 $1\%$ 的数据之后，就有 $40\%$ 的 routing 和训练完毕的 routing 一致，当训练 $40\%$ 的数据之后，这个比例提升到了 $80\%$.

作者还发现，later layers 比 early layers 饱和更快，early layer, 特别是 layer 0, 饱和的非常慢。作者认为，这是 DeepSeekMoE 放弃在第一层使用 MoE layer 的原因，因为 load balancing loss 收敛更慢。

Expert Co-activation

作者分析了 expert 之间的互相依赖程度，作者通过可视化发现，不同的 expert 之间 co-activation 的比例比较小，说明 expert redundancy 比较低

Domain Specialization

作者还探究了不同 expert 对于不同 domain 的 specialization 程度，作者发现对于 specialized domain 的数据，expert 会出现一定程度的 specialization, 但是对于通用 domain 的数据，expert 的 specialization 程度比较低。这个结论与 Mixtral MoE 的结论不同，作者认为这个原因是 Mixtral MoE 使用了 upcycling 的方式，这会限制模型的表现。因此，作者进一步强调 MoE 从零开始训练是一个更好的训练方式。

Vocabulary Specialization

作者还探究了 vocabulary 中不同 token index 与激活专家之间的关系，结果发现 later layers 的 specialization 程度更高，这与 saturation 的趋势一致

Conclusion

作者在本文中提出了 olmoe, 一个全开源的 moe 大模型系列，作者详细介绍了针对 MoE 架构和通用架构的设计，为后来的模型架构设计提供了基础。

References

arxiv

GShard

Wed, 29 Oct 2025 11:22:39 +0800

Introduction

作者认为，训练大规模的模型存在如下问题：

缺乏有效的 Model parallelism 算法
随着设备数的增加，训练时间与 model size 呈现超线性增长的关系
(tensorflow) 在 nodes 数增多时，构建所需要的时间也大幅度增长
在多个设备上 partition model 比较困难

作者在本文中构建了一个基于 sparse MoE 架构的 600B 模型。为了解决这些问题，作者作出了如下贡献：

作者提出了基于 MoE 架构的模型，来减少计算和通信开销
作者提出了 Gshard, 来自动化实现并行
作者提出了 SPMD (single program multiple data) 来减少计算表示和编译的难度

Method

架构如下图所示

其中激活专家个数为 2 个，MoE layer 和 FFN layer 交替出现。

作者对 GATE 函数进行了如下优化：

expert capacity. 每个 expert 都有一个处理 token 的上限值，超过该上限值之后 GATE 直接输出 0 进入下一层，假设 batch size 为 $N$, 专家个数为 $E$, 则该阈值定义为 $N/E$
group dispatching. 将 token 拆分为 $G$ 个 group, 每个 group 并行处理，每个 group 里每个专家处理的 token 个数上限为 $N/(GE)$.
Auxiliary loss. 作者使用了 Load Balancing loss 来实现负载均衡
Random routing. 作者选取了 top-2 的专家，当第二个专家的权重太小是，作者直接忽略第二个专家，简化为选取 top-1 的专家

算法运行如下所示

Parallel Implementation

第一步是将算法转化为线性代数的方式，算法的代码如下所示

1
2
3
4
5
6
7


gates = softmax(einsum("GSM, ME->GSE", inputs, wg))
combine_weights, dispatch_mask = Top2Gating(gates)
dispatched_expert_inputs = einsum("GSEC, GSM->EGCM", dispatch_mask, reshaped_inputs)
h = einsum("EGCM, EMH->EGCH", dispatched_expert_inputs, wi)
h = relu(h)
expert_outputs = einsum("EGCH, EHM->GECM", h, wo)
outputs = einsum("GSEC, GECM->GSM", combine_weights, expert_outputs)

第二步是通过 API 来实现并行执行

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


# Partition inputs along group (G) dim. 
inputs = split(inputs, 0, D)
# Replicate the gating weights
wg = replicate(wg)
gates = softmax(einsum("GSM, ME->GSE", inputs, wg))
combine_weights, dispatch_mask = Top2Gating(gating_logits)
dispatched_expert_inputs = einsum("GSEC, GSM->EGCM", dispatch_mask, reshaped_inputs)
# Partition dispatched inputs along expert (E) dim.
dispatched_expert_inputs = split(dispatched_expert_inputs, 0, D)
h = einsum("EGCM, EMH->EGCH", dispatched_expert_inputs, wi)

第三部就是基于 compiler infra 来基于 sharding annotation 自动化分割一个 computation graph.

Experiments

作者在机器翻译的任务上训练了若干模型，结果发现：

层数更多的模型表现更好
提高 expert capacity 有效提高模型的表现
使用更多的 expert 可以在 high-resourced 任务上提高表现
dense 模型相比于 MoE 模型拥有更强的迁移能力

从训练效率上来看

层数更多的模型的 sample efficiency 也更高
600B 的模型也可以在 4 天之内训练完毕

从内存使用效率上来看

层数相同时，weight memory 以及 activation memory 不随专家个数增加而增加
专家个数比较少（128）时，模型可以达到 roofline performance 的 $70\%$, 专家个数比较多（2048）时，模型依然可以达到 roofline performance 的 $48\%$.

Conclusion

作者探究了如何提高大规模模型的训练效率，作者提出了基于 MoE 架构的模型，然后作者还设计了 GShard, 一个自动化分割大规模模型的深度学习模块。实现结果发现，增加模型 size 可以提高模型的表现。作者还发现，SPMD 是一个更好的提高计算效率的方法。

Reference

arxiv

ST-MoE

Wed, 29 Oct 2025 11:19:37 +0800

Introduction

已有的工作如 GShard 和 Switch Transformer 均验证了 MoE 模型的有效性，但是他们的问题在于训练不稳定。

因此，在本文中，作者就提出了 ST-MoE 来解决这个问题。作者的主要贡献如下：

探究了如何平衡模型的表现与训练稳定性
提出了 router Z-loss 来解决训练的不稳定性
探究了如何设定 MoE 模型 fine tuning 时的超参数
对 MoE 的性质进行了分析

Training Stability

作者发现，直接训练 MoE 模型会面临训练不稳定性的问题，实验结果如下图所示

可以看到不同的实验，一次训练崩溃，一次训练正常。

作者接下来探究了如何提高训练的稳定性，作者有三点发现：

大多数方法都可以提高训练稳定性，但是也会是的模型表现更差
router Z-loss 可以在不损失模型表现的情况下提高训练的稳定性
multiplicative components 如 RMSNorm 等可以提高模型表现，但是训练稳定性更差

作者构建了一个 baseline 的 MoE 模型，总专家个数为 $N=32$, 激活专家个数为 $K=2$, Transformer block 按照 4 个 block 为一组，每组里包含一个 MoE layer. 为了避免初始化带来的误差，作者使用 6 个随机数种子进行训练。

Multiplicative Components

multiplicative components 指的是 GEGLU 和 RMSNorm 这些操作，其实验结果如下所示

Method	Fraction Stable	Quality
Baseline	4/6	$-1.755 \pm0.02$
Remove GEGLU	3/3	$-1.849 \pm0.02$
Remove RMS Norm. Scale Param	3/3	$-2.020 \pm0.06$

结果发现，移除掉 multiplicative components 之后，模型训练变得稳定，但是效果也变差了。

Adding Noise

接下来作者尝试了 Switch Transformer 中的 jitter noise 和 dropoout 方法，实验结果如下图所示

Method	Fraction Stable	Quality
Baseline	4/6	$-1.755 ± 0.02$
Input jitter ($10^{-2}$)	3/3	$-1.777 \pm 0.03$
Dropout (0.1)	3/3	$-1.822 \pm0.11$

实验结果显示，加入噪声对模型的表现存在负面影响。

Constraining Activations

作者接下来分析了以下 Switch Transformer 中 router 存在的问题，作者发现尽管在 router 中使用 float32 可以提高训练稳定性，但是这还不够，因此作者使用了如下的 router Z-loss:

$$ \mathcal{L}_z(x) = \frac1B\sum_{i=1}^B\left(\log\sum_{j=1}^N e^{x_j^{(i)}}\right)^2 $$

其中 $B, N$ 分别是 batch size 和专家个数，$x_j^{(i)}$ 代表了第 $j$ 个专家对 $i$ 个 token 的激活 logits. 通过增加这个 loss, 我们可以让 routing layer 的 logits 都在一个合理的范围内。

Method	Fraction Stable	Quality ($\uparrow$)
Baseline	$4/6$	$-1.755 \pm 0.02$
Update clipping (clip = 0.1)	$3/3$	$-4.206 \pm 0.17$
Router Z-Loss	$3/3$	$-1.741 \pm 0.02$

可以看到，加入 router Z-loss 之后，模型的稳定性和表现都有了提升。

因此，在本文中，作者使用的损失函数为

$$ \mathcal{L} = \mathcal{L}_{CE} + \alpha \mathcal{L}_B + \beta \mathcal{L}_z $$

其中 $\mathcal{L}_{CE}, \mathcal{L}_B$ 分别是 cross-entropy loss 和 Load Balancing loss, $\alpha,\beta$ 是对应 loss 的权重。

Numerical Precision

接下来作者探究了数值精度对训练效率和稳定性的影响，作者发现使用低精度进行训练的优势有：

通信开销更小
计算消耗更小
内存需求更小

但是低精度进行训练的问题是存在严重的 roundoff error, 其示意图如下所示

Fine-tuning

作者在本节探究了如何 fine-tune 一个 MoE 模型。

作者首先通过实验给出了 MoE 模型容易过拟合的结论，实验结果如下图所示

可以看到，在两个人呢误伤，MoE 模型相比于 dense 模型都更容易过拟合。

接下来作者探究了超参数特别是 batch size 和 learning rate 对模型表现的影响，结果如下图所示

可以看到，与 dense 模型相反，MoE 模型需要更大的 batch size 以及更小的 learning rate.

Design Sparse Models

首先作者分析了专家个数 $N$ 的设置，作者认为当专家个数超过 $256$ 之后，带来的收益就微乎其微了。并且，当专家个数增加之后，虽然算力没有变化，但是通信开销以及计算优化会变得更加困难。作者还认为，在 TPU 上，每个 core 应该进放置一个 expert

作者还对 capacity factor 进行了实验，实验结果显示，提升 capacity factor 可以提高模型的表现。但是同时也会带来计算开销，从而让模型训练更慢，实验结果如下图所示

Model	Train CF	Step Time (s) ($\downarrow$)
ST-MoE-L	1.25	$2.397$
ST-MoE-L	2.0	$2.447 (+7\%)$
ST-MoE-32B	1.25	$4.244$
ST-MoE-32B	2.0	$4.819 (+14\%)$

最终结论如下：

使用 top-2 routing, 即激活 $K=2$ 个专家，capacity factor 设置为 $1.25$, 每个 core 上最多放置一个专家

在评估时可以动态调整 capacity factor

Tracing Tokens through the Model

在这一节里，作者分析了 MoE 模型专家的性质。

首先，作者发现，每一层里，至少有一个专家对于 sentinel token 比较敏感。并且，encoder 里的专家 specialization 更强，比如某些专家会负责 punctuation, verbs, proper names 等

接下来，作者发现，专家的 specialization 在 decoder 里几乎没有，作者认为这是因为 target token distribution 不同导致的，即：

decoding 是只有一小部分 token 被一个专家处理
decoding 过程中大部分 token 都是 sentinel token

因此，每个 group 相比于 encoder 来说通常只 cover 一小部分的 semantic space

encoder 和 decoder 专家的 specialization 实验结果如下

	Layer 1	Layer 2	Layer 3	Layer 4	Layer 5	Layer 6	Uniform (32-experts)
Encoder	2.2	1.8	1.6	1.7	1.7	1.2	3.5
Decoder	3.4	3.4	3.4	3.4	3.4	3.4	3.5

可以看到，decoder 的专家分布和均匀分布差不多，这说明 decoder 里的专家 specialization 更弱。

作者还探究了不同专家对于不同语言的敏感程度，结果发现，专家对于不同语言并没有出现明显的 specialization 现象，作者分析原因认为，输入一般只含有一种语言，因此各个专家均需要去处理对应语言的 token, 这就导致了专家对于不同语种的数据处理基本上没有太大差别。

Conclusion

在本文中，作者提出了 ST-MoE 系列 MoE 大语言模型，作者通过实验优化了 MoE 模型训练的不稳定性问题。

References

Arxiv

Switch Transformer

Tue, 28 Oct 2025 09:38:12 +0800

Introduction

Kaplan scaling law 和 Chinchilla scaling law 已经给出了针对 dense transformer model 的 scaling law, 即在给定算力下如何找到最优的 model size 和 dataset size. 已有的工作如 T5 主要是扩展 dense model 来达到更高的表现, 但是他们的问题是对算力要求比较高。

为了解决这个问题，作者尝试在不增加算力的情况下提升模型的参数量，为了达到这个目的，作者构建了基于 MoE 架构的模型 Switch Transformer.

作者的主要贡献如下：

提出了基于 MoE 架构的 Switch Transformer model
探究了针对 MoE 架构的 scaling law
将 MoE model 的能力蒸馏到 small dense model 里去
若干提升训练效率和稳定性的技巧

Architecture

Switch Transformer 的架构如下图所示

MoE

MoE 的定义见 MoE tutorial, 我们假设有 $N$ 个专家，其中激活 $K$ 个专家。

之前的工作认为我们只有在激活 $>2$ 个专家时，模型才能够比较好的训练，但是在本文中，作者决定只使用 1 个 expert, 也就是 $K=1$. 作者将激活一个专家的 layer 称为 Switch layer.

作者认为 Switch Layer 有三个优势：

router computation 现在只需要将每个 token route 到 1 个 expert
每个专家的 capacity 更小，负载更加均衡
routing 的实现更简单，且通信开销也降低了

Efficient Sparse Routing

作者首先定义了expert capacity, 也就是每个 expert 处理的 token 数量，其定义如下

$$ \text{expert capacity} = \left(\frac{\text{tokens per batch}}{\text{number of experts}}\right) * \text{capacity factor} $$

其示意图如下所示

提升 capacity factor 可以减少 token overflow 的概率，但同时也会导致计算和内存的浪费。作者通过实验发现应该尽可能降低 dropped token 比例。

为了平衡每个 expert 处理 token 的个数，作者设计了 load balancing loss, 见 Load Balancing loss 来要求每个 expert 处理的 token 数基本一致。

Parallelism

作者在本节介绍了针对 MoE 模型的并行策略，如下图所示

这里，我们给定 notation 如下

Term	Description
$B$	Number of tokens in the batch.
$N$	Number of total cores.
$n$	Number of ways for data-parallelism sharding.
$m$	Number of ways for model-parallelism sharding.
$E$	Number of experts in Switch layers.
$C$	Expert capacity, the batch size of each expert.

作者将一个 Logical mesh 分为两个维度，一个是 data-parallel sharding, 用 $n$ 表示，另一个是 model-parallel sharding, 用 $m$ 表示，从而总的 cores 数为 $N=n\times m$.

Data Parallelism

对于 data prallelism 来说，每个 core 上的模型是一样的，而一个 batch 的数据被切分为了 $B/n$ 份，此时 $n=N$, $m=1$. 这种并行策略的好处是仅在 forward 和 backward 完成之后才会进行一次通信。

Model Parallelism

对于 model parallelism, 每个 core 上的数据都是全部数据（通过拷贝得到），模型被 shard 到所有 core 上，此时 $n=1,m=N$, 这种情况下，每次 forward 或者 backward，都需要进行 $N$ 次通信（假设我们使用 pipeline parallelism），这种并行策略的好处是每个 core 上的计算压力小了，但是缺点是通信开销多了。

Model and Data Parallelism

第三种总策略是同时进行 model parallelism 和 data parallelism, 这种情况下，每个 core 上保存 $B/n$ 的数据以及 shard 为 $m$ 份的 sharding weight.

Expert and Data Parallelism

第四种策略是同时进行 expert parallelism 和 data parallelism, 作者在这里假设 $n=N$ 且 $E=n=N$, 也就是每个 core 上保留 $B/n$ 的数据, 然后还有一个对应的专家。

首先，对于输入大小为 $[B,d]$ 的 tensor, 我们会进行拆分得到大小为 $[n, B/n, d]$ 的 tensor, 代表 $n$ 个设备上分别存储 $[B/n, d]$ 的数据，首先我们计算路由权重，得到

$$ [n,B/n, d]\times [d, E] \to [n, B/n, E] $$

权重的每个值 $[i,j,k]$ 代表了第 $i$ 个 core 上,第 $j$ 个 token 分配一个第 $k$ 个专家的概率，我们对其进行 softmax 与 top-K 操作之后，就得到了对应的专家（注意 Switch Transformer 中的 $K=1$, 我们通过 one-hot 编码将其输出转化为一个二进制矩阵，大小为 $[n,B/n, E, C]$, 这里的每个值 $[i,j,k,l]$ 代表了第 $i$ 个 core 上,第 $j$ 个 token 分配给第 $k$ 个专家第 $l$ 个 token, 接下来我们再基于输入 [n,B/n,d] 计算 core 到 expert 的数据，其大小为 [n,E,C,d], 计算方式为

$$ \mathrm{einsum}([n,B/n, d], [n,B/n,E,C], \mathrm{dim}=[B/n]) $$

里面的元素 [i,j,k,:] 代表了第 $i$ 个设备路由到第 $j$ 个专家的第 $k$ ($k

然后我们就可以执行 all-to-all 通信来把对应的 token 传输给对应专家所在的设备上了，通信完成之后，每个设备上的专家对收集到的输入 token 进行计算，计算完之后，再通过 all-to-all 通信来把输出传输给输入的设备。这样就完成了 MoE 模块的计算与通信

Expert, Model and Data Parallelism

目前我们只是通过增加专家个数来提高模型的表现，但是 FLOPs 并没有增加，如果我们希望通过增加 FLOPs 来提高模型表现的话，我们需要增加 expert layer 的 model size, 这就涉及到了 Model parallelism, 此时的做法和前面提到的 Expert and Data Parallelism 一样，我们在收集到对应的设备的输入之后，再执行 model parallelism 就可以了。

Results

基于以上改进，作者构建了 Switch Transformer, 模型训练的数据集为 C4, 训练的目标为 masked language modeling, 训练时，作者将 $15\%$ 的 token 替换为 [mask] token.

模型的配置如下表所示

实验结果如下图所示

实验结果发现

Switch Transformer 的表现和训练效率都超过了 Dense model
Switch transformer 的训练效率稍微优于使用 2 个 expert 的 MoE-Base 模型
Switch transformer 在 low capacity factor 的场景下效果更好

Scaling

作者对比了 MoE 模型的 scaling law, 结果如下图所示

可以看到，当我们增加专家个数的时候，模型的表现是持续提升的。并且当我们增加专家个数之后，模型的训练效率也有所提升。

作者接下来在训练时间上进行了对比，结果如下图所示

实验结果说明，switch transformer 的训练效率比 dense model 快 7 倍左右

作者进一步对比 switch transformer 和更大的 dense 模型 (T5 large, $3.5\times$ FLOPs), 实验结果证明 Switch transformer 的训练效率仍然更高。

Switch for Attention

作者还探究了在 attention layer 加入 MoE 的表现，结果发现尽管效果有提升，但是训练不稳定。

No Token Left behind

由于 tensorflow 为一个静态计算框架，tensor 的形状必须预先定义好，因此必须为每个 expert 设定 capacity, 但是这样就会导致有些 token 被 drop 掉。

因此，作者提出了 “No token left behind” 这种方法，来避免出现 token overflow 的情况。具体做法就是先按照正常的 router 逻辑进行计算，如果 router 选取出来的 top-K expert （本文中 $K=1$）都已经满载了，则选取 top-(K+1) expert 进行计算，作者发现这种方式可以保证大部分 token 都不会被 drop. 作者通过尝试之后发现，这种方法并没有带来提升。

Encouraging Exploration Across Experts

作者还探究了不同选取 top-Kexpert 的方式，作者对比了以下三种方法：

argmax
sampling from the softmax distribution
input dropout on the incoming representation
multiplicative jitter noise on the incoming representation

实验结果如下表所示

Model Quality	(Neg. Log Perp.) (↑)
Argmax	-1.471
Sample softmax	-1.570
Input dropout	-1.480
Input jitter	-1.468

作者发现，jitter 的效果最好，因此在本文中作者使用 jitter 来加入 noise.

Ablation on Few Experts

作者还使用了更少的专家进行实验，结果如下图所示

实验结果显示，即使我们只使用少数 experts, 其模型表现仍然超过了 dense model 的表现，说明了 MoE 架构的有效性。

Downstream Model Performance

作者还进一步探究了模型的预训练表现与 downstream task 任务上表现的关系，结果如下图所示

实验结果显示，不管是 baseline 还是 Switch Transformer 其预训练的表现与下游任务上的表现都是正相关的。但是，作者也发现，MoE 模型在微调之后，其表现并不总是与预训练的表现正相关。因此，作者认为进一步探究这个机制是有必要的。

Conclusion

作者在本文中提出了 Switch Transformer，一个基于 MoE 架构的 Transformer 模型。作者通过改进 MoE 算法，大幅度提高了计算和通信效率，结果发现模型比 dense model 有更高的训练效率。

作者认为，未来的工作有：

提升大规模模型的训练稳定性
解决 MoE 模型微调之后效果不如预期的问题
探究针对 MoE 模型的 scaling law
支持异构架构的 MoE 模型
在 FFN 模块意外应用 MoE 架构
将 Switch Transformer 扩展到其他的模态

References

Arxiv

Chinchilla Scaling Law

Wed, 22 Oct 2025 14:39:23 +0800

Introduction

本文中关注的研究问题为：

给定一个 FLOPs budget, 如何平衡 model size 和 dataset size 之间的关系？

即，我们希望求解如下优化问题：

$$ N_{opt}(C), D_{opt}(C) =\arg\min_{N,D,\ \mathrm{s.t.}\ FLOPs(N,D)=C} L(N,D) $$

作者通过训练 400 多个模型，构建了对应的 scaling law.

已有工作如 Kaplan scaling law 已经发现模型参数和大语言模型表现之间的关系，一个结论就是计算最优并不代表达到最优的 loss. 在本文中，作者也有相同结论，但是作者认为大模型应该使用比 Kaplan scaling law 推荐的更多的 training token. 基于这个发现，作者训练了 Chinchilla, 一个 70B 的 LLM, Chinchilla 相比 Gopher 表现有了大幅度的提升。

Scaling Law

Fix Model Size and Very Dataset Size

这个方法中，作者通过改变训练步数，来研究 FLOPs 与模型表现之间的关系，结果如下图所示

通过对实验结果进行拟合，作者发现存在关系 $N_{opt}(C)\propto C^a$ 以及 $D_{opt}(C)\propto C^b$, 拟合的结果为 $a=b=0.5$.

IsoFLOPS Profiles

这个方法中，作者使用了不同的模型大小以及算力来构建最优模型参数量与算力之间的关系。作者给定 9 个算力配置，然后选取不同参数量的模型，训练的 token 数由算力和模型参数量决定，实验结果如下图所示

结果显示，不同大小的模型的表现 (loss) 随算力上升先下降后上升。因此给定算力，存在一个最优的 model size. 作者基于拟合出来的曲线得到了 Gopher 使用的算力配置下的最优 model size 和 training tokens. 同样的，作者得到 $a=0.49,b=0.51$.

Fitting a Parametric Loss Function

这个方法中，作者对 $L(N,D)$ 进行建模，作者使用了如下的公式

$$ L(N,D) = E + \frac{A}{N^{\alpha}} + \frac{B}{D^{\beta}} $$

第一项代表了建模的误差，第二项代表了数据集充分大损失与模型参数之间的关系，第三项代表了当模型充分训练时，损失与数据集大小之间的关系。

为了求解 $(A,B,E,\alpha,\beta)$, 作者基于训练收集到的数据 $L(N_i,D_i)$, 通过 L-BFGS 算法来最小化 Huber loss 进行求解，结果得到 $(A,B,E,\alpha,\beta)=( 406.4, 410.7, 1.69, 0.34, 0.28)$.

将结果带入带上面的表达式中，然后求出梯度为 0 的点，就得到

$$ N_{opt}(C) = G\left(\frac C6\right)^a, D_{opt}(C) = G^{-1}\left(\frac C6\right)^b, \text{ where }G=\left(\frac{\alpha A}{\beta B}\right)^{1/(\alpha+\beta)}, a=\frac{\beta}{\alpha+\beta}, b=\frac{\alpha}{\alpha+\beta} $$

带入数值之后就得到 $a=0.46$, $b=0.54$. 作者对结果可视化如下图所示，左图是拟合曲线的 Contour plot, 右图对左图的一个切片

Optimal Model Scaling

作者将三种方法的结果以及 Kaplan scaling law 的结果总结放在下表中，作者假设 $N_{opt}(C)\propto C^a$ 以及 $D_{opt}(C)\propto C^b$

Approach	$a$	$b$
Kaplan	0.73	0.26
Approach 1	0.50	0.50
Approach 2	0.49	0.51
Approach 3	0.46	0.54

结果表明，三种方法的结论差不多：model size 和 dataset size 增长 debility 差不多。

作者因此给出来的不同模型大小所需要的算力以及 token, 结果如下表所示

Parameters	Approach 1		Approach 2		Approach 3
	FLOPs	Tokens	FLOPs	Tokens	FLOPs	Tokens
400 M	1.92e+19	8.0 B	1.84e+19	7.7 B	1.84e+19	7.7 B
1 B	1.21e+20	20.2 B	1.20e+20	20.0 B	1.20e+20	20.0 B
10 B	1.23e+22	205.1 B	1.32e+22	219.5 B	1.32e+22	219.5 B
67 B	5.76e+23	1.5 T	6.88e+23	1.7 T	6.88e+23	1.7 T
175 B	3.85e+24	3.7 T	4.54e+24	4.3 T	4.54e+24	4.3 T
280 B	9.90e+24	5.9 T	1.18e+25	7.1 T	1.18e+25	7.1 T
520 B	3.43e+25	11.0 T	4.19e+25	13.4 T	4.19e+25	13.4 T
1 T	1.27e+26	21.2 T	1.59e+26	26.5 T	1.59e+26	26.5 T
10 T	1.30e+28	216.2 T	1.75e+28	292.0 T	1.75e+28	292.0 T

作者基于发现的 scaling law, 对已有模型进行了探究，发现现有的大模型都存在 under-training 的现象，结果如下图所示

实验结果显示，现有的大模型的 size 应该更小（或者需要更大的算力）。作者最终的结论就是，现有的比较小的模型，需要更多的算力才能达到更好的表现。

Chinchilla

基于上一节的发现，作者提出了 Chinchilla, 一个 70B 的模型，训练使用了 1.4T token. 训练的数据集为 MassiveText 的扩展版本，训练使用的优化器为 AdamW, tokenizer 为 SentencePiece.

模型配置如下表所示

Model	Layers	Number Heads	Key/Value Size	dmodel	Max LR	Batch Size
Gopher 280B	80	128	128	16384	$4\times 10^{-5}$	$3M\to6M$
Chinchilla 70B	80	64	128	8192	$1\times 10^{-5}$	$1.5M\to3M$

Ablation Study

learning rate schedule 作者还通过 ablation study 发现，cosine learning rate cycle length 应该和训练步数差不多，当 cycle length 太长时，模型表现会下降。

Optimizer 作者对比了 Adam 和 AdamW 的表现，结果发现，AdamW 的表现优于 Adam.

High Precision 训练时，作者使用了高精度也就是 float32 来保存梯度的状态，结果显示，不管是 Adam 还是 AdamW, 使用高精度都可以提高模型的表现

Comparison with Kaplan 作者还对比了 Chinchilla 和 Kaplan 的预测结果，如下图所示

结果显示，基于 Chinchilla 预测得到的模型训练效果比 Kaplan 的更好。

Curvature of the FLOPs-frontier 作者发现，FLOP-minimal loss frontier 存在 curvature, 也就是小模型和大模型预测出来的曲线是不一样的，作者将结果展示在下图中

结果显示，从小模型拟合出来的结果比大模型拥有更高的算力使用效率，作者认为这是未来的一个研究方向。

Conclusion

本文中作者重新探究了针对 LLM 的 scaling law, 作者发现已有的大模型都存在 under-training 的现象，也就是说，模型需要更多的训练 token, 具体来讲，model size scaling 和 dataset scaling 应该处于同一水平。作者基于这个结论，提出了 Chinchilla, 一个 70B 的 LLM, 其表现超过了 280B 的 LLM.

References

arxiv

Kaplan Scaling Law

Wed, 22 Oct 2025 14:10:52 +0800

Introduction

作者首先就总结了本文的发现，主要有以下几点：

损失与模型的 scale , 即除开 embedding 的模型参数，数据集大小以及算力强相关，与 model shape 比如 depth 或者 width 关系不大
scaling law 是非常光滑的，意味着 scaling law 是一个可预测的模型
overfitting 普遍存在，当参数和数据集大小同时增加时，模型的表现会增加，但是当其中一个量固定时，提升就比较小。并且当我们将模型参数提升 8 倍时，我们只需要将数据集的大小提升 5 倍就可以避免过拟合
训练的损失函数曲线与 model size 无关，因此我们可以预测给定大小模型的表现
模型在测试集和训练集上的表现高度相关，因此我们可以基于训练集的损失来预测模型的表现
大模型比小模型拥有更高的 sample efficiency, 即更小的训练步数就可以达到相同的表现
convergence 不能说明一切，我们可以通过 early-stopping 来提高算力使用效率，避免模型花费过多的算力在较小的提升上
最优的 batch size 与 loss 呈一个 power law 的关系

Notation

Symbol	Notation
$L$	cross entropy loss
$N$	non-embedding parametters
$B$	batch size
$S$	number of training steps
$C\approx 6NBS$	estimate of total training compute
$D$	dataset size in tokens
$B_{crit}$	critical batch size
$C_{\min}$	estimate of the minimum compute to reach a given value of loss
$S_{\min}$	estimate of the minimum steps to reach a given value of loss
$\alpha_X$	power-law exponents for the scaling law of loss
$d$	hidden size of the model

为了简便，后续未经特殊说明，我们说模型参数均指的是 non-embedding 的参数

Scaling Law Overview

当数据集 $D$ 足够大时，损失与模型参数大小 $N$ 之间的关系为

$$ L(N) = \left(\frac{N_C}{N}\right)^{\alpha_N}, \alpha_N\sim 0.076, N_C\sim 8.8\times 10^{13} $$

给定模型参数大小 $N$ , 损失与数据集大小 $D$ 之间的关系为

$$ L(D) = \left(\frac{D_C}{D}\right)^{\alpha_D}, \alpha_D\sim 0.095, D_C\sim 5.4\times 10^{13} $$

给定足够大的数据集 $D$ 和最优模型大小 $N$ 时，损失与算力 $C$ 之间的关系为

$$ L(C_{\min}) = \left(\frac{C_C^{\min}}{C_{\min}}\right)^{\alpha_D}, \alpha_C^{\min}\sim 0.050, C_C^{\min}\sim 3.1\times 10^{8} $$

作者在不同大小的数据集，算力，模型大小下进行了测试，结果发现 scaling 与模型的 shape, transformer 的超参数之间的关系比较小。$\alpha_N,\alpha_D,\alpha_C^{\min}$ 等决定了当我们 scale up 数据及大小，模型大小和算力时损失的变化情况。比如当我们将模型参数提升至 2 倍时，模型的损失会降低至原来的 $0.95$.

基于发现 1 和 2, 作者发现当我们将模型的 size 提升至原来的 2 倍时，模型的数据集大小应该提升至原来的 $1.67$ 倍，具体关系为 $D\sim N^{0.74}$.

作者使用了一个统一的公式来描述损失与数据及大小和模型参数大小之间的关系

$$ L(N,D) = \left[\left(\frac{N_c}{N}\right)^{\frac{\alpha_N}{\alpha_D}}+ \frac{D_c}{D}\right]^{\alpha_D} $$

当数据集充分大时，损失与模型参数大小以及更新步数 $S$ 的关系如下

$$ L(N, S) = \left(\frac{N_C}{N}\right)^{\alpha_N} +\left(\frac{S_C}{S_{\min}(S)}\right)^{\alpha_S} $$

这里 $S_C\approx 2.1\times 10^3$, $\alpha_S\approx 0.76$, $S_{\min}(S)$ 是估计出来的最小优化步数

最优的 batch size 与损失函数之间的关系如下

$$ B_{crit}(L) = \frac{B_*}{L^{1/\alpha_B}}, B_*\sim 2*10^8 \text{ tokens}, \alpha_B\sim 0.21 $$

给定算力 $C$ 且无其他限制时，模型参数，数据及大小，batch size 和更新参数与算力之间的关系如下

$$ N\propto C^{\alpha_C^{\min}/\alpha_N}, B\propto C^{\alpha_C^{\min}/\alpha_B}, S\propto C^{\alpha_C^{\min}/\alpha_S}, D=BS $$

其中

$$ \alpha_C^{\min} = \frac{1}{\frac{1}{\alpha_S}+\frac{1}{\alpha_B}+\frac{1}{\alpha_N}} $$

实验的结果为 $N\propto C_{\min}^{0.73}$, $B\propto C_{\min}^{0.24}$ , $S\propto C_{\min}^{0.03}$. 也就是说，当我们提升算力时，提升模型的参数大小带来的收益是最高的。

Background

首先，transformer 的参数量通过计算可以得到

$$ N\approx 2dn(2d+d_{ff}) = 12nd^2 $$

这里 $d$ 是 hidden size, $n$ 是 layer 个数，$d_{ff}$ 是 MLP 的 hidden size, 这里我们假设 $d_{ff}=4d$. 计算时我们丢掉了 bias 以及 LayerNorm 的参数量。具体计算过程见 LLM parameter analysis

transformer 一次前向计算的 operations 数量大概为

$$ C_{forward}\approx 2N + 2nLd $$

这里 $L$ 是输入的 token 长度。

由于反向传播所需要的 FLOPs 是前向传播两倍，因此 transformer 的计算量为

$$ C = C_{backward} + C_{forward} = 3C_{forward}\approx 6N $$

具体计算过程见 LLM FLOPs analysis。也就是说，对于参数量为 $N$ 的 transformer model, 每个 token 所需要的 FLOPs 为 $C\approx 6N$

Empirical Results and Basic Power Laws

Transformer Shape and Hyper-parameter Independence

作者基于 $N=12nd^2$, 在保持总参数量 $N$ 不变的情况下，分别调整 $n$, $d_{ff}$ 和 number of attention heads 的个数（变化 $d$ 用于维持总参数量不变），结果如下图所示

实验结果发现，损失对于 $d_{ff}/d$, $d/n$, $d/n_h$ 都比较 robust, 说明模型的损失对模型的 shape 依赖性比较低。

Non-embedding Parameter Count

作者探究了以下 model size 对损失的影响，作者使用了不同的 $n$ 和 $d$, 然后训练得到的损失情况如下图所示

作者发现，当包含 embedding parameter 时，损失不仅依赖于模型参数量，还依赖于 layer 层数 $n$, 但是当我们排除 embedding parameter 时，模型的损失便与 layer 层数 $n$ 关系不大。这个趋势可以用以下模型来表示

$$ L(N) \approx \left(\frac{N_c}{N}\right)^{\alpha_N} $$

最终拟合的曲线如下图所示

Comparing to LSTMs and Universal Transformers

作者比较了 LSTM 和 Transformer 结构的损失，结果如下图所示

可以看到，transformer 比 LSTM 拥有更强的学习能力， LSTM 架构对于 early context 表现比较好，但是随着 context 增加，LSTM 的表现逐渐弱于 transformer. 即 transformer 的长上下文能力强于 LSTM 架构。

Generalization Among Data Distributions

模型是在 WebText2 数据集上训练的，作者进一步在其他数据集上评估了以下模型的泛化性，结果如下图所示

结果发现，模型在其他数据集上的泛化性很好。并且，模型的泛化性能仅与训练阶段的表现相关（validation loss），而与训练阶段（是否收敛）无关。

作者还评估了 model depth 对模型泛化性的影响，结果如下图所示

实验结果显示，model depth 对模型泛化性基本没有影响。

Performance with Data Size and Compute

作者探究了损失与 dataset size $D$ 之间的关系。作者固定一个模型，然后当 test loss 不再下降时停止训练，结果发现 test loss 与 dataset size $D$ 之间存在如下关系

$$ L(D) \approx \left(\frac{D_c}{D}\right)^{\alpha_D} $$

拟合结果如下图所示

接下来，基于前面计算的结果，我们有 $C\approx 6ND=6NBS$, 这里 $B$ 是 batch size, $S$ 是训练步数。给定 $C$, 作者使用不同大小的模型进行训练，batch size $B$ 保持不懂，训练步数设置为 $S=C/6BS$,实验结果显示损失与算力 $C$ 之间满足如下关系

$$ L(C) \approx \left(\frac{C_c}{C}\right)^{\alpha_C} $$

拟合结果如下图所示

作者进一步探究了 sample efficiency 与 model size 之间的关系，实验结果如下图所示

结果显示，随着 model size 增加，sample efficiency 也在增加

Charting the Infinite Data Limit and Overfitting

作者在本节探讨了同时变化 $N$ 和 $D$ 对损失变化的影响。

Proposed Equation

作者基于三个原则进行建模：

改变 vocabulary size 或者 tokenization 会 rescale loss
固定 $D$ 并且令 $N\to\infty$, 则最终损失应该接近 $L(D)$. 反之固定 $N$, 令 $D\to\infty$, 最终损失应该接近 $L(N)$
$L(N,D)$ 在 $D=\infty$ 处应该是可解析的

基于以上三条原则，将模型选择为如下形式

$$ L(N,D) = \left[\left(\frac{N_c}{N}\right)^{\frac{\alpha_N}{\alpha_D}}+ \frac{D_c}{D}\right]^{\alpha_D} $$

作者基于不同配置进行训练，基于实验结果你和得到的参数如下

Parameter	$\alpha_N$	$\alpha_D$	$N_c$	$D_c$
Value	0.076	0.103	$6.4\times 10^{13}$	$1.8\times 10^{13}$

接下来，作者探究了模型的过拟合程度，作者定义如下 metric

$$ \delta L(N,D) := \frac{L(N,D)}{L(N,\infty)} - 1 $$

带入 $L(N,D)$ 定义就得到

$$ \delta L(N,D) = \left(1 + \left(\frac{N}{N_c}\right)^{\frac{\alpha_N}{\alpha_D}}\frac{D_c}{D}\right) - 1 $$

通过测试不同的模型，作者发现 $\delta L$ 的值在 $0.02$ 左右，将实验结果带入到上面的公式就得到

$$ D \geq (5\times 10^3)N^{0.7379} $$

也就是说对于参数量为 $N$ 的模型，需要 data size $D \geq (5\times 10^3)N^{0.7379}$ 才能避免过拟合。

Scaling Laws with Model Size and Training time

作者在本节构建了损失函数与 model size $N$ 以及训练时间的 scaling law

Adjustment for Training at Critical Batch Size

已有结论说明，存在一个 critical batch size $B_{crit}$, 当 batch size 接近 $B_{crit}$ 时，增加 batch size 对计算效率影响比较小，但是当 batch size 大于 $B_{crit}$ 时，带来的提升比较小。另一方面，batch size 会影响梯度的噪声程度。因此，训练步数 $S$ 和处理的样本数 $E=BS$ 应该满足：

$$ \left(\frac{S}{S_{\min}}-1\right)\left(\frac{E}{E_{\min}}-1\right) = 1 $$

这里 $S_{\min}$ 是达到损失 $L$ 所需要的最小训练步数，而 $E_{\min}$ 是最小的训练样本数量。

作者的实验结果如下

作者将 critical batch size 定义为

$$ B_{crit}(L) := \frac{E_{\min}}{S_{\min}} $$

使用 critical batch size 进行训练可以在计算效率和算力之间达到一个平衡。

作者基于上面的实验结果探究了 critical batch size 和 model performance 之间的关系，实验结果如下图所示

可以看到，critical batch size 与 model size 的关系不大，仅与损失 $L$ 有关。作者通过以下模型拟合 critical batch size:

$$ B_{crit}(L) \approx \frac{B_*}{L^{1/\alpha_B}} $$

这里 $B_*\approx 2\times 10^8$, $\alpha_B\approx 0.21$.

给定一个 target loss $L$, 当 batch size $B>> B_{crit}$ 时，作者定义最小训练步数为

$$ S_{\min}(S) := \frac{S}{1+B_{crit}(L)/B} $$

给定 target loss $L$ 和 model size $N$, 当 batch size $B<< B_{crit}$ 时，作者定义最小算力为

$$ C_{\min}(C) := \frac{C}{1+B_{crit}(L)/B} $$

Performance with Model Size and Compute

作者使用如下公式来探究损失与 model size 和 computer 之间的关系

$$ L(N, S_{\min}) = \left(\frac{N_C}{N}\right)^{\alpha_N} +\left(\frac{S_C}{S_{\min}(S)}\right)^{\alpha_S} $$

拟合结果如下表所示

Parameter	$\alpha_N$	$\alpha_S$	$N_c$	$S_c$
Value	0.077	0.76	$6.5\times 10^{13}$	$2.1\times 10^{3}$

基于这个拟合结果，作者得到了下图的结果

作者还使用了不同的可视化方式，如下图所示

实验结果显示，上面的公式拟合的很好。

Lower Bound on Early Stopping step

作者还探究了以下 early step 与模型大小以及数据集之间的关系，作者通过分析得到如下结果

$$ S_{stop}(N,D) \gtrsim \frac{S_c}{[L(N,D)-L(N,\infty)]^{1/\alpha_S}} $$

其中 $L(N,\infty)$ 是在充分大数据集上的收敛损失。作者对实验结果进行了拟合，结果如下图所示

Optimal Allocation of the Compute Budget

作者在本节探究了最优算力与 model size $N$ 和训练数据 $2B_{crit}S_{\min}$ 之间的关系

Optimal Performance and Allocations

作者首先基于

$$ C_{\min}(C) := \frac{C}{1+B_{crit}(L)/B} $$

绘制了如下曲线图

作者发现，相比于 loss 与算力 $C$ 之间的关系，使用 $C_{\min}$ 进行拟合效果更好。

接下来，作者基于 $L(C_{\min})$ 进一步探究了给定算力如何决定最优的 model size $N(C_{\min})$. 其实验结果如下图所示

实验结果显示，model size 和算力之间有如下关系

$$ N(C_{\min}) \propto (C_{\min})^{0.73} $$

作者进一步探究了对于非最优模型大小与算力之间的关系，作者先构建了如下的关系

$$ \frac{C(N, N_{\mathrm{eff}})}{C(N_{\mathrm{eff}}, N_{\mathrm{eff}})} = \frac{N}{N_{\mathrm{eff}}} \left[ 1 + \frac{\alpha_S}{\alpha_N} \left( 1 - \left( \frac{N_{\mathrm{eff}}}{N} \right)^{\alpha_N} \right) \right]^{-\!1 / \alpha_S}. $$

对应的示意图为

实现结果发现，大小为最优模型的 $0.6\sim 2.2$ 倍只需要额外 $20\%$ 的算力。作者强调，这个实验结果对于超大模型不一定适用。

作者进一步推导了 $S_{\min}$ 和 $C_{\min}$ 之间的关系，由于 $C_{\min}=6NB_{crit}S$, 且我们前面已经有 $B\propto L^{-4.8}$ 和 $L\propto C_{\min}^{-0.05}$, 因此我们有

$$ B_{crit}\propto L^{-4.8} \propto (C_{\min})^{-0.05\times (-4.8)}\propto (C_{\min})^{0.24} $$

以及

$$ S_{\min} \propto \frac{C_{\min}}{6B_{crit}N(C_{\min})} \propto (C_{\min})^{0.03} $$

拟合的结果如下图所示

因此，基于上面的结果，当我们增加算力时，我们的主要精力应该放在增加模型大小和提高 batch size 上，而训练步数基本可以保持不变。

Another way of Derivation

作者还给出了另一种建模 $L(C_{\min})$ 的方式，即从 $L(N,S_{\min})$ 中进行推导，作者将 $B_{crit}$ 和 $S_{\min}$ 的表达式带入到 $L(N,S_{\min})$ 然后求解最小值就得到

$$ L(C_{\min})= \left(\frac{C_C^{\min}}{C_{\min}}\right)^{\alpha_C^{\min}} $$

其中

$$ \alpha_C^{\min} = \frac{1}{\frac{1}{\alpha_S}+\frac{1}{\alpha_B}+\frac{1}{\alpha_N}} \approx 0.054 $$

这和前面的结果基本是吻合的，进一步进行推导得到

$$ N(C_{\min})\propto (C_{\min})^{\alpha_C^{\min}/\alpha_N}\approx (C_{\min})^{0.71} $$

这个结果也和上面的差不多。

Contradiction and a Conjecture

作者发现，尽管拟合的 scaling law 曲线非常好，但是由于自然语言不可能达到 zero entropy, 因此该曲线最终一定会失效。作者基于更大的模型进行了实验，结果发现，模型在某一点开始就比预测的损失曲线下降的更慢。作者认为这是因为 transformer 模型已经达到了 maximal performance 导致的。

通过前面的分析，我们发现 $L(C_{\min})$ 比 $L(D)$ 下降的快，因此两者必然在某一点相交。

在前面的章节中，我们基于以下关系来决定数据集大小

$$ D\propto N^{0.74}\propto (C_{\min})^{0.74*0.73}\propto (C_{\min})^{0.54} $$

这里我们利用了 $N(C_{\min})$ 的结果

另一方面，我们有

$$ D(C_{\min}) = \frac{2C_{\min}}{6N(C_{\min})}\propto (C_{\min})^{0.26} $$

可以看到，基于训练最优导出的数据集大小相比于拟合出来的数据集大小，实际上存在过拟合。

作者进一步分析出了 $L(D(C_{\min}))$ 和 $L(C_{\min})$ 这两条曲线的交点，结果得到

$$ C^*\approx 10^4 \text{ PF-Days}, N^*\approx 10^{12}\text{ parameters}, D^*\approx 10^12\text{ tokens}, L^*\approx 1.7\text{1.7nats/token} $$

作者认为出现这种原因有以下几种情况：

$L^*$ 给出了自然语言的 entropy 的一个估计，因此当模型充分大之后，模型可能已经获取到了数据中的所有知识
$L(C_{\min})$ 可以作为数据集噪声的一个量化表现，其衡量了数据集的质量

Learning Rate Schedule

附录中，作者还探究了 learning rate 与损失之间的关系，作者使用了不同 learning rate schedule 对模型损失的影响，结果如下图所示

实验结果显示，只要 learning rate 下降的不会太快，模型的表现基本上差不太多。

作者基于实验结果得到了学习率和模型参数之间的关系如下

$$ \text{lr}(N)\approx 0.003239 - 0.0001395\log N $$

也就是说，小模型用比较大的学习率，大模型用较小的学习率。

Conclusion

作者在本文中训练了大量不同配置的大模型，然后构建了损失（损失）与模型参数，数据及大小以及算力之间的关系。实验结果发现，损失与架构和优化参数之间的关系比较小，主要由模型参数量决定，更大的模型拥有更高的采样效率。

作者认为，本文的局限在于损失函数不一定能够反应模型在其他语言任务上的表现。

References

arxiv

LLM FLOPs Computation

Wed, 15 Oct 2025 16:33:39 +0800

本文中，我们介绍如何计算基于 transformer 架构的 LLM 的 FLOPs, 计算完成之后，我们可以推导出算力 $C$ 与模型参数量 $N$，数据集大小 $D$ 之间的关系，即 $C\approx 6ND$.

Background

FLOPs

FLOPs，floating point operations，表示浮点数运算次数，一般计算 FLOPs 仅考虑加法和乘法。

假设 $A\in\mathbb{R}^{m\times p}$, $B\in\mathbb{R}^{p\times n}$, 则计算 $C=AB$ 的过程中一共需要进行 $mnp$ 次乘法运算和 $mnp$ 次加法运算，因此总的 FLOPs 为 $2mnp$.

Assumption

假设 $A\in\mathbb{R}^{m\times p}$, $B\in\mathbb{R}^{p\times n}$, $C\in\mathbb{R}^{m\times n}$ 则计算 $D=AB+C$ 的过程中一共需要进行 $mnp$ 次乘法运算和 $mnp+mn$ 次加法运算，因此总的 FLOPs 为 $2mnp+mn\approx 2mnp$.

基于上述结论，我们计算 FLOPs 时，忽略 element-wise 的操作，即我们做如下假设：

忽略 normalization 中的小常数项运算
忽略 residual connection 和 bias term 的加法
忽略注意力的 MASK 和 softmax，因为这两者都是 element-wise operation.
使用 look-up 计算 embedding layer

注，基于以上假设，我们的结果与 Chinchilla Scaling law 的结果稍有不同，但结论不变。

Notation

Math Variable	Code Variable	Description
$n$	`num_hidden_layers`	Transformer block 个数
$\vert V\vert$	`vocab_size`	词表大小
$d$	`hidden_size`	token embedding 的维度
$d_{ff}$	`intermediate_size`	MLP 的中间层的维度
$h$	`num_attention_heads`	query head 的个数
$s$	`seq_len`	length of token sequence

注：为了避免混淆，我们使用 $n$ 来表示 decode layer 的个数。

Computation

我们计算训练阶段的总 FLOPs, 记为 $C$, Kaplan scaling law 用 PF-days 作为单位，$1\text{ PF-Days}=10^{15}\times 24\times 3600\ FLOPs$. 训练阶段包括前向阶段 (forward pass) 和反向传播阶段 (backward pass). 因此

$$ C = FLOPs(\text{forward}) + FLOPs(\text{backward}) $$

Forward

decoder-only transformer 的模型架构包含三个模块：

1 层 embedding layer
$n$ 层 decoder layer
1 层 lm head layer

因此模型总的 FLOPs 为

$$ FLOPs(\text{forward}) = FLOPs(\text{embedding}) + n*FLOPs(\mathrm{decode\_layer})+FLOPs(\mathrm{lm\_head}) $$

Embedding & Lm Head

首先，对于 embedding layer, embedding layer 本质上是一个 look up table, 计算过程中不涉及浮点数运算，因此 $\boxed{FLOPs(\text{embedding})=0}$.

接下来，对于 lm_head, 这是一个 linear layer, 其权重大小为 $W\in\mathbb{R}^{d\times |V|}$, 输入为 $x\in\mathbb{R}^{s\times d}$, 因此 $\boxed{FLOPs(\mathrm{lm\_head})=2sd|V|}$.

因此，我们有

$$ FLOPs(\text{forward}) = n*FLOPs(\mathrm{decode\_layer})+ 2sd|V| $$

Decode Layer

对于 decode_layer, 其又包含了四个模块：

pre-normalization
attention
post-normalization
FFN

pre-normalization 和 post-normalization 一般是一样的，因此

$$ \begin{aligned} FLOPs(\mathrm{decode\_layer}) &= FLOPs(\mathrm{pre\_normoalization}) + FLOPs(\mathrm{Attention}) + FLOPs(\mathrm{post\_normoalization}) +FLOPs(\mathrm{FFN})\\ &= 2*FLOPs(\mathrm{normoalization}) + FLOPs(\mathrm{Attention})+FLOPs(\mathrm{FFN}) \end{aligned} $$

Normalization

现在有两种常见的 normalization, 也就是 LayerNorm 和 RMSNorm

LayerNorm 定义如下

$$ \mathrm{LayerNorm}(x) = \frac{x-\mathbb{E}[x]}{\sqrt{\mathrm{var}[x]+\epsilon}}\odot \beta + \gamma $$

其中 $\beta,\gamma\in\mathbb{R}^d$ 是可学习的参数。

对输入 $x\in\mathbb{R}^{s\times d}$, 均值需要约 $sd$ 次 FLOPs，方差 $\mathrm{var}[x]$ 只需要约 $3sd$ 次 FLOPs，这两者的计算都可以忽略。接下来就是 scaling 和 shift, 这两者都是 element-wise 操作，我们这里，因此总的 FLOPs 为

$$ \boxed{FLOPs(\mathrm{normoalization}) = 4sd} $$

RMSNorm 的作用和 LayerNorm 是一样的，但是实现上更简单

$$ \mathrm{RMSNorm}(x) = \frac{x}{\sqrt{\|x\|_2^2+\epsilon}}\odot \gamma $$

其中 $\gamma\in\mathbb{R}^d$ 是可学习的参数

对于 RMSNorm，其分析方式与 LayerNorm 基本一致，因此总的 FLOPs 为

$$ \boxed{FLOPs(\mathrm{normoalization}) = 4sd} $$

总之，不管使用哪种 normalization，其 FLOPs 都是

$$ \boxed{FLOPs(\mathrm{normoalization}) = 4sd} $$

Attention

Attention 定义如下

$$ \mathrm{Attention}(X) = \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d}}\right)V\in\mathbb{R}^{m\times d} $$

其中 $X\in\mathbb{R}^{s\times d}$, $W_Q,W_K,W_V\in\mathbb{R}^{d\times d}$

$$ Q = W_QX\in\mathbb{R}^{s\times d},\quad K =W_KX\in\mathbb{R}^{s\times d},\quad V = W_VX\in\mathbb{R}^{s\times d} $$

$Q,K,V$ 计算的 FLOPs 为 $6*sd^2$. $QK^T$ 的 FlOPs 为 $2s^2d$, $\mathrm{softmax}(\cdot)V$ 的 FLOPs 为 $2s^2d$, 最后对于 multi-head attention 还有一个 output projection layer, 其权重为 $W_O\in\mathbb{R}^{d\times d}$, 因此 FLOPs 为 $2sd^2$. 故 attention 最终的 FLOPs 为

$$ FLOPs(\mathrm{Attention})=6sd^2+2s^2d+2s^2d+2sd^2=\boxed{8sd^2+4s^2d} $$

FFN

对于 FFN，有两种常见的形式，一种基于 ReLU 激活函数，其定义如下

$$ y = \max(xW_1+b_1, 0)W_2 + b_2 $$

其中 $W_1\in\mathbb{R}^{d_{ff}\times d}$, $W_2\in\mathbb{R}^{d\times d_{ff}}$. $b_1\in\mathbb{R}^{d_{ff}}$, $b_2\in\mathbb{R}^{d}$.

对输入 $x\in\mathbb{R}^{s\times d}$, 其 FLOPs 为

$$ FLOPs(\mathrm{FFN_{ReLU}}) = 2sdd_{ff} + 2sd_{ff}d =\boxed{4sdd_{ff}} $$

其中第一项和第二项分别为为 $xW_1$ 与 $\max(xW_1+b_1, 0)W_2$ 的 FLOPs.

另一种基于 SwiGLU 激活函数，其定义为

$$ \mathrm{SwiGLU}(x) = x\odot \sigma(x) $$

其中 $\sigma(\cdot)$ 是 sigmoid 函数

FFN 的定义为

$$ y = W_2(W_3x\odot \mathrm{SwiGLU}(W_1x)) $$

其中 $W_3,W_1\in\mathbb{R}^{d_{ff}\times d}$, $W_2\in\mathbb{R}^{d\times d_{ff}}$.

对输入 $x\in\mathbb{R}^{s\times d}$, 其 FLOPs 为

$$ FLOPs(\mathrm{FFN_{SwiGLU}}) = 2sdd_{ff} + 2sdd_{ff} + 2sd_{ff}d = \boxed{6sdd_{ff}} $$

Summary

最终，decoder-only transformer 的 FLOPs 计算量为（我们假设使用 multi-head attention, 基于 SwiGLU 的 MLP）

$$ \begin{aligned} FLOPs(\text{forward}) &= FLOPs(\text{embedding}) + n*FLOPs(\mathrm{decode\_layer})+FLOPs(\mathrm{lm\_head})\\ &= n*FLOPs(\mathrm{decode\_layer})+2sd|V|\\ &= n*(2*FLOPs(\mathrm{normoalization}) + FLOPs(\mathrm{Attention})+FLOPs(\mathrm{FFN}))+2sd|V|\\ &= n*(8sd+8sd^2+4s^2d+6sdd_{ff})+2sd|V|\\ &= nsd^2\left(\frac8d + 8+\frac{4s}{d}+\frac{6d_{ff}}{d}+\frac{2|V|}{nd}\right)\\ &\approx \boxed{nsd^2\left(8+\frac{4s}{d}+\frac{6d_{ff}}{d}+\frac{2|V|}{nd}\right)} \end{aligned} $$

这里我们丢弃了 normalization 项，因为 normalization 的 FLOPs 是一个低阶项

Backward

首先，我们有如下结论：

神经网络 backward 过程的计算量（FLOPs）为 forward 过程的两倍

我们使用一个简单的例子来证明这个结论，考虑线性层 $h=Wx$, 其中 $W\in\mathbb{R}^{m\times d}$, 对于输入 $x\in\mathbb{R}^{d\times 1}$ 其 forward 过程的计算量为 $2md$.

对于反向过程，我们需要分别计算损失 $L$ 对权重和输入的梯度，即

$$ \frac{\partial L}{\partial x} = W^T\frac{\partial L}{\partial h}\in\mathbb{R}^{d\times 1}, \quad\frac{\partial L}{\partial W} = \frac{\partial L}{\partial h}\otimes x^T\in{m\times d}, $$

这里 $\frac{\partial L}{\partial h}\in\mathbb{R}^{m\times 1}$ 为损失对输出 $h$ 的梯度。因此反向传播的总计算量为

$$ 2dm + 2md = 4md $$

这里的两项分别是对 $x$ 和 $W$ 求梯度的 FLOPs.

Overall

将前向传播和反向传播的计算量汇总我们就得到一次前向传播和一次反向传播过程中，对于长度为 $s$ 的 token, 其 FLOPs 为 (multi-head attention, SwiGLU-FFN)

$$ \begin{aligned} C &= FLOPs(\text{forward}) + FLOPs(\text{backward})\\ &= 3FLOPs(\mathrm{forward}) \\ &\approx \boxed{3nsd^2\left(8+\frac{4s}{d}+\frac{6d_{ff}}{d}+\frac{2V}{nd}\right)} \end{aligned} $$

Extension

GQA

GQA 与 MHA 不同的地方在于通过共享 key 和 value 来降低 KV cache 的占用，我们假设 group number 为 $g$, 则 key 和 value 的 FLOPs 现在变成了

$$ 2sdd_h\frac{h}{g}+2sdd_h\frac{h}{g}=4sdd_h\frac{h}{g} $$

因此 attention 部分总的 FLOPs 变成了

$$ FLOPs(\mathrm{Attention})=4sd^2+2s^2d+2s^2d+4sdd_h\frac{h}{g}=4sd^2+4s^2d+4sdd_h\frac{h}{g} $$

当 $g=h$ 时，GQA 就变成了 MHA, 此时的 FLOPs 也一致。

MoE

MoE 是针对 Dense FFN 的一个改进，介绍见 MoE, 我们假设一共有 $e$ 个路由专家，其中激活 $k$ 个。

Gate layer 一般是一个 linear layer, 其权重矩阵大小为 $W_{G}\in\mathbb{R}^{d\times e}$, 因此 $FLOPs(\text{router})= 2sde$.

Expert layer 和前面提到的 FFN 一致，我们每次挑选出 $k$ 个专家进行计算，因此 expert 部分 $FLOPs(\text{expert})=6ksdd_{ff}$.

从而对于 MoE 来说，FFN 部分的 FLOPs 为

$$ FLOPs(\text{MoE}) = FLOPs(\text{router})+FLOPs(\text{expert})= \boxed{2sde+6ksdd_{ff}} $$

Simplification

我们已经得到了 transformer 的 FLOPs 计算表达式，但是其表达式比较繁琐，因此，在研究 scaling law 时，一般会进行简化。

首先，在 LLM parameter analysis 中，我们已经给出了 LLM 参数量 $N$ （基于 Qwen3）的计算结果

$$ N=n*(4d+3dd_{ff}+2hh_{d}d + 2h_{kv}h_dd) + d(2|V|+1) $$

我们这里对其进行简化，一般来说 $d_{ff}=8/3d$, $h_{kv}=h$, $h_d=d/h$, 带入就得到

$$ N = n(4d+8d^2+2d^2+2d^2) + d(2|V|+1)=n(12d^2+4d)+ d(2|V|+1) $$

我们忽略关于 $d$ 的一阶项，并且我们假设 $|V| << 12nd$, 则最终模型参数量可以近似为

$$ \boxed{N \approx 12nd^2} $$

接下来，我们基于上面的配置简化 FLOPs 表达式

$$ \begin{aligned} C &= 3nsd^2\left(8+\frac{4s}{d}+\frac{6d_{ff}}{d}+\frac{2|V|}{nd}\right) \\ &= 3nsd^2\left(24+\frac{4s}{d}+\frac{2|V|}{nd}\right)\\ &\approx 72nsd^2 \\ &= 6sN \end{aligned} $$

这里我们利用了前面的 $|V| << 12nd$ 假设，为了简便我们还舍弃了 $4s/d$.

注意到 $s$ 代表 token 序列长度，如果训练集的总 token 个数为 $D$, 则最终对于包含 $D$ tokens 的数据集和包含 $N$ 参数量的 LLM, 其训练总 FLOPs 可以近似估计为

$$ \boxed{C\approx 6ND} $$

Experiments

Setting

接下来我们定量分析一些模型的 FLOPs. 我们基于 Chinchilla scaling law 给出的实验配置 (Table A9), 我们筛掉 kv_size * n_heads != d_model 的配置，$|V|=32,000$.

各部分的 FLOPs 计算代码如下

1
2
3
4
5


def compute_flops(n, V, d, d_ff, s):
 lm_head_flops = 2 * V * d * s
 attention_flops = 8 * s * d * d + 4 * s * s * d
 feed_forward_flops = 4 * s * d * d_ff
 return lm_head_flops, attention_flops, feed_forward_flops

首先我们看一下不同大小模型的 FLOPs 分布情况

可以看到，当模型越来越大，FFN 层的算力占比越来越高，这也是为什么后来采用 MoE 架构的一个原因。

接下来，我们看一下模型 FLOPs 分布情况随上下文长度变化的情况

可以看到，随 context length 增加，attention 的算力占比逐渐上升，这符合 attention 是一个平方复杂度的算法的结论。并且，当上下文足够长之后，计算还出现了 overflow (图像右端的突然下降)。

Conclusion

在本文中，我们介绍了如何计算基于 decoder-only transformer LLM 的 FLOPs, 并推导除了 Kaplan scaling law 中使用的公式 $C=6ND$, 这为后面的 infra 和 scaling law 的学习提供了基础。

References

Notes on AdamW

Thu, 04 Sep 2025 10:27:03 +0800

作者提出了一个针对 Adam 优化器的 weight decay 方法

Introduction

作者首先回顾了动态梯度算法如 AdaGrad, RMSProp, Adam 的进展。已有工作表明动态梯度算法的泛化性要比 SGD with momentum 要差。作者在本文中探究了在 SGD 和 Adam 中使用 L2 regularization 和 weight decay 对最终模型表现的影响。结果表明，模型泛化性较差的原因在于对于 Adam, L2 regularization 的效果要比 SGD 差。

作者有如下发现：

L2 regularization 和 weight decay 不等价。在 SGD 中，L2 regularization 是等价的，但是在 Adam 中这个结论不成立。具体来说，L2 regularization 对历史参数的惩罚要小于 weight decay
L2 regularization 对 Adam 效果提升有效
weight decay 对于 SGD 和 AdamW 都很有效，在 SGD 中，weight decay 与 L2 regularization 等价
最优的 weight decay 取决于 batch, batch 越大，最优的 weight decay 越小
通过 learning rate scheduler 可以进一步提高 Adam 的表现

作者在本文中的主要贡献是通过解耦梯度更新中的 weight decay 来提高 Adam 的 regularization.

作者的主要 motivation 是提升 Adam 表现，让其可以和 SGD with momentum 相比

Method

Weight decay 的定义如下

$$ \theta_{t+1} = (1-\lambda)\theta_t - \alpha \nabla f_t(\theta_t) \tag{1} $$

其中 $\lambda$ 是 weight decay rate, $\nabla f_t(\theta_t)$ 是第 $t$ 个 batch 的梯度，$\alpha$ 是学习率。

首先，对于标准的 SGD 来说，weight decay 与 L2 regularization 等价

Proposition 1 对于标准的 SGD 来说，对损失函数 $f_t(\theta_t)$ 执行 weight decay （公式 $(1)$）与对损失函数 $f_t(\theta_t)+\lambda'/2\|\theta_t\|_2^2$ 执行梯度下降算法是等价的，这里 $\lambda'=\lambda/\alpha$。

证明比较简单，只需要写出损失函数的梯度下降更新公式即可。

基于这个结论，大部分优化算法都将 L2 regularization 和 weight decay 看做是等价的。但实际上，这个结论对于 adaptive gradient 方法来说是不成立的。结论如下

Proposition 2 令 $O$ 为一个 optimizer, 其目标函数为 $f_t(\theta)$, 当不考虑 weight decay 时，梯度更新过程为 $\theta_{t+1}\gets \theta_t-\alpha M_t\nabla f_t(\theta_t)$. 当考虑 weight decay 时，梯度更新过程为 $\theta_{t+1}\gets (1-\lambda)\theta_t-\alpha M_t\nabla f_t(\theta_t)$. 如果 $M_t\neq kI$, 则不存在 $\lambda'$, 使得 $O$ 在优化目标函数 $f_t^{reg}(\theta)=f_t(\theta)+\lambda'/2\|\theta\|_2^2$ 时，不考虑 weight decay 的梯度更新与 $O$ 在优化目标函数 $f_t(\theta)$ 时，考虑 weight decay 的梯度更新等价。

证明比较简单，只需要写出两个目标函数对应的梯度更新公式即可。

作者通过分析发现，在 adaptive gradient 方法中，对于 L2 regularization，梯度和 regularization 是打包在一起考虑的。而 weight decay 是分开考虑的。这就导致了对于梯度比较大的权重，L2 regularization 的学习率较小，从而 regularization 效应减弱。而 weight decay 中，这种效应则不存在。因此 weight decay 的 regularization 效应更强。

作者通过这个分析，给出了一个 weight decay 与 L2 regularization 相等的条件

Proposition 3 令 $O$ 为一个 optimizer, 其目标函数为 $f_t(\theta)$, 当不考虑 weight decay 时，梯度更新过程为 $\theta_{t+1}\gets \theta_t-\alpha M_t\nabla f_t(\theta_t)$. 当考虑 weight decay 时，梯度更新过程为 $\theta_{t+1}\gets (1-\lambda)\theta_t-\alpha M_t\nabla f_t(\theta_t)$. 如果 $M_t= \mathrm{diag}(s)^{-1}$ ($s_i>0,\forall i$), 则 $O$ 在优化目标函数
$$ > f_t^{reg}(\theta)=f_t(\theta)+\frac{\lambda'}{2\alpha}\|\theta\odot \sqrt{s}\|_2^2 > $$
时，不考虑 weight decay 的梯度更新与 $O$ 在优化目标函数 $f_t(\theta)$ 时，考虑 weight decay 的梯度更新等价。

上面的结论显示，对于比较大的 preconditioner $s_i$, 其在相比于 L2 regularization 被 regularized 的效应更强。

为了解耦这两个参数，作者提出了 SGDW 算法，其 weight decay 和梯度更新同时进行，算法如下图所示

在算法中，为了支持同时给 $\alpha$ 和 $\lambda$ 做 scheduling, 作者提出了一个 scaling factor $\eta_t$, $\eta_t$ 由用户定义的 scheduler SetScheduleMultiplier(t) 决定。此时，针对 SGD with momentum 的 weight decay 与 L2 regularization 是等价的

同理，我们也可以对 Adam 算法实行同样的操作，算法如下图所示

Conclusion

作者在本文中分析了 adaptive gradient 方法中 L2 regularization 与 weight decay 的不一致性。基于分析，作者提出了 SGDW 和 AdamW 两个优化算法。

References

arxiv

Notes on Adam

Thu, 04 Sep 2025 10:11:55 +0800

作者提出了 Adam, 一个一阶的优化方法，Adam 更加高效，且具有 scaling invariant 的性质。

Introduction

作者首先回顾了一下已有优化器的进展，其中主要是 SGD. 在本文中，作者提出了 Adam, 一个针对高维参数空间的一阶优化器，Adam 基于 gradient 的一阶和二阶信息为不同的参数安排不同的学习率。Adam 的来源是 adaptive moment estimation. Adam 主要是结合了 AdaGrad 和 RMSProp 两个算法的优点。

Adam 与 RMSProp 的区别在于：

RMSProp 在 rescaled gradient 上进行 momentum 的计算然后更新，而 Adam 直接使用一阶和二阶矩来进行估计
RMSProp 没有 bias-correction 项

Adam 的主要优势为：

参数更新的量级与 gradient 的 scaling 无关
步长被 stepsize 超参数限制
不要求目标函数 stationary
对于稀疏梯度 work 的比较好
优化器自带 annealing

Algorithm

Adam 的算法如下图所示

我们优化的目标函数如下

$$ \min_{\theta}\quad f(\theta) $$

这里 $f$ 一般是一个神经网络。我们记 $f(\theta)$ 在 $\theta_t$ 处的梯度为 $g_t=\nabla_{\theta}f(\theta_t)$.

算法运行时，会更新梯度 $m_t$ 以及梯度二阶矩 $v_t$ 的 exponential moving average. 超参数 $\beta_1,\beta_2$ 负责控制 exponential decay rates. 这里 $m_t$ 和 $v_t$ 分别是一阶动量（均值）和二阶动量（未中心化的 variance）的估计。由于 $m_t$ 和 $v_t$ 的初始化都是 0, 因此他们会引入 bias, 作者在后续通过修正解决了这个问题。

假设 $\epsilon=0$, 如果除了当前时刻 $t$ 之外，之前所有时刻的梯度 $g_i=0,i $$ m_t = (1-\beta_t)g_t, v_t=(1-\beta_2)g_t^2 $$

修正后的一阶和二阶矩分别为

$$ \Delta_t = \alpha \frac{(1-\beta_1)\sqrt{1-\beta_2^t}}{(1-\beta_1^t)\sqrt{1-\beta_2}} $$

当 $t$ 足够大的时候， $\beta_1^t\to0, \beta_2^t\to0$, 此时

$$ \Delta_t = \alpha \frac{1-\beta_1}{\sqrt{1-\beta_2}} $$

如果之前所有时刻的梯度不全为 0, 则依据 Cauchy-Schwarz 不等式，我们有 $(\mathbb{E}[XY])^2\leq \mathbb{E}[X^2]\mathbb{E}[Y^2]$. 令 $X=1$, $Y=g$, 则我们有

$$ (\mathbb{E}[g])^2\leq \mathbb{E}[g^2] \Rightarrow \frac{|\mathbb{E}[g]|}{\sqrt{\mathbb{E}[g^2]}}\leq 1 $$

此时，我们有

$$ \mathbb{E}[g_t] = \hat{m}_t, \mathbb{E}[g_t^2] = \hat{v}_t $$

因此，

$$ |\Delta_t| = \left|\alpha\frac{\hat{m}_t}{\sqrt{\hat{v}_t}}\right|=\alpha \left|\frac{|\mathbb{E}[g]|}{\sqrt{\mathbb{E}[g^2]}}\right|\leq\alpha $$

从而我们有

$$ |\Delta_t| \leq\begin{cases} \alpha \frac{1-\beta_1}{\sqrt{1-\beta_2}} & \text{ if }1-\beta_1>\sqrt{1-\beta_2}\\ \alpha &\text{ otherwise} \end{cases} $$

实际上，$\Delta_t$ 可以理解为一个 trust region, 可以用来保证当前更新的参数不会离原始参数太远。

作者定义 signal-noise ratio (SNR) 为

$$ SNR = \frac{\hat{m}_t}{\sqrt{\hat{v}_t}} $$

当 SNR 较小时，说明此时的不确定性比较大，因此 $\Delta_t$ 也比较小。这就避免了模型朝错误的方向更新。也就是automatic annealing.

$\Delta_t$ 还对 gradient 的 scaling 有不变的性质，这是因为，

$$ \frac{c\cdot\hat{m}_t}{\sqrt{c^2\cdot\hat{v}_t}} = \frac{\hat{m}_t}{\sqrt{\hat{v}_t}} $$

Bias Correction

上一节提到，Adam 算法的初始化是存在 bias 的，作者在本届就对齐进行了分析。令 $g$ 为目标函数 $f$ 的梯度，我们希望估计其二阶动量的期望.令 $g_1,\dots,g_T$ 分别为 $\theta_1,\dots,\theta_T$ 处的梯度估计，其中 $g_t\sim p(g_t)$ 是对应时刻梯度的分布。令 $v_0=0$, 在 $t$ 时刻，我们有

$$ v_t = \beta_2 v_{t-1} + (1-\beta_2)g_t^2 = (1-\beta_2)\sum_{i=1}^t\beta_2^{t-i} g_i^2 $$

我们希望计算 $\mathbb{E}[v_t]$ 与 $\mathbb{E}[g_t^2]$ 之间的关系，我们有

$$ \begin{aligned} \mathbb{E}[v_t] &= \left[(1-\beta_2)\sum_{i=1}^t\beta_2^{t-i} g_i^2\right]\\ &= \mathbb{E}[g_t^2]\cdot (1-\beta_2)\sum_{i=1}^t\beta_2^{t-i}+\zeta\\ &= \mathbb{E}[g_t^2](1-\beta_2^t)+\zeta \end{aligned} $$

其中当 $\mathbb{E}[g_i^2]$ 为 stationary 时，$\zeta=0$, 否则我们可以通过控制 $\beta_2$ 来让 past gradient 保持在一个较小的规模。最后，我们剩下的就是 $1-\beta_2^t$, 这也是我们在算法中进行修正的地方。

对于一阶动量 $m_t$ 的修正也是同理。

Convergence Analysis

作者在本节中使用了 online learning framework 来分写 Adam 的收敛性。给定一系列 convex cost function $f_1(\theta),\dots,f_T(\theta)$. 在 $t$ 时刻，我们的目标是基于上一个 cost function $f_t(\theta)$ 来预测 $\theta_t$.

作者在这里使用 regret 来分析，记 $f_t(\theta^*)$ 为 $t$ 时刻最优的参数对应的 cost function, regret 定义为

$$ R(T) = \sum_{t=1}^T [f_t(\theta_t) - f_t(\theta^*)] $$

其中，

$$ \theta^* = \arg\min_{\theta\in\mathcal{X}}\sum_{t=1}^Tf_t(\theta) $$

则我们有如下的结论

Theorem 1 假设

函数 $f_t$ 的梯度是有界的，即 $\|\nabla f_t(\theta)\|_2\leq G$, $\|\nabla f_t(\theta)\|_{\infty}\leq G_{\infty}$ 对任意 $\theta\in\mathbb{R}^d$ 都成立

$\{\theta_1,\dots,\theta_T\}$ 中任意两个参数的距离都是有界的，即 $\|\theta_m-\theta_m\|_2\leq D$, $\|\theta_m-\theta_n\|_{\infty}\leq D_{\infty}$ 对任意 $m,n\in\{1,\dots,T\}$ 都成立

$\beta_1,\beta_2\in[0,1)$ 满足 $\frac{\beta_1^2}{\sqrt{\beta_2}}<1$ 令 $\alpha_t=\alpha/\sqrt{t}$, $\beta_{1,t}=\beta_1\lambda^{t-1}$, $\lambda\in(0,1)$, 则我们有

$$ > R(T)\leq \frac{D^2}{2\alpha(1-\beta_1)}\sum_{i=1}^d\sqrt{T\hat{v}_{T,i}}+\frac{\alpha(1+\beta_1)G_{\infty}}{(1-\beta_1)\sqrt{1-\beta_2}(1-\gamma)^2}\sum_{i=1}^d\|g_{1:T,i}\|_2+\sum_{i=1}^d\frac{D_{\infty}^2G_{\infty}\sqrt{1-\beta_2}}{2\alpha(1-\beta_1)(1-\lambda)^2} > $$

结果说明，当我们的 data feature 稀疏且梯度有界时我们有

$$ \sum_{i=1}^d\|g_{1:T,i}\|_2<< dG_{\infty}\sqrt{T} $$

以及

$$ \sum_{i=1}^d\sqrt{T\hat{v}_{T,i}}<< dG_{\infty}\sqrt{T} $$

实际上，对于 Adam 以及 Adamgrad，这个上界可以优化到 $O(\log d\sqrt{T})$.

最终，我们可以证明 Adam 的收敛性

Corollary 1 假设

函数 $f_t$ 的梯度是有界的，即 $\|\nabla f_t(\theta)\|_2\leq G$, $\|\nabla f_t(\theta)\|_{\infty}\leq G_{\infty}$ 对任意 $\theta\in\mathbb{R}^d$ 都成立

$\{\theta_1,\dots,\theta_T\}$ 中任意两个参数的距离都是有界的，即 $\|\theta_m-\theta_m\|_2\leq D$, $\|\theta_m-\theta_n\|_{\infty}\leq D_{\infty}$ 对任意 $m,n\in\{1,\dots,T\}$ 都成立则对 $T\geq1$, 我们有

$$ > \frac{R(T)}{T}=O\left(\frac{1}{\sqrt{T}}\right) > $$

Experiment

作者在 logistic regression, MLP, CNN 等三种模型架构上进行了实验。

Conclusion

作者在本文中提出了 Adam optimizer，一个基于 AdaGrad 和 RMSProp 优点的优化器，作者通过理论验证了 Adam 的收敛性，然后通过实验验证了 Adam 的有效性。

References

arxiv

Notes on DeepSeekMoE

Fri, 29 Aug 2025 11:03:12 +0800

DeepSeek 在 2024 年 1 月发布了 DeepSeekMoE, 一个解决 MoE 模型 specialization 不足以及 redundancy 问题的大模型系列。

Introduction

作者首先回顾了已有 MoE 模型的不足，主要有两点：

knowledge hybridity: 已有 MoE 模型的专家个数比较少，这就导致每个专家需要掌握更多样化的知识，提高了训练难度
knowledge redundancy: 不同专家掌握的知识可能有重叠，从而导致了模型参数的 redundancy

为了解决这两个问题，作者提出了 DeepSeek-MoE, DeepSeek-MoE 主要做出了两点改变：

Fine-Grained Expert Segmentation: 作者使用了更多的专家，来提高每个专家的 specialization, 降低训练成本
Shared Expert Isolation: 作者在 Routing expert 的基础上，加入了 shared expert 来学习 common knowledge.

作者在 2B-A0.6B 的模型进行了实验，结果显示模型表现超过了 GShard, 说明了 DeepSeekMoE 模型架构的有效性。

作者还进一步将模型 scale 到了 16B-A2.8B 和 145B-A22B, 实验结果均验证了模型的 scaling 效果。

Method

Preliminary

作者首先回顾了 transformer 架构，transformer 的第 $\ell$ 层 decode layer 可以表示为

$$ \begin{aligned} u_{1:T}^{\ell} &= \mathrm{self\_attention}(h_{1:T}^{\ell-1}) + h_{1:T}^{\ell}\\ h_t^\ell &= \mathrm{FFN}(u_t^{\ell}) + u_t^{\ell} \end{aligned} $$

其中 $T$ 是 sequence length, $h_{1:T}^{\ell-1}$ 是第 $\ell-1$ 层 decoder layer 输出的 hidden states.

接下来，我们可以将 dense 架构转换为 MoE 架构，MoE 架构与 dense 架构不同的地方在与 $\mathrm{FFN}$ 不再是 MLP, 而是一个 MoE 模块，其表达式如下

$$ \begin{aligned} h_t^\ell &= \sum_{i=1}^N\left(g_{i,t}\mathrm{FFN}_i(u_t^{\ell})\right) + u_t^{\ell}\\ g_{i,t} &= \begin{cases} s_{i,t,}, &s_{i,t}\in\mathrm{Topk}(\{s_{j,t}\mid 1\leq j \leq N\},K)\\ 0, &\text{otherwise} \end{cases}\\ s_{i,t,} &= \mathrm{softmax}_i({u_t^{\ell}}^Te_i^{\ell}) \end{aligned} $$

这里 $N$ 是专家的总个数， $K$ 是激活专家个数， $e_{i}^{\ell}$ 是 routing layer 的权重矩阵，$\mathrm{FFN}_i$ 是每个专家对应的 FFN.

DeepSeekMoE Architecutre

DeepSeekMoE 架构如下图所示

相比于其他 MoE 架构，DeepSeekMoE 主要做了以下几点改变。

Fine-Grained Expert Segmentation

作者首先解决了每个专家学习内容过多的问题。作者的做法就是将每个 expert FFN 分割为 $m$ 个更小的专家，具体做法就是将 FFN intermediate hidden size 降低为原来的 $1/m$. 这样的话就可以在不增加模型参数量的情况下提高模型的表现。修正后的 MoE 模块为

$$ \begin{aligned} h_t^\ell &= \sum_{i=1}^{mN}\left(g_{i,t}\mathrm{FFN}_i(u_t^{\ell})\right) + u_t^{\ell}\\ g_{i,t} &= \begin{cases} s_{i,t,}, &s_{i,t}\in\mathrm{Topk}(\{s_{j,t}\mid 1\leq j \leq mN\},mK)\\ 0, &\text{otherwise} \end{cases}\\ s_{i,t,} &= \mathrm{softmax}_i({u_t^{\ell}}^Te_i^{\ell}) \end{aligned} $$

可以看到，现在我们一共有 $mN$ 个专家，激活专家个数为 $mK$ 个。作者认为，通过提高专家的粒度，我们可以有效增加专家组合的可能性，这就提高了最终的组合多样性。

Shared Expert Isolation

接下来，作者介绍了解决不同专家学习到重复知识的问题，作者的做法是在 routing Expert 的基础上加入 Shared expert. 也就是说，有固定几个专家始终都会被激活，这部分专家复杂学习通用知识，从而减少知识冗余。增加 shared expert 之后的 MoE 模块为

$$ \begin{aligned} h_t^\ell &= \sum_{i=1}^{K_s}\mathrm{FFN}_i(u_t^{\ell})+\sum_{i=K_s+1}^{mN}\left(g_{i,t}\mathrm{FFN}_i(u_t^{\ell})\right) + u_t^{\ell}\\ g_{i,t} &= \begin{cases} s_{i,t,}, &s_{i,t}\in\mathrm{Topk}(\{s_{j,t}\mid K_s+1\leq j \leq mN\},mK-K_s)\\ 0, &\text{otherwise} \end{cases}\\ s_{i,t,} &= \mathrm{softmax}_i({u_t^{\ell}}^Te_i^{\ell}) \end{aligned} $$

此时，模型中一共包含 $K_s$ 个共享专家，$mN-K_s$ 个 routing expert, 其中激活专家个数为 $mK-K_s$.

Load Balancing Loss

接下来，作者解决了训练时的 load imbalance 问题，作者提出了两个 loss 来分别解决不同层面的 load imbalance 问题。

首先，在 expert 层面，作者使用了如下的 load balancing loss:

$$ \mathcal{L} = \alpha_1\sum_{i=1}^{N'}f_iP_i $$

其中 $\alpha_1$ 是超参数，

$$ f_i = \frac{N'}{K'T}\sum_{i=1}^{N'}\mathbb{1}(\text{Token }i \text{ selects Expert }i),\quad P_i = \frac{1}{T}\sum_{t=1}^Ts_{i,t} $$

分别为以及分配给第 $i$ 个专家的 token 比例以及概率之和。$N'=mN-K_s$, $K'=mK-K_s$. $\mathbb{1}(\cdot)$ 是 indicator function.

其次，在 device 层面，作者也是用了 load balancing loss 来减少不同设备之间不必要的通信。作者将 routed experts 分为 $D$ 个 group $\mathcal{E}_1,\dots,\mathcal{E}_D$, 然后每个设备部署一个 group, group level 的 load balancing loss 定义如下：

$$ \mathcal{L} = \alpha_2\sum_{i=1}^D f_i' P_i' $$

其中 $\alpha_2$ 是超参数，

$$ f_i' = \frac{1}{|\mathcal{E}_i|}\sum_{j\in\mathcal{E}_i}f_i,\quad P_i' = \sum_{j\in\mathcal{E}_i}P_i $$

实际中，作者使用了一个较小的 $\alpha_1$ 来避免 routing collapse, 使用了一个较大的 $\alpha_2$ 来提高 Device 层面的负载均衡。

Training

作者使用了中英文数据进行训练，tokenizer 基于 BPE 算法，vocab size 为 8K. 模型训练基于 HAI-LLM.训练时使用了 TP, DP, PP, EP 等并行策略。

2B, 16B, 145B 模型的参数如下表所示

Model	2B	16B	145B
total params	2B	16.4B	144.6B
activated params	0.3B	2.8B	22.2B
hidden size	1280	2048	4096
layers	9	28	62
attention heads	10	16	32
head dimension	128	128	128
routed experts	63	64	128
activated experts	7	6	12
shared experts	1	2	4
training tokens	100B	2T	245B

Experiments

Alignment

作者针对 DeepseekMoE 16B 进行了微调，微调使用了 1.4M 的训练样本，覆盖了 math, code, QA, reasoning 等任务。

Ablation Study

作者在 2B 的模型上进行了 ablation study.

首先，作者探究了细粒度专家和共享专家的有效性，结果如下图所示

实验结果显示，与 GShard 相比，使用共享专家可以有效提高模型的表现。并且，使用更细粒度的专家也可以进一步提高模型的表现

作者还探究了共享专家与路由专家的比例，作者分别使用不同的比例进行实验，结果发现共享专家：路由专家个数为 1：3 的时候模型效果最好。

作者还探究了模型的泛化性，作者 mask 掉一部分概率最高的 routing expert, 然后从剩下的专家里进行 topK 的挑选，然后作者比较模型和 GShard 的表现，结果如下图所示

实验结果显示，DeepSeekMoE 对于 mask 操作更敏感，这说明 DeepSeekMoE 模型中专家的 specialization 更强。

作者还探究了 mask 掉共享专家对模型表现的影响，结果显示共享专家与路由专家之间的 overlap 很小，去掉共享专家之后，模型表现会变差。

作者进一步分析了共享专家与路由专家组合的有效性。作者探究 DeepSeekMoE 是否可以使用更少的路由专家来获取知识。作者通过使用不同的 activated routed experts 来进行实验，实验结果如下图所示

实验结果显示，DeepSeekMoE 仅需激活 4 个路由专家，就可以达到与 GShard 相同的表现。这说明了 DeepSeekMoE 模型中每个专家可以学习到更准确的知识。

Conclusion

作者在本文中提出了 DeepSeekMoE 架构，来提高 MoE 模型中专家的利用效率。为了达到这一点，作者首先使用了更细粒度的专家，降低每个专家的学习成本，然后，作者使用了共享专家，来降低不同专家之间的知识冗余。作者先在 2B 的模型上验证了方法的有效性，然后作者将模型 scale 到了 16B 和 125B。结果显示模型的效果均超过了以前的工作。

References

arxiv

Notes on DeepSeek-LLM

Tue, 26 Aug 2025 10:53:10 +0800

DeepSeek 在 2024 年 1 月 5 日发布了 DeepSeek LLM, 包括 7B 和 67B 两个 size, 作者主要强调了对于 scaling law 的探究

Introduction

已有的 scaling law 如 Kaplan 和 Chinchilla 介绍了 model size, dataset size, compute budget 与模型表现之间的关系。在本文中，作者进一步探究了 learning rate 和 batch size 等超参数与模型表现之间的关系。基于发现的 scaling law, 作者为不同大小的模型设置了最优的超参数。并且，作者还发现不同数据集与模型表现之间的关系。

最终，基于这些实验结果，作者提出了 DeepSeek LLM, 模型使用 2T token 进行预训练，使用 1M samples 进行后训练，后训练包括 SFT 以及 DPO.

Pre-training

Architecture

DeepSeek-LLM 的架构与 LLaMA 基本相同，作者在 67B 的模型上使用了 GQA 来提高 inference 效率。最终模型的配置如下表所示

Params	7B	67B
$n_{\text{layers}}$	$30$	$95$
$d_{\text{model}}$	$4096$	$8192$
$n_{\text{heads}}$	$32$	$64$
$n_{\text{kv\_heads}}$	$32$	$8$
Context Length	$4096$	$4096$
Sequence Batch Size	$2304$	$4608$
Learning Rate	$4.2e-4$	$3.2e-4$
Tokens	2T	2T

Data

作者主要从 Common Crawl 构建预训练数据，数据处理过程包括：去重，过滤以及 remixing 三个步骤。

对于 tokenizer, 作者使用了 BBPE 算法，tokenizer 的大小设置为 100,000, 最终的 tokenizer 大小为 102400.

Hyper Parameters

作者主要对比了一下不同 learning rate schedule 的表现：

cosine learning schedule
multi-step learning rate schedule: 包含三个 Stage, 第一个 stage 保持最大学习率，第二个 stage 将学习率降低为最大学习率的 $31.6\%$, 第三个 stage 降低为最大学习率的 $10\%$.

对比的实验结果如下图所示

实验结果显示，multi-step learning rate scheduler 的表现与 cosine learning rate 表现差不多。并且，multi-step learning rate scheduler 对于 continue pretraining 支持更好。因此在本文中作者使用了 multi-step learning rate scheduler.

Infra

作者使用了数据并行，张量并行，序列并行以及 1F1B pipeline 并行。作者还使用了 flash attention 来提高硬件利用率。

Scaling Law

本节中，作者分析了 scaling law, 主要有以下三点：

构建了针对 learning rate 和 batch size 的 scaling law
作者使用 non-embedding FLOPs/token $M$ 来表示 model scale
预训练数据的质量对最后中的 scaling 影响很大

作者首先构建了针对 batch size 和 learning rate 的 scaling law, 结果显示最优的 learning rate 和 batch size 范围都比较广，这个结论与 Kaplan 一致。

接下来，作者构建了 batch size $B$, learning rate $\eta$ 与 compute budget $C 之间的关系，实验结果如下图所示

拟合得到的曲线为

$$ \begin{aligned} \eta_{opt} &= 0.3118* C^{-0.1250}\\ B_{opt} &= 0.2920 * C^{0.3271} \end{aligned} $$

可以看到，随着 compute budget 增加，$B_{opt}$ 也逐渐增加，而 $\eta_{opt}$ 逐渐减小。并且，最优参数的范围都比较广。

接下来，作者进一步探究了 batch size 与 generalization error $L$ 之间的关系。作者希望找到 model scale $N$, data scale $D$ 与 compute budget $C$ 之间的关系，即

$$ N_{opt} \varpropto C^a,D_{opt} \varpropto C^b $$

compute budget 与 model scale, data scale 之间的关系可以近似表示为 $C=6ND$, 这个公式的推导见 LLM FLOPs computation。我们用 $N_1,N_2$ 分别表示模型的 non-embedding parameter 以及 complete parameters, 则我们可以用 $6N_1$ 或者 $6N_2$ 来近似 model scale, 但是 $6N_1$ 和 $6N_2$ 均没有考虑 attention 的计算开销，因此这两种近似的误差都比较大。

为了解决这个问题，作者提出了一个新的 model scale 表示形式，即 non-embedding FLOPS/token $M$, 其中 $M$ 包含 attention 的计算开销但是不包含 vocabulary computation. 基于这种表示，compute budget 可以近似表示为 $C=MD$. $M$ 与 $6N_1,6N_2$ 的区别表示如下所示

$$ \begin{aligned} 6N_1 &= 72nd^2\\ 6N_2 &= 72nd^2 + 6Vd\\ M &= 72nd^2+12ndl \end{aligned} $$

其中, $d$ 是 hidden size, $n$ 是 layers 个数, $V$ 是 vocabulary size, $l$ 是 sequence length. 作者在不同 scale 的模型上比较了三种表示方式，结果发现 $6N_1$ 和 $6N_2$ 要么低估，要么高估了模型的参数量。

基于 model scale 的表示方式，作者构建了如下的优化问题

$$ M_{opt}(C), D_{opt}(C) = {\arg\min}_{M,D\ s.t.\ C=MD} L(N,D) $$

作者使用了 Chinchilla 提出来的 IsoFLOP 曲线进行拟合，实验结果如下图所示

拟合的曲线为

$$ M_{opt}(C) = 0.1715*C^{0.5243}, D_{opt}(C) = 5.8316*C^{0.4757} $$

作者还进一步拟合了 compute budget 与 optimal generalization error 之间的关系，结果如下图所示

实验结果显示，作者提出的 scaling law 可以很好预测模型的表现。

最后，作者探究了以下不同数据集的 scaling law, 作者使用 early in-house data, current in-house data 以及 OpenWebText2 来将进行实验，结果如下图所示

结果显示，scaling law 与数据质量高度相关。当数据质量提升时，model scaling exponent $a$ 逐步提升，data scaling exponent $b$ 逐步下降，说明 compute budget 更多由模型参数量决定。因此，作者认为提升 compute budget 之后，我们应该优先提高模型的 model size.

Post-training

作者构建了 1.5M 的中英文指令数据。其中安全性的数据有 300K, 有帮助性的数据有 1.2M, 其中包括 $31.2\%$ 的通用数据，$46.6\%$ 的数学相关数据，$22.2\%$ 的代码数据。

post-training 包含两个阶段：

SFT:7B 的模型训练了 4 个 epoch, 67B 的模型训练了 2 个 epoch, 作者发信进一步训练 67B 的模型会导致过拟合。作者发现，模型在训练过程中会出现重复输出的情况，特别是数学 SFT 数据，为了解决这个问题，作者使用了一个两阶段的 SFT 以及 DPO.
DPO: 提高模型的能力，作者发现 DPO 可以提高模型 open-ended generation skill.

Evaluation

我们主要关注一下消融实验。

首先作者探究了分阶段 SFT 对模型表现的影响。作者发现，小模型在 math 和 code 数据集上需要训练更长时间，但是这也损害了模型的对话能力。为了解决这个问题，作者使用两阶段的训练模式，第一个阶段使用所有的数据进行训练，第二个阶段仅使用对话数据进行训练，实验结果如下表所示

Model	HumanEval	GSM8K	Repetition	IFEval
DeepSeek LLM 7B Chat Stage1	48.2	63.9	0.020	38.0
DeepSeek LLM 7B Chat Stage2	48.2	63.0	0.014	41.2

可以看到，经过第二阶段训练之后，模型的表现有所提升

接下来，作者探究了 Multi-choice question 对模型表现的影响，MCQ 要求模型不仅需要有相关的知识，还要理解每个选项的含义。作者使用 20M 中文 MCQ 来进行消融实验，结果如下表所示

Model	MMLU	CEval	CMMLU	TriviaQA	ChineseQA
DeepSeek LLM 7B Chat	49.4	47.0	49.7	57.9	75.0
DeepSeek LLM 7B Chat + MC	60.9	71.3	73.8	57.9	74.4

实验结果显示，MCQ 确实可以提高模型在上述几个 benchmark 上的表现，但是其泛化性会下降。因此，作者在 pre-training 和 fine-tuning 阶段并没有使用 MCQ 数据进行训练。

作者还探究了在 pre-training 阶段加入 instruction data, 来提高 base model 在下游 benchmark 上的表现。结果发现，base model 的表现提升优先。作者认为，尽管 instruction data 可以提高 base model 表现，但是如果 Instruction data 数量过少，则模型表现不太可能学习到有用的知识。因此，作者的做法是不在 pretraining 阶段加入 Instruction data.

最后，作者探究了 system prompt 对模型表现的影响。受 LLaMA 2 启发，作者也尝试在输入中加入 system prompt. 实验结果如下所示

Model	MT Bench
DeepSeek LLM 7B Chat	7.15
DeepSeek LLM 7B Chat + System Prompt	7.11
DeepSeek LLM 67B Chat	8.35
DeepSeek LLM 67B Chat + System Prompt	8.58

可以看到，7B 的模型加入 system prompt 之后，模型表现有所下降；67B 的模型加入 system prompt 之后，模型表现有所提升。作者认为，大模型更容易理解 system prompt 的意图，而小模型的指令跟随能力则较差，因此 system prompt 反而会影响模型表现。

Conclusion

作者在本文中提出了 DeepSeek LLM 系列大语言模型，作者详细介绍了超参数的选择以及 scaling law 等。

References

arxiv

Notes on MFA

Sat, 23 Aug 2025 16:04:34 +0800

阶跃星辰等提出了 Multi-matrix Factorization Attention (MFA), 一个新型注意力机制，用于在 KV cache 限制下最大化模型的表现。

Introduction

multi-head attention (MHA) 的问题在于，其 KV cache 的内存占用（memory footprint）随 sequence length 以及 batch size 线性增长，从而成为了 LLM 在 decoding 阶段的瓶颈。

为了解决 MHA 的内存占用过高问题，已有的工作如 MQA, GQA 等通过共享 key, value projection 来降低 KV cache size. 而 DeepSeek-V3 提出的 MLA 则是通过对 key, value projection 进行 low-rank compression, 然后只存储 latents 的方法来降低 KV cache size.

但是，已有的这些方法的问题在于，当我们设置 KV cache budget 之后，它们的表现就比标准的 MHA 要差。

基于以上这些发现，作者首先分析了已有 attention 机制的 modeling capacity, 然后使用一个统一的框架来表示这些 attention 机制。作者发现，attention heads 的个数以及 dimension 对模型表现有较大影响。

基于这个发现，作者提出了 Multi-matrix Factorization Attention (MFA), 以及其变体 MFA-Key-Reuse (MFA-KR). MFA 的主要目的是在有限的 KV cache size 下提高模型的表现。

Background

作者首先介绍了 GMHA 的概念，GMHA 由三部分组成：

QK circuit: 决定了信息之间如何交互
valueoutput (VO) circuits：决定了信息如何传递
per-head softmax attention.

接下来，作者介绍了 Fully Parameterized Bilinear Attention (FPBA), FPBA 的定义如下：

$$ O = \sum_{c=1}^d\left(\sum_{j=1}^N\phi\left(\frac{xW_cx_j}{H}\right)x_jU_c\right) $$

其中 $\phi$ 是 softmax 函数，$d$ 是模型的 hidden dimension, $N$ 是 sequence length, $W_c,U_c\in\mathbb{R}^{d\times d}$ 每个 channel 上的参数矩阵

每个 channel 都有各自的参数 $W_c, U_c$ 来获取 $x_i$ 与 $x_j$ 之间的信息
提高泛化性，所有 channel 的 $U_c$ 组合起来可以遍历 $d$ 维空间中的任意一个 permutation, 这样就避免来的信息损失
利用率高，FPBA 获取了 $x_i$ 与 $x_j$ 之间 $d$ 维空间可能的表示

基于以上这三个特点，作者认为 FPBA 是 GMHA 框架的一个 capacity upper bound. 此时每个 token 的 KV cache 占用为 $2d^2$ (key and value).

然后，作者分析了 MHA 及其变体与 GMHA 的关系，MHA 可以写作如下形式

$$ \begin{aligned} O &= \sum_{c=1}^h\left(\sum_{j=1}^N\phi\left(\frac{xQ_c(x_jK_c)^T}{\sqrt{d}}\right)x_jV_c\right)O_c^T\\ &= \sum_{c=1}^h\left(\sum_{j=1}^N\phi\left(\frac{x(Q_cK_c^T)x_j^T}{\sqrt{d}}\right)x_jV_cO_c^T\right) \end{aligned} $$

其中 $Q_c,K_c,V_c\in\mathbb{R}^{d\times h_d}$, $O_c\in\mathbb{R}^{d\times h_d}$ 分别是 query, key, value, output projection layer 对应的权重矩阵，$n$ 是 attention head 的个数，令 $h_d$ 为每个 attention 的 head dimension，则我们有 $nh_d=d$.

可以看到，MHA 实际上是一个特殊的 FPBA, 其中，$W_c$ 和 $U_c$ 分别由秩为 $h_d$ 的低秩分解 $Q_cK_c^T$ 以及 $V_cO_c^T$ 近似。此时每个 token 的 KV cache 占用为 $2d$ (key and value).

MQA 可以看作是 GQA 的一个特殊情况。对于 GQA 来说，我们有一个 group size $g\in[1, h]$, 当 $g=1$ 时，GQA 就是 MHA. 当 $g=h$ 时，GQA 就是 MQA, 通常 $g$ 满足 $h\ \%\ g=0$. GQA 的表达式与 MHA 基本相同，只是多个 head 会共享一个 $K_c$ 以及 $V_c$. 此时，每个 token 的 KV cache 占用为 $2gh_d$. 对于 MQA，其每个 token 的 KV cache 占用为 $2h_d$.

对于 MLA, 其表达式如下所示

$$ \begin{aligned} O &= \sum_{c=1}^m\left(\sum_{j=1}^N\phi\left(\frac{xS_QQ_c(x_jS_KK_c)^T}{\sqrt{d}}\right)x_jS_VV_c\right)O_c^T\\ &= \sum_{c=1}^m\left(\sum_{j=1}^N\phi\left(\frac{x(S_QQ_cK_c^TS_K^T)x_j^T}{\sqrt{d}}\right)x_jS_VV_cO_c^T\right) \end{aligned} $$

其中，$S_Q,S_K,S_V\in\mathbb{R}^{d\times C}$ 在所有的 heads 中是共享的，$Q_c,K_c,V_c\in\mathbb{R}^{C\times h_d}$ 是每个 head 的 query, key, value projection layer 的参数，是 latent factorization 的维度。与 FPBA 相比，我们可以看到，MLA 实际上是在 $d/m$ 个 head 上共享了参数，其中，$W_c$ 和 $U_c$ 分别由秩为的低秩分解 $S_QQ_cK_c^TS_K^T$ 以及 $S_VV_cO_c^T$ 近似。尽管模型中 $C>h_d$, 但是最终的 rank 仍然是 $h_d$, 因此模型的表现也就受到了限制。

Method

对已有的 attention 分析之后，作者认为，要提高模型的表现，attention 需要做到亮点：

最小化 KV cache 占用和参数量
attention 的 capacity 尽可能接近 FPBA

基于这两个原则，作者提出了 MFA, MFA 主要依赖三个策略：

提升 attention heads 的 head dimension, 通过提高 head dimension, 我们可以有效提高 attention head 的表达能力
使用矩阵分解来降低参数量
使用单一的 KV head 来降低 KV cache 内存占用

最终，MFA 的表达式如下所示

$$ \begin{aligned} O &= \sum_{c=1}^m\left(\sum_{j=1}^N\phi\left(\frac{xS_QQ_c(x_jS_K)^T}{\sqrt{d}}\right)x_jS_V\right)O_c^T\\ &= \sum_{c=1}^m\left(\sum_{j=1}^N\left(\frac{x(S_QQ_cS_K^T)x_j^T}{\sqrt{d}}\right)x_jS_VO_c^T\right) \end{aligned} $$

其中 $S_Q,S_K,S_V\in\mathbb{R}^{d\times C}$ 是所有的 attention head 所共享的，$Q_c,O_c\in\mathbb{R}^{C\times C}$ 是每个 head 的 query up projection 和 output projection, $C$ 是 latent factorization 的维度。

在 inference 的时候，由于我们只需要保存 $x_jS_K$ 和 $x_jS_V$, 因此所需要的 KV cache size 为 $2C$. 与 FPBA 相比，MFA 分别使用 $S_QQ_cS_K^T$ 和 $S_VO_c^T$ 来近似 $W_c$ 和 $U_c$, 近似矩阵的 rank 为 $C$. 由于 $C>d$, 因此其表达能力也更强，MFA 有如下优势：

scalable head count: MFA 可以支持使用更多的 attention heads, 每增加一个 heads, 所需要的额外参数为 $2C^2$. 并且，增加 attention heads 个数不会增加 KV cache 占用
enhanced head expressiveness: MFA 近似矩阵的 rank 为 $C>d$, 因此表达能力更强
Compatibility with position encodings: MFA 可以无缝集成 position encoding.

为了进一步降低 MFA 的 KV cache 占用，作者提出了 MFA-Key-Reuse (MFA-KA). 核心思想是使用 $S_K$ 来表示 $S_V$, 这样可以额外降低 $50\%$ 的 KV cache 占用，表示方法如下所示

$$ S_V = S_K + \alpha\odot NS_K = (I +\mathrm{diag}(\alpha)N)S_K $$

其中 $N\in\mathbb{R}^{N\times N}$, $\alpha\in\mathbb{R}^C$.

最终，MFA, MFA-KR 与 GQA 的对比如下图所示

不同 attention 的量化对比如下表所示

Method	KV Cache	Parameter	Heads	Factor. rank per head	Shared latent subspace Dim.	Total effec. rank
FPBA	$2d^2$	$2d^3$	$d$	$d$	$d$	$d^2$
MHA	$2d$	$4d^2$	$n$	$h_d$	$d$	$nh_d$
MQA	$2h_d$	$(2 + 2/n)d^2$	$n$	$h_d$	$h_d$	$nh_d$
GQA	$2gh_d$	$(2 + 2g/n)d^2$	$n$	$h_d$	$gh_d$	$nh_d$
MLA	$2C$	$5dC + d^2$	$m$	$h_d$	$C$	$mh_d$
MFA	$2C$	$3Cd + 2mC^2$	$m$	$C$	$C$	$mC$

Code

 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


class Step3vAttention(nn.Module):
 def __init__(self, config: Step3VLConfig, layer_idx):
 super().__init__()
 self.config = config
 self.layer_idx = layer_idx
 self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
 self.num_key_value_heads = 1
 self.total_num_kv_heads = self.num_key_value_heads
 self.num_attention_heads = config.num_attention_heads
 self.num_key_value_groups = config.num_attention_heads // self.num_key_value_heads
 self.q_size = getattr(config, "share_q_dim", self.head_dim)
 self.kv_size = self.num_key_value_heads * self.head_dim
 self.scaling = self.head_dim**-0.5
 self.is_causal = True

 self.q_proj = nn.Linear(config.hidden_size, self.q_size , bias=False)
 self.k_proj = nn.Linear(config.hidden_size, self.head_dim, bias=False)
 self.v_proj = nn.Linear(config.hidden_size, self.head_dim, bias=False)
 self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
 # query down projection normalization
 self.inter_norm = Step3vRMSNorm(self.q_size, eps=config.rms_norm_eps)
 # query up projection
 self.wq = nn.Linear(self.q_size, self.head_dim * self.num_attention_heads, bias=False)

 def forward(
 self,
 hidden_states: torch.Tensor,
 position_embeddings: Tuple[torch.Tensor, torch.Tensor],
 attention_mask: Optional[torch.Tensor],
 past_key_value: Optional[Cache] = None,
 cache_position: Optional[torch.LongTensor] = None,
 **kwargs: Unpack[FlashAttentionKwargs],
 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
 input_shape = hidden_states.shape[:-1]

 query_states = self.q_proj(hidden_states)
 key_states = self.k_proj(hidden_states).view((*input_shape, -1, self.head_dim)).transpose(1, 2)
 value_states = self.v_proj(hidden_states).view((*input_shape, -1, self.head_dim)).transpose(1, 2)

 query_states = self.inter_norm(query_states)
 query_states = self.wq(query_states).view((*input_shape, -1, self.head_dim)).transpose(1, 2)

 cos, sin = position_embeddings
 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

 ...

Conclusion

作者在本文中提出了 MFA 以及 MFA-KR, 一个在 KV cache 有限的条件下最大限度提高 attention 表达能力的 attention 机制。

References

Notes on MX-format

Thu, 21 Aug 2025 18:23:03 +0800

MX format 是一个表示数据的数据格式，在 LLM 中主要用于量化。相比于直接对整个张量进行量化，MX format 可以在更细粒度的层面控制量化，从而提高模型的表现

Microscaling

Microscaling (MS) format 如下图所示

MX format 包括三个部分：

elements $P_1,\dots,P_k$ 未 scale 的数据，要求 $P_1,\dots,P_k$ 的数据类型相同
shared scale $X$, 对 element 进行的 scale 参数，所有的 $k$ 个 bits 共享一个 $X$
block size, 决定 element block 的大小

在存储时，我们只需要存储 $X$ 以及 $P_1,\dots,P_k$, 我们假设 $X$ 需要 $w$ bits 来表示，$P_i$ 需要 $d$ bits 来表示，则我们一共需要 $w+kd$ bits 来表示这 $k$ 个元素。

Concrete MX-compliant Formats

MX-format 包含了一下几种数据格式：

Format Name	Element Data Type	Element Bits(d)	Scaling Block Size(k)	Scale Data Type	Scale Bits(w)
MXFP8	FP8 (E5M2)	8	32	E8M0	8
MXFP8	FP8 (E4M3)	8	32	E8M0	8
MXFP6	FP6 (E3M2)	6	32	E8M0	8
MXFP6	FP6 (E2M3)	6	32	E8M0	8
MXFP4	FP4 (E2M1)	4	32	E8M0	8
MXINT8	INT8	8	32	E8M0	8

GPT-oss Quantization

gpt-oss 中使用了 MXFP4 来表示 MoE 中的 down projection 以及 up projection weight matrix 的权重。

其具体操作过程如下：

我们将参数分为大小为 32 的 block
每个 block 由一个 scale $X$ 来表示，其精度为 E8M0, 即 8bits, 表示范围为 $[-127,127]$, 以及 $32$ 个元素 $P_i$ 来表示，每个元素的精度为 E2M1, 即 4bits, 表示范围为 $[-6.0,6.0]$.
由于每个元素由 4bits 来表示，因此我们将两个元素合并在一起来表示

在加载时，我们可以用如下代码来恢复 $P_i$ 的值到 FP8

 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


FP4_VALUES = [+0.0,+0.5,+1.0,+1.5,+2.0,+3.0,+4.0,+6.0,-0.0,-0.5,-1.0,-1.5,-2.0,-3.0,-4.0,-6.0]


def convert_moe_packed_tensors(
 blocks,
 scales,
 *,
 dtype: torch.dtype = torch.bfloat16,
 rows_per_chunk: int = 32768 * 1024,
) -> torch.Tensor:
 import math
 # scales are represented with uini8
 scales = scales.to(torch.int32) - 127

 assert blocks.shape[:-1] == scales.shape, f"{blocks.shape=} does not match {scales.shape=}"

 lut = torch.tensor(FP4_VALUES, dtype=dtype, device=blocks.device)

 *prefix_shape, G, B = blocks.shape
 rows_total = math.prod(prefix_shape) * G

 blocks = blocks.reshape(rows_total, B)
 scales = scales.reshape(rows_total, 1)

 # each byte representing 2 elements and represented with unit8
 out = torch.empty(rows_total, B * 2, dtype=dtype, device=blocks.device)

 for r0 in range(0, rows_total, rows_per_chunk):
 r1 = min(r0 + rows_per_chunk, rows_total)

 blk = blocks[r0:r1]
 exp = scales[r0:r1]

 # nibble indices -> int64
 idx_lo = (blk & 0x0F).to(torch.long)
 idx_hi = (blk >> 4).to(torch.long)

 sub = out[r0:r1]
 sub[:, 0::2] = lut[idx_lo]
 sub[:, 1::2] = lut[idx_hi]

 torch.ldexp(sub, exp, out=sub)
 del idx_lo, idx_hi, blk, exp

 out = out.reshape(*prefix_shape, G, B * 2).view(*prefix_shape, G * B * 2)
 # to match for now existing implementation
 return out.to(torch.float8_e5m2)

References

Notes on flashattention

Thu, 21 Aug 2025 11:32:53 +0800

作者提出了 flashattention, 一个通过降低 multi head attention 内存访问开销来提高 attention 计算效率的方法

Introduction

Transformer 的 attention 是一个平方度复杂度的算法，这个平方复杂度既体现在时间复杂度上（矩阵乘法），也体现在空间复杂度上（需要存储中间结果）。因此，要降低 attention 的复杂度，我们有两种思路：

从时间复杂度上入手，比如使用稀疏 attention 机制或者线性注意力机制
从空间复杂度上入手，比如使用 GQA, MQA 等减少内存的访问开销

本文提出的 flashattention 就属于降低空间复杂度的一种做法。作者认为，我们应该设计一种 IO-aware 的 attention 算法，来减少 attention 计算式的内存访问开销，进而提高 attention 的计算效率。

作者首先提到，一个未解决的问题就是：

降低 attention 的内存访问开销是否可以提高 attention 的计算效率？

作者发现，已有的一些工作虽然在理论上降低了 attention 的计算效率，但是在实际中，他们的效果并没有提升太多。作者分析原因认为，已有工作主要关注于降低 FLOPs, 但是忽略了内存访问开销。

因此，作者在本文中就提出了 flashattention, 一个 IO-aware 的 attention 算法，作者通过尽可能降低内存访问开销来提高模型的计算效率。具体做法就是，避免从内存中读写 attention matrix, 作者认为这个目标有两个挑战：

计算 softmax 的时候不访问所有的输入
在反向传播时不存储中间的 attention matrix

作者提出了两个方法来分别解决这两个问题：

作者使用了 tiling 技巧，将 input 分成多个 block, 然后分别进行处理，进而降低 softmax 的内存访问开销
作者使用了 recompute 技巧，在反向传播时，重新计算 softmax normalization factor

通过这些改进，我们可以让 attention 运行更快，并且降低内存访问开销。

作者还从理论上分析了 flashattention 的复杂度，提供了理论基础。

作者通过实验验证了 flashattention 的有效性，主要是三点：

训练效率更高：相比于 Huggingface 和 Megatron, flashattention 的训练效率提升了 2-3 倍
模型的表现更好：相比于 GPT-2, 模型的 perplexity 提升了 0.7 个点左右
速度更快：flashattention 比标准的 attention 实现快 3 倍以上

Background

Hardware Performance

作者首先介绍了以下 GPU 的内存架构，如下图所示

可以看到，GPU 内存可以分为三个层级：

SRAM: GPU 的寄存器，容量小，但是访问速度极快
High bandwith memory (HBM): GPU 的高速内存，访问速度较快，容量中等
DRAM: CPU 内存，容量最大，但是访问速度较慢

接下来作者介绍了 Execution model 的概念，GPU 有多个线程来执行同一个操作（SPMD），这个操作也被称为 kernel, kernel 会从 HBM 中加载输入到 SRAM 中进行计算，然后写回 HBM.

对一个算法，我们可以将其归类为 compute-bound 和 memory-bound 两类，我们可以用 arithmetic intensity 来进行区分，arithmetic intensity 定义为 arithmetic operations 与 memory access 的比率。

compute bound: 算法的瓶颈在于算力，由于算力不足导致运行时间慢，比如矩阵乘法
memory-bound: 算法的瓶颈在于内存访问效率，比如 element-wise 操作或者是 reduction

为了提高 memory-bound 类型算法的效率，我们进行 kernel fusion, 即把多个访问同一片内存的操作放在一起处理，避免多次读写内存

Standard Attention Implementation

作者还回顾了一下标准化的 attention 实现。

Forward Pass

给定 $Q,K,V\in\mathbb{R}^{N\times d}$, 其中 $N$ 是序列长度, $d$ 是 head dimension, attention 的定义如下

$$ S = QK^T\in\mathbb{R}^{N\times N},\quad P = \mathrm{softmax}(S)\in\mathbb{R}^{N\times N},\quad O = PV\in\mathbb{R}^{N\times d} $$

这里 $\mathrm{softmax}$ 是逐行计算的。

算法的执行过程如下

我们有第一个结论

Proposition 1 标准化 attention 前向传播时访问 HBM 的内存访问开销为 $\mathcal{O}(Nd+N^2)$.

证明：对于 attention, 我们需要从 HBM 中加载 $Q,K,V\in\mathbb{R}^{N\times d}$, 然后输出 $O\in\mathbb{R}^{N\times d}$ 并保存到内存中。

首先我们需要计算 $S = QK^T$, 这一步需要加载 $Q,K$ 并将 $S$ 保存到 HBM 中，内存访问量为 $\mathcal{O}(Nd+N^2)$.

接下来，我们需要计算 $P = \mathrm{softmax}(S)$, 这一步需要加载 $S$ 然后将 $P$ 保存到 HBM 中，内存访问量为 $\mathcal{O}(N^2)$.

最后，我们需要计算 $O = PV$, 这一步需要加载 $P$ 和 $V$ 然后将 $O$ 保存到 HBM 中，内存访问量为 $\mathcal{O}(Nd+N^2)$.

总的来说，标准化 attention 的内存访问开销为 $\mathcal{O}(Nd+N^2)$.

Backward Pass

标准 attention 反向传播过程如下图所示

Proposition 2 标准化 attention 反向传播时访问 HBM 的内存访问开销为 $\mathcal{O}(Nd+N^2)$.

证明：对于标准化 attention 的反向传播，我们需要从 HBM 中加载 $Q,K,V,dO\in\mathbb{R}^{N\times d}$ , 然后输出 $dQ,dK,dV$ 并保存到 HBM 中。

首先我们计算 $dV=P^TdO$, 这一步需要加载 $P,dO$ 并将 $dV$ 保存到 HBM 中，内存访问开销为 $\mathcal{O}(Nd+N^2)$.

接下来我们计算 $dP=dOV^T$, 这一步需要加载 $dO, V$ 并将 $dP$ 保存到 HBM 中，内存访问开销为 $\mathcal{O}(Nd)$.

然后我们计算 $dS$, 这一步需要加载 $P$ 并将 $dS$ 保存到 HBM 中，内存访问开销为 $\mathcal{O}(N^2)$.

对于 $dQ$ 和 $dK$ 的计算，内存访问开销都是 $\mathcal{O}(Nd+N^2)$.

因此，标准化 attention 的内存访问开销为 $\mathcal{O}(Nd+N^2)$.

Method

作者在本节首先介绍了 flashattention 算法，然后作者证明了 flashattention 的正确性以及分析了复杂度。最后作者对 flashattention 进行扩展得到了 Block-sparse Flashattention.

Flashattention

attention 模块的输入是 $Q,K,V\in\mathbb{R}^{N\times d}$, 输出是 $O\in\mathbb{R}^{N\times d}$, 作者的目标是减少计算过程中的 HBM 访问次数

作者分别使用了 tiling 和 recomputation 来解决 attention 前向传播和反向传播中的内存访问开销。flashattention 的核心思想是，我们将 $Q,K,V$ 分割成 block, 然后在 block 层面进行加载和计算。

Tiling

首先作者介绍了一下如何使用 tiling 来计算 softmax.

给定一个向量 $x\in\mathbb{R}^{B}$, 其 softmax 计算方式如下

$$ m(x) = \max_i x_i,\ f(x) = [e^{x_1-m(x)},\dots,e^{x_B-m(x)}], \ \ell(x)=\sum_if(x)_i, \ \mathrm{softmax}(x) = \frac{f(x)}{\ell(x)} $$

如果我们现在有两个向量 $x^{(1)}, x^{(2)}\in\mathbb{R}^{B}$, 记 $x=[x^{(1)}, x^{(2)}]^T\in\mathbb{R}^{2B}$, 我们可以将 $\mathrm{softmax}(x)$ 的计算分解为

$$ \begin{aligned} m(x) &= \max(m(x^{(1)}), m(x^{(2)}))， f(x) = [e^{m(x^{(1)})-m(x)}f(x^{(1)}),e^{m(x^{(2)})-m(x)}f(x^{(2)})]\\ \ell(x) &= e^{m(x^{(1)})-m(x)}\ell(x^{(1)}) + e^{m(x^{(2)})-m(x)}\ell(x^{(2)}), \mathrm{softmax}(x) = \frac{f(x)}{\ell(x)} \end{aligned} $$

因此，如果我们额外记录 $m(x)$ 以及 $\ell(x)$ 这两个量，那么我们可以每次仅计算 softmax 的一个 block. 具体细节见softmax.

Recomputation

在反向传播过程中，一般我们需要存储 $S,P\in\mathbb{R}^{N\times N}$, 需要的空间复杂度为 $\mathcal{O}(N^2)$. 但是，通过存储 $O\in\mathbb{R}^{N\times d}$ 以及 $(m,\ell)$, 我们可以避免重新计算 $S,P$,这可以看做是 gradient checkpointing. 但是与 checkpointing 相比，因为 flashattention 减少了内存访问开销，因此其反向过程并没有变得更慢。

Algorithm

最终，flashattention 的算法如下图所示

Analysis

Correctness

算法的正确性由定理 1 给出

Theorem 1 flashattention (即算法 1) 输出 $O=\mathrm{softmax}(QK^T)V$, 其时间复杂度为 $\mathcal{O}(N^2d)$, 空间复杂度为 $\mathcal{O}(N)$.

证明：时间复杂度主要由矩阵乘法决定。在计算 $S_{ij}=Q_iK_j^T$ 时，所花费的 FLOPS 为 $\mathcal{O}(B_rB_cd)$. 在计算 $\tilde{P}_{ij}V_j$ 时，所花费的 FLOPS 为 $\mathcal{O}(B_rB_cd)$. 循环一共执行了

$$ T_cT_r = \left\lceil\frac{N}{B_c}\right\rceil\left\lceil\frac{N}{B_r}\right\rceil $$

从而总的 FLOPS 为

$$ \mathcal{O}\left(\frac{N^2}{B_rB_c}B_rB_cd\right) = \mathcal{O}(N^2d) $$

在 flashattention 的计算过程中，我们只需要保存 $(\ell, m)$ 即可，因此需要的额外内存空间为 $\mathcal{O}(N)$.

接下来，我们可以证明 flashattention 的正确性，我们使用归纳法来证明。令 $j$ 满足 $0\leq j\leq T_c$, $K_{:j}\in\mathbb{R}^{jB_c\times d}$, $V_{:j}\in\mathbb{R}^{jB_c\times d}$ 分别为 $K$ 和 $V$ 的前 $jB_c$ 行。 $S_{:, :j}=QK_{:j}^T\in\mathbb{R}^{N\times jB_c}$, $P_{:,:j}=\mathrm{softmax}(S_{:,:j})\in\mathbb{R}^{N\times jB_c}$, $m^{(j)}, \ell^{(j)}, O^{(j)}$ 分别为 $m,\ell, O$ 的第 $j$ 个元素。我们证明经过第 $j$ 次迭代后，HBM 中保存的是

$$ m^{(j)}=\mathrm{rowmax}(S_{:,:j})\in\mathbb{R}^N, \ell^{(j)}=\mathrm{rowsum}(\exp(S_{:,:j}-m^{(j)}))\in\mathbb{R}^N, O^{(j)} = P_{:,:j}V_{:j}\in\mathbb{R}^{N\times d} $$

当 $j=0$ 时，上面的结果显然成立。现在我们假设对某个 $j=0,\dots, T_c-1$ 上面的结果成立，我们需要证明对 $j+1$ 也成立。

首先

$$ m^{(j+1)}=\max(m^{(j)}， \tilde{m}) = \max(\mathrm{rowmax}(S_{:,:j}), \mathrm{rowmax}(S_{:,j:j+1}))=\mathrm{rowmax}(S_{:,:j+1}) $$

接下来

$$ \begin{aligned} \ell^{(j+1)} &= \exp(m^{(j)}-m^{(j+1)})\ell^{(j)} + \exp(\tilde{m}-m^{(j+1)})\tilde{\ell}\\ &=\exp(m^{(j)}-m^{(j+1)})\mathrm{rowsum}(\exp(S_{:,:j}-m^{(j)})) + \exp(\tilde{m}-m^{(j+1)})\mathrm{rowsum}(\exp(S_{:,j:j+1}-\tilde{m}))\\ &= \mathrm{rowsum}(\exp(S_{:,:j}-m^{(j+1)})) + \mathrm{rowsum}(\exp(S_{:,j:j+1}-m^{(j+1)}))\\ &= \mathrm{rowsum}(\exp(S_{:,:j+1}-m^{(j+1)})) \end{aligned} $$

最后，我们计算 $O^{(j+1)}$ 得到：

$$ \begin{aligned} O^{(j+1)} &= \mathrm{diag}(\ell^{(j+1)})^{-1}(\mathrm{diag}(\ell^{(j)})\exp(m^{(j)}-m^{(j+1)})O^{(j)}+\exp(\tilde{m}-m^{(j+1)})\exp(S_{:,j:j+1}-\tilde{m})V_{:,j:j+1})\\ &= \mathrm{diag}(\ell^{(j+1)})^{-1}(\mathrm{diag}(\ell^{(j)})\exp(m^{(j)}-m^{(j+1)})P_{:,:j}V_{:,:j}+\exp(-m^{(j+1)})\exp(S_{:,j:j+1})V_{:,j:j+1})\\ &= \mathrm{diag}(\ell^{(j+1)})^{-1}(\mathrm{diag}(\ell^{(j)})\exp(m^{(j)}-m^{(j+1)})\mathrm{diag}(\ell^{(j)})^{-1}\exp(S_{:,:j}-m^{(j)})V_{:,:j}+\exp(-m^{(j+1)})\exp(S_{:,j:j+1})V_{:,j:j+1})\\ &= \mathrm{diag}(\ell^{(j+1)})^{-1}(\exp(-m^{(j+1)})\exp(S_{:,:j}))V_{:,:j}+\exp(-m^{(j+1)})\exp(S_{:,j:j+1})V_{:,j:j+1})\\ &= \mathrm{diag}(\ell^{(j+1)})^{-1}( \begin{bmatrix} \exp(S_{:,:j}-m^{(j+1)}) & \exp(S_{:,:j}-m^{(j+1)}) \end{bmatrix}\begin{bmatrix} V_{:,:j} \\ V_{:,j:j+1} \end{bmatrix}\\ &= \mathrm{softmax}(S_{:,:j+1})V_{:,:j+1} \end{aligned} $$

因此上面的结果对 $j+1$ 也成立，从而 flashattention 的结果对 $j=0,\dots,T_c$ 都成立。

Forward Pass of flashattention

第一个问题是如何提高 softmax 计算的效率，作者的做法先先计算 normalization constant 然后再分别计算不同的 column.

给定 $Q,K,V\in\mathbb{R}^{N\times d}$, 其中 $N$ 是序列长度, $d$ 是 head dimension, attention 的定义如下

$$ S = QK^T\in\mathbb{R}^{N\times N},\quad P = \mathrm{softmax}(S)\in\mathbb{R}^{N\times N},\quad O = PV\in\mathbb{R}^{N\times d} $$

我们有 $S_{ij}=q_i^Tk_j$, 这里 $q_i$ 和 $k_j$ 分别是 $Q$ 和 $K$ 的第 $i$ 列以及第 $j$ 列， normalization constant 定义为：

$$ L_i = \sum_{j=1}^N \exp\left(q_i^Tk_j\right) $$

对任意 $i$, 计算 $L_i$ 只需要 $\mathcal{O}(N)$ 的空间复杂度。

令 $v_j$ 是 $V$ 的第 $i$ 列，则输出 $O$ 的第 $i$ 列 $o_i$ 为

$$ o_i = P_{i:}V = \sum_{j=1}^N P_{ij}v_j = \sum_{j=1}^N\frac{\exp(q_i^Tk_j)}{L_i}v_j $$

这个过程中，对任意 $i$, 计算 $o_i$ 也只需要 $\mathcal{O}(N)$ 的空间复杂度。

因此，在 $L_i$ 已经计算好的情况下，我们可以在 $\mathcal{O}(N)$ 的空间复杂度下计算 $o_i$.

最终，flashattention 的 forward pass 过程如下图所示

接下来，作者分析了 flashattention 的内存访问开销。结论如下

Theorem 2 令 $N$ 为 sequence length, $d$ 为 head dimension, $M$ 是 SRAM 的 size, 且满足 $d\leq M\leq Nd$. 则 flashattention 前向传播的内存访问开销为 $\Theta(N^2d^2M^{-1})$.

证明：由 Algorithm 1（或者 Algorithm 2）可以知道，$K$ 和 $V$ 的每一个元素都只需要从 HBM 中加载一次，而每一次外层循环都会从 HBM 中加载一次 $O$ 和 $Q$, 因此总的 HBM 访问次数为 $\mathcal{O}(Nd+NdT_c)=\mathcal{O}(NdT_c)$.

接下来，我们给出每一次内层循环的内存访问开销，这是由 SRAM 的大小决定的。由于我们需要 SRAM 可以存储 $K_j\in\mathbb{R}^{B_c\times d}$ 以及 $V_j\in\mathbb{R}^{B_c\times d}$ ，我们的 block size 需要满足

$$ B_cd = \mathcal{O}(M) \Rightarrow B_c = \mathcal{O}\left(\frac{M}{d}\right) $$

同理，对于 $O$ 和 $Q$, 我们有

$$ B_rd = \mathcal{O}(M) \Rightarrow B_r = \mathcal{O}\left(\frac{M}{d}\right) $$

最后，我们还需要 SRAM 可以存储 $S_{ij}\in\mathbb{R}^{B_r\times B_c}$, 因此

$$ B_rB_c=\mathcal{O}(M) $$

这样，

$$ B_c = \mathcal{O}\left(\frac{M}{d}\right), B_r=\mathcal{O}\left(\min\left(\frac{M}{d},\frac{M}{B_c}\right)\right)=\mathcal{O}\left(\min\left(\frac{M}{d},d\right)\right) $$

从而

$$ T_c = \frac{N}{B_c} = \mathcal{O}\left(\frac{Nd}{M}\right) $$

最终，总的内存访问开销为

$$ \mathcal{O}(NdT_c) = \mathcal{O}\left(\frac{N^2d^2}{M}\right) $$

一般来说, $d$ 的大小为 $64-128$, $M$ 的大小为 $100 KB$ 左右, $d^2« M, 因此 flashattention 的内存访问开销远小于标准化 attention 的内存访问开销。

作者还证明 flashattention 的内存访问开销是一个下界，即

Proposition 3 令 $N$ 为 sequence length, $d$ 为 head dimension, $M$ 是 SRAM 的 size, 且满足 $d\leq M\leq Nd$. 则不存在一个对任意 $M\in[d,Nd]$ 都可以在内存访问开销为 $\Theta(N^2d^2M^{-1})$ 的条件下完成 attention 计算的算法。

证明可以用反证法，基本思想是加载 $Q,K,V$ 的 HBM 访问次数至少为 $\mathcal{O}(Nd)$.

Backward Pass of flashattention

第二个问题是能否在线性空间复杂度下计算 attention 的反向传播过程。

首先我们记损失函数为 $\phi$, 然后令 $\phi$ 对 $O,Q,K,V$ 的梯度分别为 $dO,dQ,dK, dV\in\mathbb{R}^{N\times d}$, 我们的目标是计算 $dQ, dK, dV$.

$dV$ 的计算是最容易的，我们有 $dV=P^TdO$, 因此

$$ dv_j = \sum_{i=1}^N P_{ij}do_i = \sum_{i=1}^N\frac{\exp(q_i^Tk_j)}{L_i}do_i $$

由于我们已经计算了 $L_i$, 因此，$dv_j$ 只需要 $\mathcal{O}(d)$ 的空间复杂度。

接下来，注意到 $dP=dOV^T$, 因此我们有

$$ dP_{ij} = do_i^Tv_j $$

计算的空间复杂度也是要 $\mathcal{O}(N)$ 的

注意到 $P_{i:}=\mathrm{softmax}(s_{i:})$, 且 $y=\mathrm{softmax}(x)$ 的 Jacobian 是 $\mathrm{diag}(y)-yy^T$ (推导过程见 softmax), 我们有

$$ dS_{i:} = (\mathrm{diag}(P_{i:})-P_{i:}P_{i:}^T)dP_{i:} = P_{i:} \odot dP_{i:} - (P_{i:}^TdP_{i:})P_{i:} $$

我们定义

$$ D_i = P_{i:}^TdP_{i:}= \sum_{j=1}^N\frac{\exp(q_i^Tk_j)}{L_i}do_i^Tv_j = do_i^T\sum_{j=1}^N\frac{\exp(q_i^Tk_j)}{L_i}v_j = do_i^To_i $$

$D_i$ 的空间复杂度也只需要 $\mathcal{O}(N)$.

则

$$ dS_{i:} =P_{i:} \odot dP_{i:} - D_iP_{i:} $$

我们有

$$ dS_{ij} = P_{ij}dP_{ij} - D_iP_{ij} = P_{ij}(dP_{ij}-D_i) $$

注意到 $S_{ij}=q_i^Tk_j$, 我们有

$$ dq_i = \sum_{j=1}^N dS_{ij}k_j = \sum_{j=1}^NP_{ij}(dP_{ij}-D_i)k_j = \sum_{j=1}^N\frac{\exp(q_i^Tk_j)}{L_i}(do_i^Tv_j-D_i)k_j $$

因此计算 $dq_i$ 的空间复杂度为 $\mathcal{O}(d)$.

同样的，

$$ dk_j = \sum_{j=1}^N dS_{ij}q_i = \sum_{j=1}^NP_{ij}(dP_{ij}-D_i)q_i = \sum_{j=1}^N\frac{\exp(q_i^Tk_j)}{L_i}(do_i^Tv_j-D_i)q_i $$

其空间复杂度为 $\mathcal{O}(N)$.

总之，attention 的反向传播过程所需要的空间复杂度为 $\mathcal{O}(N)$.

作者发现有两点可以改进：

attention mask 不需要存储，我们只需要保存 forward pass 时的输入，然后在 backward pass 时重新生成即可，这样只需要 $\mathcal{O}(N)$ 的空间复杂度。
计算 softmax 的梯度是，如果使用公式 $D_i=P_{i:}^TdP_{i:}$ 来计算的话，由于 $P_{i:}\in\mathbb{R}^N$, 可能会导致超过 SRAM 的内存使用限制，因此，我们可以使用 $D_i=do_i^To_i$ 来避免这个问题，其中 $o_i\in\mathbb{R}^d$.

最终，flashattention 的 backward pass 过程如下图所示

经过前面的分析，flashattention 的反向传播的时间复杂度为 $\mathcal{O}(N^2)$, 空间复杂度为 $\mathcal{O}(N)$.

Theorem 5 令 $N$ 为 sequence length, $d$ 为 head dimension, $M$ 是 SRAM 的 size, 且满足 $d\leq M\leq Nd$. 则 flashattention 反向传播的内存访问开销为 $\Theta(N^2d^2M^{-1})$.

定理的证明与 Theorem 2 基本一致，我们此处不再赘述。

Block-sparse Flashattention

当 attention 具有 block sparsity 的性质时，作者提出了 blck-sparse flashattention 来进一步提高 attention 的计算效率。

给定 $Q,K,V\in\mathbb{R}^{N\times d}$, 以及一个 mask $M\in\{0,1\}^{N\times N}$, 我们需要计算

$$ S = QK^T\in\mathbb{R}^{N\times N},\quad P = \mathrm{softmax}(S\odot \mathbb{1}_{M})\in\mathbb{R}^{N\times N},\quad O = PV\in\mathbb{R}^{N\times d} $$

其中当 $M_{kl}=1$ 时， $(S\odot \mathbb{1}_ {M})_ {kl}=S_ {kl}$, 否则 $(S\odot \mathbb{1}_ {M})_{kl}=0$.

Block-sparse attention 的算法如下所示

Proposition 4 令 $N$ 为 sequence length, $d$ 为 head dimension, $M$ 是 SRAM 的 size, 且满足 $d\leq M\leq Nd$. 则 block-sparse attention 的内存访问开销为 $\Theta(Nd+N^2d^2M^{-1}s)$, 其中 $s$ 是 block-sparse mask 中的非零 block 的比例

证明与 Theorem 2 的证明是类似的，总的内存访问开销为 $\mathcal{O}(Nd+NdT_c)$, 但是在计算的过程中，由于 mask 矩阵的 block-sparsity, 我们实际上只需要计算一小部分 $M_{ij}\neq0$ 的情况，因此最终的内存访问开销为

$$ \mathcal{O}\left(Nd+\frac{N^2d^2}{M}s\right) $$

可以看到，attention mask 的 sparsity 越高，block-sparse flashattention 的效率也就越高。当 $N$ 非常大时，$s 通常为 $1/\sqrt{N}$ 或者 $N^{-1}\log N$, 从而最终的内存访问开销为 $\mathcal{O}(N\sqrt{N})$ 或者 $\mathcal{O}(N\log N)$.

作者对比了以下 block-sparse flashattention 和 flashattention 的效率对比，结果如下图所示

Experiment

作者通过实验验证了 flashattention 的有效性，如下表所示

Attention	Standard	FlashAttention
GFLOPs	66.6	75.2
HBM R/W (GB)	40.3	4.4
Runtime (ms)	41.7	7.3

可以看到，尽管 flashattention 相比于标准化 attention 需要更多的算力，但是由于其内存访问开销更少，所以最终的运行时间大有了大幅度降低

作者还探究了 block size 对 flashattention 性能对的影响，实验结果如下图所示

可以看到，随着 block size 增加，循环次数降低，内存访问开销也逐渐降低。但是当 block size 充分大 ( $> 256$) 之后，运行时间就会被别的因素所限制，并且过大的 block size 可能会导致 SRAM 的内存溢出

作者首先在 BERT 和 GPT-2 上验证了 flashattention 的表现，BERT 的实验结果如下表所示

BERT Implementation	Training time (minutes)
Nvidia MLPerf 1.1	$20.0\pm1.5$
FlashAttention (ours)	$17.4\pm1.4$

GPT-2 的实验结果如下表所示

Model implementations	OpenWebText (ppl)	Training time (speedup)
GPT-2 small - Huggingface	18.2	9.5 days (1.0 )
GPT-2 small - Megatron-LM	18.2	4.7 days (2.0 )
GPT-2 small - FlashAttention	18.2	2.7 days (3.5 )
GPT-2 medium - Huggingface	14.2	21.0 days (1.0 )
GPT-2 medium - Megatron-LM	14.3	11.5 days (1.8 )
GPT-2 medium - FlashAttention	14.3	6.9 days (3.0 )

实验结果显示，flashattention 比 Huggingface 快 3 倍左右，比 Megatron 快 1.7 倍左右

训练速度：实验显示，flashattention 在 BERT 上，比 MLPerf 1.1 快 $15\%$, 在 GPT-2 上比 HuggingFace 快 3 倍，比 Megatron 快 1.8 倍
准确率：flashattention 是第一个在 Path-X 上比随机表现更好的 transformer 模型；block-sparse flashattention 是第一个在 Path-256 上比随机表现更好的的 sequence model

Conclusion

作者提出了 flashattention, 一个通过优化标准 attention 内存访问效率来提高 attention 计算效率的方法，作者详细介绍了算法设计的原理与证明，并通过实验证明了结果的有效性。

References

arxiv

Notes on StreamingLLM

Wed, 20 Aug 2025 10:16:35 +0800

作者提出了 StreamingLLM, 一个基于 attention sink 来提高 sliding window attention 在超长上下文场景下表现的方法。

Introduction

已有的基于 softmax attention 的架构的问题在于很难扩展到长上下文的场景，主要原因有两点：

KV cache 会随着序列长度增加而商城，从而提高 decoding 的 latency
序列长度超过预训练的 context length 之后，模型表现会急剧下降

为了解决这个问题，已有的方法可以分为三类：

length extrapolation: 使用 RoPE 或者 AliBi 等方法来扩展 LLM 的 context length, 这类方法的问题是扩展的上下文长度仍然有限，对于 streaming 的场景作用有限
context window attention: 扩展 LLM 的上下文长度，如 flash attention 等来降低 attention 的计算和内存开销。这类方法也是只在有限的上下文场景下 work
Improving LLMs’ Utilization of Long Text: 更好利用长上下文的数据

基于已有的工作的发现，作者提出了本文研究的核心问题：

如何在不损失模型表现和效率的情况下，提高模型在无限长上下文场景下的表现。

为了解决这个问题，作者首先分析了 sliding window attention 的不足，作者发现，sliding window attention 在超过 KV cache size 之后，表现也会急剧下降。作者通过实验发现，sliding window attention 表现急剧下降的原因在于 attention sink, 也就是模型损失了对于初始 token 的关注，从而导致模型表现下降。

基于 attention sink, 作者设计了 StreamingLLM, 用于提高 sliding window attention 在长上下文场景下的表现，结果发现，模型的表现有了大幅度的提升。

作者还进一步在预训练阶段加入了 sink token 充当初始 token, 进一步提高模型的表现。

Method

Attention Sink

作者首先探究了一下 softmax attention 以及 sliding window attention 性能下降的节点，实验结果如下图所示

可以看到，softmax attention 性能急剧下降的节点为 pre-training 的 context length; 而 sliding window attention 性能急剧下降的节点为 KV cache size.

接下来，作者分析了一下不同 layer 的 attention 分布情况，如下图所示

可以看到，初始的 2 层 layer 里 attention logits 的分布比较均匀。但是在后续的 layer 里，第一个 token 的权重都大幅度上升。

作者分析原因认为，sliding window attention 在超过 KV cache size 之后性能急剧下降的主要原因是初始 token 不再参与 softmax 的计算，这导致了 softmax 的计算出现了比较大的变化，从而模型的表现开始下降。

为了探究初始 token 对最终模型表现的影响因素是语义层面还是位置层面的，作者将初始的 token 替换为 \n, 并比较了模型的表现，结果如下表所示

Llama-2-13B	PPL (↓)
0 + 1024(Window)	5158.07
4 + 1020	5.40
4"\n"+1020	5.60

可以看到，把初始的四个 token 替换为 \n, 并不影响模型最终的表现，这说明是初始 token 的位置信息在发挥作用。

作者接下来探究了一下模型架构的影响，实验结果如下表所示

Cache Config	0+2048	1+2047	2+2046	4+2044	8+2040
Falcon-7B	17.90	12.12	12.12	12.12	12.12
MPT-7B	460.29	14.99	15.00	14.99	14.98
Pythia-12B	21.62	11.95	12.09	12.09	12.02
Cache Config	0+4096	1+4095	2+4094	4+4092	8+4088
Llama-2-7B	3359.95	11.88	10.51	9.59	9.54

可以看到，不同的模型架构都存在这个问题，这说明 sliding window attention 的影响与架构无关。并且，作者认为，使用初始 4 个 token 就可以有效的避免模型的性能下降，进一步增加初始 token 的数量不会有进一步提升。

作者分析 attention sink 出现的原因在于，

初始的 token 对于后续所有的 token 都是可见的，因此其会携带一些信息
在预训练阶段，模型并没有一个一致的初始 token 来标注起始信息，这导致模型会默认使用第一个 token 来储存一些信息。

为了解决这个问题，作者就提出了缓存初始 token 的方法，具体做法就是，在 sliding window attention 的基础上，我们还会加上初始 token 的信息，作者展示示意图如下所示

也就是说，我们初始 token 始终会参与计算（论文中初始 token 数量为 4），然后我们会维持一个大小为 3 的 KV cache 队列来进行最终 sliding window attention 的计算，这样，每次计算 attention 的时候，我们就会使用 $\# \text{iniital token} + \# \text{sliding window token}$ 这么多的 token 来计算 attention. 作者对比了不同 attention 的计算方式，如下图所示

前面是在 inference 阶段进行优化的，作者现在进一步探究在 pre-training 阶段加入 attention sink 参与训练对模型表现的影响。

[[softmax-off-by-one]] 提出了我们应该加入一个 zero sink token, 其计算公式如下

$$ \mathrm{softmax}_1(x)_i = \frac{\exp(x_i)}{1 + \sum_{j=1}^N \exp(x_j)} $$

这里 $x\in\mathbb{R}^N$ 是输入的序列。我们可以将 sink token 视为一个 key 以及 value 都是 0 向量的特殊 token.

在本文中，作者使用了一个可学习的 sink token. 作者对比了原始 softmax attention, 使用 zero sink attention, learnable sink attention 三种方法的表现，结果如下表所示

Cache Config	0+1024	1+1023	2+1022	4+1020
Vanilla	27.87	18.49	18.05	18.05
Zero Sink	29214	19.90	18.27	18.01
Learnable Sink	1235	18.01	18.01	18.02

可以看到 zero sink 仍然需要一部分初始 token 来维持模型的表现。作者在论文中推荐使用 learnable sink.

Experiments

作者首先验证了 StreamlingLLM 在不同架构上的表现，结果如下图所示

实验结果显示，StreamingLLM 可以扩展到 4M 的上下文

接下来，作者探究了以下在 Pretraining 阶段加入 learnable sink token 对模型表现的影响，结果如下图所示

可以看到，加入 sink token 之后对模型的表现没有显著影响。并且，模型在下游任务上的表现与标准的 softmax attention 表现差不多。

作者还对 StreamlingLLM 进行了可视化，结果如下图所示

作者进一步评估了 StreamingLLM 在下游任务上的表现，我们主要关注一下 ARC 上的表现，结果如下图所示

可以看到，full attention 出现了 OOM error, 而 sliding window attention 虽然避免了 OOM 的问题，但是其表现非常差。而 StreamingLLM 则进一步提高了 Sliding Window attention 的表现。

Conclusion

作者在本文中提出了 StreamingLLM, 一个在 Sliding window attention 中加入 sink token 来避免超过 cache size 之后模型表现急剧下降的问题。作者详细介绍了 attention sink 现象以及解决方法。

References

Efficient Streaming Language Models with Attention Sinks

Notes on gpt-oss

Tue, 19 Aug 2025 16:14:56 +0800

openAI 发布了 gpt-oss 大语言模型，包含 120B-A5.1B 以及 20.9B-A3.6B 两个 size, 作者强调了模型的 instruction following, tool use, 以及 adaptive thinking 能力

Method

Architecture

gpt-oss 系列是一个基于 MoE transformer 架构的 LLM. 架构中交替使用 sliding window attention 和 full attention, sliding window size 为 128 token, 架构图如下所示

模型的配置如下表所示

Model	120B	20B
Date	2025/8/5	2025/8/5
# Total Parameters	116B	20B
# Activated Parameters	5.13B	3.61B
# MoE Layers	36	24
Hidden Dim	2880	2880
MoE Intermediate Dim	2880	2880
Attention	GQA+sliding window	GQA+sliding window
Attention Head Dim	64	64
Attention bias	True	True
# Attention Heads	64	64
# Key-Value Heads	8	8
# Experts (total)	128	32
# Experts Active Per Token	4	4
# Shared Experts	0	0
QK-Norm	No	No

在架构上，gpt-oss 做的主要改变有：

Q, K, V projection layer, expert layer, routing layer 都使用了 bias
修改了 expert layer 中 SwiGLU 的定义
attention 中额外使用了一个 attention sink

SwiGLU

大多数模型使用的基于 SwiGLU 的 MLP 定义如下

$$ y = W_2(W_3x \odot \mathrm{SwiGLU}(W_1x)) $$

其中 $\mathrm{SwiGLU}(x)=x\odot\mathrm{sigmoid}(x)$, 在 gpt-oss 模型中，作者首先定义了两个常数 $\alpha=1.702$, $\mathrm{limit}=7.0$, 然后 SwiGLU MLP 的定义如下

$$ \begin{aligned} o_1&=W_1x+b_1,\\ o_3&=W_3x+b_3\\ o_1&=\mathrm{clamp}(o_1,\max=\mathrm{limit})\\ o_3&=\mathrm{clamp}(o_3,\min=-\mathrm{limit},\max=\mathrm{limit})\\ o_3&= o_3\odot \mathrm{sigmoid}(\alpha\cdot o_3)\\ o_3&= (o_1+1)\odot o_3\\ y &= W_2o_3 \end{aligned} $$

Attention Sink

作者使用了 attention sink 来避免 window attention 在超过 kv cache size 之后，表现大幅度下降的问题。

Quantization

为了降低模型的内存占用量，作者使用了 PTQ 来训练 MoE 的权重，使用的精度为 MXFP4, 这样每个参数由 4.25 bits 来表示。最终，模型的参数存储格式如下：

通过这个流程，gpt-oss-120B 可以部署在 80GB 内存的 GPU 上，gpt-oss-20B 可以部署在 16GB 内存的 GPU 上。模型各部分参数量如下表所示

Component	120b	20b
MLP	114.71B	19.12B
Attention	0.96B	0.64B
Embed + Unembed	1.16B	1.16B
Active Parameters	5.13B	3.61B
Total Parameters	116.83B	20.91B
Checkpoint Size	60.8GiB	12.8GiB

这里在计算激活参数的时候，没有没有考虑 embedding 的参数量。

Pre-training

预训练细节不多，主要是使用了 flash attention 进行加速计算，使用了 Triton 进行了 kernel 的优化，gpt-oss-120b 训练了 120M H100-hours, gpt-oss-20B 的训练时间是 gpt-oss-120Bd 的十分之一左右

Post-training

post-training 的数据包括 coding, math 以及 science 等，主要使用 RL 进行训练

作者介绍了以下 post-training 使用的格式，即 harmony chat format. 角色的优先级如下所示

1

System > Developer > User > Assistant > Tool

作者还加入了 channels 来限制可以使用的信息，比如使用 analysis 来表示 CoT tokens, 使用 commentary 来表示 function calling 等，一个具体的例子如下：

Evaluation

gpt-oss 系列的表现如下表所示

Conclusion

作者提出了 gpt-oss 系列大语言模型，gpt-oss 在架构上与已有的主流模型架构如 Qwen3, DeepSeek-V3 等都有一定区别

References

technical report

Notes on QK-Norm

Wed, 13 Aug 2025 16:12:11 +0800

作者提出了 QK norm, 一个解决 softmax 注意力权重不稳定的 scaling 算法。

Problem Definition

Softmax 可以用于将 logits 转化为一个概率分布，但是 softmax 问题是输入的微小差别会对输出产生巨大影响，甚至会 mask 掉其他信号。因此我们需要选取合适的缩放因子，来解决 softmax 的极端值问题。SDPA 的可视化结果如下图所示

Method

作者首先提出了 QKNorm, 其使用了一个可学习的 scaling 参数来控制 $QK^T$ 的范围，进而让 attention 的 pattern 更加分散。

作者首先回顾了以下已有的进展，主要是三点：

FixNorm: 将 word embedding 限制为单位长度
PerNorm: 使用 Pre-norm 替换 Post-norm
ScaleNorm: 使用 $\ell_2$ normalization 替换 LayerNorm, 并乘以一个可学习的 scaling 参数。

作者基于这三点进行了改进，改进后的 attention 定义如下

$$ \mathrm{softmax}\left(g\cdot \hat{Q}\hat{K}^T\right)V $$

其中,

$$ \hat{Q} = [\frac{q_1}{\|q_1\|_2},\dots,\frac{q_m}{\|q_m\|_2}], \hat{K} = [\frac{k_1}{\|k_1\|_2},\dots,\frac{k_n}{\|k_n\|_2}] $$

是对原始的 $Q, K$ 按列进行 $\ell_2$ normalization 得到的结果， $g$ 是一个可学习的参数，其初始化值为

$$ g_0 = \log_2(L^2-L) $$

这里 $L$ 是训练数据 $97.5$ 分位。

使用这种动态缩放之后，attention 的分布变得更加分散了，结果如下图所示：

References

Query-Key Normalization for Transformers

Notes on GLM-4.5

Wed, 13 Aug 2025 12:27:48 +0800

智谱 AI 提出了 GLM4.5, 包含 GLM4.5 和 GLM-4.5-Air,两个 MoE LLM. 模型大小分别为 355B-A22B 和 106B-A12B, GLM4.5 主要关注 agentic, reasoning 以及 coding 三个领域。

Introduction

作者认为，通用模型有三个关键能力，即 ARC：

Agent: 与外部工具以及真实世界进行交互
Reasoning: 解决数学和科学领域的复杂问题
Coding: 解决真实世界软件工程相关问题

已有的商业模型如 o1/o3, Claude Sonnet 4 已经在 ARC 上达到了非常好的表现，但是开源模型仍然比较稀缺

基于这个目标，作者就提出了 GLM4.5 和 GLM-4.5-Air, 来统一完成三个不同的目标。

Method

Pre-training

Architecture

GLM-4.5 是一个基于 MoE 架构的 LLM, 架构与 DeepSeek-MoE 相似，作者做了如下几点改变：

在 MoE layer 中，使用了 loss-free balance routing, 然后使用了 sigmoid function 作为 routing score 的 normalization.
与 Kimi-k2 和 DeepSeek-V3 相比，作者降低了 head dimension, 提升了 number of layers. 作者认为更深的模型更有利于提高模型的 Reasoning 表现
attention 上，作者使用了 GQA, 对于 #RoPE, 作者使用了 partial RoPE, 只旋转每个 token 的前半部分，作者还将 attention heads 的个数增加到了 2.5 倍，作者发现增加 attention heads 可以提高模型的 Reasoning 表现
作者还使用了 QK-Norm 来防止 attention logits 爆炸
作者还使用了一个 MoE layer 作为 MTP layer 来支持 speculative decoding.

模型与 DeepSeek-V3 和 Kimi-k2 的对比如下

Model	GLM-4.5	GLM-4.5-Air	Step 3	Kimi K2
Date	2025/8/8	2025/8/8	2025/7/25	2025/7/28
# Total Parameters	355B	106B	316B	1043B
# Activated Parameters	32B	12B	38B	32B
# Dense Layers	3	1	5	1
# MoE Layers	89	45	56	60
# MTP Layers	1	1	0	0
Hidden Dim	5120	4096	7168	7168
Dense Intermediate Dim	12288	10944	18432	18432
MoE Intermediate Dim	1536	1408	5120	2048
Attention	GQA	GQA	MFA	MLA
Attention Head Dim	128	128	256	192
# Attention Heads	96	96	64	64
# Key-Value Heads	8	8	1	64
scoring	sigmoid	sigmoid	softmax	softmax
# Experts (total)	160	128	48	384
# Experts Active Per Token	8	8	3	8
# Shared Experts	1	1	1	1
QK-Norm	Yes	No	No	No

Pre-training Data

预训练数据包括四个方面

Web: 过滤低质量数据和使用模版产生的数据
Multilingual: 基于 webpages 和 Fineweb-2
Code: 基于 GitHub 和其他代码平台，作者使用了 [[Fill in the middle]] 来训练模型。
Math & Scirence: 训练一个 classifier 来给数据进行打分。

最终，预训练数据一共包括 23T token.

Pre-training Recipe

预训练包括 2 个阶段:

Pre-training: 使用网页数据进行训练
Mid-training: 加入 code, math, science 数据进行训练，在这个阶段，作者使用了 repo-level 的 code 数据，合成的 reasoning 数据以及长上下文数据。作者将模型上下文从 4K 扩展到 32K，然后在扩展到 128K.

作者在 pre-training 的时候使用了 random truncation, 在 mid-training 的时候使用了 best-fit packing 技巧

训练时，与 Kimi-k2 一样，作者使用了 Muon 作为优化器。作者使用了 cosine decay schedule. batch size 从 16M token 到 64M token.

Post-training

Post-training 分为两个阶段：

Stage 1, Expert Training. 构建 agent, reasoning, General chat 三个 domain 的专家模型
Stage 2, Unified Training. 使用 self-distillation 来汇总多个模型的能力

训练框架如下图所示

SFT

两个 stage 都由 SFT 开始，

在 Stage 1 里，SFT 的目标是让 expert model 掌握初步的 chat, reasoning 以及 tool-use 的能力。作者使用了一小部分包含 CoT 的 SFT 数据进行训练
在 Stage 2 中，SFT 的目标是将不同的 expert model 蒸馏到一个模型中，作者使用了百万级的数据，包含 reasoning 任务和通用的 chat 数据，来训练模型的 hybrid reasoning 能力

在训练模型的 tool-use 能力是，作者发现，function call 在 code 场景下会出现混淆，提高了模型的学习成本。因此，作者的解决方法是使用了类似 XML 的 special token tags

Recall 与之相反，Kimi-K2 认为模板应该尽可能简洁，因此 Kimi 采取了 TypeScript 作为 function call 的语言

从专家模型进行采样是，作者进行了数据过滤。还对数据进行了分级结果发现，使用难题进行训练可以提升模型 $2\%\sim4\%$ 的表现，多次采样也可以提高模型的表现

Agentic SFT 数据的构建包括四个步骤：

Agentic Framework and Tool Collection: 收集 MCP 和 tool API
Task Synthesis: 合成不同的 agentic 任务
Trajectory Generation: 采样生成的 rollout
Quality Filtering: 过滤低质量的数据

RL

Reasoning RL

这个阶段使用了 GRPO 算法进行训练，与 DAPO 一样，作者去除了损失函数中的 KL divergence。

首先，作者探究了课程学习对模型表现的影响，结果发现，课程学习可以有效提高模型的性能。因此，作者构建了一个 2 阶段的课程学习框架。实验结果如下图所示

可以看到，在第二个阶段，模型可以进一步通过更难的题目获得提升。

其次，作者探究了以下渐进式扩展模型上下文对模型表现的影响。DeepScaleR 认为，逐步提高模型的上下文长度，可以有效提高模型的表现。但是，本文确认为这种方法会损害模型的性能，原因在于，模型在 SFT 阶段的上下文长度就是 64K, 如果我们降低模型的上下文长度，这会导致训练数据分布不一致，从而影响模型的长上下文表现。因此作者直接在 64K 的上下文上进行训练。

接下来，作者探究了以下采样温度对模型表现的影响，温度太低会导致模型探索能力下降，太高的话会导致输出质量下降。因此作者动态调整采样温度来平衡模型的性能以及探索能力。

[!tip] Kimi-K2 认为随着 RL 训练的进行，我们应该逐步降低采样温度来稳定模型的表现

最后，作者分析了以下 code 以及 Science RL 中的一些问题。对于 code RL, 作者发现，我们应该在 sequence 层面而不是 token 层面进行平均。对于 Science RL, 作者强调了高质量数据的重要性。实验结果如下图所示

Agent RL

作者主要关注 web-search 以及 code generation 两个任务。对于 web-search, 作者构建了一个数据合成 pipeline, 用于生成 multi-step reasoning 的 QA 数据。构建过程包括基于知识图谱的 multi-hop reasoning 和 human-in-the-loop 的内容提取。对于 code generation, 作者基于 GitHub 的 PR 以及 issues 构建了 benchmark

RL 的训练目标如下

$$ \mathcal{L}(\theta) = \mathbb{E}_{x\sim\mathcal{D}}\left[\frac1K\sum_{i=1}^K(r(x,y_i) - \bar{r}(x))\right] $$

其中 $(x,y_i)$ 是基于 $\pi_{\mathrm{old}}$ 采样的 trace, $\bar{r}(x) = 1/k\sum_{i=1}^Kr(x,y_i)$ 是平均的 reward. 计算损失时，只有模型的回答参与计算。

作者发现，通过训练模型的 web-search 以及 code generation 能分，模型在 tool-use 以及 coding 任务上的表现也有了提升。作者还是用了 format penalty 来保证模型输出格式的正确性。如果格式不对的话，模型获得的奖励是 0

Recall 在 GLM-4.V-Thinking 中，作者认为应该在 cold-start SFT 阶段让模型学会格式要求，而不是在 RL 阶段加入 format reward

由于 agent RL 的训练比较耗时，为了提高训练效率。作者首先基于 SFT 模型进行 agent RL 训练，训练到一定步数之后，作者使用 self-distillation 来将能力蒸馏回 SFT model, 接下来再基于 Self-distillation 后的 SFT 模型来进行 agent RL 训练

Recall Seed1.5-VL 和 LlaMA3.2 都提到了使用 multi-round SFT-RL 的形式来提高模型的表现

作者还发现，随着交互轮数的提升，模型的表现也有相应提升。实验结果如下图所示

General RL

General RL 用于提高模型的整体表现，解决潜在的问题以及提升关键能力，作者主要使用了 RLHF 和 RLAIF 两种方法

对于 Holistic RL, 作者收集了 5000 条 prompt, reward 基于人类反馈和 AI 反馈。人类反馈用于训练一个 reward model, 对于 AI 反馈，作者构建了 scoring rubrics. 然后作者将两种反馈结合在一起

对于 Instruction following RL, 作者构建了基于规则的奖励，reward model 的奖励以及 critical model 的奖励。实验结果显示，这种奖励方式可以有效降低模型的 reward hacking

对于 function calling RL, 作者使用了 step-wise rule-based RL 来提高模型表现。对于 end-to-end multi-turn RL, 作者训练了一个 expert model 来蒸馏专家到模型。

最后，对于 Pathology RL, 作者希望通过 RL 来解决潜在的问题，比如语言混合输出，重复输出以及格式错误等。作者构建了一批模型容易出错的数据，然后来训练模型。

Infra

作者针对不同任务分别构建了不同的 scheduling 模式：

对于通用 RL 任务，作者将 training engine 和 inference engine 放在一个 worker 来提高效率
对于 agentic RL 任务，作者将 training 和 inference engine 分开，来提高 data throughput

在训练时，作者使用了 BF16 精度，在推理时，作者使用了 FP8 精度来提高推理效率。

针对 agentic RL 任务，作者还进行了优化。与 Kimi-k2 类似，作者让 inference engine 持续产出 rollout, 然后让 training engine 来更新模型权重，最后同步到 inference engine 上

Experiments

整体表现如下图所示，GLm4.5 在 ARC benchmark 上的平均表现达到了第三名。

具体来看，

在 agentic benchmark 上, GLM4.5 仅次于 o3 的表现
在 coding benchmark 上，GLM4.5 次于 Claude Opus 4 和 Claude Sonnet 4, 排第三名
在通用能力上，GLM

人工对比 coding agent 能力的结果如下图所示

Conclusion

作者提出了 GLM4.5，一个基于 MoE 架构的大语言模型系列，包含 GLM4.5(355B-A22B) 和 GLM4.5-Air(106B-A12B) 两个模型，作者详细介绍了模型的架构，训练，数据和评估。

References

arxiv

Notes on GQA

Thu, 07 Aug 2025 18:08:36 +0800

Google Research 在 23 年 12 月份提出了 Group Query Attention (GQA), 一个提升 multi-head attention 效率的方法。GQA 自 Qwen2 系列开始被应用。

Introduction

Multi-head attention (MHA) 的问题在于 inference 阶段，每次 decoding，都需要重新加载 attention 模块中 query layer, key layer 和 value layer 的权重，而加载权重会受带宽限制。

已有的工作有 MQA, 也就是我们把多个 head 的 key layer 以及 value layer 压缩成一个，这样对于 $h$ 个 head 的 attention，我们有 $h$ 个 query layer，$1$ 个 key layer 以及 1 个 value layer. 但是 MQA 的问题在于其会导致性能下降，而且训练过程会不稳定。

因此，在本文中作者就作出了两点贡献：

如何将一个 MHA 模型转化为一个一个 MQA 模型
提出了 Group Query Attention (GQA)，在保持模型性能的同时，提高计算效率

Method

Uptraining

将 MHA 模型转化为 MQA 模型分为两步：

将 MHA 权重转化为 MQA 权重
额外的预训练

具体来讲，作者使用了一个 mean pooling 的方法，来将不同 head 的 query layer 以及 key layer 的权重转化为 MQA 对应 layer 的权重。然后作者 pre-training 若干步来让模型适应新的结构。

GQA

GQA 的思路在于在 MHA 和 MQA 之间达到一个平衡，也就是说我们将 key layer 和 value layer 进行分组，每个组内共享一个 key layer 和 value layer, 我们假设有 $h$ 个 head，$G$ 个 group，那么

$G=1$ 时，所有的 head 共享一个 key layer 和一个 value layer, 此时 GQA 等价于 MQA
$G=H$ 时，每个 head 都有一个 key layer 和一个 value layer, 此时 GQA 等价于 MHA
$1

三者的示意图如下所示

Code

MQA 的代码也比较好理解，我们首先定义 group size，即 num_key_value_heads, 然后基于 group size 定义对应的 key layer self.k_proj 和 value layer self.v_proj.

计算得到 key_states 和 value_states 之后，在计算 attention，即 eager_attention_forward 的时候，我们对 key_states 和 value_states 进行复制，即 repeat_kv

 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


def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
 batch, num_key_value_heads, slen, head_dim = hidden_states.shape
 if n_rep == 1:
 return hidden_states
 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)

def eager_attention_forward(
 module: nn.Module,
 query: torch.Tensor,
 key: torch.Tensor,
 value: torch.Tensor,
 attention_mask: Optional[torch.Tensor],
 scaling: float,
 dropout: float = 0.0,
 **kwargs: Unpack[TransformersKwargs],
):
 key_states = repeat_kv(key, module.num_key_value_groups)
 value_states = repeat_kv(value, module.num_key_value_groups)

 attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
 if attention_mask is not None:
 causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
 attn_weights = attn_weights + causal_mask

 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
 attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
 attn_output = torch.matmul(attn_weights, value_states)
 attn_output = attn_output.transpose(1, 2).contiguous()

 return attn_output, attn_weights

class Qwen3Attention(nn.Module):
 def __init__(self, config: Qwen3Config, layer_idx: int):
 super().__init__()
 self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
 self.scaling = self.head_dim**-0.5

 self.q_proj = nn.Linear(
 config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
 )
 self.k_proj = nn.Linear(
 config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
 )
 self.v_proj = nn.Linear(
 config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
 )
 self.o_proj = nn.Linear(
 config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
 )

 def forward(
 self,
 hidden_states: torch.Tensor,
 position_embeddings: tuple[torch.Tensor, torch.Tensor],
 attention_mask: Optional[torch.Tensor],
 past_key_value: Optional[Cache] = None,
 cache_position: Optional[torch.LongTensor] = None,
 **kwargs: Unpack[FlashAttentionKwargs],
 ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:
 input_shape = hidden_states.shape[:-1]
 hidden_shape = (*input_shape, -1, self.head_dim)

 query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
 key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
 value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)

 cos, sin = position_embeddings
 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

 attn_output, attn_weights = eager_attention_forward(
 self,
 query_states,
 key_states,
 value_states,
 attention_mask,
 dropout=0.0,
 scaling=self.scaling,
 sliding_window=self.sliding_window, # diff with Llama
 **kwargs,
 )

 attn_output = attn_output.reshape(*input_shape, -1).contiguous()
 attn_output = self.o_proj(attn_output)
 return attn_output, attn_weights

Conclusion

本文中，作者提出了一个解决 multi-query attention 的 uptraining 方法，以及提出了 GQA，一个结合 MHA 表现和 MQA 效率的新型注意力机制。

References

arxiv

Notes on MQA

Thu, 07 Aug 2025 18:06:37 +0800

Google 在 2019 年提出了 multi-query attention (MQA), 用于解决 MQA 内存带宽瓶颈问题。

Method

Background

对于 multi-head attention, 我们假设其 hidden size 为 $d$, 有 $h$ 个 heads, 每个 head 的 size 为 $d_h=d/h$, 输入 sequence 长度为 $n$, batch size 为 $d$. 则总的 arithmetic operations 为 $O(bnd^2)$. 总的内存访问量为 $O(bnd + bhn^2+d^2)$, 第一项是 $Q,K,V$ 的内存占用（$Q,K,V$ 分别是 query, key 和 value layer 的输出），第二项是 attention score 的占用，第三项是 query, key 和 value layer 的权重。

因此，其 Memory Access Ratio (MAR), 也就是内存访问量与 arithmetic operations 之比为

$$ O\left(\frac1k + \frac{1}{bn}\right) $$

对于现代的 GPU 来说，其一般算力比较强，但是内存访问带宽相对较慢，因此我们希望 MAR 越低越好，以充分发挥 GPU 的算力。

MHA Analysis

在训练的时候，由于我们知道 ground truth sequence, 因此我们可以并行计算。但是在 inference 的时候，我们只能 token-by-token 进行计算，因此我们分析一下 token-by-token 场景下的 MAR

我们整体的 arithmetic operations 还是 $O(bnd^2)$.

但是，现在我们要调用 $n$ 次 multi-head attention, 因此我们总的内存访问量为 $O(bn^2d + nd^2)$, 第一项是 $K$ 和 $V$ , 第二项是 query, key 和 value layer 的权重。

这种情况下，MAR 就变成了

$$ O\left(\frac{n}{d} + \frac{1}{b}\right) $$

当 $n\approx d$ 或者 $b\approx 1$ 时，MAR 就非常接近于 1，意味着内存带宽成了一个主要的瓶颈。为了解决这个问题，我们有两种做法：

提升 batch size $b$, 也就是同时 inference 多次
降低 $K$ 和 $V$ 的大小

MQA

MQA 的做法就是第二种，也就是降低 $K$ 和 $V$ 的大小，但是 $K,V$ 分别是 key 和 value layer 的输出，要降低输出大小，我们就必须改变 key 和 value layer 的 size。基于这个考虑，作者在所有的 head 上共享了一个 key 和 value layer，也就是说，原来

1
2


self.k_proj = nn.Linear(hidden_size, num_heads * head_dim) # (d, n*d_h)
self.v_proj = nn.Linear(hidden_size, num_heads * head_dim) # (d, n*d_h)

现在在 MQA 里，其变成了

1
2


self.k_proj = nn.Linear(hidden_size, head_dim) # (d, n*d_h)
self.v_proj = nn.Linear(hidden_size, head_dim) # (d, n*d_h)

MQA Analysis

我们还是在 token-by-token 的场景下进行分析。

我们整体的 arithmetic operations 还是 $O(bnd^2)$.

调用 $n$ 次 multi-query attention 的总的内存访问量为 $O(bnd +bn^2d_h+ nd^2)$, 第一项是 $q$ , 第二项是 $K$ 和 $V$ , 第三项是是 query, key 和 value layer 的权重。

此时，MAR 变成了

$$ O\left(\frac{1}{d} + \frac{n}{dh}+\frac{1}{b}\right) $$

现在，我们就将 $n/d$ 这一项给降低了 $h$ 倍。如果我们的 batch size 足够大的话，理论上 MQA 应该能极大提高整体的计算效率。

Conclusion

MQA 为了追求极致的内存带宽占用，选择使用单一的 key 和 value, 来极大提高 inference 的 decoding 效率，但是后来在 GQA 中验证发现，MQA 虽然非常高效，但是其表现比较差，这也是后来没有得以应用的原因。

References

Arxiv

Notes on Moonlight

Thu, 07 Aug 2025 10:49:32 +0800

Kimi 提出了 Moonlight, 一个基于 Muon optimizer 训练得到的 16B-A3B MoE LLM. 作者详细介绍了如何 scale up muon optimizer.

Introduction

Muon 验证了 Muon optimizer 在小语言模型 nanoGPT 上的表现，但是对于更大规模 LLM 的表现，尚未有人探究。因此 Kimi 就希望在大规模 LLM 上验证 Muon optimizer 的表现。作者主要进行了两点改进：

加入 weight decay
调整了不同参数更新的 scale

基于改进后的 Muon optimizer, 其训练效率相比于 AdamW 提升了 2 倍。作者基于 Muon Optimizer 训练得到了 Moonlight, 一个 16B-A3B 的 MoE LLM.

作者主要作出了三点贡献：

探究了 weight decay 在 scaling Muon 时的作用
分布式 Muon optimizer 的实现
验证了 Muon optimizer 的 scaling law

Method

Background

作者首先介绍了一下 Muon optimizer, 给定步数 $t$, 参数矩阵 $W_{t-1}$, momentum $\mu$, 学习率 $\eta_t$ 以及目标函数 $\mathcal{L}_t$, Muon optimizer 的更新方式如下：

$$ \begin{aligned} M_t &= \mu M_{t-1} + \nabla\mathcal{L}_t(W_{t-1})\\ O_t &= \mathrm{Newton-Schulz}(M_t)\\ W_t &= W_{t-1} - \eta_t O_t \end{aligned} $$

这里 $M_t$ 是 gradient 的 momentum, 初始化为 $M_0=0$. 在上面的更新公式中，Newton-Schulz 的作用是求解 $(M_tM_t^T)^{-1/2}M_t$. 令 $M_t=U\Sigma V^T$ 为 SVD 分解，我们有

$$ (M_tM_t^T)^{-1/2}M_t = UV^T $$

这是一个半正交矩阵，即 $(UV^T)^T(UV^T)=I$.

Newton-Schulz 迭代的具体公式如下：

$$ X_0 = \frac{M_t}{\|M_t\|_F},\quad X_k = aX_{k-1} + b(X_{k-1}X_{k-1}^T)X_{k-1} + c(X_{k-1}X_{k-1}^T)^2X_{k-1} $$

其中，normalization 是为了保证 Newton-Schulz 的收敛性。 $a,b,c$ 是三个超参数，在 Muon 中设置为 $(a,b,c)=(3.4445, 4.7750, 2.0315)$.

Scaling up Muon

作者发现，尽管 Muon 在小规模场景下 work 的很好，但是大规模性场景下的收益就非常有限了。作者发现，这是因为模型的参数以及每一层输出的 RMS 变得很大，这可能会影响模型的性能。因此，作者就和 AdamW 一样使用 weight dacay 来避免这个问题，即

$$ W_t =W_{t-1} - \eta_t(O_t + \lambda W_{t-1}) $$

作者通过实验对比了 AdamW, vanilla Muon 和 Muon w/ weigth decay 三者的表现，实验结果如下图所示

实验结果显示，尽管 vanilla Muon 手链最快，但是由于其权重增长很快，因此最后模型的表现不如 AdamW 和 Muon w/ weigth decay.

接下来，作者分析了以下更新矩阵的 Root Mean Square (RMS), 结论是 Muon optimizer 的 RMS 与参数矩阵的形状相关：

Lemma For a full-rank matrix parameter of shape $[A, B]$, its theoretical Muon update RMS is $\sqrt{1/\max(A, B)}$.

证明如下：通过 Newton-Schulz 迭代，我们得到 $O_t=UV^T$, 其中 $M_t=U\Sigma V^T$ 是 SVD 分解，我们有

$$ \mathrm{RMS}(O_t) = \sqrt{\frac{\sum_{i=1}^A\sum_{j=1}^BO_{t,i,j}^2}{AB}}=\sqrt{\frac{r}{AB}} $$

其中, $r=\mathrm{rank}(M_t)$ , 这样就完成了证明。

而 Adam 和 AdamW 的 RMS 都在 $1$ 附近。作者认为 RMS 也会影响模型表现：

当 $\max(A,B)$ 过大时，如 dense MLP matrix, 其更新就会变得很小，限制了模型的表现
当 $\max(A,B)$ 过小时，如 GQA 中的 KV head 或者 DeepSeek-V3 中的 MLA, 更新又会变得很大，导致训练不稳定。

因此，作者就提出了一个 rescaling 的技巧，来消除 Muon optimizer 的影响。

作者通过实验发现，AdamW 的 RMS 通常在 $0.2\sim0.4$ 左右，因此，作者将 Muon optimizer 的更新设置如下

$$ W_t = W_{t-1} - \eta_t(0.2\cdot O_t\cdot \sqrt{\max(A,B)} + \lambda W_{t-1}) $$

基于这个改变， Muon 和 AdamW 可以共享学习率以及 weight decay 参数。

Distributed Muon

ZeRO-1 天然适合 AdamW, 因为 AdamW 都是 element-wise 进行计算的。但是 Muon 则需要梯度矩阵的全部信息。因此，作者就针对 ZeRO-1 进行适配，提出了 Distributed Muon, 分布式版本将优化器的状态进行切分，然后加入了两个额外的操作：

DP gather: 将 ZeRO-1 切分的梯度矩阵 gather 为一个完整的矩阵
Calculate Full Update: 对完整的梯度矩阵执行 Newton-Schulz 迭代

最终，Distributed Muon 的算法如下图所示

最后，作者分析了一下 distributed Muon 和 distributed AdamW 的内存和算力占用：

内存开销：Muon 只有一阶矩，而 AdamW 有二阶矩，因此 Muon 的额外内存开销为 AdamW 的一半。
通信开销：对于 ZeRO-1，通信开销来源于三个过程：All-Gather 参数 $P$ 用于前向传播, Reduce-Scatter 梯度 $G$ 用于反向传播, All-Gather 更新后的参数 $P$ 用于下一轮的前向传播。AdamW 不引入额外通信，所以其每个参数的通信量为 $4+4=8$, 分别代表 $G$ 和 $P$ 的通信量。而 Muon 则需要额外的一次通信来得到 full matrix, 因此每个参数通信量为 $4+4+2=10$, 分别代表 $P, G$ 和 full matrix. 也就是说，分布式 Muon 的通信量最高为 AdamW 的 $1.25$ 倍。实际上由于我们使用 multiple DP, 这个比例会更接近于 $1.0$.
latency：Distributed Muon 相比于 AdamW latency 更高，这是因为 Muon 需要进行 DP gather 以及计算 Newton-Schulz 迭代。但实际上，latency 很小，因为 Newton-Schulz 迭代只需要迭代 5 次，并且 optimizer 的 end-to-end latency 相比于 forward-backward 过程是可以忽略的。一些额外的技巧也可以降低 latency.

实际在训练的过程中，作者发现 Distributed Muon 相比于 AdamW 并没有太明显的 latency.

Experiments

Scaling Law of Muon

作者分析了一下 Muon Optimizer 的 scaling law, 实验结果如下图所示

实验结果表明，在最优设置下，Muon Optimizer 只需要 $52\%$ 的 FLOPs 就可以达到 AdamW 的表现

Pretraining with Muon

作者分贝使用 AdamW 和 Muon 训练模型，然后评测了以下模型在不同 benchmark 上的表现，结果如下图所示

可以看到，在相同的设置下，Muon optimizer 的表现更好。

Dynamics of Singular Spectrum

Muon optimizer 的核心思想就是让比较难更新的方向也能被更新到，本节作者就探究了 Muon 是否满足这个性质，作者对参数矩阵进行 SVD 分解，然后定义 SVD entropy 如下

$$ H(\sigma) = -\frac{1}{\log n}\sum_{i=1}^n\frac{\sigma_i^2}{\sum_{j=1}^n\sigma_j^2}\log\frac{\sigma_i^2}{\sum_{j=1}^n\sigma_j^2} $$

作者对 SVD entropy 可视化如下

可以看到，Muon optimizer 的 SVD entropy 比 AdamW 更大，这说明 AdamW 的更新方向更多更广，验证了 Muon optimizer 的核心思想

SFT with Muon

作者还在 SFT 阶段验证了 Muon optimizer 的有效性。实验结果如下图所示

结论主要有两个：

预训练阶段与 SFT 阶段使用不同的优化器时，模型表现没有明显区别
SFT 阶段使用 Muon 可以达到与 AdamW 差不多的表现，但是最好还是在 pre-training 阶段使用 Muon

Conclusion

作者探究了如何 scale up Muon Optimizer. 通过改进，作者在 16B-A3B 的 MoE LLM 上验证了 Muon Optimizer 的性能。实验结果发现，Muon Optimizer 的训练效率比 AdamW 提升了 2 倍左右。

作者提出了三个未来可行的研究方向：

目前 Muon 只能针对 2D 参数进行优化，其他参数仍然依赖于 AdamW 优化器，是否可以使用 Muon 优化所有参数？
Muon optimizer 可以理解是 spectral norm 下的 steepest descent 方法，如何将其扩展到 Schatten norm 是一个可以研究的方向
实验里提到，预训练和 SFT 阶段使用不同的 optimizer, 表现不是最优的，如何解决这个因为不同 optimizer 导致的性能差距是一个需要解决的问题。

References

Muon is Scalable for LLM Training

Notes on Hunyuan-Large

Wed, 06 Aug 2025 16:46:32 +0800

腾讯混元提出了 Hunyuan-Large, 一个 389B-A52B 的 MoE LLM, 上下文长度为 256K.

Introduction

Hunyuan-Large 主要在三个方向进行了改进：

使用了更高质量的合成数据：模型使用了 7T 的预训练数据，其中包含了 1.5T 的合成数据
优化了模型的架构：作者提出了 KV cache compression, recycle routing, expert-specific learning rate scaling 策略来提高模型的表现
探究了 MoE 模型的 scaling law: 作者探究了 MoE 模型的 scaling law

Pre-training

Architecture

Hunyuan-Large 是一个基于 MoE 的 transformer 架构，attention 部分使用了 GQA, position encoding 使用了 RoPE, MLP 的激活函数为 SwiGLU. 在 MoE layer 中，Hunyuan-Large 使用了 shared experts. 最终，模型的配置如下图所示

KV Cache Compression

为了减少 KV cache 的内存开销，作者使用了两个技巧：

GQA: 通过共享 KV projection 的参数，来减少内存访问次数
[[CLA]]: 在相邻的 layer 中共享 KV cache, 来进一步压缩 KV cache

在 Hunyuan-Large 中，作者将 GQA 的 group size 设置为 8, 然后相邻的 2 层 layer 共享 KV cache.

假设输入的 batch size 为 $B$, sequence 长度为 $L$, layers 个数为 $\ell$, attention heads 个数为 $h$, KV heads 个数为 $h_{kv}$, 每个 head 的 hidden size 为 $d_h$, 则每一层的 GQA 需要缓存 $K,V\in\mathbb{R}^{B\times _{kv}\times L\times d_h}$， KV cache 的总占用为

$$ 2\times B\times h_{kv}\times L\times d_h \times \ell \times 2=4BLh_{kv}d_h\ell $$

第一个 $2$ 是因为同时缓存 K 和 V, 第二个 $2$ 是因为一般使用 bfloat16 数据格式。

对于 CLA, 因为连续两层共享相同的 KV cache，因此结果除以 2; 对于 MHA, $h_{kv}=h$; 对于 MQA, $h_{kv}=1$. 最后，KV cache 的内存占用如下表所示

Attention Mechanism	KV Cache Memory
MHA	$4BLhd_h\ell$
GQA	$4BLh_{kv}d_h\ell$
MQA	$4BLd_h\ell$
CLA	$2BLhd_h\ell$
GQA+CLA	$2BLh_{kv}d_h\ell$

可以看到，使用 GQA+CLA 之后，模型的 kv cache 占用相比于 MHA 变成了

$$ \frac{2BLh_{kv}d_h\ell}{4BLhd_h\ell}=\frac{1}{16} $$

也就是说，Hunyuan-Large 的 KV cache 内存占用下降到了 MHA 的 1/16.

Expert Routing Strategy

作者采用了 shared expert + activated expert 的形式，其中包含 1 个 shared expert, 然后从 16 个专家里激活 1 个专家。

为了解决 MoE 中 expert capacity 难以设定的问题，作者提出了一个 recycle routing 的策略，基本思想就是，当 activated expert 的容量超出限制时，会从其他没有超出容量限制的专家里重新进行激活。

Expert Specific Learning Rate Scaling

作者使用 AdamW 作为优化器，作者探讨了如何设定学习率。基于之前的工作，最优的学习率与 batch size 相关：

$$ \epsilon_{\mathrm{opt}}(B) = \frac{2\epsilon_{\max}}{\sqrt{\frac{\mathcal{B}_{\mathrm{noise}}}{B}}+\sqrt{\frac{B}{\mathcal{B}_{\mathrm{noise}}}}} $$

这里 $\epsilon_{\max}$ 是 AdamW 的学习率, $\mathcal{B}_{\mathrm{noise}}$ 是训练速度与数据使用效率的一个平衡因子。

但是，在 MoE 模型中，不同专家处理的 token 是不一样的。基于 load balancing loss, shared expert 和 activated expert 处理的 token 个数比例大概是 $n :1$, 其中 $n=16$ 是总的专家个数。因此，对于 shared expert, 作者使用 $\epsilon_{\mathrm{opt}}(B)$ 作为学习率，然后对于 activated expert, 作者使用 $\epsilon_{\mathrm{opt}}(B/n)$ 作为学习率。

Data

预训练数据包括收集和合成。收集的数据主要来自互联网，覆盖中英文两种语言。

合成数据包括 4 个步骤：

instruction generation: 作者使用高质量的语料作为 seed, 然后生成多样的 instruction 覆盖不同的 domain
Instruction evolution: refine 上一步生成的 instruction
Response generation: 使用 specialized model 来生成回答
response filtering: 对生成的回答进行过滤

数据合成的流程如下图所示

tokenizer 大小为 128K, 由 tittoken tokenizer 和额外的 28K token 组成。

Pre-training Recipe

作者首先探究了一个针对 MoE 模型的 scaling law. 结果发现，最优的激活参数量为 58.1B, training token 个数为 5.6T. 经过平滑之后，作者最终将模型的激活参数两定为 52B, 训练 token 数定为 $7T$.

在训练时，作者将学习率分为了 3 个 stage:

warmup phase
gradual decay phase
concise annealing phase

上面的三个 stage 结束之后，作者加入了两个 stage 来扩展模型的上下文长度从 32K 扩展到 256K. 训练的数据包括 75% 的短文本和 25% 的长文本。两个 stage 训练的 token 数均为 $10B$ 左右。

Post-training

post-training 分为 SFT 和 RLHF 两个阶段。

SFT

SFT 数据副高 math, coding, logical reasoning 等 domain, 包含超过 1M 的数据。

SFT 训练了 3 个 epoch, 学习率从 2e-5 降低到 2e-6, 为了避免 overfitting, 作者使用了 0.1 的 attention dropout 和 0.2 的 hidden dropout.

[!tip] 作者发现，MoE 模型可以从 dropout 中学习到更多

RLHF

作者使用 DPO 来进行 RLHF, 作者同时使用了 offline 和 online 的数据来进行训练，前者是收集的数据，后者是当前 policy 生成的数据。与 LLaMA 3 和 Nemotron-4 一样，为了提高训练稳定性，对于 chosen reponse, 作者使用了 SFT loss.

作者还是用了 exponential moving average 策略来减少 reward hacking 现象，以及降低 alignment tax.

Experiment

对于 base 版本，作者对比了 LLaMA 3, Mixtral, DeepSeek-V2, 实验结果如下图所示

Instruction 版本的表现如下图所示

Conclusion

作者提出了 Hunyuan-Large, 一个 389B-A52B 的 LLM, 上下文长度为 256K. 作者详细介绍了模型的架构，数据和训练方式。

References

arxiv

Notes on GSPO

Wed, 06 Aug 2025 11:26:26 +0800

Qwen 提出了 Group Sequence Policy Optimization (GSPO), 一个针对 GRPO 进行改进的 RL 算法。GSPO 在 sequence 层面计算 importance ratio, 避免了 token-level 计算带来的训练不稳定性。

Introduction

GRPO 的问题在于训练超大规模的 LLM 时，会出现 model collapse, Qwen3 和 MiniMax-01 均对 GRPO 算法进行了改进。

在本文中，作者认为 GRPO 算法的问题在于 Importance sampling weight 的计算是有问题的，这导致了误差随生成回答长度累积，最后被 clipping 机制放大，从而导致模型崩溃。

为了解决这个问题，作者提出了 GSPO, 一个在 sequence 层面进行 Importance sampling 的 RL 算法，GSPO 还对 reward 进行 normalization 来保持训练的稳定性

Method

Preliminary

对于一个大语言模型 $\pi_{\theta}$, 我们将 $x$ 记为 query, 将 $y$ 记为 $\pi_{\theta}$ 针对 $x$ 的 response, 即

$$ \pi_{\theta}(y\mid x) = \prod_{t=1}^{|y|}\pi_{\theta}(y_t\mid x, y_{其中 $|y|$ 代表 response $y$ 的长度, 一般我们还有一个 reward model $r$, 用于生成奖励 $r(x,y)\in[0, 1]$.

[!tip] Note 本文中没有提及 KL divergence, 作者认为这不是重点，所以没有在公式中标出。

PPO 算法包含四个模型：reference model, 也就是 $\pi_{old}$, policy model, 也就是 $\pi_{\theta}$, value model, 用于计算 advantage, 以及 reward model, 用于计算奖励。 PPO 算法如下面公式所示

$$ \mathcal{J}_{\mathrm{PPO}}(\theta) = \mathbb{E}_{(x,y)\sim\mathcal{D},y_{\leq t}\sim \pi_{\theta_{old}}(\cdot\mid q)}\left[ \min\left(r_t(\theta)\hat{A}_t,\mathrm{clip}\left(r_t(\theta), 1-\epsilon, 1+\epsilon\right)\hat{A}_t\right) \right] $$

其中

$$ r_t(\theta) = \frac{\pi_{\theta}(y_t\mid x, y_{< t})}{\pi_{\theta_{old}}(y_t\mid x, y_{< t})} $$

是 token $y_t$ 的 importance ratio, $\hat{A}_t$ 是 $y_t$ 的 advantage estimation.

PPO 算法的问题在于其严重依赖 value model, 一般 value model 与 policy model 的大小相当，以保持内存和算力的负载均衡。

为了解决 PPO 的这个问题，GRPO 通过多次采样然后计算 relative advantage 来避免使用 value model. 其目标函数为

$$ \mathcal{J}_{\mathrm{GRPO}}(\theta) = \mathbb{E}_{(x, y)\sim\mathcal{D},\{y_i\}_{i=1}^G\sim \pi_{\theta_{old}}(\cdot\mid q)}\left[ \frac{1}{G}\sum_{i=1}^G\frac{1}{|y_i|}\sum_{t=1}^{|y_i|}\min\left(r_{i,t}(\theta)\hat{A}_{i,t},\mathrm{clip}\left(r_{i,t}(\theta), 1-\epsilon, 1+\epsilon\right)\hat{A}_{i,t}\right) \right] $$

其中，$G$ 是针对 query $x$ 的多次采样 response,

$$ r_{i,t}(\theta) = \frac{\pi_{\theta}(y_{i,t}\mid x, y_{i,< t})}{\pi_{\theta_{old}}(y_{i,t}\mid x, y_{i,< t})} $$

是 importance ratio,

$$ \hat{A}_{i,t} = \hat{A}_{i} = \frac{r(x,y_i) - \mathrm{mean}(\{r(x,y_i)\}_{i=1}^G)}{\mathrm{std}(\{r(x,y_i)\}_{i=1}^G)} $$

是使用 group response 估计得到的 advantage.

Motivation

在 Qwen3 中已经提到，对于稀疏的 MoE 模型，我们必须使用更大的 batch size 来最大化内存和算力使用效率。但是，更大的 batch 意味着有一些数据必须是 off-policy 的，因此 PPO 和 GRPO 就需要使用 clipping 来降低 off-policy sample 对模型表现产生过大影响。

基于 clipping 机制，作者认为 GRPO 的目标函数是 “ill-pose” 的，其原因在于 GRPO 计算的 Importance ratio 是与 RL 训练目标不匹配的。

通常，importance sampling 的公式如下：

$$ \mathbb{E}_{z\sim\pi_{\mathrm{tar}}}[f(z)] = \mathbb{E}_{z\sim\pi_{\mathrm{beh}}}\left[\frac{\pi_{\mathrm{tar}}(z)}{\pi_{\mathrm{beh}}(z)}f(z)\right] $$

其中 $\pi_{\mathrm{tar}}$ 是目标分布， $\pi_{\mathrm{beh}}$ 是采样分布, importance ratio $\pi_{\mathrm{tar}}(z)/\pi_{\mathrm{beh}}(z)$ 负责对采样进行修正。

对于 GRPO, 其 importance ratio 是在 token 层面定义的，而一次采样是在 sequence 层面定义的，因此这种区别就导致了 GRPO 存在 high-variance noise.

因此，作者认为，Importance ratio 的关键在于优化目标应该与 reward 的粒度是一致的。

GSPO

基于上面的想法，作者就提出了 GSPO, 作者首先针对 LLM 改写了上面的 importance sampling 的公式

$$ \mathbb{E}_{x\sim \mathcal{D}, y\sim\pi_{\theta}(\cdot\mid x)}[r(x,y)] = \mathbb{E}_{x\sim\mathcal{D}, y\sim\pi_{\theta_{old}}(\cdot\mid x)}\left[\frac{\pi_{\theta}(y\mid x)}{\pi_{\theta_{old}}(y\mid x)}r(x,y)\right] $$

这里的 Importance ratio 表现了采样的回答 $\pi_{old}(y\mid x)$ 与目标回答 $\pi_{\theta}(y\mid x)$ 之间的差距，这是从 sequence 层面上体现的。

因此，作者基于上面的式子，构建出了 GSPO 的目标函数

$$ \mathcal{J}_{\mathrm{GSPO}}(\theta) = \mathbb{E}_{(x, y)\sim\mathcal{D},\{y_i\}_{i=1}^G\sim \pi_{\theta_{old}}(\cdot\mid q)}\left[ \frac{1}{G}\sum_{i=1}^G\min\left(s_{i}(\theta)\hat{A}_{i},\mathrm{clip}\left(s_{i}(\theta), 1-\epsilon, 1+\epsilon\right)\hat{A}_{i}\right) \right] $$

这里 $\hat{A}_{i}$ 与 GRPO 一致， $s_i(\theta)$ 是 normalize 之后的 importance ratio, 定义如下

$$ s_i(\theta) = \left(\frac{\pi_{\theta}(y_{i}\mid x)}{\pi_{\theta_{old}}(y_{i}\mid x)}\right)^{\frac{1}{|y_i|}} = \exp\left(\frac{1}{|y_i|}\sum_{i=1}^{|y_i|}\frac{\pi_{\theta}(y_{i,t}\mid x, y_{i,< t})}{\pi_{\theta_{old}}(y_{i,t}\mid x, y_{i,< t})}\right) $$

这里作者使用了 length normalization 来控制 variance.

作者对比了以下 GSPO 和 GRPO 两者的梯度，我们忽略期望的计算，直接写出内部的梯度，有

$$ \begin{aligned} \nabla_\theta \mathcal{J}_{\mathrm{GSPO}}(\theta) &= \frac{1}{G}\sum_{i=1}^G\left(\frac{\pi_{\theta}(y_{i}\mid x)}{\pi_{\theta_{old}}(y_{i}\mid x)}\right)^{\frac{1}{|y_i|}}\hat{A}_i\cdot \frac{1}{|y_i|}\sum_{t=1}^{|y_i|} \nabla_{\theta} \log\pi_{\theta}(y_{i,t}\mid x, y_{可以看到，两者不同的地方在于 token 的 reweight 方式，GRPO 中先加权再求和；而 GSPO 中则是先求和，然后在 sequence 层面进行加权

GSPO-token

为了支持 multi-turn RL 等需要细粒度 advantage 的场景，作者对 GSPO 的目标函数进行了改进，得到了 GSPO-token, 其目标函数为

$$ \mathcal{J}_{\mathrm{GSPO-token}}(\theta) = \mathbb{E}_{(x, y)\sim\mathcal{D},\{y_i\}_{i=1}^G\sim \pi_{\theta_{old}}(\cdot\mid q)}\left[ \frac{1}{G}\sum_{i=1}^G\frac{1}{|y_i|}\sum_{t=1}^{|y_i|}\min\left(s_{i,t}(\theta)\hat{A}_{i},\mathrm{clip}\left(s_{i,t}(\theta), 1-\epsilon, 1+\epsilon\right)\hat{A}_{i}\right) \right] $$

其中，

$$ s_{i,t}(\theta) = \mathrm{sg}[s_i(\theta)]\frac{\pi_{\theta_{old}}(y_{i,t}\mid x, y_{i,< t})}{\mathrm{sg}[\pi_{\theta_{old}}(y_{i,t}\mid x, y_{i,< t})]} $$

这里 $\mathrm{sg}$ 是 stop gradient operation. 通过这种改写方式，GSPO-token 与 GSPO 的优化目标一致，但是可以在更细粒度的 token 层面进行优化。

Evaluation

作者使用 Qwen3-30B-A3B 进行测试。对于 GRPO，作者发现必须使用 Routing Replay training strategy 才能提升 MoE RL 的收敛性， Routing Replay 的具体做法就是保留 $\pi_{\theta_{old}}$ 的激活专家，在 $\pi_{\theta}$ 中使用相同的专家来保持稳定性。但是，GSPO 则不需要这个技巧。

实验结果如下图所示。

从实验结果可以看到，GSPO 比 GRPO 更加高效，效果也更好。

作者带对比了 GSPO 与 GRPO 的 clipping 比例，结果发现，GSPO clipping 的 token 比例比 GRPO 高几个数量级，作者认为，尽管 GSPO clip 了更多 token, 但是 GSPO 更加高效，这也说明 GRPOtoken 层面的梯度估计是存在噪声的。

作者还介绍了以下 MoE 模型中 RL 训练的 expert-activation volatility 现象，也就是说，MoE 模型参数更新之后，对于同一个 response, 激活的专家可能飞铲个不同。作者举例说明，对于 Qwen3-30B-A3B, 更新一次梯度之后，有 $10\%$ 左右的激活专家变得不一样了。

作者最后介绍了以下 GSPO 的两个优势，一个是解决了 GRPO 依赖于 Routing Replay 的问题；第二个是支持 SGLang 和 vLLM 等推理框架。

Conclusion

作者提出了 GSPO，一个在 sequence 层面计算 importance ratio 的 RL 算法，解决了 GRPO 在训练大规模 MoE LLM 时出现的训练不稳定性以及 mode collapse 现象。

References

Group Sequence Policy Optimization

Notes on Muon blog

Tue, 05 Aug 2025 11:10:51 +0800

Muon (MomentUm Orthogonalized by Newton-Schulz) 是一个针对二维神经网络的优化器，它基于 SGD-momentum 改进，增加了一个 Newton-Schulz 的后处理步骤

Method

Newton-Schulz (NS) 的目的是用一个正交矩阵近似一个给定矩阵，即

$$ \mathrm{Ortho}(G) = \arg\min_{O} \{\|O-G\|_F: \text{either } O^TO=I\text{ or } OO^T=I\} $$

也就是说，NS iteration 将 SDG-moment 的更新矩阵替换为了“最近的” semi-orthogonal matrix. 这等价于将更新矩阵替换为 $UV^T$, 其中 $USV^T$ 是更新矩阵的 SVD 分解。

[!tip] 作者观察到，对于 SGD-momentum 和 Adam 来说，其在基于 transformer 的神经网络里有非常高的 condition number, 也就是 optimizer 仅在少数几个方向上进行优化。作者认为，通过正交化，可以有效提高模型在其他方向上的更新速度，进而提高模型表现

Newton-Schulz

作者提到，正交化矩阵的方法有很多，比如 SVD 分解，但是其问题是非常慢，还有 Coupled Newton iteration, 但是其精度要求非常高，必须要在 float32 以上。

作者因此使用了 Newton-Schulz iteration.

令 $G=USV^T$ 是 SGD-momentum 更新矩阵的 SVD 分解，则基于系数 $(a,b,c)$ 的 NS iteration 定义如下：

$$ \begin{aligned} G' &= aG + b(GG^T)G + c(GG^T)^2G\\ &= (aI+b(GG^T)+c(GG^T)^2)G\\ &= (aI+bUS^2U^T+cUS^4U^T)USV^T\\ &= U(aS+bS^3+cS^5)V^T \end{aligned} $$

也就是说，如果我们定义五次多项式函数 $\phi(x)=ax+bx^3+cx^5$, 然后执行 $N$ 次 NS iteration, 则我们得到 $U\phi^N(S)V^T$, 其中 $\phi^N$ 代表 $\phi$ 复合 $N$ 次。

为了保证 NS iteration 收敛到 $\mathrm{Ortho}(G) = UV^T$, 我们必须保证两点：

$S$ 的值，也就是 $G$ 的奇异值必须在区间 $[0,1]$ 上
$\phi$ 必须满足 $\phi^N\to 1$, $N\to\infty$, $\forall x\in[0,1]$.

为了满足第一个条件，我们可以对 $G$ 进行 rescale, 即 $G\gets G/\|G\|_F$, rescale 不影响最终的结果，即 $\mathrm{Ortho}(G) = \mathrm{Ortho}(cG)$.

对于 $\phi(x)$, 我们有很多选择，比如我们定义 $(a,b,c):=(2,-1.5,0.5)$ 就得到如下结果

Coefficient Optimization

尽管 $(a,b,c):=(2,-1.5,0.5)$ 已经满足了第二个条件，但是我们还是想进一步优化，优化的方向主要有两个：

让 $a$ 尽可能大，这是因为 $\phi'(0)=a$ 控制了较小奇异值的收敛速率。
对于所有的 $x\in[0,1]$, 我们希望 $\phi^N(x)\in[1-\epsilon, 1+\epsilon]$, $N\to\infty$. 这样 NS iteration 的结果与 $\mathrm{Ortho}(G)$ 不会相差太远。

作者发现， $\epsilon$ 可以设置为 $0.3$ 而不影响 Muon optimizer 的收敛性。因此，作者的目标现在是

$$ \begin{aligned} \max\quad &a\\ \mathrm{s.t.}\quad &\lim_{N\to\infty}\phi^N(x)\in[0.7, 1.3] \end{aligned} $$

作者通过 ad-hoc gradient 方法求解得到一组数值解为 $(a,b,c)=(3.4445, 4.7750, 2.0315)$, 作者将这组数值应用于 Muon optimizer 中。迭代结果如下图，可以看到，当 $x\approx0$ 时，函数变得更加陡峭。

实验中，作者发现，仅需迭代五次，最终的结果就 work 的很好。作者还尝试了不同的多项式，结果发现并没有太大的提升。

Algorithm

最终，Muon Optimizer 的算法如下

其中, NewtonSchulz5 算法伪代码定义如下

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14


def newtonschulz5(G, steps=5, eps=1e-7):
 assert G.ndim=2
 a, b, c = (3.4445, -4.7750, 2.0315)
 X = G.bfloat16()
 X /= (X.norm() + eps)
 if G.size(0) > G.size(1):
 X = X.T
 for _ in range(steps):
 A = X @ X.T
 B = b * A + c * A @ A
 X = a * X + B @ X
 if G.size(0) > G.size(1):
 X = X.T
 return X

Analysis

本节作者分析了以下 Muon 的内存占用和算力开销。

在 NS iteration 之前，Muon optimizer 和 SGD-moment 是一样的。

对于 $n\times m$ 的矩阵（假设 $m\leq n$），首先 NS iteration 会进行转置，NS iteration 的每一步需要 $2(2nm^2+m^3)$ FLOPs, 其中括号前面的系数 $2$ 代表精度。因此，Muon 相比于 SGD momentum 需要的额外 FLOPs 为 $2T(2nm^2+m^3)$, 其中 $T$ 是迭代次数。

使用 baseline 进行一次训练（前向 + 后向），所需要的 FLOPS 为 $6nmB$, 其中 $B$ 是 batch size. 因此，Muon 的 FLOP 开销至多为 $Tm/B$, 其中 $m$ 是模型的 hidden size, $B$ 是 batch size, $T$ 是 NS iteration 的步数。

作者分别基于 nanoGPT 和 LLaMA-405B 进行验证，结果发现，Muon optimizer 带来的额外开销不足 $1\%$.

作者发信啊，使用 Nesterov-style momentum 可以比普通的 SGD-momentum 效果更好，因此作者在 muon 中使用了前者。

作者还发现，对于 QKV layer，分别进行优化效果会更好。

Experiments

Limitation and Future Work

Muon 仅被设计用于优化 2D 参数（因为涉及矩阵计算），其余的参数仍然需要 AdamW 等优化器参与。

作者认为未来的工作有：

能否 scale up Muon Optimizer
分布式优化
在 fine-tuning 和 RL 阶段使用 Muon Optimizer

Conclusion

作者提出了 Muon optimizer，该优化器在 nanoGPT speedrun 上取得了 SOTA 的结果，作者详细介绍了优化器的工作原理。

References

Muon: An optimizer for hidden layers in neural networks

Notes on Kimi-k2

Thu, 24 Jul 2025 10:56:50 +0800

Kimi-k2 是一个总参数为 1T, 激活参数为 32B 的 MoE 大语言模型，模型使用 15.5T token 进行训练，optimizer 使用了 MuonClip. 作者主要关注模型的 agent 能力

Introduction

作者首先强调现在 LLM 发展主要是 Agentic Intelligence, 也就是让模型自主感知，规划，思考和与环境交互。

基于这个目标，作者就提出了 Kimi K2, 一个 1.04T 总参数，32B 激活参数的 MoE 模型，用于解决实现 agent Intelligence 中遇到的问题。作者主要进行了三点改进：

MuonClip, 一个基于 Muon 的优化算法，来提高 Kimi K2 对 token 的利用效率以及提高训练的稳定性
大规模的 agentic 数据合成 pipeline: 作者构建了一个用于合成工具调用，agent 数据的 pipeline
通用的 RL 框架，作者将 RLVR 和 self-critic rubric reward mechanism 结合起来，用于提升模型的表现

Pre-training

Architecture

Kimi-K2 的架构与 DeepSeek-V3 相似，配置如下表所示

指标	DeepSeek-V3	Kimi K2	$\Delta$
# Layers	61	61	=
Total Parameters	$671\text{B}$	$1.04\text{T}$	$\uparrow 54\%$
Activated Parameters	$37\text{B}$	$32.6\text{B}$	$\downarrow 13\%$
Experts (total)	256	384	$\uparrow 50\%$
Experts Active per Token	8	8	=
Shared Experts	1	1	=
Attention Heads	128	64	$\downarrow 50\%$
Number of Dense Layers	3	1	$\downarrow 67\%$
Expert Grouping	Yes	No	-

与 DeepSeek-V3 相比，模型主要进行了以下改动：

作者认为提高专家的稀疏性，可以有效提高模型表现。因此作者将专家个数提升了 50%.
为了降低模型在 inference 阶段的算力开销，作者将 attention heads 的个数降低了 50%.

Sparsity Scaling Law 作者首先构建了基于 Muon 的 sparsity law, 这里 sparsity 定义为激活专家个数与总专家个数之比，即：

$$ \mathrm{sparsity} = \frac{\# \mathrm{activated\ experts}}{\# \mathrm{total\ experts}} $$

作者在小规模上的实验结果如下图所示

Observation 实验结果表明，在相同算力（激活专家个数）下，模型的表现随 sparsity 提高而增加。

但是，进一步提高 sparsity, 会让 infra 变的难以优化，因此在 Kimi-K2 里，将 sparsity 定义为 $48$.

Number of attention heads 作者还分析了探究了 attention heads 的最优配置。DeepSeek-V3 将 attention heads 的个数设置为 layers 层数的 2 倍，来提高带宽利用率以及提高计算效率。但是，当上下文长度增加之后，attention head 变成了 computational bound. 作者对比了不同配置模型的表现，实验结果如下：

Observation 实验结果表明，使用双倍 attention head 带来的收益是比较小的，只有 $0.5\%$ 到 $1.2\%$ 左右

因此，作者在 Kimi-K2 中，奖 attention head 的个数设置为 $64$.

Data

Kimi-K2 主要强调了 token efficiency, 即每个 token 对模型表现提升的贡献。token efficiency 越高，说明每个 token 对最终模型的贡献也就越高

相比于 Kimi-k1.5, Kimi-K2 使用了一个 rephrasing pipeline, 来提高高质量 token 的利用率，作者在 knowledge 和 mathematics domain 上进行了实验。

最终，Kimi-K2 的预训练数据包括了 15.5T token, 主要覆盖 Web text, code, mathmatics 以及 Knowledge 四个 domain. 大部分的数据处理与 Kimi-k1.5 相同。

Knowledge Data Rephrasing 作者构建了一个 synthetic rephrasing 框架来提高 knowledge token 的利用率，框架主要包含以下几个模块：

Style and perspective-diverse prompting: 作者构建了一系列 prompt, 来让 LLM 从多个角度和风格来重新表述原始文本
Chunk-wise auto-regressive generation: 为了保持模型在长文档中的 coherence 以及避免信息损失，作者使用了一个 chunk-based rewriting 策略，也就是对每个 chunk 分别进行改写，然后将它门汇总在一起
Fidelity verification: 作者对原始文本和改写文本进行了 fidelity 检查来保证语义的一致性。

作者对比了以下三个设置对模型在 SimpleQA benchmark 上表现的影响，这三个设置分别是：

原始数据集训练 10 epoch
改写数据一次，然后训练 10 epoch
改写数据一次，训练 1 epoch

实验结果如下表所示

# Rephrasings	# Epochs SimpleQA	Accuracy
0 (raw wiki-text)	10	23.76
1	10	27.39
10	1	28.94

可以看到，改写的策略可以有效提高模型在 SimpleQA benchmark 上的表现

Math Data Rephrasing 对于数学相关数据，作者基于 SwallowMath, 将数据改写成了学习笔记 (learning-note) 的风格。然后，作者还将其他语言的高质量文本翻译为英文。

Recall 个人认为，数据改写在一定程度上也算是合成数据，对于合成数据这一块，微软提出的 phi 系列是一个比较好的参考

MuonClip Optimizer

Kimi-K2 使用了 Muon optimizer, 之前的实验结果说明了在相同的 compute budget 下，Muon optimizer 的表现超过了 AdamW. 也就是说，Muon optimizer 更加高效。

但是，Muon optimizer 的问题是训练不稳定。为了解决这个问题，作者提出了 QK-clip, 一个 weight clipping 机制，来显示约束 attention logits. QK-clip 的机制是当输出的 logits 超过某一个阈值之后，就对齐进行截断。

每个 head 的 attention 的计算公式如下

$$ O = \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d}}\right)V $$

其中, $d$ 是 hidden size, $Q, K,V$ 分别是 query, key, value, 定义如下：

$$ Q = XW_Q, K=XW_K, V=XW_V $$

这里 $W_Q,W_K,W_V$ 是模型可学习的参数。

作者定义每个 head 的 max logit 如下：

$$ S_{\max}^h = \frac{1}{\sqrt{d}} \max_{X\in\mathcal{B}}\max_{i,j} [QK^T]_{ij} $$

最简单的做法就是直接进行截断，也就是

$$ W_Q\gets \gamma^\alpha W_q, W_K\gets \gamma^{1-\alpha}W_K $$

其中 $\gamma=\min(1, \tau S_{\max})$, 这里 $S_{\max}=\max_h S_{\max}^h$ 是所有 head 对应 $S_{\max}^h$ 的最大值。

但是，实际中，作者发现只有少部分 head 会出现 logits 爆炸现象。为了提高计算效率，作者针对每个 head 单独进行 scaling, 也就是 $\gamma_{h}=\min(1, \tau S_{\max}^h)$. 对于 MLA 架构，作者仅在 unshared 模块使用 clipping:

$q^c$ 以及 $k^c$, scaling factor 为 $\sqrt{\gamma_h}$
$q^R$, scaling factor 为 $\gamma_h$

最后，作者将 Muon, weight decay, RMS matching 和 QK-clip 汇总在一起，得到 Kimi-k2 使用的 MuonClip optimizer, 算法如下所示：

接下来，作者对比了以下 Muon 和 MuonClip 两个优化器的表现，作者分别使用这两个优化器训练 53B-A9B 的 MoE 模型，实验结果如下

实验结果表明，Muon 优化器在 1000 步左右就出现了 logits 爆炸的现象，但是 MuonClip 通过优化可以避免这个问题。作者还发现，MuonClip 可以减少 loss spike 的产生，整体的 loss 变化情况如下图所示

作者还在附录中说明 QK-clip 并不影响最终的收敛性，并且，为了尽可能减少对模型训练的干扰，作者发现，仅在训练初期 QK-clip 被激活:

在初始的 70, 000 步里，有 12.7% 的 attention heads 至少触发了一次 QK-clip, 并且将它们的 $S_\max$ 降到了 100 以下
接下来的 70, 000 步里，QK-clip 就不再被激活

Infra

Kimi-k2 使用了 16-way 的 PP, 16-way 的 EP 以及 ZeRO-1 Data Parallelism. 具体流程如下

作者发现通过增加 warm-up mircro-batches, 我们可以有效重叠 EP 的 all-to-all communication 以及 computation. 但是，对于 [[DualPipe]], 其要求双倍的存储来保存参数和梯度，并且提高 PP 粒度会引入更多的 bubble, 因此 Kimi-K2 没有使用 DualPipe.

为了减少 [[1F1B]] 的 PP 通信开销，作者在反向传播的过程中同时进行 PP 通信，来进一步重合通信与计算。

作者还发现，使用更小的 EP group 可以提高整体的训练速度，因此作者将 EP 设置为 16.

作者发现，剩余的 GPU 内存不足以存放 MoE 的 activation, 因此作者采取了三个办法解决这个问题：

Selective recomputation: 作者对 LayerNorm, SwiGLU, MLA 的 up-projection 和 MoE 的 down-projections 等进行重新计算。
FP8 storage for insentive activations: 作者使用了 FP8 精度来存储 MoE up-projections 以及 SwiGLU. 作者发现使用 FP8 进行计算可能会导致性能下降，因此作者并没有使用 FP8 进行计算。
Activation CPU offload: 作者将其余的 activation 放在 CPU RAM 上，在 1F1B 的过程中，作者再进行加载

Training Recipe

模型上下文长度为 4096, 使用 MuonClip 进行优化，使用 WSD lr Scheduler.weight decay 为 0.1, global batch size 为 67M tokens.

预训练结束之后，作者加入了一个 long-context activation stage. 这个阶段，作者使用了 400B 的 4K 上下文的 token 和 60B 的 32K 上下文的 60B token. 最后作者使用 YARN 将模型上下文长度扩展到 128K.

Post-training

SFT

数据的构建主要是基于：

prompt 的多样性
Response 的质量

作者构建了一系列的数据生成 pipeline, 然后使用 Kimi-k1.5 来生成多样化的回答，最后再进行质量评估和过滤。这里，作者主要介绍了一下 tool-use 数据的构建

受 ACEBench 启发，作者构建了一个模仿真实世界 tool-use 的数据合成 pipeline. pipeline 如下图所示

pipeline 主要包含三个阶段：

tool spec generation: 作者基于 MCP tools 和 self-evolved 的数据合成策略来构建 tool repository. MCP tools 包括 3000+ 的 tools,合成了 20，000 个 tools
Agent and task generation: 作者还用不同的 system prompt 以及不同的 tools combination 来构建对应的 agent. 然后对不同的 agent, 作者构建了对应的成功标准均，工具调用模式以及评估方式
Trajectory generation: 作者首先构建不同风格的 LLM, 然后作者使用一个 simulator 来执行 tool call 并给予反馈。

最后，作者对数据进行了过滤。

作者还加入了 coding 以及软件工程等任务相关数据来进一步提高模型的 agent 能力

RL

RL 与 Kimi-K1.5 差不多。作者进一步提高了 RL 阶段的算力并做出了亮点改进：

作者构建了类似 Gym 的框架，用于扩展 RL 的能力
作者加入了更多 RLVR 的任务

Data 数据主要包括以下几类：

Math, STEM and logical tasks: 数据构建的原则为多样化和中等难度
Instruction following: 作者基于 hybrid rule verification 和 multi-source instruction generation 来合成复杂的指令跟随数据
Faithfulness: 作者训练了一个 judge model 来提供 reward
Coding & Software Engineering: 作者从开源数据收集并合成了代码相关数据
Safety. 提高模型的安全性，防止 jailbreak

Reward 作者使用了 self-critique rubric reward 的奖励机制。

首先，对于 Kimi-k2 的回答，作者会使用另一个 kimi-k2 作为 critic，基于三个方面进行排序：

core rubric: AI 的核心价值观
prescriptive rubric: 避免 reward hacking
human-annotated rubric: 特定的上下文

在训练的过程中，critic 也会基于 verifiable signals 进行 refine

RL training RL 的训练目标与 Kimi-k1.5 相同

$$ \mathcal{L}(\pi_\theta) = \mathbb{E}_{x\sim \mathcal{D}}\left[\frac1K\sum_{i=1}^K \left(r(x,y_i)-\bar{r}(x) -\tau\log\frac{\pi_{\mathrm{\theta}}(y_i\mid x)}{\pi_{\mathrm{old}}(y_i\mid x)}\right)^2\right] $$

其中 $\bar{r}(x)=1/K\sum_{i=1}^Kr(x,y_i)$ 是 sample response 的平均奖励。

作者做了以下几点改进来提高模型在不同 domain 上的表现：

Budget control: 作者针对不同任务分别设置了最大 token 限制，模型超过这个限制会受到惩罚。猜测应该是 length penalty 或者类似 Qwen3 一样，直接中断思考过程输出最终答案
PTX loss: 作者使用了 PTX loss 来提高模型对于高质量数据的利用率以及降低模型的过拟合
Temperature Decay: 作者发现，训练后期保持较高的采样温度会影响模型的表现，因此作者设置了一个 schedule, 来逐步降低采样温度。

RL Infra

RL infra 与 Kimi-k1.5 类似。主要包含三个模块。其中 train engine 和 inference engine 两者互相切换，一个 engine 进行 training 的时候，另一个 engine 就进行 inference, 为了提高 engine 切换的效率，作者使用了 checkpoint engine 来传输更新后的参数。

Evaluation

模型评估结果如下图所示

评估结果显示，模型在 coding 和通用任务上的表现仅次于 Claude 4, 大部分 benchmark 上都是第二名甚至是 sota.

Conclusion

作者在本文中提出了 Kimi-K2, 一个总参数为 1B 的 MoE 大语言模型。作者在架构，数据和优化器上进行了创新。并且通过 post-training 显著提高了模型的 agent 以及 tool-use 能力

作者发现模型主要存在的问题有：

reasoning 任务过难或者 tool 的定义不清晰的时候，模型会使用很多 token
有时候工具调用可能会降低模型的表现
模型在 agentic coding 任务上的能力需要进一步提升

References

Github

LLM Parameter Computation

Tue, 22 Jul 2025 10:50:47 +0800

本文中，我们介绍一下如何计算 LLM 的参数量。我们将基于 Qwen3 模型架构出发，对模型架构进行拆解，然后给出 LLM 参数量计算公式。

Dense Model

我们首先来看一下 Qwen3 的架构，如下图所示

这里，Qwen3ForCausalLM 就是我们的的 LLM, 其关键代码如下

1
2
3
4


class Qwen3ForCausalLM:
 def __init__(self, config):
 self.model = Qwen3Model(config)
 self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

我们假设 vocab_size, 也就是词表大小为 $|V|$ (我们用 $V$ 表示词表), hidden_size 为 $d$, 则总参数量为

$$ \mathrm{parameter}(\texttt{Qwen3ForCausalLM}) =d|V| + \mathrm{parameter}(\texttt{Qwen3Model}) $$

Qwen3Model 包含三个（含参数的）模块，分别是 nn.Embedding, Qwen3DecodeLayer 以及 Qwen3RMSNorm, 分别代表了输入 token 的 embedding layer, Transformer block 和对输出的 normalization. 其关键代码如下：

1
2
3
4
5
6
7


class Qwen3Model:
 def __init__(self, config: Qwen3Config):
 self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
 self.layers = nn.ModuleList(
 [Qwen3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
 )
 self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

其中，nn.Embedding 参数量与 lm_head 一样，都是 $d|V|$.

对于 normalization, 现在大部分 LLM 用的都是 RMSNorm, 其定义如下：

$$ \mathrm{RMSNorm}(x) = \frac{x}{\sqrt{\|x\|_2^2+\epsilon}}\odot \gamma $$

其参数量为：$d$.

如果说我们使用的是 LayerNorm, 则其定义如下：

$$ \mathrm{LayerNorm}(x) = \frac{x-\mathbb{E}[x]}{\sqrt{\mathrm{var}[x]+\epsilon}}\odot \beta + \gamma $$

其参数量为： $2d$.

因此，Qwen3Model 的参数量为

$$ \mathrm{parameter}(\texttt{Qwen3Model})=d|V| + N*\mathrm{parameter}(\texttt{Qwen3DecoderLayer}) + d $$

这里第一项为 nn.Embedding, 第三项为 Qwen3RMSNorm，第二项里，$N$ 代表 decode layer 的个数，也就是 config.num_hidden_layers.

Qwen3DecoderLayer 包含了四个模块，其关键代码如下

1
2
3
4
5
6


class Qwen3DecoderLayer:
 def __init__(self, config: Qwen3Config):
 self.self_attn = Qwen3Attention(config=config, layer_idx=layer_idx)
 self.mlp = Qwen3MLP(config)
 self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 self.post_attention_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

因此，Qwen3DecoderLayer 的参数量为

$$ \mathrm{parameter}(\texttt{Qwen3DecoderLayer}) = 2d + \mathrm{parameter}(\texttt{Qwen3MLP}) + \mathrm{parameter}(\texttt{Qwen3Attention}) $$

其中，第一项是两个 Qwen3RMSNorm 的参数。

对于 Qwen3MLP, 其定义如下：

$$ y = W_2(W_3x\odot \mathrm{SwiGLU}(W_1x)) $$

这里 $W_3,W_1\in\mathbb{R}^{d_{ff}\times d}$, $W_2\in\mathbb{R}^{d\times d_{ff}}$, $d_{ff}$ 是 MLP 的 hidden size, 代码中用 intermediate_size 来表示，因此

$$ \mathrm{parameter}(\texttt{Qwen3MLP}) = dd_{ff} + dd_{ff} + d_{ff}d=3dd_{ff} $$

这里三项分别代表 $W_1,W_3,W_2$ 的参数量。

如果说，我们使用原始 transformer 的 MLP, 也就是

$$ y = W_2\max(0,W_1x+b_1)+b_2 $$

其中 $W_1\in\mathbb{R}^{d_{ff}\times d}$, $W_2\in\mathbb{R}^{d\times d_{ff}}$, $b_1\in\mathbb{R}^{d_{ff}}$, $b_2\in\mathbb{R}^d$, 则总参数为

$$ \mathrm{parameter}(\texttt{TransformerMLP}) = d_{ff}d + dd_{ff} + d_{ff} +d = 2d_{ff}d + d_{ff} + d $$

这里的四项分别代表了 $W_1,W_2,b_1,b_2$.

Qwen3Attention

接下来，就是 Attention 部分的参数，Qwen3Attention 的关键代码如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


class Qwen3Attention:
 def __init__(self, config: Qwen3Config, layer_idx: int):
 self.q_proj = nn.Linear(
 config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
 )
 self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias)
 self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias)
 self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias)
 self.q_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps)
 self.k_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps)

这里，我们先定义几个量：

我们将 head 的个数记为 $h$, 即 num_attention_heads
我们将每个 head 的 hidden size 记为 $h_d$, 即 head_dim
我们将 key 和 value head 的个数记为 $h_{kv}$ , 即 num_key_value_heads

Qwen3Attention 的参数由以下几个部分组成：

Query projection: $W_{Q}\in\mathbb{R}^{hh_d\times d}$, $b_Q\in\mathbb{R}^{hh_d}$
Key projection: $W_K\in\mathbb{R}^{h_{kv}h_d\times d}$, $b_K\in\mathbb{R}^{h_{kv}h_d}$
Value projection: $W_V\in\mathbb{R}^{h_{kv}h_d\times d}$, $b_V\in\mathbb{R}^{h_{kv}h_d}$
output projection: $W_O\in\mathbb{R}^{d\times hh_d}$
RMSNorm：前文已经提到过，两个 normalization (query norm 以及 key norm) 的总参数量为 $2h_d$.

因此， Qwen3Attention 部分的总参数量为

$$ \mathrm{parameter}(\texttt{Qwen3Attention}) = hh_{d}d + 2h_{kv}h_dd +dhh_d + 2h_d = 2hh_dd + 2h_{kv}h_dd + 2h_d $$

分别代表 $W_Q, W_O, W_K,W_V$ 和两个 normalization layer 的参数量。

注意，这里我们没有加入 bias, 这是因为 QKV bias 在 Qwen3 中被取消，取而代之的是两个 normalization.

如果我们查看 Qwen2Attention 的代码，我们可以得到 Qwen2Attention 的总参数量为

$$ \mathrm{parameter}(\texttt{Qwen2Attention}) = 2hh_{d}d + 2h_{kv}h_dd +h_{kv}h_d+hh_d+h_{kv}h_d $$

分别代表 $W_Q, W_K, W_V,W_O$ 和 $b_Q,b_K, b_V$ 的参数量。

我们将计算结果汇总在一起就得到：

$$ \begin{aligned} \mathrm{parameter}(\texttt{Qwen3ForCausalLM}) &=d|V| + \mathrm{parameter}(\texttt{Qwen3Model})\\ &=2d|V| + N*\mathrm{parameter}(\texttt{Qwen3DecoderLayer}) + d\\ &= N*(2d + \mathrm{parameter}(\texttt{Qwen3MLP}) + \mathrm{parameter}(\texttt{Qwen3Attention})) + d(2|V|+1)\\ &= N*(2d+3dd_{ff}+2hh_{d}d + 2h_{kv}h_dd + 2h_d) + d(2|V|+1)\\ \end{aligned} $$

这是针对 Qwen3ForCausalLM 的参数量计算。这里的变量定义如下：

Math Variable	Code Variable	Description
$N$	`num_hidden_layers`	Transformer block 个数
$\vert V\vert$	`vocab_size`	词表大小
$d$	`hidden_size`	token embedding 的维度
$d_{ff}$	`intermediate_size`	MLP 的中间层的维度
$h_d$	`head_dim`	每个 head 的维度
$h$	`num_attention_heads`	query head 的个数
$h_{kv}$	`num_key_value_heads`	key 和 value head 的个数

Verification

接下来，我们就可以基于 Qwen3 的模型来验证了，比如，Qwen3-32B 的配置如下:

Field	Value
$N$	64
$\vert V\vert$	151936
$d$	5120
$d_{ff}$	25600
$h_d$	128
$h$	64
$h_{kv}$	8

根据上式，最终的参数量为

$$ \mathrm{parameter}(\texttt{Qwen3-32B}) = 32.762123264*10^9\approx 32.8B $$

我们使用 Qwen3-32B 的 index.json 可以得到其真实的参数量为

$$ \texttt{total\_size}/\texttt{precision} = 65524246528/2 = 32762123264 $$

与我们计算的结果一致（这里除以 2 的原因是其表示模型权重文件的总大小，以 bytes 为单位，一般模型都是 bfloat16, 大小为 2 个 bytes, 因此总参数量为总大小除以权重的精度）

MoE Model

MoE model 与 Dense model 不同的地方在于每一层的 FFN, 因此，其总参数计算方式为：

$$ \mathrm{parameter}(\texttt{Qwen3MoeForCausalLM})=N*(2d+\mathrm{parameter}(\texttt{Qwen3MoE})+hh_{d}d + 2h_{kv}h_dd +dhh_d + 2d) + d(2|V|+1) $$

对于 MoE layer, 其关键代码为：

1
2
3
4
5
6


class Qwen3MoeSparseMoeBlock:
 def __init__(self, config):
 self.num_experts = config.num_experts
 self.top_k = config.num_experts_per_tok
 self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False)
 self.experts = nn.ModuleList([Qwen3MoeMLP(config, intermediate_size=config.moe_intermediate_size) for _ in range(self.num_experts)])

我们记 $n$ 为总专家个数，即 num_experts，记 $k$ 为激活专家个数，即 num_experts_per_tok 或者 top_k,

首先 gate 的参数量为 $dn$, 接下来每个 expert 都是一个 Qwen3MLP, 因此 experts 总参数量为 $n * 3dd_{ff}$. 这样 MoE layer 的总参数量为

$$ \mathrm{parameter}(\texttt{Qwen3MoE}) = nd + 3ndd_{ff} $$

在推理时，只有一部分专家，也就是 $k$ 个专家会参与计算，此时激活参数量为

$$ \mathrm{parameter\_activated}(\texttt{Qwen3MoE}) = nd + 3kdd_{ff} $$

我们带入到 Qwen3MoeForCausalLM 中就得到：

模型总参数量为：

$$ \mathrm{parameter}(\texttt{Qwen3MoeForCausalLM})=N*(2d+nd + 3ndd_{ff}+hh_{d}d + 2h_{kv}h_dd +dhh_d + 2h_d) + d(2|V|+1) $$

模型激活参数量为：

$$ \mathrm{parameter\_activated}(\texttt{Qwen3MoeForCausalLM})=N*(2d+nd + 3kdd_{ff} + 2h_{kv}h_dd +dhh_d + 2h_d) + d(2|V|+1) $$

这里的变量定义如下：

Math Variable	Code Variable	Description
$N$	`num_hidden_layers`	Transformer block 个数
$\vert V\vert$	`vocab_size`	词表大小
$d$	`hidden_size`	token embedding 的维度
$d_{ff}$	`intermediate_size`	MLP 的中间层的维度
$h_d$	`head_dim`	每个 head 的维度
$h$	`num_attention_heads`	query head 的个数
$h_{kv}$	`num_key_value_heads`	key 和 value head 的个数
$n$	`num_experts`	总专家个数
$k$	`top_k` (`num_experts_per_tok`)	激活专家个数

MoE Verification

我们用 Qwen3-235B-A22B 来验证，其配置如下：

Field	Value
$N$	94
$\vert V\vert$	151936
$d$	4096
$d_{ff}$	1536
$h_d$	128
$h$	64
$h_{kv}$	4
$n$	128
$k$	8

计算得到总参数量为

$$ \mathrm{parameter}(\texttt{Qwen3-235B-A22B}) = 235.09363456*10^9\approx 235B $$

计算得到激活参数量为

$$ \mathrm{parameter\_activated}(\texttt{Qwen3-235B-A22B}) = 22.14456064 * 10^9 $$

实际总参数量为

$$ 470187269120/2 = 235093634560.0\approx 235B $$

可以看到，计算结果与实际相符。

Extension

通过以上计算过程，我们可以很轻松将上述公式扩展到混合架构或者是 MLA 上，对于混合架构，我们分别计算不同 attention 的 layer 个数，然后分别计算。对于 MLA, 我们可以替换 Attention 的计算逻辑。

对于小语言模型系列，比如 0.6B 和 1.7B, 模型的大部分参数集中在 embedding 上，因此 Qwen3 采取了 tie embedding 的方式来减少参数量，具体做法就是 nn.Embedding 和 lm_head 共享参数。

Visualization

接下来，我们来可视化一下不同大小模型不同模块的参数量占比。计算的代码如下：

 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


def compute_param_distribution( N: int, V: int, d: int, d_ff: int, h_d: int, h: int, h_kv: int, tie_word_embeddings: bool) -> dict:
 """
 Compute parameter distribution across different components of the model.

 Args:
 N: Number of layers
 V: Vocabulary size
 d: Hidden dimension
 d_ff: Feed-forward dimension
 h_d: Head dimension
 h: Number of attention heads
 h_kv: Number of key/value heads
 tie_word_embeddings: Whether input and output embeddings are tied

 Returns:
 dict: Parameter counts for each component
 """
 # Embedding parameters
 embedding = V * d
 if not tie_word_embeddings:
 embedding *= 2

 # Attention parameters per layer
 attention_per_layer = h * h_d * d + 2 * h_kv * h_d * d + d * h * h_d + 2 * h_d
 attention_total = attention_per_layer * N

 # FFN parameters per layer
 ffn_per_layer = 3 * d * d_ff
 ffn_total = ffn_per_layer * N

 # RMSNorm (2d), QK-Norm (2d), output normalization (d)
 others = N * (4 * d) + d

 return {
 "Embedding": embedding,
 "Attention": attention_total,
 "FFN": ffn_total,
 "Others": others,
 }

对于 dense 模型，每部分的参数量可视化如下：

可以看到，随着模型 size 增加，模型大部分参数量都集中在 FFN 上。

接下来，我们可视化一下 MoE 模型的参数分布，核心计算代码如下：

 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


def compute_moe_param_distribution(N, V, d, d_ff, h_d, h, h_kv, n, k, activate=False) -> dict:
 """
 Compute parameter distribution across different components of the MoE model.

 Args:
 N: Number of layers
 V: Vocabulary size
 d: Hidden dimension
 d_ff: Feed-forward dimension
 h_d: Head dimension
 h: Number of attention heads
 h_kv: Number of key/value heads
 n: Number of experts
 k: Number of experts per token
 activate: Whether to compute activated parameters

 Returns:
 dict: Parameter counts for each component
 """
 # Embedding parameters
 embedding = 2 * V * d

 # Attention parameters per layer
 attention_per_layer = h * h_d * d + 2 * h_kv * h_d * d + d * h * h_d + 2 * h_d
 attention_total = attention_per_layer * N

 # FFN parameters per layer
 if activate:
 moe_per_layer = n * d + 3 * k * d * d_ff
 else:
 moe_per_layer = n * d + 3 * n * d * d_ff
 ffn_total = moe_per_layer * N

 # RMSNorm (2d), QK-Norm (2d), output normalization (d)
 others = N * (2 * d + 2 * h_d) + d

 return {
 "Embedding": embedding,
 "Attention": attention_total,
 "MoE": ffn_total,
 "Others": others,
 }

结果如下

可以看到，MoE 模型的大部分参数还是集中在 MoE 模块上，但是由于其稀疏机制，在激活的参数里，MoE 占比从 95% 以上降低到了 60% 左右。

Conclusion

在本文中，我们基于 Qwen3 大语言模型系列，介绍了如何计算 dense 模型和 MoE 模型的参数量。模型的参数量计算为后面的显存占用以及优化提供了基础。

References

Qwen3 Collection

Notes on Magistral

Wed, 16 Jul 2025 11:04:04 +0800

Magistral 是 Mistral 提出的一个 reasoning model 系列，主要针对 math 和 code 两个 domain

Introduction

Mistral 在 2025 年 6 月 12 日发布了 Magistral ，一个 reasoning model, 包含两个模型，一个是由纯 RL 训练得到的 Magistral Medium, 另一个是由 SFT 和蒸馏 Magistral Medium 得到的 Magistral Small

作者首先介绍了一下本文的贡献：

介绍了如何仅使用 RL (而不是用蒸馏) 来训练 Magistral Medium
infra 上的改进，主要使用最新的权重来更新 generator
多语种能力，支持多种语言
系统性探究了 RLVR 的能力边界
开源了 Magistral small (24B)

Method

RL

RL 算法基于 GRPO 改进，主要有以下几点：

去掉了 KL Divergence loss, 这一点跟 DAPO 是一致的，提升模型的探索能力
Loss Normalization，在 sample 层面做平均，还是跟 DAPO 一致，减少不同长度输出对训练造成的影响
Advantage Normalization, 作者首先将每个 token 的 Advantage 定义为 $\hat{A}_i=R_i-\mu$, 其中 $\mu$ 是每个 group 的 reward 均值，然后在每个 mini-batch 里对 advantage 进行 Normalization (这里 $\hat{A}_{mean}$ 和 $\hat{A}_{std}$ 分别为 mini-batch advantage 的均值和方差)：

$$ \hat{A}_{i,t}^{norm}=\frac{\hat{A}_i-\hat{A}_{mean}}{\hat{A}_{std}} $$

CLIP-Higher, 跟 DAPO 一致，提高稀有 token 的被采样概率
Eliminating non-divsere groups, 跟 DAPO 一致，去掉过于简单和过于难的题目

最终 RL 阶段训练的损失函数为：

$$ \mathcal{L}(\theta) = \mathbb{E}_{q\sim P(Q),\{o_i\}_{i=1}^G\sim \pi_{old}(\cdot\mid q)}\frac{1}{\sum_{i=1}^G|o_i|}\sum_{i=1}^G\sum_{t=1}^{|o_i|}\min\left[r_{i,t}(\theta)\hat{A}_{i,t}^{norm}, \mathrm{clip}(r_{i,t}(\theta), 1-\epsilon_{low}, 1+\epsilon_{high})\hat{A}_{i,t}^{norm}\right], \mathrm{s.t.}, \exists 1\leq m < n \leq G, r_m\neq r_n. $$

其中

$$ r_{i,t}(\theta) = \frac{\pi_{\theta}(o_{i,t}\mid q, o_{i,Reward Shaping

作者还基于四个维度来构建 reward: formatting, correctness, length, 以及 language consistency

Formatting 针对数学和代码问题，作者要求模型输出符合特定的格式

Tag requirements: 思考过程用 <think></think> 包含，且只能包含一个 tag
mathematical responses: 对于数学问题，结果用 \boxed{} 包含
code response：包含一个 markdown block

[!tip] Future 最新的 GLM-4.1V-Thinking 认为，不应该在 RL 训练阶段加入 format reward

correctness 基于答案的正确性分配奖励

math correctness：使用 rule-based verifier 进行打分，使用 parser 和 Sympy 来比较模型输出以及 ground truth
code correctness: 构建单元测试，评估输出代码是否能通过所有的单元测试

Length penalty 与 DAPO 一致，使用 Length penalty 来惩罚过长的回答。

Language consistency 减少模型混合语言输出的问题。作者的做法是将 10% 的问题从英语转化为其他语种，然后使用 fastText 进行分类，确保内容都是一个语言。作者发现，通过这个简单的修改，就提高了模型的语言跟随能力。

[!tip] Recall MiMo 中也遇到了混合语言输出的问题，但是其并没有给出解决办法。

作者还在 system prompt 中规定了输出的格式以及语言。

[!tip] Observation 作者发现，RL 的训练对于 system prompt 非常的敏感，因为 system prompt 会提高 model 的 entropy, 然后提高模型的探索能力。

Data

数据包括 math problems 以及 code problems

Math

作者首先收集了 700K 样本，然后作者通过预处理以及过滤来保证数据的质量：

format filtering 要求问题完整，答案准确且可验证；去掉证明题和 multi-part 的问题；改写多项选择题为解答题，提高难度防止 reward hacking

difficulty filtering 使用两阶段的过滤策略。

第一阶段使用 LLM 进行多次采样，然后去除掉比较简单或者比较复杂的
第二阶段使用 RL checkpoint 进行多次采样，去除掉标准答案可能会存在问题的题目

最终一共得到 38K 的样本

Code

首先去掉没有 solution 和 test 的问题；然后去除掉标准答案不能通过所有 test 的问题。最终一共得到 35K 的样本，包含 Python 和 C++ 两种语言

Training

训练框架如下图所示

对于两个模型，作者采用了不同的方式进行训练

Magistral Medium: 使用 pure RL 进行训练
Magistral small: 使用 SFT + RL 进行训练

Magistral Medium 训练时满足的要求：

dataset is not too easy: 太简单的题目对模型提升没有帮助
Generation length does not stop growing: 逐步提升模型的最大输出长度
KV-cache memory burden is not too large: 降低 batch size 来减少 KV-cache 的内存占用

Magistral small 训练

SFT 数据集包括两部分，一部分是 Magistral Medium 回答正确的这部分数据，第二部分是公开数据集，包括 [[OpenThoughts]] 和 [[OpenR1]] 两个数据集，作者使用 Magistral Medium 来生成回答。作者还加入了 10% 的 instruction tuning 数据来保持模型的通用能力。

RL RL 的训练与 Magistral Medium 一致。

Infra

作者首先介绍了一下 RL 训练的 infra, infra 主要包括三个模块：

Trainers: 用于更新模型的权重
Generators: 用于采样，生成 roll-out
Verifiers: 对模型输出的结果进行打分

分布式 RL 训练的主要问题在于，不同长度的 roll-out 花费的时间不一致，作者发现，最长的 roll-out 和最短的 roll-out 的时间相差超过 5 倍以上。

因此，作者就提出了异步生成这个方法。具体的做法就是

首先由 Generator 生成多条 roll-out
当 roll-out 完成之后，立马用 Verifiers 对轨迹进行打分
收集 roll-out 以及对应的 reward, 直到达到给定的 batch 大小
使用 Trainer 更新 Generator 的权重, 将更新后的权重同步给 Generator，这样其他 generator 在生成新的 token 时用的就是新的权重

[!tip] Recall MiMo 里的做法是，当我们收集到给定数量的 roll-out 之后，我们就基于这些 roll-out 更新权重，然后进行下一次采样

训练时，对于每个 rank, 只要其收集到足够的 roll-out, 就会进行一次梯度更新。

Evaluation

Ablation

RL 的泛化性 作者探究了 RL 的 cross-domain generalization 能力，实验结果如下

Model	AIME’24	LiveCodeBench v5
Starting Checkpoint	32.2	22.7
RL (Math only)	62.5	38.3 (+15.6)
RL (Code only)	49.7 (+17.5)	42.7

可以看到，不管是使用 math 还是 code 数据单独进行训练，模型在另一个 domain 上的表现都有所提升。

[!tip] Future 这个结论与最新的 GLM-4.1V-Thinking 结论一致

Distillation v.s. RL for small models 作者探究了对于小语言模型，使用 RL 进行训练的效果更好，还是使用蒸馏的效果更好。

[!tip] Observation 实验结果发现，仅使用 RL 的效果与蒸馏差不多，甚至更好

Batch size 作者还探究了 batch size $n_{batch}$ 以及 mini-batch size $n_{mini}$ 的影响。这里 batch size 指的是用于更新梯度的 roll-out 数量，mini-batch size 指的是计算梯度的 roll-out 数量，作者还定义了并行生成 roll-out 的数量 $n_{async}$. 当 $n_{async}>> n_{batch}$ 时，生成的 sequence 就很可能是 off-policy 的。作者固定 $n_{async}=4096$, 然后对比了不同的 $n_{batch}$ 和 $n_{mini}$ 对模型表现的影响，结果如下：

![tip] Observation 当 $n_{batch}=n_{mini}$ 时，模型表现差不太多（左图）；当 $n_{batch}$ 为常数，而 $n_{mini}$ 逐渐减小时，模型表现会逐渐变差。

Advantage normalization 作者对比了三种针对 advantage 的 normalization 方式：

mini-batch: 在每个 mini-batch 里对模型进行 normalization
Group normalization: 在每个 group 里进行 normalization
no normalization: 没有 normalization

实验结果如下

作者发现，三种方式的区别并不是很大。

Analysis

length dimension 作者首先保存模型的 weight, 然后使用 PCA 进行降维，并在 2 维上进行可视化，作者对权重进行扰动，然后记录模型的 reward 以及输出长度。结果发现，模型存在一个 length dimension. 可视化结果如下图所示

multimodal extension 由于 Magistral medium 和 Magistral small 都是基于 MLLM 中的 LLM 开发得到的，作者还探究了更换 LLM 的 checkpoint 之后，原始 MLLM 的表现，结果发现，模型在多模态 reasoning 相关任务上的表现也得到了提升，结果如下图所示：

Impact on other capabilities 作者还探究了 RL 训练对模型其他表现的影响，结果发现，RL 训练可以提高模型的 tool calling 和指令跟随能力，实验结果如下表所示

Category	Benchmark	Mistral Medium 3	Magistral Medium
Function calling	Internal bench	87.2	87.4
Instruction following	IFEval	86.8	87.4

failed approaches 作者还介绍了一些尝试失败的做法：

partial Reward: 对于 coding 任务，作者使用 test 的通过率作为奖励，结果发现效果并不好，这是因为一些错误的解法的 test 通过率也很高
entropy bonus loss: 作者发现在损失函数中加入 entropy bonus loss 之后，模型的训练变得不稳定，而且效果不如使用更高的 $\epsilon_{high}$
作者还进一步验证在 PPO loss 中加入 KL divergence loss, 结果发现效果并不好，这与 DAPO 的结论一致
作者还尝试先 SFT Magistral Medium, 再进行 RL, 结果发现 RL 可以大幅度提高 SFT checkpoint 的表现。

Conclusion

本文中，作者提出了 Magistral, 一个针对 math 和 code 的 reasoning model, 作者介绍了训练细节。但是，从方法层面来看，和 DAPO 区别不是很大。关键点应该是作者详细介绍了各种消融实验，为后来相关探索提供了经验。

References

arxiv

Notes on SmolLM3

Tue, 15 Jul 2025 11:01:13 +0800

Hugging Face 在 2025 年 7 月 8 号发布了 SmolLM3, 一个 3B 的，128K 上下文，支持 6 种语言，支持 dual mode reasoning 的小语言模型。

Pre-training

Architecture

SmolLM3 是一个基于 transformer 的小语言模型，其模型架构与 SmolLM2 相似，SmolLM3 做了以下改进：

使用 GQA 代替 multi-head attention. 作者通过消融实验发现 GQA 和 MHA 的效果差不多，并且还可以减少 KV cache 的 size
NoPE: 作者使用了NoPE, 来选择性（每 4 个 layer）移除 position embedding. 这个方法可以在不损害模型短文本能力的同时，提高模型长上下文的表现
Intra-Document Masking: 不同的文档之间使用 attention masking 隔开
Training stability: 与 olmo 2 一样，作者移除了 embedding layer 的 weight decay, 来提升训练的稳定性。

Data

与 SmolLM2 一样，作者使用了11.2T token 进行训练，训练包括 3 个 stage。作者针对数据混合策略进行了消融实验，实验配置如下：

Stage	Stable phase	Stable phase	Decay phase
Description	Base training	High quality injection	LR Decay
Tokens	8T	2T	1.1T
Web	85%	75%	63%
Code	12%	15%	24%
Math	3%	10%	13%
datasets	web: FineWeb-Edu, DCLM, FineWeb2, Fineweb2-HQ code: The Stack v2, StarCoder2 PRS, Jupyter and Kaggle notebooks, Github issues, StackExchange math: FineMath3, InfiWebMath3+	Adding Stack-Edu, FineMath4+, InfiWebMath4+, MegaMath	upsampling of high-quality code data upsampling math data and introducing instruction and reasoning datasets such as OpenMathReasoning

其中，Web data 包含 12% 的多语种数据。

Training Recipe

训练过程中，作者使用了 2.36M tokens 的 batch size, 上下文长度为 4096. 优化器为 AdamW

作者使用了 nanotron 进行训练， datatrove 来处理数据， lighteval 来评估模型的表现。

模型在 384 张 H100 上训练了 24 天。分布式训练的配置如下

Mid-training

Mid-training 的主要目标为扩展模型的上下文以及提升模型的 reasoning 能力。

Long Context Extension

在预训练阶段结束之后，作者使用了额外的 100B tokens 来扩展模型的上下文。作者将扩展过程分为两个 stage:

Stage 1: 将模型的上下文从 4K 提升到 32K. 具体做法是将 RoPE 的 base frequency 提升到 1.5M
Stage 2: 将，模型的上下文从 32K 提升到 64K. 具体做法是将 RoPE 的 base frequency 提升到 5M

训练过程中，作者对 math, code 和 reasoning data 进行了上采样。作者发现，对长文本数据进行上采样并不会提高模型在 RULER 和 HELMET 上的表现。

[!tip] Recall Qwen2.5-1M 里也分析了长文本数据的问题，也就是大部分长文本数据依然是局部相关性强，而全局相关性弱

与 Qwen2.5 一样，作者还是用了 YARN 来在推理阶段进一步提高模型的上下文长度，作者发现模型可以处理 128K 上下文长度的文本。

Reasoning Mid-training

扩展模型上下文长度之后，作者还额外增加了一个 mid-training stage 来提高模型的 reasoning 能力。这个阶段的目标在于提升模型的通用能力。作者希望模型不针对特定的 domain, 如 math 或者 code 等。

训练过程中，作者使用了 35B 的 token. 数据来源包括 Open-Thoughts3-1.2M 以及 NVIDIA 的 Llama-Nemetron-Post-Training-Dataset-v1.1. Reasoning trace 由 DeepSeek-R1 生成。作者使用了 ChatML 的格式，还使用了 Packing 来提升训练效率。训练持续了 4 个 epoch.

Post-training

如何构建 dual instruction model 来同时支持 reasoning 和 non-reasoning 任务并没有一个共识，大部分模型的数据都是非公开的。因此，作者就构建了一个 training pipeline, 用于提升模型在两种模式下的能力。

Chat Template

作者首先构建了一个 chat template, 用于支持 reasoning 和 non-reasoning 两种模式。该 chat template 支持用户使用 /think 和 /no_think flag 来控制模型的思考模式。在 non-reasoning 模式下，作者还在模型输出中 prefill 了 <think>\n\n</think>, 这一点与 Qwen3 一致。

SmolLM3 还支持工具调用，其在 chat template 中加入了两种描述方式：XLM tools 和 Python tools.

SmolLM3 还在 system prompt 中加入了 metadata, 如知识的截止时间，当前的 reasoning mode 等。

SFT

作者针对 math, code, general reasoning, instruction following, 以及 multilinguality 这几个领域来提升模型的表现。

作者首先使用 reasoning mode 的 Qwen3-32B 来合成数据，合成数据的过程如上图所示。

最终，SFT 阶段的数据包括 1.8B token, 其中 1B 为 non-reasoning mode 的 token, 覆盖 12 个数据集， 0.8B 为 reasoning token, 覆盖 10 个数据集。作者训练了 4 个 epoch, 使用了 Packing 的技巧。

APO

SFT 阶段之后，作者使用了 Tulu3 的 preference dataset 用于 non-reasoning mode 的训练，然后合成了一批数据用于 reasoning mode 的训练，这批合成数据使用 Qwen3 32B 和 Qwen3 0.6B 生成得到，具体做法就是 Qwen3 32B 输出的结果定义为正样本，Qwen3 0.6B 输出的结果定义为负样本。

作者使用了 APO 算法来进行训练，APO 是 DPO 的一个变体， DPO 的目标函数为

$$ \mathcal{L}_{DPO}(x,y_w,y_{l}; \theta) = -\log \sigma(r_\theta(x,y_w) - r_\theta(x, y_l)) $$

其中

$$ r_\theta(x,y) = \beta\log \frac{\pi_{\theta}(y\mid x)}{\pi_{ref}(y\mid x)} $$

$\beta>0$ 是一个超参数。APO 的目标函数如下

$$ \mathcal{L}_{APO}(x,y_w,y_l;\theta) = -\sigma(r_\theta(x,y_w)) + \sigma(r_\theta(x,y_l)) $$

作者发现，模型的 reasoning 能力提升之后，其长上下文能力反而下降了。作者认为这是因为在 reasoning mid-training stage 的训练中，提升 reasoning 能力损害了模型的 long context 能力。并且，APO 的训练数据上下文长度大多都在 24K 左右。为了解决这个问题，作者提出了 Model merging 的方法

Model Merging

作者使用 MergeKit 来完成 model merging 的任务。merge 的过程包括两步：

构造一个 model soup, 包括 APO 的每个 checkpoint
将 model soup 与 mid-training 的一个拥有强上下文能力的 checkpoint 结合起来，作者使用了 linear merge, APO model soup 和 mid-training checkpoint 的权重分别为 0.9 和 0.1.

[!tip] Recall Kimi-k1.5 也使用了 model merging 的方法，来将 long CoT 模型的能力迁移到 short-CoT 模型上去

Evaluation

作者首先评估了一下 base model 的表现，评估使用了 12 个 benchmark, 对比的模型包括 Qwen3, Gemma3, LLaMA 3.2. 结果如下图所示

模型的多语种表现如下

Instruction (non-reasoning) 版本的评估结果如下

Reasoning 版本的评估结果如下

可以看到，Qwen3-4B 的表现是最好的。而 SmolLM3 的表现在 3B 左右也是非常强劲的

Conclusion

作者提出了 SmolLM3, 一个拥有长上下，文多语种以及 dual mode reasoning 能力的大语言模型，作者详细介绍了数据，训练以及 model merging 的技巧，来提高模型的表现。

References

blog

Notes on Qwen2.5-1M

Sat, 12 Jul 2025 11:00:47 +0800

Qwen 在 2025 年 1 月提出了 Qwen2.5-1M，一个拥有 1M 上下文长度的大语言模型系列。包含 7B，14B 两个开源模型以及 API 模型 Qwen2.5-Turbo. 主要改进方法包括长上下文数据合成，渐进式预训练以及多阶段 post-training 等。作者还对 inference 进行了优化，提高了 inference 的效率。

Method

架构上，Qwen2.5-1M 与 Qwen2.5 的架构一致，Qwen2.5-1M 包括 7B，14B 两个 size，还包括一个基于 MoE 的 API 模型 Qwen2.5-Turbo，不同的点在于，Qwen2.5-1M 的上下文长度为 1M，最大生成长度为 8K

Pretraining

Data 作者首先从 CC, arxiv, book, code repositories 等 domain 收集了原始数据。但是，作者发现，原始数据的局部相关性强，但是全局相关性弱。因此，作者基于原始数据进行了增广，来提高数据的长上下文依赖关系。具体有三个任务：

Fill in the middle: FIM 是 openAI 提出来的一个做法，核心思想就是将填空类问题转化为 next-token-prediction 的问题。通过这种方式，作者希望提升模型理解长上下文依赖的能力
Keyword-based and Position-based retrieval: 基于 keywords 或者 position 来找到对应的 paragraph，这个任务的目的是提高模型识别并连接相关信息的能力
Paragraph Reordering: 对输入的 paragraphs 进行随机打乱，然后要求模型重新组织段落的关系

Training 作者将训练拆分为了 5 个 stage：

stage 1 和 stage 2 与 Qwen2.5 的训练过程一致，stage 1 的上下文长度为 4096，stage 2 的上下文长度为 32768, 训练时，作者使用了 ABF 技巧来将 RoPE 的 base frequency 从 10,000 调整到了 1,000,000.
stage 3, stage 4 和 stage 5 分别将模型的上下文长度扩展到了 65,536 tokens, 131,072 tokens 以及 262,144 tokens, 对应的 RoPE base frequency 分别为 1M, 5M 和 10M. 训练时，作者使用了 75% 的长文本和 25% 的短文本，这样可以保证模型在短文本任务上的表现

最后，作者在评估了一下每个 stage 的表现，结果如下表所示

Training Length				RULER
	Avg.	4K	8K	16K	32K	64K	128K
32,768 Tokens	82.3	96.8	94.7	95.9	92.2	76.4	37.6
65,536 Tokens	86.8	96.5	95.5	93.6	92.5	86.7	56.0
131,072 Tokens	92.5	96.5	95.9	93.0	92.6	93.0	83.8
262,144 Tokens	92.7	95.6	93.8	93.1	94.1	91.9	87.6

可以看到，随着训练的上下文长度的提升，模型在更长上下文下的能力也有提升，说明模型具有一定的泛化性。

Post-training

Post-training 阶段与 Qwen2.5 一样，也分为了 SFT 和 RL 两个阶段。

在 SFT 阶段，作者从预训练预料中选择了一部分长文档的片段，然后让 Qwen2.5 来生成对应的 query，query 类型包括 summarization, information retrieval, multi-hop QA 等任务。接下来，作者使用 Qwen-Agent 框架基于全文来回答这些问题。最后，作者基于生成的 query，全文，以及模型产生的回答作为训练数据。

SFT 训练时，作者拆分为了两个 stage。 stage 1 作者在 32768 的上下文上进行训练，来提高模型短文本回答能力。第二个阶段，作者混合了 262,144 和 32768 上下文长度的训练数据。

RL 训练时，与 Qwen2.5 不一样的是，作者进使用了 offline RL，也就是 DPO。作者仅在 8192 的上下文长度上面进行训练。作者认为，长上下文的 RL 训练是非常耗时的，并且作者发现，短文本上进行 RL 的训练之后，模型在长文本上的表现也能得到提升。结果如下表所示

Model	Before RL	After RL
Qwen2.5-7B-Instruct-1M	7.32	8.08 (+0.75)
Qwen2.5-14B-Instruct-1M	8.56	8.76 (+0.20)
Qwen2.5-Turbo	7.60	8.34 (+0.74)

Inference

前面是训练部分的优化，主要是提升模型的上下文能力。接下来，作者详细介绍了如何在 Inference 阶段提升整体的推理效率和减少内存占用。

Length Extrapolation

与 Qwen2.5 一样，Qwen2.5-1M 也是用了 Dual Chunk Attention 和 YARN 来在推理阶段扩展模型的上下文长度，作者做了如下实验，来对比 Qwen2.5, Qwen2.5-1M 加上 DCA 之后的影响

结果显示，Qwen2.5-1M 的表现比 Qwen2.5 更好，并且加上 DCA 之后，两者的表现都有进一步的提升。

Sparse Attention

为了进一步提高计算效率，作者基于 MInference 来加速 perfilling phase. 并结合了前面的技巧来防止模型性能下降。

MInference MInference 的主要思想就是在长上下文中，有一些 critical token 对最终结果的影响是更大的。因此我们可以识别出这些 critical token 并只计算这些 token 对应的 attention score. 这些 critical token 对应的 pattern 被称为 Vertical-Slash pattern.

为了识别出这个 pattern，作者首先进行离线搜索，来决定最优的 configuration。这个 configuration 决定了 attention 应该如何计算。在 Inference 阶段，MInference 首先计算最后一个 query 和前面所有 key 的 attention，然后基于 configuration 来动态选择 pattern。通过 MInference，我们可以降低 10 倍以上的内存和算力消耗。

Integrating with Chunked prefill 但是 MInference 的问题在于，整个 sequence 是并行处理的，这会导致内存占用持续上升。为了解决这个问题，作者提出了 chunked prefilling 的技巧，来降低 VRAM 的消耗。具体做法就是，将整个 sequence 分为若干个 chunk，然后每个 chunk 里，选取最后 64 个 token 作为 query，在每个 chunk 中分别识别出 critical token，这样就降低了 MInference 的内存占用

接下来，作者在集成 DCA 的时候，发现性能有所下降。作者认为，这是由于 DCA 的 position id 信息不连续所导致的，为了解决这个问题，作者在选择 critical token 的时候，使用了连续版的 position id 信息。在最终推理的时候，还是使用 DCA 本身的位置信息。

Sparsity refinement 前面提到，MInference 需要先进行离线搜索决定最优的 configuration，但是对于 1M token 的上下文，这个过程还是非常耗时的。因此，作者构建了一个加速离线搜索的方法，具体做法就是定义两个 attention score，一个是 full attention, 另一个是 sparse attention，然后计算两者的差值，如果说相差比较小，则说明 critical token 抓住了全局信息，这个配置是有效的。其公式定义如下：

$$ \mathrm{Attention\_Recall} = \exp\left(\log\sum_{0\leq j\leq i}\exp \left(\frac{q^Tk_j}{\sqrt {d}}\right) - \log\sum_{j\in\mathcal{critical}}\exp \left(\frac{q^Tk_j}{\sqrt {d}}\right)\right) $$

Attention Recall 越高，说明选取的 critical token 越好，其 configuration 也就越好。

作者进一步分析了 sparse attention 对 accuracy 的影响，结果如下

可以看到，仅使用 MInference 会导致模型性能下降，但是加入 refinement 之后，模型的表现基本上和 full attention 差不太多。

Inference Engine

Kernel Optimization 作者还对 inference engine 进行了优化，作者使用 BladeLLM 作为 Qwen2.5-1M 的推理引擎。

作者主要做了两点优化，第一是对 sparse attention kernel 进行了优化，提高了 sparse attention 的计算效率，结果发现，在 1M 的上下文下，BladeLLM 比 flash attention 要快 27.8 倍。

第二是针对 MoE kernel 的优化。作者发现，decoding 的表现是与 memory access speed 相关的。具体来讲，当 batch size 超过 32 之后，获取模型参数成了效率的瓶颈。因此，作者使用了一系列技巧来提高 memory access 的效率

Pipeline parallelism 作者还对 Chunked pipeline parallelism 进行了优化，Chunked pipeline parallelism 的问题在于，在长上下文的场景下，不同长度的 chunk 会对 attention 的计算时间产生很大影响。不同的计算时间会产生 pipeline bubbles.

BladeLLm 使用了 Dynamic Chunked pipeline parallelism 来解决这个问题，该方法通过计算复杂度来调整每个 chunk 的大小，进而使得最终的处理时间尽可能一致

Scheduling 作者还在 Scheduling 上进行了优化，已有的推理引擎主要分为四个模块：API server, scheduler, model runner 以及 decoder

已有方法的问题在于，non-GPU 的操作会占用大量时间，导致 GPU 利用率非常低。因此，作者在 BladeLLM 中进行了改进，使用了 Totally Asynchronous Generator (TAG) 的架构，主要有：

Scheduler：动态分配 KV cache，类似于 speculative sampling, 而不必等前面的结果完成
Runner: 基于 Scheduler 分配的任务直接进行处理，处理完之后直接处理下一个任务
Decoder：基于 token id，进行解码，然后发送给前端的 API server

Evaluation

作者主要在三个 benchmark 上进行了评测：

RULER: RULER 是 Needle-in-ahaystack 任务的一个扩展笨笨，其要求模型从不相关的上下文中找到多个 “needles” 或者回答多个问题，数据最长为 128K tokens.
LV-Eval: LV-Eval 要求模型从上文本中同时理解多个 evidence fragments，数据最长为 256K tokens
Longbench-Chat: 评估模型在长上下文下与人类偏好对齐的程度，数据最长为 100K tokens

Qwen2.5-1M 与 Qwen2.5 的对比表现如下

可以看到，相比于 Qwen2.5，Qwen2.5 模型的表现有了大幅度的提升。

Conclusion

在本文中，作者提出了 Qwen2.5-1M 系列大语言模型，包括 7B，14B 两个 size，以及一个 MoE 架构的 API 模型 Qwen2.5-Turbo。作者在训练和推理两方面进行了改进，最终将模型的上下文长度扩展到了 1M。从现在的角度来看，不管是 Reasoning model 还是 agent 的训练都依赖 long Context 作为基础能力。

References

Qwen2.5-1M Technical Report

Notes on Qwen2.5

Sat, 12 Jul 2025 10:51:42 +0800

2024 年 12 月 Qwen 发布了 Qwen 2.5 系列大语言模型，包括 7 个 dense 模型以及两个 MoE 模型，Qwen2.5 在 pre-training 阶段使用了 18T token 进行。在 post-training 阶段使用了 1M 的样本，还使用了 DPO 以及 GRPO 来进行 RL 的训练

Qwen2.5 主要在以下方面进行了改进

模型方面，提供了更多的 size，Qwen2 中只有 0.5B, 1.5B, 7B, 72B 四个 size, 在 Qwen2.5 中，加入了 3B, 14B 和 32B 三个 size 的模型
数据方面，pre-training 阶段使用了 18T 的 token， post-training 阶段使用了 1M 的样本
功能方面，Qwen2.5 支持更长的上下文长度（8K），支持结构化输入和输出，拥有更强的工具调用能力。

Method

Architecture

模型架构这方面，Qwen2.5 和 Qwen2 的模型架构是一致的，tokenizer 页没有太大变化。为了支持工具调用，作者额外增加了 18 个 control token

Pre-training

data Qwen2.5 从以下方面提高了预训练数据的质量

Better data filtering: 使用 Qwen2-Instruct 来过滤掉质量的数据，然后从多维度对训练数据进行打分，从而提高数据的质量
Better math and code data: 加入了 Qwen2.5 Math 以及 Qwen2.5 Coder 的训练数据来提高模型的数学和代码能力
Better synthetic data: 作者使用 Qwen2-72B-Instruct 以及 Qwen2-Math-72B-Instruct 来合成 math, code, knowledge domain 的数据，然后通过过滤以及 Qwen2-Math-RM-72B 来提高数据的质量
Better data mixture: 作者使用 Qwen2-Instruct 来分类，然后平衡不同 domain 的数据分布。作者发现 e-commerce, social media 以及 entertainment 的数据重复性高，且大多都是机器生成的。而 technology, science 以及 academic research 等 domain 的数据质量更高。作者对不同 domain 的数据进行了上采样或者下采样。

基于这个过程，作者一共收集了18T tokens

Hyper-parameters 作者构建了针对超参数的 scaling law，即决定最优的训练超参数如 batch size, learning rate 等

作者通过实验得到了 model size $N$ 以及 pre-training data size $D$ 与 learning rate $\mu_{opt}$ 和 batch size $B_{opt}$ 之间的关系。

Long context pre-training 为了提升模型的上下文长度，作者将 pre-training 拆分为两个 stage，第一个 stage 的上下文长度为 4096，第二个 stage，作者将上下文长度从 4096 扩展到 32768.

在提升模型上下文过程中，作者使用 ABF 技巧将 Position Encoding 的 base frequency 从 10,000 提升到了 1,000,000.

对于 Qwen2-5-Turbo，作者实现了渐进式上下文长度扩展策略，模型上下文长度扩展经历四个阶段：32768, 65536, 131072 到最终的 262,144. 此时，RoPE 的 base frequency 为 10,000,000. 在训练的每个阶段，作者都使用了 40% 的长文本以及 60% 的短文本，以保证在扩展模型上下文长度的同时，还能保持模型在不同上下文长度下的表现。

为了提高模型在 inference 时的长上下文表现，作者使用了 Dual Chunk Attention 和 YARN 两个技巧。通过这两个技巧，作者将 Qwen2.5-Turbo 的上锈阿文扩展到了 1M，将其他模型的上下文长度扩展到了 131072.

Post-training

Qwen2.5 的 post-training 分为两个大的 stage: SFT 和 RL，其中 RL 又分为两个小的 stage，分别是 offline RL 和 online RL

在 SFT 阶段，作者主要做了以下改进：

Long-sequence generation: 作者将 Qwen2.5 的输出长度提升到了 8192, 为了扩展模型输出的长度，作者构建了 Long-response 数据集，然后基于 back-translation 来生成对应的 query，最后使用 Qwen2 来过滤低质量的数据
Math: 作者在 SFT 阶段加入了 Qwen2.5-Math 的 CoT 数据，包括公开数据集，K12 问题集一集合成数据等。作者通过 rejection sampling 以及 annotated answers 来生成 CoT 过程
Code: 作者加入了 Qwen2.5-Coder 的 SFT 数据，作者基于多个 agent 来生成多样化高质量的 Instruction, 然后还从 code-related QA website 以及 Github 上获取数据来扩展数据集。对于最终的数据，作者使用了 sandbox 来保证代码的质量
Instruction following: 作者构建了一个基于 code 的验证框架，让 LLM 同时生成 Instruction 和对应的验证代码，验证的单元测试。最后，通过 rejection sampling 来得到最终的数据集
Structured Data Understanding: 作者还构建了针对 tabular QA, fact verification, error correction 以及 structured understanding 等数据集。作者在回答中加入 CoT，作者提高了模型对 structured data 的理解能力
Logical Reasoning: 作者构建了 70,000 个不同 domain 的 query，有多种格式，覆盖了 analogical reasoning, causal reasoning 等 domain
Cross-Lingual Transfer: 作者使用了一个翻译模型，来将 Instruction 转换到 low-resource language 上，进而提高模型在对应语种上的表现
Robust System Instruction: 作者构建了不同的 system prompt 用于提升 system prompt 的多样性。作者发现，使用不同的 system prompt 可以减少模型的 variance, 提高模型的 robustness.
Response Filtering: 作者使用了多种自动化标注方法来保证最终 response 的质量

最终，作者一共收集到 1M 的 SFT 样本，模型训练了两个 epoch

在 RL 阶段，作者首先基于 SFT model 来进行采样，然后将高质量的回答作为正样本，低质量的回答作为负样本，通过这个过程，一共采集到了150K的样本。最后，作者使用 DPO 来进行训练。

然后，作者进行了 online stage 的 RL 训练，这一阶段主要是对齐模型与人类的价值观。这一阶段的数据包括公开数据集，私有数据集。作者使用不同的 checkpoint 来进行采样，然后作者使用 GRPO 来进行训练.

Long Cotnext Fine-tuning

作者还针对 Qwen2.5-Turbo 做了额外的 post-training, 来进一步提高其在长上下文下的表现。

在 SFT 阶段，作者使用了一个两阶段方法，第一阶段仅在短文本上进行训练（上下文长度为 32768），这一阶段的训练数据与其他 Qwen2.5 的模型训练数据相同。第二个阶段，作者混合了短文本和长文本（262144）来进行训练，来提高模型在长上下文情景下的指令跟随能力

在 RL 阶段，作者使用了和其他 Qwen2.5 模型相同的训练策略。作者认为：

长上下文下训练 RL 代价很大
reward model 更偏向于长文本
RL 尽管只在短文本上进行训练，其还是可以提高模型在长上下文下的表现。

Evaluation

我们仅关注 instruction 版本的 72B,32B 和 7B 模型

可以看到，Qwen2.5 72B 模型表现和 LLaMA3.1 405B 表现差不多，其他两个 size 的模型基本上达到了 SOTA

最后，作者评估了一下 DCA+YaRN v.s. Full attention 的表现，结果如下图所示

可以看到，使用 DCA+YaRN 之后，模型的推理效率比 full attention 要快 3-4 倍。

Conclusion

在本文中，作者提出了 Qwen2.5 系列大语言模型，包括 7 个 dense 模型以及两个 MoE 模型，作者详细介绍了模型的 pre-training 和 post-training. 评测结果发现 Qwen2.5 模型基本上达到了 SOTA.

作者认为，未来工作有：

使用更多更多样化的 pre-training 和 post-training 数据
多模态大模型的构建，特别是 omni-modal
提高模型的 Reasoning 能力

References

Arxiv

Dual Chunk Attention

Sat, 12 Jul 2025 10:41:12 +0800

Dual Chunk Attention (DCA) 由阿里巴巴在 2024 年 9 月份提出，DCA 是一个无需训练的，扩展 LLM 上下文长度的方法，后续，DCA 被应用于 Qwen2, Qwen2.5, Qwen2.5-1M 以及 Qwen3 中，与 YARN 一起作为扩展模型上下文的有效手段

Introduction

提升 LLM 上下文长度的方法可以分为两类：一类是 training-free 的，包括 LM-infinite 和 StreamingLLM 等，这些方法以损失 long range dependency 为代价来保持较低的 perplexity。另一类为了保留全局信息，则是通过外插来扩展模型的上下文，主要工作我们在 YARN 中已经回顾了。

第二类方法的问题在于，其依赖训练，在 training-free 的 setting 下，这些方法也会导致 perplexity 的上升

因此，在本文中，作者就提出了 Dual Chunk Attention (DCA) ，一个无需训练的，扩展 LLM 上下文长度的方法。DCA 的主要做法是将 attention 的计算进行分块，这样就可以提高计算效率。

通过实验，作者给出了三点关键发现：

Extrapolation： DCA 可以在无需训练的情况下，将 LLM 的上下文提升到 32K，而不导致 Perplexity 大幅度增加
Orthogonality： DCA 可以和其他方法一起使用，如 YARN, 这一点已经在 Qwen2.5-1M 以及 Qwen3 中得到了应用
Long Context Understanding: DCA 可以在无需训练的情况下，在长上下文设置下，达到已有 SOTA 模型的表现

Preliminary

对于一个长度为 $L$ 的 token 序列，我们首先定义对应的 position id 如下

$$ P_{\mathbf{q}} = [0,1,\dots,L-1],\quad P_{\mathbf{k}} = [0,1,\dots,L-1] $$

然后，对于第 $i$ 个位置和第 $j$ 个位置的 token，其 attention score 定义为：

$$ \langle f(\mathbf{q}, i), f(\mathbf{k}, j)\rangle =\langle R_{\theta,i}\mathbf{q}, R_{\theta,j}\mathbf{k}\rangle =\mathbf{q}^TR_{\theta, i-j}\mathbf{k} $$

具体细节参考 Position Encoding 中的 RoPE 部分介绍。这里面的关键在于，最后的结果只与相对位置 $i-j$ 相关，而与绝对位置 $i$ 和 $j$ 无关。因此，我们可以用一个相对位置矩阵 $M\in\mathbb{R}^{L\times L}$ 来表示这个信息，其中 $M_{ij}=P_{\mathbf{q},i}- P_{\mathbf{k},j}$ 代表了第 $i$ 个位置的 query $\mathbf{q}$ 和第 $j$ 个位置的 key $\mathbf{k}$ 的相对位置信息，其示意图如下所示

原始版本的 RoPE 的问题在于，在训练时，模型没有见过更长的上下文，因此其泛化性也最差，这一点在 YARN 已经得到了验证

Method

DCA 的关键在于将 sequence 分割为若干个 Chunk，然后将 attention 的计算拆分为三个部分：

intra-chunk：负责计算每个 chunk 内部的 attention
inter-chunk ：负责计算 chunk 之间的 attention
successive-chunk：负责计算相邻两个 chunk 之间的 attention

为了更好理解 DCA，我们接下来假设 $L=12$, 这时我们有

$$ P_{\mathbf{q}} = [0,1,\dots,11],\quad P_{\mathbf{k}} = [0,1,\dots,11] $$

Intra-Chunk Attention

我们首先定义个超参数 chunk size $s>0$, 然后我们将我们的 sequence 分割成 $L/s$ 个 chunk，然后每个 chunk 重新进行编号，就得到了如下的 position id

$$ P_{\mathbf{q}}^{Intra} = [0,1,\dots,L-1]\mod s,\quad P_{\mathbf{k}}^{Intra} = [0,1,\dots,L-1]\mod s $$

接下来，我们定义 intra-chunk 的相对位置矩阵 $M$，此时，我们仅在每个 chunk 内部进行计算 attention，即

$$ M[i][j] = \begin{cases} P_{\mathbf{q},i}^{Intra} - P_{\mathbf{k},j}^{Intra},&\text{if}\lfloor P_{\mathbf{q},i}/s \rfloor = \lfloor P_{\mathbf{k},j} /s\rfloor ,\\ 0,&\text{otherwise} \end{cases} $$

在上面的例子中，假设 $s=6$, 那么我们新的 position id 就变成了

$$ \begin{aligned} P_{\mathbf{q}}^{Intra} = [0,1,2,3,4,5,0,1,2,3,4,5]\\ P_{\mathbf{k}}^{Intra} = [\underbrace{0,1,2,3,4,5}_{\text{Chunk 0}},\underbrace{0,1,2,3,4,5}_{\text{Chunk 1}}] \end{aligned} $$

对其进行可视化，我们就得到

Inter-Chunk Attention

接下来，我们来看一下不同 chunk 之间如何计算彼此的 attention。在 Intra-chunk attention 计算中，我们忽略了跨 chunk 的信息，而且，由于现在的 position id 不再是单调递增的了，我们直接使用 $P_{\mathbf{q}}^{Intra}$ 和 $P_{\mathbf{k}}^{Intra}$ 给出的位置信息不对，这也是为什么我们在 Intra-chunk attention 中要求 query 和 key 在同一个 chunk 中才能计算的原因。

为了解决这个问题，作者构建了一个新的 position id。首先我们引入一个新的超参数 $c>\max_i P_{\mathbf{q},i}$, $c$ 代表了模型预训练时的上下文长度，如 4096。

接下来，基于 $c$, 我们定义新的 position id 如下：

$$ P_{\mathbf{q}}^{Inter} = [c-1,c-1,\dots,c-1]\in\mathbb{R}^s,\quad P_{\mathbf{k}}^{Inter} = P_{\mathbf{k}}^{Intra} $$

注：这里的 $P_{\mathbf{q}}^{Inter}$ 指的是某一个 chunk 中的 position id，每个 chunk 中的 position id 都相同，请参考例子理解，后面不再赘述。

也就是说，在计算跨 chunk 的 attention 的时候，我们直接把 query 的 position id 设置为最大值，然后 key 的 position id 依然使用 intra-chunk 的位置信息。由于 $\max_i P_{\mathbf{k},i}=s-1$, 因此我们有

$$ M[i][j] = P_{\mathbf{q},i}^{Inter} - P_{\mathbf{k},j}^{Inter} = c - 1 - P_{\mathbf{k},j}^{Inter}\geq c - 1 - (s- 1) \geq c-s. $$

最后，我们对于 inter chunk 的位置矩阵 $M$ 定义如下：

$$ M[i][j] = \begin{cases} P_{\mathbf{q},i}^{Inter} - P_{\mathbf{k},j}^{Inter},&\text{if}\lfloor P_{\mathbf{q},i}/s \rfloor \neq \lfloor P_{\mathbf{k},j} /s\rfloor ,\\ 0,&\text{otherwise} \end{cases} $$

在上面的例子中，当 $c=10$ 时，$c-1=9$, 我们有

$$ P_{\mathbf{q}}^{Inter}=[\underbrace{9,9,9,9,9,9}_{\text{Chunk 0}},\underbrace{9,9,9,9,9,9}_{\text{Chunk 1}}] $$

对其进行可视化，得到

Successive-Chunk Attention

现在我们既可以计算 intra-chunk，也可以计算 inter-block 的 attention，但是问题是对于相邻的 chunk，其位置信息不对了，从上面的可视化中，我们可以看到，当 $P_{\mathbf{q},i}=6$ , $P_{\mathbf{k},j}=5$ 时，我们有

$$ P_{\mathbf{q},i}^{Inter} - P_{\mathbf{k},j}^{Inter}=9-5=4\neq 1 = P_{\mathbf{q},i}-P_{\mathbf{k},j} $$

也就是说，inter-block 的 attention 会破坏原有的相对位置信息，因此我们就通过 successive chunk attention 来解决这个问题，使得 $P_{\mathbf{q},i}^{Inter} - P_{\mathbf{k},j}^{Inter}\approx P_{\mathbf{q},i}-P_{\mathbf{k},j}$.

作者发现，这个问题不是所有的 chunk 都有，而是只存在于相邻的 chunk 中，因此，作者又加入了一个超参数 $w>0$ 代表了 local window size，我们可以直接将其设置为 $c-s$, 通过这个 local window，我们调整对应的 position id 如下：

$$ P_{\mathbf{q}}^{Succ} = [\overbrace{s,s+1,\dots,s+w-1}^{w \text{ elements}},c-1, \dots,c-1]\in\mathbb{R}^s,\quad P_{\mathbf{k}}^{Succ} = P_{\mathbf{k}}^{Inter} $$

对于 successive chunk 的位置矩阵 $M$ 定义如下：

$$ M[i][j] = \begin{cases} P_{\mathbf{q},i}^{Inter} - P_{\mathbf{k},j}^{Inter},&\text{if}\lfloor P_{\mathbf{q},i}/s \rfloor - \lfloor P_{\mathbf{k},j} /s\rfloor=1 ,\\ 0,&\text{otherwise} \end{cases} $$

在上面的例子中，我们设置 $w=4$, 就得到

$$ P_{\mathbf{q}}^{Succ}=[\underbrace{6,7,8,9,9,9}_{\text{Chunk 0}},\underbrace{6,7,8,9,9,9}_{\text{Chunk 1}}] $$

对其进行可视化，得到

Computation

接下来，我们把所有的改进放在一起，就得到

$$ M[i][j] = \begin{cases} P_{\mathbf{q},i}^{Intra} - P_{\mathbf{k},j},&\text{if}\lfloor P_{\mathbf{q},i}/s \rfloor - \lfloor P_{\mathbf{k},j} /s\rfloor=0 ,\\ P_{\mathbf{q},i}^{Succ} - P_{\mathbf{k},j},&\text{if}\lfloor P_{\mathbf{q},i}/s \rfloor - \lfloor P_{\mathbf{k},j} /s\rfloor=1 ,\\ P_{\mathbf{q},i}^{Inter} - P_{\mathbf{k},j},&\text{if}\lfloor P_{\mathbf{q},i}/s \rfloor - \lfloor P_{\mathbf{k},j} /s\rfloor>1. \end{cases} $$

基于上面的位置矩阵 $M$，我们再依次计算对应的 attention score

$$ \langle f(\mathbf{q}, i), f(\mathbf{k}, j)\rangle = \begin{cases} \langle f(\mathbf{q}, P_{\mathbf{q},i}^{Intra}) , f(\mathbf{k}, P_{\mathbf{k},j})\rangle,&\text{if}\lfloor P_{\mathbf{q},i}/s \rfloor - \lfloor P_{\mathbf{k},j} /s\rfloor=0 ,\\ \langle f(\mathbf{q}, P_{\mathbf{q},i}^{Succ}) , f(\mathbf{k}, P_{\mathbf{k},j})\rangle,&\text{if}\lfloor P_{\mathbf{q},i}/s \rfloor - \lfloor P_{\mathbf{k},j} /s\rfloor=1 ,\\ \langle f(\mathbf{q}, P_{\mathbf{q},i}^{Inter}) , f(\mathbf{k}, P_{\mathbf{k},j})\rangle,&\text{if}\lfloor P_{\mathbf{q},i}/s \rfloor - \lfloor P_{\mathbf{k},j} /s\rfloor>1. \end{cases} $$

Code

首先是 RotaryEmbedding 部分的修改

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22


class DCARotaryEmbedding(nn.Module):
 def __init__(self, max_len, chunk_size, local_window):
 self.max_len = max_len
 self.chunk_size = chunk_size
 self.local_window = local_window

 self.inv_freq = ...

 def forward(self, x):
 q_t = torch.arange(self.chunk_size)
 qc_t = (q_t + self.chunk_size).clamp(max=self.chunk_size)
 k_t = torch.arange(seq_len) % self.chunk_size

 q_freqs = torch.outer(q_t, self.inv_freq) # seq_len x dim/2
 qc_freqs = torch.outer(qc_t, self.inv_freq)
 k_freqs = torch.outer(k_t, self.inv_freq) # seq_len x dim/2

 q_emb = torch.cat((q_freqs, q_freqs), dim=-1) # seq_len x dim
 qc_emb = torch.cat((qc_freqs, qc_freqs), dim=-1)
 k_emb = torch.cat((k_freqs, k_freqs), dim=-1) # seq_len x dim
 # compute related sin, cos
 return q_sin, q_cos, qc_sin, qc_cos, k_sin, k_cos

attention 计算时的逻辑

 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


class Attention(nn.Module):
 def forward(...):

 key_states = apply_rotary_pos_emb(key_states, k_cos, k_sin, position_ids)
 q_states_intra = apply_rotary_pos_emb(query_states[:, :, :chunk_len, :], q_cos, q_sin,
 position_ids[:, :chunk_len])
 k_states_prev = key_states[:, :, :chunk_len, :]
 v_states_prev = value_states[:, :, :chunk_len, :]
 # first chunk
 flash_result = do_flash_attn(q_states_intra, k_states_prev, v_states_prev)
 flash_results.append(flash_result)
 remain_len = kv_seq_len - chunk_len

 while remain_len > 0:
 flash_per_chunk = []
 begin = kv_seq_len - remain_len
 curr_chunk_len = min(chunk_len, remain_len)
 end = begin + curr_chunk_len
 # current chunk, intra-chunk attention
 q_states_intra = apply_rotary_pos_emb(query_states[:, :, begin:end, :], q_cos, q_sin,
 position_ids[:, begin:end])

 k_states_intra = key_states[:, :, begin:end, :]
 v_states_intra = value_states[:, :, begin:end, :]
 flash_result = do_flash_attn(q_states_intra, k_states_intra, v_states_intra)
 flash_per_chunk.append(flash_result)
 # successive chunk attention
 q_states_succ = apply_rotary_pos_emb(query_states[:, :, begin:end, :], qc_cos, qc_sin,
 position_ids[:, begin:end])
 flash_result = do_flash_attn(q_states_succ, k_states_prev, v_states_prev, False)
 flash_per_chunk.append(flash_result)
 # inter chunk attention
 if begin - (k_states_prev.size(-2)) > 0:
 prev_len = k_states_prev.size(-2)
 q_states_inter = apply_rotary_pos_emb(query_states[:, :, begin:end, :], qc_cos, qc_sin,
 position_ids[:, chunk_len - 1][:, None].repeat(1, curr_chunk_len))
 k_states_inter = key_states[:, :, :begin - prev_len, :]
 v_states_inter = value_states[:, :, :begin - prev_len, :]
 flash_result = do_flash_attn(q_states_inter, k_states_inter, v_states_inter, False)
 flash_per_chunk.append(flash_result)

 flash_results.append(flash_per_chunk)
 k_states_prev = k_states_intra
 v_states_prev = v_states_intra
 remain_len = remain_len - chunk_len
 # merge the final results
 attn_output = merge_attn_outputs(flash_results)

Evaluation

作者还分析了一下 DCA 的效率，结果如下

可以看到，在 flash attention 的基础上加上 DCA 之后，内存占用和推理时间并没有发生太大变化

作者还分析了三种 attention 对结果的贡献，如下图所示

结果显示，intra block 的 perplexity 是最低的，但是其在下游任务上表现是最差的。当三者结合在一起之后，perplexity 和下游任务上的表现都是最好的。

Conclusion

本文中，我们回顾了 Qwen 系列扩展大模型上下文的方法 Dual Chunk Attention （DCA）通过将 attention 切分成更小的 chunk，然后将 attention 的计算分为 intra-chunk，inter-chunk 和 successive-chunk，分别处理 chunk 内部，chunk 之间以及相邻 chunk 的 attention，通过这种方式，在无需训练的情况下，我们可以有效将模型上下文长度扩展 4 倍以上。

References

Arxiv
Github

Notes on Qwen2

Sat, 12 Jul 2025 10:36:43 +0800

Introduction

2024 年 9 月 Qwen 发布了 Qwen2 系列技术报告，Qwen2 系列包括 4 个 dense 模型（0.5B, 1.5B, 7B, 72B）和一个 MoE 模型（总参数 57B，激活参数 14B），作者主要在架构，数据和长上下文上进行了改进。

Method

Model

对于 dense 模型，Qwen2 在 Qwen-LLM 的基础上做了如下改动：

使用 Group Query Attention (GQA) 替换 MHA，来优化 KV cache，提高 throughput
使用 Dual Chunk Attention 和 YARN 来提高模型上下文长度和训练效率

其余与 Qwen 一致，包括 SwiGLU，RoPE，RMSNorm 和 pre-normalization

对于 MoE 模型，Qwen2-MoE 基于 Qwen1.5 进行了改进，主要是 3 点：

作者使用了更细粒度的专家个数，作者认为细粒度的专家可以提供更丰富的 combination，这一点与 olmoe 的结论相同
与 DeepSeek-MoE 一样，作者使用了共享专家和路由专家
作者使用了类似 upcycling 的方法来初始化模型。假设一共有 $n$ 个专家，每个专家的维度为 $h_E$, 原始 dense 模型的维度为 $h_{FFN}$, 那么我们会把 dense 模型的参数复制 $[nh_E/h_{FFN}]$ 次，这样就可以扩展到任意个数的 MoE 模型上。作者还对参数进行 shuffle，来提高 diversity。最后，作者还对 50% 的参数进行随机初始化，来提高模型的 capacity。

模型配置如下

Configuration	0.5B	1.5B	7B	72B	57B-A14B
Hidden Size	896	1,536	3,584	8,192	3,584
# Layers	24	28	28	80	28
# Query Heads	14	12	28	64	28
# KV Heads	2	2	4	8	4
Head Size	64	128	128	128	128
Intermediate Size	4,864	8,960	18,944	29,568	2,560
# Routed Experts	-	-	-	-	64
# Activated Experts	-	-	-	-	8
# Shared Experts	-	-	-	-	8
Embedding Tying	True	True	False	False	False
Vocabulary Size	151,646	151,646	151,646	151,646	151,646
# Trained Tokens	12T	7T	7T	7T	4.5T

Pre-training

预训练阶段的数据基于 Qwen 和 Qwen1.5，数据处理策略如下：

使用基于 heuristic 和 model-based 方法来过滤掉低质量的数据
加入了 code， math 和 multilingual 的数据
平衡了各个类别的数据分布

初始数据包括 12T token，经过过滤得到 7T token。作者发现，使用 12T token 进行训练，模型的表现不如使用 7B token 训练得到的模型效果好。因此除了 0.5B 的模型，其他模型使用的都是 7T 的 token

对于 MoE 模型，作者使用了额外的 4.5T token 来进行预训练。

在训练过程中，作者还加入了 multi-task instruction 数据，来提高模型的上下文学习能力和指令跟随能力。

作者还将 Qwen2 模型系列的上下文长度从 4096 扩展到 32768，扩展过程中作了三个改动：

加入了更多高质量的长上下文数据
将 RoPE 的 frequency 从 10,000 提升到了 1,000,000
使用了 YARN 来扩展上下文长度
使用了 Dual Chunk Attention 来优化 attention 的计算

Post-training

post-training 包括 SFT 和 RLHF 两个阶段

数据包括 SFT 数据和 RLHF 使用的偏好数据

数据标注过程有：

使用 InsTag 对数据进行打标
选取高质量的 instruction
构建了一个 self-evolution 策略，来扩展 instruction 数据
请人类来标注数据

作者还合成了一些数据，合成数据的过程如下：

rejection sampling：对 LLM 进行多次采样，然后保留结论正确的数据作为 SFT 数据，以正确和错误的数据对作为偏好数据
Execution feedback：对于代码任务，使用 Python 来验证答案的正确性
Data Repurposing：对于写作类任务，以文档为输入，让 LLM 生成对应的 instruction
Constitutional Feeback：基于预设的 principle 来生成回答

最终，SFT 数据包括 500, 000 条样本

RLHF 的训练包括 offline stage 和 online stage，offline stage 就是用收集到的偏好数据。在 online stage，作者使用 reward model 来给输出的回答进行打分，然后再使用 DPO 进行训练。

与 Qwen 不同，Qwen2 中作者使用了 Online Merging Optimizer 来解决因为 alignment 导致的性能降低

Conclusion

本文提出了 Qwen2 系列，在 Qwen2 中，首次使用了 GQA 代替 MHA，Qwen2 在上下文上做出了初步探索

References

Qwen2 tech report

Notes on Qwen1.5

Thu, 03 Jul 2025 17:37:39 +0800

Qwen 在 24 年 1 月份发布了 Qwen1.5，包含 0.5B, 1.8B, 4B, 7B, 14B, 32B, 72B, 以及 110B 6 个 size，还有一个 MoE 模型。

介绍

Qwen1.5 的主要特点：

支持 12 中语言
统一支持 32768 tokens 上下文长度。
提供量化版本（Int4、Int8、AWQ、GGUF）以适应低资源环境或部署需求。

训练过程使用了 DPO 以及 PPO 来进行对齐

Qwen1.5-MoE

Qwen1.5-MoE 的激活参数为 2.7B，一共包含 64 个专家，其中激活 4 个专家，共享 4 个专家

相比于 Qwen1.5-7B，去训练的 FLOPS 降低了 75%，inference 的速度提高了 174%

Qwen1.5-MoE 采用了改进的 MoE 架构，主要优化包括：

细粒度专家（Fine-grained experts） ：通过将 FFN 层划分为多个片段，构建更多专家而不增加参数总量。
初始化策略（Upcycling） ：基于 Qwen-1.8B 初始化模型，并引入随机性以加速收敛。
路由机制（Routing Mechanism） ：在每个 MoE 层中使用 64 个专家，其中 4 个共享专家始终激活，60 个路由专家中有 4 个被激活，提高了灵活性和效率。

效率对比

作者对比了 throughput (requests processed per second) 以及 tokens per second (TPS):

Model	Throughput	TPS
Qwen1.5-7B-Chat	1.15	2298.89
Qwen1.5-MoE-A2.7B-Chat	2.01	4010.27

References

Notes on YaRN

Thu, 03 Jul 2025 14:40:49 +0800

Introduction

YaRN (Yet Another RoPE extentionN method) 时23年9月EleutherAI等提出来的一个扩展LLM上下文长度的方法，后来被Qwen系列模型所应用。

Preliminary

作者首先回顾了一下RoPE, 具体内容请参见上一篇blog。并使用了 $f_{W}(x_m, m, \theta_d)$ 来表示RoPE：

$$ f_{W}(x_m, m, \theta_{d}) = \Theta_{\theta, m}^d W x_m $$

其中 $\Theta_{\theta, m}^d\in\mathbb{R}^{d\times d}$ 是多维旋转矩阵, $\theta_d=[\theta_{0,d},\dots,\theta_{(d-2)/2,d}]\in\mathbb{R}^{d/2}$, $\theta_{i,d}=\theta^{2i/d}$, $\theta>0$ 是一个超参数，RoPE中设置为 $\theta=10000$, $W\in\{W_q,W_k\}\subset\mathbb{R}^{d\times d}$ 是对应的query/key projection layer的权重矩阵， $x\in\mathbb{R}^{d}$ 是输入的hidden states.

接下来，作者定义了两个新的变量：

scaling factor 假设预训练的上下文长度为 $L$, 扩展的上下文长度为 $L'>L$, 则我们定义scaling factor 为

$$ s = \frac{L'}{L} $$

易知 $s>1$.

wavelength 我们将 $\lambda_d$ 定义为 $i$-th 维的RoPE embedding对应的 wavelength：

$$ \lambda_{i,d} = \frac{2\pi}{\theta_{i,d}} = 2\pi\theta^{2i/d}, \ i=0,\dots,(d-2)/2 $$

wavelength描述了对于第$i$ 个维度，RoPE旋转一周 ($2\pi$) 所需要的上下文长度。

基于 $f_{W}(x_m, m, \theta_d)$, 作者统一了已有的扩展上下文长度的方法，作者将不同的扩展方法使用一个通用函数 $f_W(x_m,g(m), h(\theta_d))$ 来表示，这里 $g(m)$ 和 $h(\theta_d)$ 分别代表了不同的长度外推方法所使用的变换。

Position Interpolation

Position Interpolation (PI) 的核心思想在于，我们可以通过Interpolation，将超过预训练长度的文本给压缩到当前最大长度，以此来避免RoPE外推产生的问题，其对应的公式为

$$ f'_W(x_m,m,\theta_d) = f_W(x_m, \frac{mL}{L'}, \theta_d) $$

其中 $L'>L$ 为我们扩展之后的上下文长度， $L$ 为我们预训练的上下文长度。使用通用函数表示的话，我们有

$$ g(m) = \frac{m}{s},\quad h(\theta_d) = \theta_d $$

NTK-aware Interpolation

PI的问题是，并没有考虑不同维度的wavelength。基于NTK理论，DNN当输入维度比较低，且embedding缺少高频内容时，模型就会很难学习到高频信息。对应到RoPE里面，输入的token position id是低位信息（1维），而输出的RoPE是一个 $d$ 维的复杂向量。因此，当输入token非常相似却距离非常近时，RoPE就会丢失高频的细节信息

因此，为了解决这个问题，作者对不同的维度使用了不同的缩放策略：维度比较小时，其缩放的更多，维度比较大是，其缩放的更少。

基于这个策略，作者提出了NTK-aware interpolation，其定义如下：

$$ g(m) = m, \quad h(\theta_{i,d}) = \theta'^{-2i/d}, i=0,\dots,(d-2)/2 $$

其中

$$ \theta' = \theta\cdot s^{\frac{d}{d-2}} $$

实际中，这种方法会产生out-of-bound的值，因此最终结果会比PI要差一点，为了解决这个问题，一般会使用比 $s$ 更大的scaling factor.

上式的推导基于一个简单的假设：我们希望最后一个维度的wavelength在scaling之后，是线性变化的，即 $\theta'_{(d-2)/2,d}=s\theta_{(d-2)/2,d}$, 求解之后，我们就得到了上面的定义

PI 和NTK-aware interpolation的问题在于，我们对不同的维度的处理都是一样的。这类不在乎wavelength的方法被称为blind interpolation methods, 接下来要介绍的就是基于wavelength的方法，即target interpolation methods.

NTK-by-parts Interpolation

与NTK-aware interpolation， NTK-by-parts interpolation基于wavelength来考虑不同维度上所做的变换。

对于低维度，其 $\theta_{i,d}$ 非常大，因此旋转一周所需要的上下文长度也非常大。实际上就导致某些维度的embedding并不是均匀分布的，（比如说只有 $0\sim\pi$ 这个区间的embedding），这个时候，模型就只能访问到绝对位置信息，而访问不到相对位置信息。另外，当我们对所有的维度都进行scale的时候，所有的token都会与彼此更加靠近，这损害了模型对于局部信息的获取能力。基于这些认知，作者基于wavelength，对不同的维度分别进行处理：

如果wavelength远小于上下文长度 $L$，则我们不做任何处理
如果wavelength等于或者大于上下文长度 $L$，则我们使用NTK-aware interpolation 进行处理
对于中间的其他维度，我们进行了一个trade off

作者定义了一个ratio $r$ 来描述原始上下文长度 $L$ 和 wavelength 之间的关系

$$ r(i,d) = \frac{L}{\lambda_{i,d}} = \frac{L}{2\pi\theta'^{2i/d}} $$

基于这个ratio，我们可以定义上面的三种处理方式对应的权重

$$ \gamma(r) = \begin{cases} 0, &\text{if } r < \alpha\\ 1, &\text{if } r > \beta\\ \frac{r-\alpha}{\beta-\alpha}, &\text{otherwise} \end{cases} $$

其中， $\beta>\alpha>0$ 是超参数， $r<\alpha$, $r<\beta$ 分别代表了上面的第1种，第2种情况。

最后，NTK-by-parts interpolation的定义如下

$$ g(m) = m, \quad h(\theta_i) = \left(1-\gamma(r(i,d)\right)\frac{\theta_i}{s} + \gamma(r(i,d))\theta_i $$

作者通过实验发现，在LLaMA上，$\alpha=1$ 和 $\beta=32$ 是一个比较好的选择

Dynamic NTK Interpolation

在实际中，一个经常遇到的场景就是sequence length会从1逐步上升到最大上下文长度，比如说inference的时候。对于这种情况，我们有两种解决方法：

在整个inference周期中，RoPE的scaling factor都设置为 $s=K'/L$, 其中$L'$ 是扩展后的上下文长度
在每次foward的过程汇总，都更新sclaing factor $s=\max(1, \ell'/L)$, 这里 $\ell'$ 是当前sequence的长度

作者发现，方案1在sequence长度小于 $L$的时候性能会下降，并且当上下文长度超过 $L'$ 时，性能下降的更快。但是，方案2可以让模型的性能下降曲线更平缓。因此，作者将这种inference-time方法称为 Dynamic Scaling method, 当其与NTK-aware方法结合时，就得到了 Dynamic NTK interpolation

作者通过实验发现，Dynamic NTK interpolation在$L'=L$ 时，效果非常好

YaRN

在YaRN中，作者针对Dynamic NTK interpolation做了进一步改进，也就是在计算attention softmax时，加入了一个温度参数 $t>0$, 这样attention的计算就变成了

$$ \mathrm{Attn}(Q, K, V) = \mathrm{softmax}\left(\frac{QK^T}{t\sqrt{d}}\right)V $$

作者发现，通过这种scaling的方式，YaRN可以在不改变代码的前提下，更改attention的机制。并且，其不增加训练和推理的cost

作者将YaRN定义为结合了NTK-by-parts interpolation和上述scaling技巧的方法

对于LLaMA，作者推荐使用如下参数：

$$ \sqrt{\frac{1}{t}} = 0.1\ln(s) + 1. $$

实验结果如下

作者发现：

对于合适的 $t$, 扩展上下文之后，perplexity会变的更好
最好的$t$ 对于不同的位置和样本提升都是一样的

Evaluation

实验结果如下

Extension Method	Trained Tokens	Context Window			Evaluation Context Window Size
			2048	4096	6144	8192	10240
PI (s = 2)	1B	8k	3.92	3.51	3.51	3.34	8.07
NTK ($\theta$ = 20k)	1B	8k	4.20	3.75	3.74	3.59	6.24
YaRN (s = 2)	400M	8k	3.91	3.50	3.51	3.35	6.04

可以看到，YaRN使用的数据更少，并且当模型扩展到10240的时候，其表现下降的最慢，这说明了YaRN在扩展上下文长度时的有效性

原始RoPE，dynamic-PI和dynamic-YaRN的对比

可以看到，RoPE的上下文扩展能力很差，Dynamic-YaRN的表现最好。

Code

YaRN的实现在HuggingFace/src/transformers/modeling_rope_utils.py 里的 _compute_yarn_parameters 函数里，其返回 inv_freq 以及 attention_factor 两个量，前者代表了 $\theta_d$, 后者代表 $t\sqrt{d}$.

 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



def get_mscale(scale):
 if scale <= 1:
 return 1.0
 return 0.1 * math.log(scale) + 1.0

def find_correction_dim(num_rotations, dim, base, max_position_embeddings):
 """Inverse dimension formula to find the dimension based on the number of rotations"""
 return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / (2 * math.log(base))

def find_correction_range(low_rot, high_rot, dim, base, max_position_embeddings):
 """Find dimension range bounds based on rotations"""
 low = math.floor(find_correction_dim(low_rot, dim, base, max_position_embeddings))
 high = math.ceil(find_correction_dim(high_rot, dim, base, max_position_embeddings))
 return max(low, 0), min(high, dim - 1)

def linear_ramp_factor(min, max, dim):
 if min == max:
 max += 0.001 # Prevent singularity

 linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min)
 ramp_func = torch.clamp(linear_func, 0, 1)
 return ramp_func


class YaRNRotaryEmbedding(nn.Module):
 def __init__(self, config):
 beta_fast = config.get("beta_fast") or 32
 beta_slow = config.get("beta_slow") or 1
 dim = config["head_dim"]
 factor = config.max_position_embeddings / original_max_position_embeddings
 pos_freqs = base ** (torch.arange(0, dim, 2) / dim)
 inv_freq_extrapolation = 1.0 / pos_freqs
 inv_freq_interpolation = 1.0 / (factor * pos_freqs)

 low, high = find_correction_range(beta_fast, beta_slow, dim, base, original_max_position_embeddings)

 # Get n-dimensional rotational scaling corrected for extrapolation
 inv_freq_extrapolation_factor = 1 - linear_ramp_factor(low, high, dim // 2)
 inv_freq = (inv_freq_interpolation * (1 - inv_freq_extrapolation_factor) + inv_freq_extrapolation * inv_freq_extrapolation_factor)

 attention_factor = get_mscale(factor)

 def forward(self, x, position_ids):
 ...
 return cos, sin

实际的transformers代码中，Qwen使用的还是默认的RoPE，在inference时如果我们需要扩展上下文，可以通过修改config的形式：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18


from transformers import pipeline

model_name_or_path = "Qwen/Qwen3-8B"

generator = pipeline(
 "text-generation",
 model_name_or_path,
 torch_dtype="auto",
 device_map="auto",
 model_kwargs={
 "max_position_embeddings": 131072,
 "rope_scaling": {
 "rope_type": "yarn",
 "factor": 4.0,
 "original_max_position_embeddings": 32768,
 },
 }
)

Conclusion

在本文中，作者首先构建了一个统一的表征不同上下文长度扩展的形式，接下来作者分析了不同上下文长度扩展的不足，并提出了YaRN这种上下文长度扩展方式，结果发现，YaRN不仅在短上下文长度下面表现很好，当上下文长度扩展之后，其表现依然非常优秀。

References

Notes on Qwen-LLM

Thu, 03 Jul 2025 10:47:27 +0800

Qwen 在 23 年 9 月份发布了 Qwen 系列大语言模型，包括 1.8B， 7B，14B 三个 size，训练过程使用了 3T token. 作者还基于 Qwen，构建了 Code-Qwen-Chat，Math-Qwen-Chat 等系列领域大语言模型。

Pre-training

Data

数据一共使用了 3T token，主要是 public web documents, encyclopedia, books, codes, etc，覆盖了中文和英文两种语言

数据处理：

语言识别
去重，包括 MinHash 和 LSH 算法
质量过滤，包括基于规则和和基于 ML 的方法
上采样，特定数据会进行上采样
加入指令数据，提高模型的 zero-shot 和 few-shot 表现

Tokenization

BPE tokenizer，最终的 tokenizer 大小为 152K

Architecture

模型架构基于 LLaMA，改动：

tie embdding: input embdding 和 output embdding 使用的权重相同
position encoding:RoPE, inverse frequency 的精度为 FP32
bias: 取消了大部分的 bias，增加了 QKV bias，来提高模型的外推能力
Pre-Norm & RMSNorm
Activation function: SwiGLU

Training

上下文长度：2048
attention：flash attention
optimizer：AdamW， $\beta_1=0.9$, $\beta_2=0.95$, $\epsilon=10^{-8}$.
data type: BF16

Context Extention

使用了三个技巧：

NTK-aware position interpolation
log-N scaling
window attention

后续前两个统一成了 YARN.

observation: lower layer 对上下文长度扩展更敏感, 因此作者动态调整了 window size

Post-training

包括 SFT 和 RLHF 两个阶段

SFT

data：使用了 ChatML 格式

RLHF

PPO 算法

reward model 构建：基于 Qwen-base model

RL 训练：先更新 value model 50 steps

发现：top-p 设置为 0.9 比设置为 1.0 更好

Tool-use and Agent

作者使用了 self-instruct 来进行 SFT，基于 ReAct 构建数据，数据包括 2000 条高质量数据

Specialization

Code-Qwen

code-qwen 基于 qwen continue Pretraining 得到，然后基于 code-qwen 进行 sft 得到 code-qwen-chat，包括 7B 和 14B 两个 size

Math-Qwen

基于 qwen 直接 SFT 得到，包括 7B 和 14B 两个 size

Conclusion

作者在本文中介绍了 Qwen 系列大语言模型，模型使用了 3T token，作者介绍了训练的细节以及如何扩展到领域大语言模型 Code-Qwen 和 Math-Qwen

References

Hands on LLM(2) Transformer

Sun, 29 Jun 2025 11:40:39 +0800

Transformer 实现

我们采用 top-down 的形式构建 transformer 的代码

架构

我们以 Qwen3 的代码为例子讲解 Assignment1 的代码实现

我们通过在 transformer 架构上加上一个 linear layer 就可以完成不同的下游任务，比如：

Qwen3ForQuestionAnswering
Qwen3ForCausalLM
Qwen3ForSequenceClassification

因此，大语言模型是 transformer 的一个附加产物

CausalLM

编写大语言模型的第一步为定义 Qwen3ForCausalLM

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


class CausalLM(nn.Module):
 def __init__(self, config):
 self.model = Transformer(config)
 self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

 def forward(self, ...):
 outputs = self.model(***)
 logits = self.lm_head(outputs)

 return logits

这里 lm_head 的作用就是构建 embedding space 到 vocabulary 的映射，即 $\mathbb{R}^d\to\mathbb{R}^{|V|}$

Transformer

transformer 部分包括四个部分：

Embedding Layer：将 token 映射到 embedding space
layers：Transformer 的主体部分，由 $n$ 个 DecodeLayer 组成
Norm：在输出之前，进行一次 Normalization
Position Embedding：由于输入的 sequence 长度是固定的，因此我们提前计算好每一层的 position embedding

Transformer 部分的代码

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22


class Transformer(nn.Module):
 def __init__(self, config):
 self.embedding = nn.Embedding(config.vocab_size, config.hidden_size, config.pad_token_id)
 self.layers = nn.ModuleList(
 [DecodeLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
 )
 self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 self.rotary_emb = RotaryEmbedding(config)

 def forward(self, input_ids,...):
 input_embeds = self.embedding(input_ids)

 hidden_states = input_embeds
 position_embeddings = self.rotary_emb(hidden_states, position_ids)

 for decode_layer in self.layers:
 layer_outputs = decode_layer(hidden_states, position_ids, position_embeddings)
 hidden_states = layer_outputs[0]

 hidden_states = self.norm(hidden_states)

 return logits

DecodeLayer

DecodeLayer 就是 transformer 的核心部分，里面包含四个模块：

Pre-Normalization：一般是 RMSNorm 或者 LayerNorm
Attention：self-attention
Post-Normalization：与 Pre-Normalization 一致
MLP：FFN，SwiGLU 或者 MoE

DecodeLayer 还会使用 residual connection 来防止梯度消失

DecodeLayer 部分的代码

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21


class DecodeLayer(nn.Module):
 def __init__(self, config, layer_idx):
 self.attn = Attention(config, layer_idx)
 self.mlp = MLP(config)
 self.pre_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 self.post_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

 def forward(self, hidden_states, position_ids, position_embeddings):
 residual = hidden_states
 hidden_states = self.pre_norm(hidden_states)
 hidden_states = self.attn(hidden_states, position_ids, position_embeddings)
 # residual
 hidden_states = hidden_states + residual

 residual = hidden_states
 hidden_states = self.post_norm(hidden_states)
 hidden_states = self.mlp(hidden_states)
 # residual
 hidden_states = hidden_states + residual

 return hidden_states

我们接下来按照

Normalization
MLP
Attention
Position embedding

的顺序来介绍

RMSNorm

RMSNorm 的作用和 LayerNorm 是一样的，但是实现上更简单

$$ \mathrm{LayerNorm}(x) = \frac{x-\mathbb{E}[x]}{\sqrt{\mathrm{var}[x]+\epsilon}}\odot \beta + \gamma $$

其中 $\beta,\gamma\in\mathbb{R}^d$ 是可学习的参数

$$ \mathrm{RMSNorm}(x) = \frac{x}{\sqrt{\|x\|_2^2+\epsilon}}\odot \gamma $$

其中 $\gamma\in\mathbb{R}^d$ 是可学习的参数

RMSNorm 代码实现

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11


class RMSNorm(nn.Module):
 def __init__(self, d, eps):
 self.weight = nn.Parameter(torch.ones(d))
 self.eps = eps

 def forward(self, x):
 input_dtype = x.dtype
 x = x.to(torch.float32)
 variance = x.pow(2).mean(-1, keepdim=True)
 x = x * torch.rsqrt(variance + self.eps)
 return self.weight * x.to(input_dtype)

MLP

现在大语言模型的 MLP 使用的激活函数一般都是 SwiGLU, 其定义为

$$ \mathrm{SwiGLU}(x) = x\odot \sigma(x) $$

其中 $\sigma(\cdot)$ 是 sigmoid 函数

MLP 的定义为

$$ y = W_2(W_3x\odot \mathrm{SwiGLU}(W_1x)) $$

其中 $W_3,W_1\in\mathbb{R}^{d_{ff}\times d}$, $W_2\in\mathbb{R}^{d\times d_{ff}}$

一般地，由于 FFN 只有两个权重矩阵，且 $d_{ff}=4d$, 在 SwiGLU 中，为了保证参数量一致，其隐藏层大小设置为 $d_{ff}'=\frac23d_{ff}=\frac83 d$.

MLP 的代码如下所示

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11


def SwiGLU(x):
 return x * torch.sigmoid(x)

class MLP(nn.Module):
 def __init__(self, d, d_ff):
 self.gate_proj = nn.Linear(d, d_ff, bias=False)
 self.up_proj = nn.Linear(d, d_ff, bias=False)
 self.down_proj = nn.Linear(d_ff, d, bias=False)

 def forward(self, x):
 return self.down_proj(SwiGLU(self.gate_proj(x)) * self.up_proj(x))

Attention

我们先不考虑 position embedding，直接看 attention，attention 定义为

$$ \mathrm{Attention}(X) = \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d}}\right)V\in\mathbb{R}^{m\times d} $$

其中 $X\in\mathbb{R}^{m\times d}$,

$$ Q = W_QX\in\mathbb{R}^{m\times d},\quad K =W_KX\in\mathbb{R}^{n\times d},\quad V = W_VX\in\mathbb{R}^{n\times d} $$

在自回归模型里，我们还会加上 mask, 让每个 token 只能看见前面的 token 的信息

$$ \mathrm{Attention}(X) = \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d}}\odot M\right)V $$

其中

$$ M = [M_{ij}] = \begin{cases} 1, &\text{ if } i < j\\ 0, &\text{ otherwise} \end{cases} $$

self-attention 的代码如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


def scaled_dot_product_attention(Q, K, V, mask) -> torch.Tensor:
 d_k = Q.shape[-1] # d_k
 scaled_factor = 1 / d_k**0.5

 scores = torch.einsum("... s_q d_k, ... s_k d_k -> ... s_q s_k", Q, K)
 scores *= scaled_factor

 if mask is not None:
 scores = scores.masked_fill(mask == 0, float("-inf"))

 scores = scores.softmax(dim=-1)

 return torch.einsum("... s_q s_k, ... s_k d_v -> ... s_q d_v", scores, V)

Multi-Head Attention

Multi-Head Attention 定义如下

$$ \mathrm{MultiHeadAttention}(X) = [\mathrm{Attention}_1(X),\dots,\mathrm{Attention}_h(X)]W_o\in\mathbb{R}^{m\times d} $$

其中 $W_o\in\mathbb{R}^{d\times d}$, 且每一个 Attention heads 的维度会从 $d\to d/h$.

Multi-Head Attention 的主要作用为：

让不同的 head 关注不同的信息
并行计算，提高计算效率

MHA 代码

 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


class MultiHeadAttention(nn.Module):
 def __init__(self, d_model: int, num_heads: int) -> None:
 self.q_proj = Linear(d_model, d_model)
 self.k_proj = Linear(d_model, d_model)
 self.v_proj = Linear(d_model, d_model)
 self.output_proj = Linear(d_model, d_model,)

 def forward(self, x, position_embeddings, mask):
 Q = rearrange(self.q_proj(x), "... seq_len (num_heads head_dim) -> ... num_heads seq_len head_dim",
 num_heads=self.num_heads, head_dim=self.head_dim)
 K = rearrange( self.k_proj(x), "... seq_len (num_heads head_dim) -> ... num_heads seq_len head_dim",
 num_heads=self.num_heads, head_dim=self.head_dim)

 if mask is None:
 mask = torch.ones(Q.shape[-2], K.shape[-2])
 mask = torch.tril(mask)

 if position_embeddings is not None:
 sin, cos = position_embeddings
 Q = apply_rotary_pos_emb(Q, sin, cos)
 K = apply_rotary_pos_emb(K, sin, cos)

 V = rearrange(self.v_proj(x), "... seq_len (num_heads head_dim) -> ... num_heads seq_len head_dim",
 num_heads=self.num_heads, head_dim=self.head_dim)

 output = scaled_dot_product_attention(Q, K, V, mask=mask)
 output = rearrange(output, "... num_heads seq_len head_dim -> ... seq_len (num_heads head_dim)")
 return self.output_proj(output)

Position Encoding

Attention 对于输入的顺序是不敏感的，也就是

$$ \mathrm{Attention}(Q, \Pi K, \Pi V) = \mathrm{Attention}(Q, K, V) $$

这里 $\Pi\in \{0,1\}^{d\times d}$ 是一个置换矩阵 (permutation matrix)

Transformer 的解决方法是在 query 和 key 上加上位置信息：

$$ Q' = Q + PE(Q),\ K'=K + PE(K) $$

这样

$$ \mathrm{Attention}(Q, K, V) = \mathrm{softmax}\left(\frac{(Q + PE(Q))(K + PE(K))^T}{\sqrt{d}}\right)V\in\mathbb{R}^{m\times d} $$

就包含了位置信息

绝对位置编码

Transformer 的使用的位置编码如下所示

$$ PE(pos, 2i) = \sin\left(\frac{pos}{10000^{2i/d}}\right) $$$$ PE(pos, 2i+1) = \cos\left(\frac{pos}{10000^{2i/d}}\right) $$

RoPE

苏剑林老师提出了 Position Encoding，现在已经被广泛使用

$$ q' = R_{\theta,m}^dq, k' = R_{\theta,n}^d k $$

这样 $\langle q, k\rangle$ 就仅包含两者的相对位置信息

$$ \langle q_m, k_n\rangle = x^TW_qR_{\theta, n-m}^d W_kx_n $$

RoPE 的矩阵定义如下

$$ R_{\theta,m}^d = \mathrm{diag}(M_1,\dots,M_{d/2}) $$

其中

$$ M_i = \begin{bmatrix} \cos m\theta_i & -\sin m\theta_i\\ \sin m\theta_i & \cos m\theta_i \end{bmatrix} $$

这里

$$ \theta_i = \frac{1}{10000^{2(i-1)/d}}, i\in\{1,2,\dots,d/2\} $$

简化后得到

$$ R_{\theta,m}^dq = \begin{bmatrix} \cos m\theta_0\\ \cos m\theta_0\\ \vdots\\ \cos m\theta_{d/2}\\ \cos m\theta_{d/2}\\ \end{bmatrix}\odot \begin{bmatrix} x1\\ x2\\ \vdots\\ x_{d-1}\\ x_d\\ \end{bmatrix} + \begin{bmatrix} \sin m\theta_0\\ \sin m\theta_0\\ \vdots\\ \sin m\theta_{d/2}\\ \sin m\theta_{d/2}\\ \end{bmatrix}\odot \begin{bmatrix} -\ x_2\\ x_1\\ \vdots\\ -x_d\\ x_{d-1}\\ \end{bmatrix} $$

RoPE 代码 Naive 实现

RotaryEmbedding 代码

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12


class RotaryEmbedding(nn.Module):
 def __init__(self, dim):
 self.inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))

 def forward(self, x, position_ids):
 freqs = einsum(torch.arange(self.max_seq_len),
 self.inv_freq, "seq_len, d_k_half -> seq_len d_k_half")[token_positions]

 sin = torch.sin(freqs)
 cos = torch.cos(freqs)

 return sin, cos

计算部分代码

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


def apply_rotary_pos_emb(x: torch.Tensor, sin: torch.Tensor, cos: torch.Tensor) -> torch.Tensor:
 x_even = x[..., ::2] # (seq_len, d_k_half)
 x_odd = x[..., 1::2] # (seq_len, d_k_half)
 odds = cos * x_even - sin * x_odd # (...,seq_len, d_k_half)
 evens = sin * x_even + cos * x_odd # (...,seq_len, d_k_half)
 stacked = torch.stack((odds, evens), -2) # (...,seq_len, 2, d_k_half)
 stacked_trans = rearrange(
 stacked, "... seq_len double d_k_half -> ... seq_len d_k_half double"
 ) # (...,seq_len, d_k_half, 2)
 out = rearrange(
 stacked_trans, "... seq_len d_k_half double -> ... seq_len (d_k_half double)"
 ) # (..., seq_len, d_k)
 return out

RoPE 标准实现

RotaryEmbedding 代码 (LLaMA)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16


class LlamaRotaryEmbedding(nn.Module):
 def __init__(self, config: LlamaConfig, device=None):
 inv_freq = inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64) / dim)) # d_k_half

 def forward(self, x, position_ids):
 # (bsz, d_k_half, 1)
 inv_freq_expanded = self.inv_freq[None, :, None].expand(position_ids.shape[0], -1, 1)
 # (bsz, 1, seq_len)
 position_ids_expanded = position_ids[:, None, :]
 # (bsz, seq_len, d_k_half)
 freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2)
 emb = torch.cat((freqs, freqs), dim=-1) # (..., seq_len, d_k)
 cos = emb.cos() # (..., seq_len, d_k)
 sin = emb.sin() # (..., seq_len, d_k)

 return cos, sin

计算部分代码

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


def rotate_half(x):
 x1 = x[..., : x.shape[-1] // 2]
 x2 = x[..., x.shape[-1] // 2 :]
 return torch.cat((-x2, x1), dim=-1)



def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
 cos = cos.unsqueeze(unsqueeze_dim)
 sin = sin.unsqueeze(unsqueeze_dim)
 q_embed = (q * cos) + (rotate_half(q) * sin)
 k_embed = (k * cos) + (rotate_half(k) * sin)
 return q_embed, k_embed

References

Unified perspective on dLLM and LLM

Sat, 28 Jun 2025 15:02:09 +0800

Introduction

在上一篇blog里，我们介绍了MLE和KL minimization的等价性。在这篇blog里，我们将要基于这个等价性，推导masked diffusion LLM, autoregressive LLM, any-order diffusion LLM之间的等价性。最终，我们发现，这几种建模方式本质上都是一致的。

Preliminary

对于 $\bm{x}=(x_1,\dots,x_L)\in\mathbb{R}^L$, 基于概率的链式法则，我们有

$$ p(\bm{x}) = \prod_{i=1}^Lp(x_i\mid x_{1D的类别分布，对于随机变量 $X\in\{1,\dots, K\}$以及概率分布 $\bm{p}=(p_1,\dots,p_K)$ 我们定义类别分布为 $\mathrm{Cat}(\bm{x};\bm{p})$, 其中 $\bm{x}\in\mathbb{R}^K$是一个向量，其PMF定义为：

$$ f(\bm{x}; \bm{p}) = \prod_{i=1}^Kp_i^{[x=i]} $$

这里 $[x=i]$ 当 $x=i$时为$1$，否则为 $0$.

一个比较好理解的例子是骰子，骰子有$K=6$种可能性，每一面出现的可能性都是 $p_i=1/6$. 当 $\bm{x}$是one-hot向量是，其代表了出现某一面的概率。

Discrete-Time Discrete Markov Chain. 接下来我们考虑离散时间离散Markov链，我们现在的随机变量不仅与状态 $\{1,\dots,K\}$ 还与时间有关，我们记 $X_n$ 为 $n$ 时刻的状态分布。这样，我们可以规定不同时刻之间的状态转换矩阵：

$$ Q_{ij} = \mathrm{Pr}(x_{n+1}=j\mid x_n=i) $$

Categorical distribution

autoregressive LLM

首先是我们最常见的自回归大语言模型，给定文本语料 $\bm{x}=(x_1,\dots,x_L)\sim p_{data}$，大语言模型的目标在于求解最大似然估计

$$ \begin{aligned} \theta^* &= \arg\max_{\theta}\mathbb{E}_{\bm{x}\sim p_{data}}[\log p(\bm{x}\mid \theta))]\\ &= \arg\max_{\theta}\mathbb{E}_{\bm{x}\sim p_{data}}\left[\log \left(\prod_{i=1}^Lp(x_i\mid x_{这里我们使用了链式法则.

Any order autoregressive LLM

接下来，我们证明Any order autoregressive LLM的目标函数等价于求解MLE。令 $S_D$ 为在集合 $\{1,\dots, D\}$ 上所有可能的permutation，令 $\sigma\in S_D$ 为一个其中的一个permutation，则我们有

$$ p(\sigma) = \frac{1}{D!} $$

因此

$$ p(x) = \sum_{\sigma\in S_d}p(x,\sigma) = \sum_{\sigma\in S_D}p(x\mid \sigma)p(\sigma) = \frac{1}{D!}\sum_{\sigma\in S_D}p(x\mid \sigma)=\mathbb{E}_{\sigma\sim U(S_D)}[p(x\mid \sigma)] $$

这里 $U(\cdot)$ 代表均匀分布，从而我们有

$$ \log p(x) = \log \mathbb{E}_{\sigma\sim U(S_D)}[p(x\mid \sigma)] \geq \mathbb{E}_{\sigma\sim U(S_D)}[\log p(x\mid \sigma)] $$

这里我们使用了Jensen不等式

那么对于any-order autoregressive model, 我们有：

$$ \log p(x) \geq \mathbb{E}_{\sigma\sim U(S_D)}\sum_{t=1}^D \log p(x_{\sigma(t)}\mid x_{\sigma(也就是说，对于any-order autoregressive model，我们的目标函数是MLE的一个下界。

Masked Language modeling

Discrete Diffusion Model

Any-order diffusion LLM

Masked diffusion LLM

Reference

Structured Denoising Diffusion Models in Discrete State-Spaces

Hands on LLM(1) Tokenizer

Sat, 24 May 2025 19:56:34 +0800

Introduction

在自然语言处理中，tokenizer的作用是将一个文本序列通过一个字典转化为一个token id的序列。我们回顾图片分类任务，模型在预测的时候，实际上预测的是类别对应的id，而不是类别本身。tokenizer做的事情就是提供一个类似于从类别到对应id的字典。

一般来说，一个tokenizer处理文本序列的过程有两步：

pre-tokenize，也就是预处理，我们需要将文本序列分割成合适大小的chunks (words)
tokenize，构建chunks (words)到token id的映射

实际上, huggingface的tokenizer包括四个步骤, 其中第二第三个步骤与上述一致. 在pre-tokenize之前, 我们有一个normalization过程, 该过程会对文本序列进行处理, 如将文本序列变为小写, 删掉声调符号等, 如下面例子所示：

1
2
3
4
5
6


from tokenizers import normalizers
from tokenizers.normalizers import NFD, StripAccents

normalizer = normalizers.Sequence([NFD(), StripAccents()])
normalizer.normalize_str("Héllò hôw are ü?")
# "Hello how are u?"

在tokenize之后, 我们会有一个post-processing过程, 比如BERT会在生成的token系列前后加入 [CLS] token 和 [SEP] token, 例子如下:

1
2
3
4
5
6
7


from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
token_ids = tokenizer.encode("I love NLP.")
print(token_ids)
# [101, 146, 1567, 21239, 2101, 119, 102]
# represents [[CLS], "I", "love", "NL", "##P", ".", [SEP]]

其完整流程如下图所示 (图源: huggingface llm-course)

构建好tokenizer之后, 我们还要保证tokenizer提供两个接口：

encoding, 给定文本序列, 将其映射到字典中去得到token id序列
decoding, 给定token id序列, 将其解码成文本序列

接下来, 我们将简单介绍一下word tokenizer, character tokenizer以及byte tokenizer, 并分析它们各自的不足。然后, 我们介绍现代大语言模型中使用最多的BPE tokenizer。最后, 我们介绍一些sub-word tokenizer。

Training-free tokenizer

本节我们将要介绍word tokenizer, character tokenizer以及byte tokenizer, 它们的特点就是简单易懂, 不需要额外的规则和学习.

Word tokenizer

给定一个文本序列, 我们现在需要将其转化为一个token序列。一个比较自然的想法是, 我们按照空格将序列拆分成若干个单词, 这样每个单词的语义都能比较好地保留。下面是一个例子：

1
2
3
4
5
6


import tiktoken

tokenizer = tiktoken.get_encoding("gpt2")
indices = tokenizer.encode("hello world")
# indices = [31373, 995]
# decode = ["hello", " world"]

接下来我们基于一个预定义好的词典, 将其转化为一个token id的序列。

word tokenizer的优点是能够保留语义信息，且压缩率比较高（每个token包含的bytes数），其问题是不能处理预定义好的词典之外的词 (out of vocabulary, OOV)。现有的处理方法是使用 <UNK> token来表示这些OOV的词。但这样显然会丢失语义信息, 因为我们编码成 <UNK> token之后, 就没办法再解码回原有的语义信息了。

word tokenizer的缺点为：

单词数量很大, 很多罕见单词的出现频率很低, 降低了tokenizer的利用率
对于不在词典内的单词只能用<UNK> token表示, 损害了语义信息

既然基于word的tokenizer有OOV的问题，我们能否想办法解决这个问题呢？答案是可以的, 我们可以使用 character tokenizer。

Character tokenizer

Character tokenizer的基本思想是使用字符而不是单词来编码文本序列。其实现方式如下：

1
2
3
4
5
6


class CharacterTokenizer:
 def encode(self, s: str) -> list[int]:
 return list(ord(c) for c in s)

 def decode(self, token_ids: list[int]) -> str:
 return "".join(chr(token_id) for token_id in token_ids)

character tokenizer的词表大小取决于我们的编码方式，UTF-8的编码大概有110K code points。character tokenizer的缺点总结如下：

character tokenizer会导致我们的词表非常大
和word tokenizer一样, 很多character非常罕见, 会降低词表的利用率
token序列的上下文语义信息较差

Byte tokenizer

我们发现, character tokenizer和word tokenizer的词表都很大, 我们能否想办法降低词表大小, 提升每个token的利用率呢？答案是使用Byte tokenizer.

Byte tokenizer的基本思想是, 所有的字符(character)都是由byte组成的, 比如对于UTF-8编码来说, 每个字符由1-4个byte组成。因此, 所有满足UTF-8编码的文本, 我们都可以将它们转换为基于byte的token序列。由于现在的大部分文本都是基于UTF-8的, 因此, 我们只讨论UTF-8编码的文本。

Byte tokenizer的实现如下：

1
2
3
4
5
6


class ByteTokenizer:
 def encode(self, s: str) -> list[int]:
 return list(s.encode("utf-8"))

 def decode(self, token_ids: list[int]) -> str:
 return bytes(token_ids).decode("utf-8")

byte tokenizer的词表很小, 其词表大小为 256, 这是因为一个byte可以有256中可能的值.

尽管byte tokenizer实现简单，并且词表也很小，可以说byte tokenizer解决了character tokenizer和word tokenizer的问题。但是，byte tokenizer的问题在于，其encode的到的token序列可能会非常长！我们知道，transformer计算量与token序列的长度是平方级关系的，也就是说token序列长度增加10倍，整体的计算量就会增加100倍，因此我们势必需要考虑token序列的长度。

总之，byte tokenizer的问题为：

产生的token序列过长, 增加了transformer的计算量
没有上下文语义信息

总结

我们总结一下word tokenizer, character tokenizer以及byte tokenizer三者各自的特点:

Feature	Word Tokenizer	Character Tokenizer	Byte Tokenizer
Granularity	Coarse	Medium	Fine
Vocabulary	Yes	No	No
Support OOV	Bad	Good	Best
#Tokens	Small	Large	Very Large
Chinese	Yes	Yes	Yes
Support Spell Error	Bad	Yes	Yes
Context	Good	Bad	Worst

因此, 这三种tokenizer尽管实现起来很简单, 但是其都有各自的问题. 为了解决这些问题, 我们的做法就是折衷, 使用sub-word tokenizer, 也就是介于word tokenizer和byte tokenizer之间的方法.

BPE

基本原理与实现

实际生活中，对于出现频率比较高的词，我们会有一个简写的方式，也就是我们使用一个新的单词来表示这个词。比如在英语中，我们会使用plz 来代替 please 以及使用how r u 来代替how are you。 BPE，即byte pair tokenizer的原理也是类似的，对于出现频率比较高的byte pair或者character pair, 我们会使用一个新的token来表示这个pair，这样就压缩了sequence的长度。

BPE算法包括以下几个步骤:

对文本序列进行pre-tokenize, 分割成不同的单词
当len(vocab)<vocab_size时, 重复以下步骤:
1. 对所有单词, 统计其相邻character或者byte pair的频率
2. 计算出现频率最高的pair, 使用一个新的token来表示这个pair
3. 将新的token和其对应的token_id加入到vocab中, 并更新单词的分割表示

算法如下图所示 (参考文献2)

注意：实际上，我们实现的是BBPE (byte BPE算法)，BBPE与BPE的区别在于我们的最小单元是character还是bytes。本质上原理是一致的

实现代码见 Github naive BPE

高效实现

BPE的原理很简单, 我们也实现了其naive版本, 但是naive版本的问题是太慢了。因此我们将要优化naive版本的效率。

首先我们发现, 我们不需要遍历所有的word, 只有含有best_pair的word我们才会进行处理, 因此, 我们的第一个改进就是使用 pair_to_word 来记录每个pair的来源, 比如：

1
2
3
4


pair_to_word = {
 (b' ', b't'): [b' the', b' it'],
 (b't', b'h'): [b'the']
}

这样, 我们在merge的时候, 直接使用 pair_to_word[best_pair] 来获取需要被更新的token序列就可以了。

其次, 注意到每次merge之后, 我们都需要重新计算一次 pair_freq, 而实际上, 只有被merge的token序列才需要被重新计数, 其他大部分token序列都是不需要重新计数的。因此, 一个改进点就是我们在merge的过程中就更新 pair_freq, 而不是重新计算。为了达到这个目标, 我们其实只需要两个操作。我们用(b'x', b'a', b'b', b'y') 和 best_pair=(b'a', b'b')来说明, merge之前, 这个序列贡献的pair_freq为：

1
2
3
4
5


{
 (b'x', b'a'): 1,
 (b'a', b'b'): 1,
 (b'b', b'y'): 1,
}

merge之后, token序列变成了(b'x', b'z', b'y') (假设best_pair对应的新的token为b'z'), 这时候的计数为:

1
2
3
4
5
6
7


{
 (b'x', b'a'): 0,
 (b'a', b'b'): 0,
 (b'b', b'y'): 0,
 (b'x', b'z'): 1,
 (b'z', b'y'): 1,
}

也就是说, merge之后, 三个pair的计数减少了1, 分别是(token_seq[i-1], merge_pair[0]),merge_pair 和 (merge_pair[1], token_seq[i+2])。两个pair的个数增加了1, 分别是 (token_seq[i-1], new_token)和(new_token, token_seq[i+2]) (这里我们假设merge_pair=(token_seq[i], token_seq[i+1]))。

基于这个结论，我们就可以优化BPE算法了，具体逻辑就是：

pretokenize, 将 text 切分为若干个 word
计算word_count, pair_freq, pair_to_word, 使用splits记录每个word对应的token分布
重复以下过程：
1. 挑选频率最高的pair将其merge为一个新的token, 基于pair_to_words更新对应的pair_freq
2. 对每个split, 按照上述方式更新pair_freq和split

其具体实现见 Github

Other subword tokenizers

WordPiece

WordPiece是Google在预训练BERT时采用的tokenizer，WordPiece的基本思想和BPE差不多，都是从一个较小的vocab开始的。

首先，WordPiece会通过加上prefix ##来把单词进行切分，比如 word 会被拆分为：

1

['w', '##o', '##r', '##d']

接下来，对于pair $(a, b)$, WordPiece定义了merge pair的规则如下：

$$ \mathrm{score}((a, b)) = \frac{\#(a, b)}{\#a \times \#b} $$

其中 $\#(a, b)$, $\#a$, $\#b$ 分别代表 pair $(a, b)$, 元素 $a$ 和元素 $b$ 的频率。通过这个方式，我们会给予包含元素出现频率较低的pair更高的优先级。通过这个方式，我们选取score最高的pair，然后将其用一个新的token表示，然后和BPE算法一样，继续这一过程直到我们的vocab size达到指定大小。

在tokenize的时候，WordPiece会找出现在vocab中的最长的subword, 比如对于'hugs', 假设从左向右在词典中的最长subword是'hug', 那么'hugs' 就会被拆分为 ['hug', '##s']。如果我们在词表中找不到对应的subword，这个时候我们就会使用'[UNK]'来表示。

具体实现见 Github wordpiece (基于huggingface llm course)。代码实现除了选择最优pair的方式不同之外，和BPE基本一致。

Unigram

Unigram也是由Google提出来的tokenizer，与BPE和wordpiece不同，unigram从一个非常大的vocab开始，然后merge token来降低vocab的size，直到达到指定大小。初始的vocab可以基于BPE算法或者使用prefix subword来构建。并且，初始vocab还包含所有的base characters来保证所有的word都可以被tokenize。

算法的描述如下:

我们来看一下算法的细节, 首先对于一个word, 我们有多种切割方式, 比如'bug'可以被切分为如下三种形式:

1

[['b', 'u', 'g'], ['b', 'ug'], ['bu', 'g']]

unigram 假设每个 word 出现的概率是其 subword 出现概率的乘积, 即对于包含 $n$个subword的单词 $\bm{x}=(x_1,\dots,x_n)$, 我们有:

$$ p(\bm{x}) = \prod_{i=1}^n p(x_i) $$

其中，对于给定的vocab $\mathcal{V}$, 我们有：

$$\sum_{v\in\mathcal{V}} p(x)=1$$

unigram的目的就是选择合适的切分 $\bm{x}\in S(\bf{x})$ (这里我们用 $\bf{x}$ 表示单词本身, 用 $\bm{x}$ 表示 $\bf{x}$ 的一个切分), 使得 $p(\bm{x})$的概率最大. 这样我们就可以写出unigram的损失函数了:

$$ \mathcal{L} = \sum_{i=1}^{N} \log\left(\sum_{\bm{x}\in S(\bf{x})}p(\bm{x})\right) $$

其本质就是: 我们希望对每个单词找到一种合适的切分, 切分得到的subword的概率分布满足其求和为1, 并且使得每个单词的概率最大.

但是直接对上面概率最大化的问题就是我们每个subword的概率是未知的, unigram的做法是使用EM算法求解这个问题.

当我们求解完成之后, 对每个subword, 我们都尝试将其从 $\mathcal{V}$中移除, 然后计算移除后的损失 $loss_i$, 我们依照$loss_i$对subword进行排序, 然后我们去掉 $\eta \%$ 比例的subword.

unigram的伪代码逻辑如下:

1
2
3
4
5
6
7
8
9


while len(model) > vocab_size:
 scores = compute_scores(model)
 sorted_scores = sorted(scores.items(), key=lambda x: x[1])
 # Remove percent_to_remove tokens with the lowest scores.
 for i in range(int(len(model) * percent_to_remove)):
 _ = token_freqs.pop(sorted_scores[i][0])

 total_sum = sum([freq for token, freq in token_freqs.items()])
 model = {token: -log(freq / total_sum) for token, freq in token_freqs.items()}

其中 compute_scores 用于计算最优分割以及从model中去掉每个token之后的loss.

具体实现见 Github wordpiece (基于huggingface llm course)。代码实现的关键在于为每个word选取最优分割，huggingface是采取了动态规划的方法，也就是我们使用 dp[i] 来表示 word[:i] 的最优score，这样我们有：

$$ dp[i] = \max_{0 \leq j < i} dp[j]* p(word[j:i]),\quad \mathrm{s.t.}\ word[j:i]\in \mathcal{V} $$

这里的乘法代表 $p(\bm{x}) = \prod_{i=1}^n p(x_i)$, 在实现的时候我们会取log变成加法，然后概率会由频率来代替。

Subword tokenizer总结

sub-word tokenizer的对比 (来自huggingface llm course)：

Method	BPE	WordPiece	Unigram
Start Vocabulary	Small	Small	Large
Train	Merge tokens	Merge tokens	Remove tokens
Training Step	Merge with most frequent pair	Merge with best score	Remove all tokens minimized the loss
Learns	Merge rules and a vocab	A vocab	A vocab with a score for each token
Encoding	Splits into words and applies merge rules	Find the longest subword from the beginning that is in the vocab	Finds the most likely split into tokens with learned scores
Model	GPT	BERT	T5

实践

tiktoken

tiktoken是openAI提出来的一个BPE tokenizer, openAI的模型都基于这个tokenizer, 其主要用于调用GPT系列模型是对token进行计数, 我们可以在tokenizer 这个网站查看其分词情况.

SentencePiece

SentencePiece是google开源的一个无监督的text tokenizer，其实现了BPE和unigram两种算法，SentencePiece还是一个语言无关的tokenizer，使其更适合多语种大语言模型的开发。

Tokenizer

tokenizer 是huggingface推出的为基于transformer服务的tokenizer库, 其支持BPE, wordpiece和unigram等分词算法, 使用简便. 并且, huggingface的tokenizer包括两种:

fast tokenizer, 即Tokenizer库, 这个库是基于Rust开发的
slow tokenizer, 这个是transformer库里模型自带的, 比如ChatGLM就有自己开发的tokenizer

huggingface比较了并行处理时两者的区别:

Setting	Fast Tokenizer	Slow Tokenizer
`batched=True`	10.8s	4min41s
`batched=False`	59.2s	5min3s

huggingface提供的tokenizer库已经非常齐全了, 如果我们要训练新的基于transformer的模型的话，建议直接使用Huggingface的AutoTokenizer。

总结

特性	SentencePiece	Tokenizer	tiktoken
是否适合中文	√	√	×
是否适合英文	√	√	√
是否适合训练	√	√	×
是否快速	√	√	fast
是否用于 GPT 系列	×	×	√
是否可解码	√	√	√
是否支持多语言	√	√	×

结论

本文中, 我们介绍了大语言模型中的tokenizer, 我们从byte level, word level到sub-word level, 再到现代大语言模型最常使用的BPE tokenizer, 并给出了其（高效版本）实现。最后, 我们介绍了一下tokenizer-free的大语言模型和huggingface的tokenizer库。在未来, 我们将继续深入了解大语言模型的基本原理和实现细节。

参考文献

Notes on attention bias

Thu, 22 May 2025 15:25:07 +0800

Introduction

我们知道，transformer使用position encoding的一个原因就是，attention layer具有置换不变性，也就是说，我们随机打乱输入token的顺序，并不影响其最终结果 (我们后面会证明，实际上只对key和value具有置换不变性，对query具有置换等变性，也就是改变query的顺序之后，结果的顺序也相应改变)。因此为了让模型学习到正确的上下文知识，我们需要加上position encoding。

已有的工作大部分都在讨论如何构建更好的position encoding，但是鲜有工作探究为什么attention layer具有置换不变性. 因此，本文将从这一点出发，抽丝剥茧探究其内在原因，最后通过数学公式证明原始transformer是如何具有置换不变性的。

attention layer介绍

原始transformer layer的架构比较简单，其结构具有attention-LayerNorm-FFN-LayerNorm的形式。给定输入 $X\in\mathbb{R}^{d\times m}$ 和上下文 $Y\in\mathbb{R}^{d\times n}$. 其中，attention的定义为

$$ \mathrm{Attn}(X, Y, Y) = V\mathrm{softmax}\left(\frac{K^TQ}{\sqrt{d}}\right)\in\mathbb{R}^{d\times m} $$

其中 $d$是模型的hidden_size, $Q=W_QX\in\mathbb{R}^{d\times m}$, $K=W_KY\in\mathbb{R}^{d\times n}$, $V=W_VY\in\mathbb{R}^{d\times n}$, $W_Q, W_K, W_V\in\mathbb{R}^{d\times d}$ 分别是QKV projection layer的参数.

LayerNorm的定义为：

$$ \mathrm{LayerNorm}(x) = \frac{x-\mathbb{E}[x]}{\sqrt{\mathrm{var}[x]+\epsilon}}\odot \gamma + \beta $$

其中 $\epsilon>0$是一个超参数， $\gamma, \beta\in\mathbb{R}^d$ 是可学习的参数.

FFN的定义为：

$$ \mathrm{FFN}(x) = W_2\max(0, W_1x+b_1)+b_2 $$

其中 $W_1\in\mathbb{R}^{d_{ff}\times d}$, $W_2\in\mathbb{R}^{d\times d_{ff}}$, $b_1\in\mathbb{R}^{d_{ff}}$, $b_2\in\mathbb{R}^d$ 是可学习的参数。

最后，一个attention layer的的结构可以表达为：

$$ X = X + \mathrm{LayerNorm}(\mathrm{Attn}(X, Y, Y))\\ X = X + \mathrm{LayerNorm}(\mathrm{FFN}(X))\\ $$

置换不变性的定义

置换不变性(permutation invariant)的定义：假设 $f:\mathbb{R}^n\to\mathbb{R}^n$，如果

$$ f(\sigma(\bm{x})) = (f(\bm{x})) $$

则我们说 $f$是置换不变的. 这里 $\sigma:\mathbb{R}^n\to\mathbb{R}^n$ 是一个置换函数 (permutation function). 当输入的是一个矩阵时，我们默认置换其列，即对 $X=[X_1,\dots,X_n]\in\mathbb{R}^{d\times n}$, 我们有 $\sigma(X)=[X_{\sigma_1},\dots, X_{\sigma_n}]=Y\Pi $, 其中 $\Pi\in\mathbb{R}^{n\times n}\in \{0,1\}^{n\times n}$ 是一个置换矩阵 (permutation matrix)。

置换等变性 (permutation equivariant)的定义：假设 $f:\mathbb{R}^n\to\mathbb{R}^n$，如果

$$ f(\sigma(\bm{x})) = \sigma(f(\bm{x})) $$

则我们说 $f$是置换等变的.

attention的置换不变性与置换等变性

我们首先证明attention 对于key和value是置换不变的，即

$$ \boxed{\mathrm{Attn}(X, \sigma(Y),\sigma(Y)) = \mathrm{Attn}(X, Y, Y)} $$

证明: 我们直接计算即可得到：

$$ \begin{aligned} \mathrm{Attn}(X, \sigma(Y),\sigma(Y)) &= V\Pi\mathrm{softmax}\left(\frac{(K\Pi)^TQ}{\sqrt{d}}\right)\\ &=V\Pi\mathrm{softmax}\left(\frac{\Pi^TK^TQ}{\sqrt{d}}\right)\\ \end{aligned} $$

由于softmax是按列计算的，置换只是改变了元素的顺序，因此我们自然有

$$ \mathrm{Attn}(X, \sigma(Y),\sigma(Y)) = V\Pi\Pi^T\mathrm{softmax}\left(\frac{K^TQ}{\sqrt{d}}\right)=V\mathrm{softmax}\left(\frac{K^TQ}{\sqrt{d}}\right)=\mathrm{Attn}(X, Y, Y) $$

这里我们使用了性质 $\Pi\Pi^T=\mathbf{I}$.

接下来我们证明，attention对于query是置换等变的，即

$$ \boxed{\mathrm{Attn}(\sigma(X), Y, Y) = \sigma(\mathrm{Attn}(X,Y,Y))} $$

证明:

$$ \begin{aligned} \mathrm{Attn}(\sigma(X), Y, Y) &= V\mathrm{softmax}\left(\frac{K^TQ\Pi}{\sqrt{d}}\right)\\ &= V\mathrm{softmax}\left(\frac{K^TQ\Pi}{\sqrt{d}}\right)\\ &= V\mathrm{softmax}\left(\frac{K^TQ}{\sqrt{d}}\right)\Pi\\ &= \mathrm{Attn}(X,Y,Y)\Pi\\ &= \sigma(\mathrm{Attn}(X,Y,Y)) \end{aligned} $$

从以上的证明可以看到，attention layer对于key和value具有置换不变性，也就是说，我们改变文字顺序不影响最终的输出结果。但是，我们发现，尽管我们证明了attention具有置换不变性，我们却忽略了一件事：那就是我们计算query, key和value的时候，没有加上bias! 为什么bias如此重要呢？这是因为，$W\sigma(x) = \sigma(Wx)$, 但是 $W\sigma(X)\neq \sigma(Wx+b)$. 因此，我们就会思考，难道是transformer实际上可以通过增加bias的方式来让模型学习到上下文知识？事实上并非如此，我们将要通过分析表明，我们计算query, key和value时，增加的query bias和key bias会被softmax操作给消除掉，而key bias则会被LayerNorm消除掉。因此，我们加与加bias，对attention的置换不变性没有任何影响。

Bias对attention layer的影响

接下来，我们考虑在计算query, key和value时加入bias。为了简化，我们只考虑query为一个向量的情况，即 $X=\bm{x}\in\mathbb{R}^d$, 我们计算query, key和value如下：

$$ \bm{q} = W_Q\bm{x}+\bm{b}_Q\in\mathbb{R}^{d}\\ K = W_KY + \bm{b}_K\mathbf{1}^T\in\mathbb{R}^{d\times n}\\ V = W_VY + \bm{b}_V\mathbf{1}^T\in\mathbb{R}^{d\times n} $$

这里 $\mathbf{1}^T\in\mathbb{R}^{n}$. 我们这里简化了scaling的操作，因为其不对结果产生影响。

注：以下证明参考了【参考文献2】

我们首先展开attention中的 $V$:

$$ \begin{aligned} \mathrm{Attn}(\bm{x}, Y, Y) &= V\mathrm{softmax}\left(K^T\bm{q}\right)\\ &= \left(W_VY + \bm{b}_V\mathbb{1}^T\right)\mathrm{softmax}\left(K^T\bm{q}\right)\\ &= W_VY\mathrm{softmax}\left(K^T\bm{q}\right) + \bm{b}_V\mathbb{1}^T \mathrm{softmax}\left(K^T\bm{q}\right) \end{aligned} $$

由于 $\mathrm{softmax}\left(K^T\bm{q}\right)\in\mathbb{R}^{n}$的列求和为$1$, 因此，$\mathbb{1}^T\mathrm{softmax}\left(K^T\bm{q}\right)=1$, 我们有

$$ \mathrm{Attn}(\bm{x}, Y, Y) = W_VY\mathrm{softmax}\left(K^T\bm{q}\right) + \bm{b}_V $$

接下来，我们展开 $K$:

$$ \begin{aligned} \mathrm{Attn}(\bm{x}, Y, Y) &= W_VY\mathrm{softmax}\left(K^T\bm{q}\right) + \bm{b}_V\\ &= W_VY\mathrm{softmax}\left((W_KY + \bm{b}_K\mathbf{1}^T)^T\bm{q}\right) + \bm{b}_V\\ &= W_VY\mathrm{softmax}\left(Y^TW_K^Tq + \mathbf{1}\bm{b}_K^T\bm{q}\right) + \bm{b}_V\\ \end{aligned} $$

这里，我们需要用到softmax函数的平移不变性，即 $\mathrm{softmax}(\bm{x}+\delta\mathbf{1})=\mathrm{softmax}(\bm{x})$, 这里 $\delta\in\mathbb{R}$ 是一个常数，证明起来很简单：

$$ \mathrm{softmax}(\bm{x}+\delta)_i = \frac{e^{x_i+\delta}}{\sum_{j}e^{x_j+\delta}} = \frac{e^{x_i} * e^{\delta}}{\sum_{j}e^{x_j} * e^{\delta}} = \mathrm{softmax}(\bm{x})_i $$

而这里 $\bm{b}_K^T\bm{q}\in\mathbb{R}$，因此我们可以将这一项给去掉，我们得到：

$$ \mathrm{Attn}(\bm{x}, Y, Y) = W_VY\mathrm{softmax}\left(Y^TW_K^T\bm{q}\right) + \bm{b}_V $$

接下来，我们展开 $\bm{q}$,

$$ \boxed{ \begin{aligned} \mathrm{Attn}(\bm{x}, Y, Y) &= W_VY\mathrm{softmax}\left(Y^TW_K^T\bm{q}\right) + \bm{b}_V\\ &= W_VY\mathrm{softmax}\left(Y^TW_K^T(W_Q\bm{x}+\bm{b}_Q)\right) + \bm{b}_V\\ \end{aligned}} $$

因此，我们最终的结论为： key bias对attention输出没有任何贡献，query bias和key bias会影响结果。

到这里，看了参考文献3，我本以为可以进一步简化。但实际上并不行。参考文献3关于“transformer block is equivariant"的结果是错的，因为在attention layer之后还有一个LayerNorm，而LayerNorm不是置换不变的，这也是LayerNorm和BatchNorm之间的区别。也就是如果我们在nn.Linear后加一个BatchNorm，那么nn.Linear的bias是无效的，反之如果是LayerNorm的话，则bias是有效的.

为什么没有bias

实际上这个问题并没有定论。特别是加入position encoding之后，就更难探究bias对最终结果的影响了。但是，我认为一个原因就是bias其实就是某种先验知识，假设输入满足高斯分布，那么我们有

$$ \mathbb{E}[W\bm{x}+b] = b $$

加上先验知识后，当训练数据出现distribution shift之后，模型在训练过程中可能就会不稳定(PaLM). 而后来将LayerNorm替换为RMSNorm，使用RoPE而不是其他的additive position encoding, 我认为也是避免模型学习到先验知识，从而影响其泛化性。在未来，我认为transformer里应该是没有bias的，尽管这样效果可能会差一些，但是其稳定性更好，泛化性应该也会更好。

结论

在本文中，我们分析了attention的性质，我们发现，在原始transformer架构中，attention对于key和value有置换不变性，对于query有置换等变性。然后，我们给出了一些猜测，也就是bias会让模型产生先验知识，而这种先验知识很可能会影响训练的稳定性和模型的泛化性。

参考文献

附录

下面是测试上面结论的python代码

 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


import torch
import torch.nn as nn
import torch.nn.functional as F

# 设置随机种子，确保可复现性
torch.manual_seed(42)

# 输入参数
batch_size = 1
seq_len = 16
embed_dim = 1024 # 嵌入维度
num_heads = 32 # 多头注意力头数
head_dim = embed_dim // num_heads

# 输入张量 (batch_size, seq_len, embed_dim)
x = torch.randn(batch_size, seq_len, embed_dim)


# 有 bias 的 QKV 线性层
class Attention(nn.Module):
 def __init__(self, embed_dim, q_bias=False, k_bias=False, v_bias=False):
 super().__init__()
 self.q = nn.Linear(embed_dim, embed_dim, bias=q_bias)
 self.k = nn.Linear(embed_dim, embed_dim, bias=k_bias)
 self.v = nn.Linear(embed_dim, embed_dim, bias=v_bias)

 def forward(self, x):
 B, N, C = x.shape
 q = self.q(x).reshape(B, N, num_heads, head_dim).transpose(1, 2)
 k = self.k(x).reshape(B, N, num_heads, head_dim).transpose(1, 2)
 v = self.v(x).reshape(B, N, num_heads, head_dim).transpose(1, 2)

 attn = (q @ k.transpose(-2, -1)) * (1.0 / (head_dim**0.5))
 attn = attn.softmax(dim=-1)
 x = (attn @ v).transpose(1, 2).reshape(B, N, C)
 return x, attn


# 初始化模型
model_no_bias = Attention(embed_dim, q_bias=False, k_bias=False, v_bias=False)
model_with_bias = Attention(embed_dim, q_bias=False, k_bias=True, v_bias=False)

model_with_bias.q.weight.data = model_no_bias.q.weight.data
model_with_bias.k.weight.data = model_no_bias.k.weight.data
model_with_bias.v.weight.data = model_no_bias.v.weight.data

# 推理
out_no_bias, attn_no_bias = model_no_bias(x)
out_with_bias, attn_with_bias = model_with_bias(x)


# 比较差异
diff_output = torch.abs(out_no_bias - out_with_bias).mean()
diff_variance = torch.abs(out_no_bias - out_with_bias).var()
diff_attn = torch.abs(attn_no_bias - attn_with_bias).mean()

print("\nMean difference in output:", diff_output.item())
print("Mean difference in variance:", diff_variance.item())
print("Mean difference in attention weights:", diff_attn.item())

# Mean difference in output: 1.2734082233123445e-08
# Mean difference in variance: 1.7173628739783402e-16
# Mean difference in attention weights: 3.949708116124384e-09

Notes on Position encoding

Mon, 19 May 2025 10:46:39 +0800

本文前半部分参考参考文献1，推荐大家看博客原文。

Position encoding总结

在上一篇blog 中, 我们介绍了 Attention 的两个性质，也就是在不加 position encoding 的情况下，Attention 对于 query 是 permutation equivariant 的，对于 key 和 value 是 permutation invariant 的。

但是“我爱你”和“你爱我”这两句话所表示的含义应该是不一样的，我们将这两句话作为 key 和 value 的时候，我们发现模型的输出是一致的，这显然是不能接受的。因此，我们就需要加入 position encoding，让模型学习到语序信息，从而明白不同的语序有不同的含义。

下面是测试代码（来自参考文献1）

 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


import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel

model_id = "meta-llama/Llama-3.2-1B"
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModel.from_pretrained(model_id)

text = "The dog chased another dog"
tokens = tok(text, return_tensors="pt")["input_ids"]
embeddings = model.embed_tokens(tokens)
hdim = embeddings.shape[-1]

W_q = nn.Linear(hdim, hdim, bias=False)
W_k = nn.Linear(hdim, hdim, bias=False)
W_v = nn.Linear(hdim, hdim, bias=False)
mha = nn.MultiheadAttention(embed_dim=hdim, num_heads=4, batch_first=True)

with torch.no_grad():
 for param in mha.parameters():
 nn.init.normal_(param, std=0.1) # Initialize weights to be non-negligible

output, _ = mha(W_q(embeddings), W_k(embeddings), W_v(embeddings))

dog1_out = output[0, 2]
dog2_out = output[0, 5]
print(f"Dog output identical?: {torch.allclose(dog1_out, dog2_out, atol=1e-6)}") #True

Position encoding 可以分为绝对位置编码 (absolute position encoding, APE)，相对位置编码 (relative position encoding, RPE) 以及可学习的位置编码。可学习位置编码主要是 BERT 类的模型在使用，其训练成本比较高，本文不做讨论。绝对位置编码是原始 transformer 里提出的编码模式，现在的大多数基于 transformer 模型使用的都是相对位置编码。

本文中，我们先介绍位置编码应该具有的性质，然后我们分别介绍绝对位置编码和相对位置编码，我们将着重关注苏剑林老师提出来的 RoPE。

位置编码

在介绍位置编码之前，我们首先应该关注位置编码的性质，位置编码的目标是为输入的 token embedding 增加位置信息，那么理想的位置编码应该是怎么样的呢？

我们这里直接引用参考文献1 中给定的性质：

性质 1: token sequence 中每个位置的位置编码都是唯一的。这个很好理解，如果不唯一的话，那么根据前面推导的性质，这两个位置的 attention 输出就完全一致了
性质 2: 线性相关性。也就是说，如果我们知道了位置 $p$ 处的位置编码，那么理想情况下，我们应该能比较简单地得到 $p+k$ 处的位置编码，理想情况下，我们应该有 $PE(p+k)=W_kPE(p)$.
性质 3: 泛化到长上下文中去。我们希望位置编码不仅在 8K 的上下文起作用，还希望位置编码能够泛化到 32K 的上下文
性质 4: 生成模式是固定的。固定的模式有助于模型更好地学习位置相关的信息
性质 5: 可以扩展到多维。我们希望位置编码可以从文本扩展到图片再到视频，也就是从 $1D$ 到 $nD$.

绝对位置编码

绝对位置编码依照其名称，其思想就是为每个位置的 token 分配一个固定的位置信息，也就是对于输入的 hidden states $\bm{x}=[\bm{x}_1,\dots,\ \bm{x}_m]\in\mathbb{R}^{m\times d}$, 我们有

$$ \bm{x}_i' = \bm{x}_i + p_i, i=1,\dots, m $$

这里，$p_i\in\mathbb{R}^d$. 我们的 attention 就变成了

$$ \mathrm{Attn}(X) = \mathrm{softmax}\left(\frac{(Q+P)(K+P)^T}{\sqrt{d}}\right)V\in\mathbb{R}^{m\times d} $$

这里

$$ P = [p_1,\dots,p_m]\in\mathbb{R}^{m\times d}， Q= W_QX\in\mathbb{R}^{m\times d}, K=W_KX, V=W_VX\in\mathbb{R}^{n\times d} $$

整数位置编码

一个最简单的想法就是我们使用正整数来标记 token 所在的位置，也就是

$$ PE(i) = [i, \dots, i]=i\mathbf{1}_{d\times 1}\in\mathbb{R}^d,\ i=1,\dots,m $$

可以看到，这个简单的设计满足性质 1，性质 2，性质 3，性质 4.

但是，注意到 attention 的输入 $X$ 通常是经过 Layer Normalization 处理过后的，因此其按列符合正态分布，并且均值和方差一般较小。当我们加上整数位置编码之后，其 token 本身的信息就会被污染，也就是信噪比非常低。一个解决方法就是我们对 $PE(i)$ 进行 normalization，即

$$ PE(i)' = \frac{1}{m}PE(i) = \frac{i}{m}\mathbf{1}_{d\times 1} $$

现在所有的位置编码的值都比较小，但是我们发现新的位置编码不满足性质 2 了，这是因为现在位置编码还和 sequence 长度有关，我们从位置 $p$ 到位置 $p+k$ 不仅取决于 $k$ 还取决于 sequence 长度 $m$

二进制位置编码

既然整数位置编码的主要问题是对输入影响太大，我们能否找一个不影响输入的整数位置编码方式呢？参考文献1 提出了二进制位置编码，因为每个 token 是 $d$ 维的，因此我们可以使用 $d$ 位二进制来表示 $i$. 比如说，当 $d=3$, $m=4$ 时，我们的位置编码分别为

$$ PE(0) =p_{(000)_2} = [0, 0, 0],\ PE(1) =p_{(001)_2}= [0, 0, 1],\ PE(2) =p_{(010)_2} = [0, 1, 0],\ PE(3) =p_{(011)_2} = [0, 1, 1] $$

现在，我们二进制位置编码满足性质 1，性质 2. 对于性质 3，由于 $d$ 位二进制的表示范围为 $[0, 2^d-1]$，因此其泛化性受到 $d$ 的影响。

参考文献1 画出了不同位置的值的变化情况。我们这里也模仿绘制出类似的曲线图

我们发现，二进制位置编码高位，也就是 $PE(i)_{d}$ 的变化很慢，而低位，也就是 $PE(i)_{0}$ 变化很快，

二进制位置编码解决了整数位置编码的信噪比过低和线性相关性。但是其问题是其对不同位置的 token embedding 产生的影响是不一样的。比如位置 1 和位置 2 的相同的 token embedding 之间的区别是：

$$ (\bm{x}_2 + PE(2)) - (\bm{x}_1 + PE(1)) = (\bm{x}_2-\bm{x}_1)+ [0, 1, -1] $$

一般来说, $\bm{x}_2-\bm{x}_1$ 比较小，因此使用二进制位置编码的问题是输入位置的微小变化（增加一个 token 或减少一个 token）都会对最终结果产生巨大影响。因此，我们需要想办法解决这个问题。

Sinusoidal

前面提到二进制位置编码的问题是相邻 token 之间变化太大，不够光滑。因此我们想要增加一个光滑性质，也就是说我们希望：

位置编码值在 $[-1, 1]$ 之间，防止对 token embedding 产生影响
相邻 token 的位置编码尽可能相近，即 $|PE(k+p)-PE(p)| \leq \delta |k|$, 其中 $\delta>0$ 是一个比较小的数。
与二进制一样，高位的变化比较慢，低位的变化比较快

一个想法就是利用三角函数 $\sin$ 或者 $\cos$，三角函数满足前两个性质，对于第三个性质，我们可以通过控制频率来满足。这样我们得到的位置编码就具有如下形式：

$$ PE(p, i) = \sin\left(\frac{p}{\theta^{i/d}}\right) $$

其中 $\theta$ 是我们的超参数。

我们现在来推导一下上面位置编码的线性相关性：

$$ PE(p+k) = \sin\left(\frac{p+k}{\theta^{i/d}}\right)=PE(p)\cos\left(\frac{k}{\theta^{i/d}}\right) + \cos\left(\frac{p}{\theta^{i/d}}\right)\sin\left(\frac{k}{\theta^{i/d}}\right) $$

我们发现，$\sin$ 位置编码不满足线性相关性。但是出现的 $\cos$ 给了我们启发，也就是我们可以同时使用 $\sin$ 和 $\cos$ 来完成位置编码，这也是原始 transformer 里提出来的 Sinusoidal 位置编码，其形式为：

$$ \begin{aligned} PE(p, 2i) &= \sin\left(\frac{p}{\theta^{2i/d}}\right)\\ PE(p, 2i+1) &= \cos\left(\frac{p}{\theta^{2i/d}}\right) \end{aligned} $$

现在，记 $\omega_i=1/\theta^{2i/d}$, 我们再推导一下线性相关性，就得到：

$$ \begin{aligned} \begin{bmatrix} PE(p+k, 2i)\\ PE(p+k, 2i+1)\\ \end{bmatrix}&=\begin{bmatrix} \sin \omega_i(p+k)\\ \cos \omega_i(p+k) \end{bmatrix}\\ &=\begin{bmatrix} \sin \omega_i(\omega_ip)\cos(\omega_ik)+\cos w_i(\omega_ip)\sin(\omega_ik)\\ \cos \omega_i(\omega_ip)\cos(\omega_ik)-\sin w_i(\omega_ip)\sin(\omega_ik) \end{bmatrix}\\ &= \begin{bmatrix} \cos(\omega_ik)& \sin(\omega_ik)\\ -\sin(\omega_ik)& \cos(\omega_ik) \end{bmatrix}\begin{bmatrix} PE(p, 2i)\\ PE(p, 2i+1)\\ \end{bmatrix} \end{aligned} $$

也就是说，Sinusoidal 位置编码满足线性相关性。对于 Sinusoidal 位置编码我们也可以进行可视化：

相对位置编码

前面介绍了绝对位置编码，每个位置的位置编码是固定的。但是绝对位置编码的问题是，模型比较难以学习相对位置关系。

举个例子，我们提到上下文时，通常会使用“上一节”，“上一章”这些表示相对位置关系的词。

因此，我们希望让模型学习相对位置关系而不是绝对位置关系，因为相对关系更符合我们的认知。

RoPE

RoPE 由苏剑林老师提出，最早应用于 LLaMA 架构（没有确认），后续被大多数模型所采用。

之前的 PE 大多数关注于加性位置编码，也就是假设位置编码的形式为 $\bm{x}+\bm{p}$, 基于这种假设，已有的工作基本都集中于优化下面的 Q 和 K 的内积

$$ \langle f_q(\bm{x}_q, m), f_k(\bm{x}_k, n) \rangle $$

这里 $f_q(\bm{x}_q, m)=W_q(\bm{x}_q+\bm{p}_m)$, $f_k(\bm{x}_k, n)=W_k(\bm{x}_k+ \bm{p}_n)$.

而 RoPE 里面，作者使用了一个不同的假设： 假设内积应该仅包含两者的相对信息，也就是

$$ \langle f_q(\bm{x}_q, m), f_q(\bm{x}_k, n)\rangle := g(\bm{x}_q,\bm{x}_k, m-n) $$

这里的 $f$ 和 $g$ 都是未知函数。我们的目标就是从这个公式中推导出一个合适的位置编码出来。

不失一般性，我们可以假设

$$ f_q(\bm{x}_m,0) = \bm{x}_q,\quad f_q(\bm{x}_n, 0) = \bm{x}_k $$

这个假设代表初始条件下，我们不对输入做任何改变，也就是不增加位置信息。

2D 推导

与 RoPE 一样，我们直接使用复平面来进行推导。

我们假设 $d=2$, 注意到二维平面上的每个点都可以表示为如下形式

$$ \bm{z} = (x,y) = re^{i\theta} $$

其中 ($\mathrm{atan2}$ 定义参考维基百科)

$$ r = \|\bm{z}\|_2 = \sqrt{x^2+y^2}\in\mathbb{R},\quad \theta = \mathrm{atan2}(y, x)\in\mathbb{R}, $$

现在，对于三个向量 $f_q(\bm{x}_q, m)$, $f_q(\bm{x}_k, n)$, $g(\bm{x}_q,\bm{x}_k, m-n)$ 我们可以写出其极坐标形式：

$$ \begin{aligned} f_q(\bm{x}_q,m) &:= r_q(\bm{x}_q,m)e^{i\theta_q(\bm{x}_q,m)}\\ f_k(\bm{x}_k, n) &:= r_k(\bm{x}_k, n)e^{i\theta_k(\bm{x}_k, n)}\\ g(\bm{x}_q,\bm{x}_k, m-n) &:= r_g(\bm{x}_q,\bm{x}_k, m-n)e^{i\theta_g(\bm{x}_q,\bm{x}_k, m-n)} \end{aligned} $$

我们计算内积并比较同类项得到：

$$ \begin{aligned} r_g(\bm{x}_q,\bm{x}_k, m-n) &:= r_q(\bm{x}_q,m)r_k(\bm{x}_k, n)\\ \theta_g(\bm{x}_q,\bm{x}_k, m-n) &:= \theta_q(\bm{x}_q,m)-\theta_k(\bm{x}_k, n) \end{aligned}\tag{3} $$

我们接下来分别推导 $r_g(\bm{x}_q,\bm{x}_k, m-n)$ 和 $\theta_g(\bm{x}_q,\bm{x}_k, m-n)$ 的形式

$r_g(\bm{x}_q,\bm{x}_k, m-n)$

我们令 $m=n=0$ 可以得到初始条件

$$ r_g(\bm{x}_q,\bm{x}_k, 0) = r_q(\bm{x}_q,0)r_k(\bm{x}_k, 0)=\|\bm{q}\|_2\|\bm{k}\|_2 $$

我们再令 $n=0$,得到

$$ r_g(\bm{x}_q,\bm{x}_k, m) = r_q(\bm{x}_q,m)r_k(\bm{x}_k, 0)=r_q(\bm{x}_q,m)\|\bm{k}\|_2=\frac{r_g(\bm{x}_q,\bm{x}_k, m-n)}{r_k(\bm{x}_k, n)}\|\bm{k}\|_2 $$

这里最后一个等式带入了原始等式 (3)，注意到左侧与 $n$ 无关，因此右侧我们选取 $n=1$, 得到

$$ r_g(\bm{x}_q,\bm{x}_k, m) = \frac{r_g(\bm{x}_q,\bm{x}_k, m-1)}{r_k(\bm{x}_k, 1)}\|\bm{k}\|_2 =\cdots= r_g(\bm{x}_q,\bm{x}_k, 0)\left(\frac{\|\bm{k}\|_2 }{r_k(\bm{x}_k, 1)}\right)^{m+1} $$

令 $m=0$ 我们有

$$ r_k(\bm{x}_k, 1) = \|\bm{k}\|_2. $$

因此我们最终的表达式为：

$$ r_g(\bm{x}_q,\bm{x}_k, m) = r_g(\bm{x}_q,\bm{x}_k, 0) = \|\bm{q}\|_2\|\bm{k}\|_2. $$

并且，通过分别设置 $m=0$ 以及 $n=0$ 我们还可以得到

$$ r_q(\bm{x}_q,m) = \|\bm{q}\|_2,\quad r_k(\bm{x}_k, n) = \|\bm{k}\|_2 $$

$\theta_g(\bm{x}_q,\bm{x}_k, m-n)$

令 $m=n=0$, 我们得到初始条件

$$ \theta_g(\bm{x}_q,\bm{x}_k, 0) = \theta_q(\bm{x}_q,0)-\theta_k(\bm{x}_k, 0)=\theta_q-\theta_k $$

令 $n=1$, 我们有

$$ \begin{aligned} \theta_g(\bm{x}_q,\bm{x}_k, m-1) &= \theta_q(\bm{x}_q,m)-\theta_k(\bm{x}_k, 1)\\ &=\theta_g(\bm{x}_q,\bm{x}_k, m-n) + \theta_k(\bm{x}_k, n)-\theta_k(\bm{x}_k, 1) \end{aligned} $$

这里我们带入了公式 (3)，注意到公式左边与 $n$ 无关，因此在公式右侧我们令 $n=0$, 得到

$$ \theta_g(\bm{x}_q,\bm{x}_k, m-1) = \theta_g(\bm{x}_q,\bm{x}_k, m)+ \theta_k(\bm{x}_k, 0)-\theta_k(\bm{x}_k, 1) $$

分别令 $m=1,2,\dots$ 并相加这些等式，我们得到

$$ \theta_g(\bm{x}_q,\bm{x}_k, 0) = \theta_g(\bm{x}_q,\bm{x}_k, m) + m(\theta_k(\bm{x}_k, 0)-\theta_k(\bm{x}_k, 1)) $$

即

$$ \theta_g(\bm{x}_q,\bm{x}_k, m) = m(\theta_k(\bm{x}_k, 1)-\theta_k(\bm{x}_k, 0))+(\theta_q-\theta_k)\tag{4} $$

注意到

$$ \theta_g(\bm{x}_q,\bm{x}_k, m) = \theta_q(\bm{x}_q,m)-\theta_k(\bm{x}_k, 0)=\theta_q(\bm{x}_q,m)-\theta_k $$

带入上式我们就得到

$$ \theta_q(\bm{x}_q,m) = m(\theta_k(\bm{x}_k, 1)-\theta_k(\bm{x}_k, 0))+\theta_q $$

在 (4) 式中再令 $m=m-n$，并带入 $\theta_q(\bm{x}_q,m)$ 就有

$$ \theta_k(\bm{x}_k,n) = n(\theta_k(\bm{x}_k, 1)-\theta_k(\bm{x}_k, 0))+\theta_k $$

汇总

最后，我们将以上结果放在一起，就得到

$$ f_q(\bm{x}_q,m) = \bm{q}e^{im\theta}, f_v(\bm{x}_k,m) = \bm{k}e^{in\theta} $$

这里 $\theta=\theta_k(\bm{x}_k, 1)-\theta_k(\bm{x}_k, 0)$ 是一个超参数，用于控制频率。

我们记

$$ R_{\theta,m} = \begin{bmatrix} \cos m\theta & -\sin m\theta\\ \sin m\theta & \cos m\theta \end{bmatrix} $$

则我们有：

$$ f_q(\bm{x}_q,m) = R_{\theta,m}\bm{q}, f_v(\bm{x}_k,m) = R_{\theta,n}\bm{k}. $$

并且

$$ \langle f_q(\bm{x}_q,m), f_v(\bm{x}_k,m)\rangle = \bm{q}^TR_{\theta,m-n}\bm{k} \tag{5} $$

多维扩展

上面是 2D 的情况，对于多维情况，苏剑林老师通过将两个元素组对，然后分别进行处理，得到了多维的情形：

$$ R_{\theta,m}^d = \begin{bmatrix} R_{\theta_1,m} & & && & \\ & & R_{\theta_2,m} & & & \\ &&&& \ddots & \\ &&&& & R_{\theta_{d/2},m} \end{bmatrix}\in\mathbb{R}^{d\times d} $$

我们可以验证公式 (5) 仍然是成立的。

RoPE 的远程衰减性质

我们接下来看一下结果与相对距离 $m-n$ 之间的关系，注意到

$$ \langle f_q(\bm{x}_q,m), f_v(\bm{x}_k,m)\rangle = \bm{q}^TR_{\theta,m-n}\bm{k} = \sum_{i=1}^{d/2} \bm{q}_i^TR_{\theta, m-n}\bm{k}_i $$

这里 $\bm{q}_i=[q_{2i},q_{2i+1}]^T$, $\bm{k}_i=[k_{2i},k_{2i+1}]^T$ 分别是对应的 pair，我们考虑其中一个分量，不妨假设 $\|\bm{q}\|_2=\|\bm{k}\|_2=1$, 我们有

$$ \begin{aligned} \bm{q}_i^TR_{\theta, m-n}\bm{k}_i &\leq \bm{q}_i^TR_{\theta, m-n}\bm{q}_i\\ &= \bm{q}_i^T\left(\frac{R_{\theta, m-n}+R_{\theta, m-n}^T}{2}\right)\bm{q}_i\\ &\leq \lambda_{\max}\left(\frac{R_{\theta, m-n}+R_{\theta, m-n}^T}{2}\right) \\ &= \cos (m-n)\theta_i \end{aligned} $$

其中，第一个不等式是因为两个向量相等时其内积最大，第二个不等式是由于二次型最大值为矩阵的特征值。

这样我们就有

$$ \langle f_q(\bm{x}_q,m), f_v(\bm{x}_k,m)\rangle \leq \sum_{i=1}^{d/2}\cos (m-n)\theta_i. $$

我们可以简单画出对应的曲线：

这里针对不同的配置，结果也会略微不同，具体分析可以参加知乎回答 RoPE的远距离衰减

RoPE 代码实现与理解

Naive 实现

我们接下来看一下如何实现 RoPE

$$ \Theta_m\bm{x}=\begin{bmatrix} \cos m\theta_0 & -\sin m\theta_0 & &&\cdots &\cdots &\cdots \\ \sin m\theta_0 & \cos m\theta_0 & &&\cdots &\cdots &\cdots\\ & & \cos m\theta_1 & -\sin m\theta_1 & \cdots &\cdots &\cdots\\ & & \sin m\theta_1 & \cos m\theta_1 & \cdots &\cdots &\cdots \\ &&&& \ddots &\vdots &\vdots\\ &&&& & \cos m\theta_{d/2} & -\sin m\theta_{d/2}\\ &&&& & \sin m\theta_{d/2} & \cos m\theta_{d/2} \end{bmatrix}\begin{bmatrix} x_1\\ x_2\\ \vdots \\ x_d \end{bmatrix} $$

在实现的时候，我们一般根据 $\sin$ 和 $\cos$ 进行分组，也就是

$$ \Theta_m\bm{x}=\begin{bmatrix} \cos m\theta_0\\ \cos m\theta_0\\ \vdots\\ \cos m\theta_{d/2}\\ \cos m\theta_{d/2}\\ \end{bmatrix}\odot \begin{bmatrix} x1\\ x2\\ \vdots\\ x_{d-1}\\ x_d\\ \end{bmatrix} + \begin{bmatrix} \sin m\theta_0\\ \sin m\theta_0\\ \vdots\\ \sin m\theta_{d/2}\\ \sin m\theta_{d/2}\\ \end{bmatrix}\odot \begin{bmatrix} -\ x_2\\ x_1\\ \vdots\\ -x_d\\ x_{d-1}\\ \end{bmatrix} $$

我们通常按照奇偶 index 来分别计算，然后通过重排序来得到最终的结果，实现代码如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


def apply_rotary_pos_emb(x: torch.Tensor, sin: torch.Tensor, cos: torch.Tensor) -> torch.Tensor:
 x_even = x[..., ::2] # (seq_len, d_k_half)
 x_odd = x[..., 1::2] # (seq_len, d_k_half)
 odds = cos * x_even - sin * x_odd # (...,seq_len, d_k_half)
 evens = sin * x_even + cos * x_odd # (...,seq_len, d_k_half)
 stacked = torch.stack((odds, evens), -2) # (...,seq_len, 2, d_k_half)
 stacked_trans = rearrange(
 stacked, "... seq_len double d_k_half -> ... seq_len d_k_half double"
 ) # (...,seq_len, d_k_half, 2)
 out = rearrange(
 stacked_trans, "... seq_len d_k_half double -> ... seq_len (d_k_half double)"
 ) # (..., seq_len, d_k)
 return out

LLaMA 实现

 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


def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
 freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
 t = torch.arange(end, device=freqs.device) # type: ignore
 freqs = torch.outer(t, freqs).float() # type: ignore
 freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64
 return freqs_cis

def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
 ndim = x.ndim
 assert 0 <= 1 < ndim
 assert freqs_cis.shape == (x.shape[1], x.shape[-1])
 shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
 return freqs_cis.view(*shape)


def apply_rotary_emb(
 xq: torch.Tensor,
 xk: torch.Tensor,
 freqs_cis: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
 xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
 xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
 freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
 xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
 xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
 return xq_out.type_as(xq), xk_out.type_as(xk)

在 LLaMA 中，我们首先还是计算 $\theta_i$, 然后在计算的过程中，我们将 $(x_i,x_{i+1})$ 视作一个复数，然后乘以 $\exp(im\theta)$, 最后再取实部得到最终的结果

通用实现

实际上，naive 版本的实现与现在大语言模型所采用的实现并不一致，我们先看一下现有的大语言模型的 RoPE 实现，这里我们将 LLaMA的transformer代码放在下面，

 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


def rotate_half(x):
 """Rotates half the hidden dims of the input."""
 x1 = x[..., : x.shape[-1] // 2]
 x2 = x[..., x.shape[-1] // 2 :]
 return torch.cat((-x2, x1), dim=-1)

def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
 cos = cos.unsqueeze(unsqueeze_dim)
 sin = sin.unsqueeze(unsqueeze_dim)
 q_embed = (q * cos) + (rotate_half(q) * sin)
 k_embed = (k * cos) + (rotate_half(k) * sin)
 return q_embed, k_embed

class LlamaRotaryEmbedding(torch.nn.Module):
 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
 super().__init__()

 self.dim = dim
 self.max_position_embeddings = max_position_embeddings
 self.base = base
 inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
 self.register_buffer("inv_freq", inv_freq)

 self.max_seq_len_cached = seq_len
 t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)

 freqs = torch.einsum("i,j->ij", t, self.inv_freq)
 # Different from paper, but it uses a different permutation in order to obtain the same calculation
 emb = torch.cat((freqs, freqs), dim=-1)
 self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)
 self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)

 def forward(self, x, seq_len=None):
 # x: [bs, num_attention_heads, seq_len, head_size]
 if seq_len > self.max_seq_len_cached:
 self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)

 return (
 self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
 self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
 )

我们将上述代码翻译成公式，现在我们的 $\Theta$ 变成了 (对应 emb = torch.cat((freqs, freqs), dim=-1))

$$ \Theta = [\theta_0,\dots,\theta_{d/2},\theta_0,\dots,\theta_{d/2}]^T $$

实际上 $\sin$ 部分对应的向量现在变成了

$$ [-x_{d/2+1}, -x_{d/2+2}, \dots, -x_{d}, x_1, \dots, x_{d/2}]^T $$

我们带回到原始公式，可以得到对应的 RoPE 操作变成了

$$ R_{\theta,m}^d=\begin{bmatrix} \cos m\theta_0 & & & -\sin m\theta_0 & \cdots &\cdots &\cdots \\ & & \cos m\theta_1 & &-\sin m\theta_1 & \cdots &\cdots \\ & & & \cos m\theta_2 & &-\sin m\theta_2 & \cdots \\ \vdots&\vdots&\vdots&\vdots& \vdots &\vdots &\vdots\\ & & & &\cos m\theta_{d/2 - 1} & & -\sin m\theta_{d/2 - 1} \\ \sin m\theta_0 && & \cos m\theta_0 &&\cdots &\cdots \\ & & \sin m\theta_1 & &\cos m\theta_1 & \cdots &\cdots \\ & & & \sin m\theta_2 & &\cos m\theta_2 & \cdots \\ \vdots&\vdots&\vdots&\vdots& \vdots &\vdots &\vdots\\ &&&& \sin m\theta_{d/2 - 1}& & \cos m\theta_{d/2 - 1} \end{bmatrix} $$

这列每一行的 $\cos$ 和 $\sin$ 都相差了 $d/2$ 列.

因此，这里的区别在于，原始 RoPE 计算的 pair 为 $(x_{i}, x_{i+1})$, 而 LLaMA 里的 RoPE 计算的 pair 为 $(x_{i}, x_{i+d/2})$. transformers library 使用这种方式，可以减少计算量，提高整体的计算效率。

为了适应使用 LLaMA 中实现的 RoPE 的，Huggingface 对权重进行了转换，使得基于原始 RoPE 实现的模型也可以获得加速.

假设 $d=8$，原始 RoPE 的 pair 为 [(q_0, q_1), (q_2, q_3), (q_4, q_5), (q_6, q_7)], 新的 pair 为 [(q_0, q_4), (q_1, q_5), (q_2, q_6), (q_3, q_7)]. 我们希望对 index 进行 remap，我们发现一个满足条件的 permutation 为 [0, 2, 4, 6, 1, 3, 5, 7], 也就是 q_0->q_0, q_2->q_1, …, q_7->q_7.

但是，如果我们在推理时这样做，就会降低整体速度，因此 Huggingface 的做法是改变 $W_Q$ 和 $W_K$ 的权重，具体来说，就是 $\Pi q=(\Pi W_Q)x$，左边是在线转换，右侧离线转换。转换好 $W_Q$ 之后，正常计算就可以了。具体代码为

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


# permute for sliced rotary
def permute(w, n_heads=n_heads, dim1=dim, dim2=dim):
 return w.view(n_heads, dim1 // n_heads // 2, 2, dim2).transpose(1, 2).reshape(dim1, dim2)

state_dict = {
 f"model.layers.{layer_i}.self_attn.q_proj.weight": permute(
 loaded[f"layers.{layer_i}.attention.wq.weight"]
 ),
 f"model.layers.{layer_i}.self_attn.k_proj.weight": permute(
 loaded[f"layers.{layer_i}.attention.wk.weight"]
 ),
 ...
}

结论

本文中，我们回顾了位置编码，包括绝对位置编码和相对位置编码，我们着重介绍了 RoPE 的原理，推导以及代码实现。

参考文献

Notes on Qwen3

Thu, 15 May 2025 14:48:11 +0800

Qwen 在 2025 年 5 月发布了 Qwen3 系列大语言模型，Qwen3 包括 6 个 dense 模型和 2 个 MoE 模型，主要亮点为多语种能力，自适应快慢思考能力以及支持用户设置 thinking budget.

Qwen3 包括 6 个 dense 模型和 2 个 MoE 模型，其旗舰模型是一个 235B 的 MoE 模型，激活参数为 22B. Qwen3 系列的主要亮点如下:

快慢思考融合，模型原生支持在 reasoning/non-reasoning 模式之间切换
Reasoning budget, 用户可以指定思考需要的 budget，来平衡 latency 和 performance
Distillation, 使用蒸馏的方法训练小模型，大幅度提高模型的表现
多语种支持，相比于 Qwen2.5，Qwen3 支持 119 中语言和方言

Method

Model

Qwen3 的 dense 模型的架构与 Qwen2.5 基本一致，包括使用 GQA , SwiGLU, RoPE, RMSNorm 和 pre-normalization. Qwen3 进一步移除了 QKV bias, 然后加入了 QK-Norm 来提高训练的稳定性。

Qwen3 的 MoE 架构使用了 128 个专家，激活专家个数为 8 个。与 Qwen2.5-MoE 不同，Qwen3 里没有使用 shard experts。并且，Qwen3 加入了 global-batch load balancing loss,来提高 expert 的特化程度。

在 tokenizer 方面，Qwen 系列的 tokenizer 一直都是一样的，这也是 Qwen 系列领先的一点。

模型的具体参数如下两张表所示。

MoE 架构：上下文长度为 128K，128 个专家，每个 token 由 8 个专家负责处理

Qwen3-235B-A22B, 总参数 235B，激活参数 22B
Qwen3-30B-A3B, 总参数 30B，激活参数 3B

Models	Layers	Heads (Q / KV)	# Experts (Total / Activated)	Context Length
Qwen3-30B-A3B	48	32 / 4	128 / 8	128K
Qwen3-235B-A22B	94	64 / 4	128 / 8	128K

dense 架构: Qwen3-32B, Qwen3-14B, Qwen3-8B, Qwen3-4B, Qwen3-1.7B, and Qwen3-0.6B

Models	Layers	Heads (Q / KV)	Tie Embedding	Context Length
Qwen3-0.6B	28	16 / 8	Yes	32K
Qwen3-1.7B	28	16 / 8	Yes	32K
Qwen3-4B	36	32 / 8	Yes	32K
Qwen3-8B	36	32 / 8	No	128K
Qwen3-14B	40	40 / 8	No	128K
Qwen3-32B	64	64 / 8	No	128K

Pre-training

预训练数据一共包括 36T token，覆盖了 119 种语言。数据包括 coding, STEM, reasoning, books, multilingual texts 以及合成数据。

为了扩展训练数据，作者微调了 Qwen2.5-VL 来从 PDF 文档中提取文字，然后使用 Qwen2.5 来进行修正。最终收集到了几 T 的 token。另外，作者还使用 Qwen2.5, Qwen2.5-Math, Qwen2.5-Coder 来合成不同格式的数据，包括教科书，QA，指令以及代码片段等。最后，作者加入了更多的多语种数据。

作者从 educational value, fields, domains 以及 safety 对数据进行了标注。在数据混合时，Qwen3 在 instance 层面进行操作。

预训练阶段包括 3 个 stage：

General Stage (S1): 这一阶段的目的是让模型掌握世界知识，使用了 30T 的 token，模型上下文长度为 4096
Reasoning Stage (S2): 这一阶段的目的是提高模型的推理能力，使用了 5T 的高质量 token，模型上下文长度为 4096，数据包括 STEM, coding, reasoning 以及合成数据
Long Context Stage (S3): 这一阶段的目的是提升模型的长上下文能力，使用了几百 B的 token，模型上下文长度为 32768.训练时数据混合 75% 的长文档数据，25% 的短文本数据。作者将 RoPE 的 frequency 从 10000 提升到了 1,000,000. 作者还是用 YARN 以及 Dual Chunk Attention 来提高 inference 效率

对 pre-training 的 base model 进行评测之后，作者发现：

Qwen3-235B-A22B-Base 超过了其他 base 模型的表现，包括 DeepSeek-V3 Base, Llama-4-Maverick Base, Qwen2.5-72B Base
Qwen3-MoE 模型与相同大小的 Qwen3-Dense 模型参数相比，其只需要 1/5 的参数就可以达到相同的表现
Qwen3-MoE 模型与 2 倍参数量的 Qwen2.5-MoE 模型表现差不多
Qwen3-Dense 模型与大一个量级的 Qwen2.5-Dense 模型表现差不多

Post-training

Qwen3 的 post-training 如下图所示：

对于旗舰模型 (Qwen3-235B-A22B, Qwen3-32B) 的训练，Qwen3 使用了一个四阶段的训练 pipeline。对于轻量化模型（其他模型）的训练，Qwen3 使用了知识蒸馏。

旗舰模型的训练包括四个阶段，前两个阶段用于提升模型的 reasoning 能力，后两个阶段用于将 reasoning 和 non-reasoning 能力结合起来。

Flashship Model

Stage 1 (Long CoT Cold Start) 这个阶段的目的是让模型掌握 reasoning 的基础。这个阶段使用了数学，代码，逻辑推理和通用的 STEM 相关问题。每个问题都有参考答案或者 test-cases. 作者使用了 Qwen2.5-72B 来过滤数据，包括 non-verifiable prompts 以及太简单的 prompt. 作者认为，这一阶段应该减少训练使用的样本和训练步数。

Stage 2 (Reasoning RL) 这个阶段的目的是提升模型的 reasoning 能力。该阶段使用了 3,995 条过滤得到的样本，算法为 GRPO. 作者发现提高 batch size 和每个 query 的 rollouts 可以提高模型的表现。作者通过调整模型的 entropy 来控制 exploration 和 exploitation 的平衡

Stage 3 (Thinking Mode Fusion) 这一阶段的目的是将 non-reasoning 能力加入到之前的 reasoning 模型中。作者在第二阶段的 model 上进行了 continual SFT，然后构建了一个 chat template 用于融合两种模式。

reasoning 数据来源于 stage1 的 rejection sampling 和 stage 2 的模型. non-reasoning 数据来源于各种任务，如 coding, math, multilingual 等。为了保证模型的多语种能力，作者还加入了一些翻译相关的数据。

作者还构建了一个 chat template, 用于统一数据格式。chat template 如下图所示

作者使用 /think 和 /no_think 来标记两种模式，对于 non-reasoning mode, 其 <think></think> 会被置空。模型在默认情况下处于 reasoning mode, 因此作者加入了一些不包含 /think 的 reasoning 数据。

作者发现，通过这种 Think mode fusion, 模型可以学会在 reasoning mode 和 non-reasoning mode 下进行回答，因此，模型也可以基于中间结果来给出最终的答案。当超出 budget 之后，作者使用以下 Instruction

Considering the limit time by the user. I have to give the solution based on the thinking directly now. \n</think>.\n\n

来让模型直接终止思考二给出最终的答案。

Stage 4 (General RL) 这个阶段的目的是提升模型在不同场景下的能力。作者构建了一个 reward system 来覆盖 20 多种不同的任务。这些任务包括：instruction following, format following, preference alignment, agent ability 以及 abilities for specialized scenarios.

作者构建了三种不同的 rewards:

Rule-based rewards: 覆盖的任务包括 instruction following 和 format following
Model-based rewards: 作者使用 Qwen2.5-72B 来判别答案的正确性
Model-based Reward without reference answer: 作者训练一个 reward model 来给模型的回答进行打分

Lightweight Model

对于轻量化的模型，作者发现直接通过蒸馏可以有效提高学生模型的表现，并且训练效率也更高。蒸馏训练包括两个阶段：

Off-policy Distillation: 这个阶段的目的是让模型拥有基本的 reasoning 能力并且可以在不同的模式中进行切换。作者使用了教师模型的 reasoning 输出和 non-reasoning 输出来蒸馏学生模型
On-policy Distillation: 在这个阶段，学生模型生成回答，然后基于教师模型的输出，使用 KL-divergence 来更新学生模型的参数

Evaluation

Thinking budget. 作者发现当我们提高 Thinking budget 之后，模型的表现是可以持续提升的。结果如下图

Efficiency of distillation. 作者发现使用 distillation 可以大幅度提高模型的表现和训练效率。下面是结果如下图所示

Effects of Thinking mode fusion and RL 作者进一步探究了三个 stage 对模型表现的影响，为此，作者构建了 in-house benchmarks 来评估模型的表现，这些 benchmarks 包括：

CounterFactrQA. 问题是不符合事实的，用于评估模型的幻觉
LengthCtrl. 有长度要求的写作任务，评估生成内容长度和给定长度之间的差别
ThinkFollow. 多轮对话，每轮对话随机插入 /think 和 /no_think flag，评估模型是否能在两种模式之间切换
Tooluse. 评估模型的工具调用能力

结果如下

结论如下：

Stage3 可以提高模型在两种 reasoning mode 切换的能力，并且 stage3 还可以提高模型的通用以及 instruction following 能力
Stage4 进一步提高模型在两种模式下的通用，instruction following 和 agent 能力
Stage3 和 stage4 并没有显著提高模型在 knowledge, STEM, math 和 coding 相关任务上的表现。甚至在一些竞赛如 AIME24 上模型的表现还有所下降，作者认为这是由于我们提升了模型的通用能力而导致其特化能力下降导致的，作者认为作为一个通用模型，这是可以接受的。

Conclusion

在本文中，作者提出了 Qwen3 系列大语言模型，包括 6 个 Dense 模型和 2 个 MoE 模型。Qwen3 模型标志了一个新的 SOTA，其特点主要是快慢思考结合，thinking budget，以及多语种。

作者认为后续工作有以下几点：

使用更高质量的数据来进行预训练
优化模型架构和训练方式，提升模型的上下文
提高针对 RL 的计算资源，来进一步提高模型的 agent 能力

References

Arxiv
Github

Notes on LLaMA4 blog

Wed, 30 Apr 2025 10:44:19 +0800

介绍

Meta在2025年4月10号发布了LLaMA4系列，包含三个模型：Llama 4 Scout, Llama 4 Maverick 以及Llama 4 Behemoth, 三个模型都基于MoE架构，且支持多模态

Model	Layers	Heads (Q / KV)	Context Length	#Parameters (activated/total)	#Experts (activated/total)	#Tokens
LLaMA 4 Behemoth	-	-	-	288B / 2T	(1 shared + 1 routed) / 16	-
LLaMA 4 Maverick	48	40/8	1M	17B / 109B	(1 shared + 1 routed) / 128	~22T
LLaMA 4 Scout	48	40/8	10M	17B / 400B	(1 shared + 1 routed) / 16	~40T

训练数据截止到2024年8月。LLaMa4支持200多种语言，其中100多种语言的训练token数超过了1B

亮点

原生多模态 LLaMA 4是一个原生多模态架构
超长上下文 LLaMA 4的上下文超过了1M
iRoPE 通过交替dense和MoE MLP来提高整体推理效率
基于MetaCLIP的vision encoder
MetaP 使用MetaP来调整超参数
FP8精度训练

Pre-training

LLaMA 4 仍然是一个基于transformer的架构，但是引入了MoE，其示意图如下所示

MoE架构中包含1个shared expert以及1个routed expert. 并且，与其他LLM不同，LLaMA 4使用了一个交替MLP个MoE的架构，即iRoPE，即特定的transformer layer是MoE架构，其余的是MLP架构，其核心代码如下：

1
2
3
4
5


self.is_moe_layer = layer_idx in config.moe_layers
if self.is_moe_layer: # the 128E model interleaves dense / sparse
 self.feed_forward = Llama4TextMoe(config)
else:
 self.feed_forward = Llama4TextMLP(config, intermediate_size=config.intermediate_size_mlp)

early fusion. LLaMA 4称其一个原生多模态大模型，但是其架构仍然是 Vision Encoder-MLP-LLM 的形式，其不同点在于patch embedding没有使用convolution, 而是使用 nn.Unfold直接进行展平，然后使用一个线性层与vision encoder进行对齐。代码如下

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18


class Llama4UnfoldConvolution(nn.Module):
 def __init__(self, config):
 super().__init__()
 kernel_size = config.patch_size
 if isinstance(kernel_size, int):
 kernel_size = (kernel_size, kernel_size)
 self.unfold = torch.nn.Unfold(kernel_size=kernel_size, stride=config.patch_size)
 self.linear = nn.Linear(
 config.num_channels * kernel_size[0] * kernel_size[1],
 config.hidden_size,
 bias=False,
 )

 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 hidden_states = self.unfold(hidden_states)
 hidden_states = hidden_states.permute(0, 2, 1)
 hidden_states = self.linear(hidden_states)
 return hidden_states

其他训练优化技巧如下：

MetaP：用于选择超参数
FP8 precision：与 DeepSeek-V3 一样，使用FP8精度进行训练
mid-training：在预训练阶段之后，额外增加了一个训练阶段，来提高模型的长上下文等关键能力

Post-training

post-training包括三个阶段：

SFT
online RL
DPO

作者发现SFT和DPO会限制模型的探索能力，特别是在math, coding等domain。为了解决这个问题，作者使用LlaMA对问题进行难度分级，然后移除了50%的简单数据。

在online RL阶段，作者设计了一个continuous online RL策略，让模型在训练和筛选问题两种模式之间进行切换，以平衡效率和准确率。

DPO的目的是为了提升模型输出的质量

评测

benchmark	LLaMA 4 Maverick	LLaMA 4 Maverick	LLaMA 4 Scout	Gemeni 2.0 Flash	GPT-4o
MMMU	76.1	73.4	69.4	71.7	69.1
Math Vista	73.7	70.7	73.1	63.8
ChartQA	-	90.0	88.8	88.3	85.7
DocVQA	-	94.4	94.4	-	92.8
LiveCodeBench	49.4	43.4	32.8	34.5	32.3
MMLU Pro	82.2	80.5	74.3	77.6	-
GPQA Diamond	73.7	69.8	57.2	60.1	53.6

结论

LLaMA 4 采用了MoE架构，是一个原生的多模态大模型系列。在架构上，与DeepSeek-MoE, aria和OLMoE不同，LLaMA4并没有增加expert granularity，OLMoE分析认为，增加granularity可以提高模型的flexibility，下面总结了一下相关模型的参数

Model	Layers	Heads (Q / KV)	Context Length	#Parameters (activated/total)	#Experts (activated/total)	#Tokens
DeepSeek-MoE(144.6B)	62	32/32	2048	22.2B/144.6B	(1 shared + 7 routed)/64	245B
DeepSeek-V3	61	128(MLA)	128K	37B/671B	(1 shared + 8 routed)/257	14.8T
Aria	28	20/20	64K	3.5B/24.9B	(2 shared+ 6 routed)/66	6.4T(text)
OLMoE	16	16/16	4096	1.3B/6.9B	8/64	5T
LLaMA 4 Maverick	48	40/8	1M	17B / 109B	(1 shared + 1 routed) / 128	~22T
LLaMA 4 Scout	48	40/8	10M	17B / 400B	(1 shared + 1 routed) / 16	~40T

Reference

Notes on Qwen3 blog

Tue, 29 Apr 2025 11:23:04 +0800

介绍

Qwen3发布，包含两种架构的模型，每种架构均包含对应的base model和post-trained model.

MoE架构：上下文长度为128K，128个专家，每个token由8个专家负责处理

Qwen3-235B-A22B, 总参数235B，激活参数22B，
Qwen3-30B-A3B, 总参数30B，激活参数3B

Models	Layers	Heads (Q / KV)	# Experts (Total / Activated)	Context Length
Qwen3-30B-A3B	48	32 / 4	128 / 8	128K
Qwen3-235B-A22B	94	64 / 4	128 / 8	128K

dense架构: Qwen3-32B, Qwen3-14B, Qwen3-8B, Qwen3-4B, Qwen3-1.7B, and Qwen3-0.6B

Models	Layers	Heads (Q / KV)	Tie Embedding	Context Length
Qwen3-0.6B	28	16 / 8	Yes	32K
Qwen3-1.7B	28	16 / 8	Yes	32K
Qwen3-4B	36	32 / 8	Yes	32K
Qwen3-8B	36	32 / 8	No	128K
Qwen3-14B	40	40 / 8	No	128K
Qwen3-32B	64	64 / 8	No	128K

亮点

Hybrid Thinking modes qwen3支持两种思考模式： thinking mode 和 non-thinking mode，前者用于解决复杂的问题，后者用于解决简单的问题
multilingual support qwen3支持119中语言和方言
Improved agentic capabilities 提升了qwen3的coding和agentic能力，并支持MCP

训练

Pre-training

Qwen3使用了36T token进行训练（与之相比，Qwen2.5使用方的token数量为18T），数据来源于互联网和PDF，作者使用Qwen2.5-VL来提取内容，然后使用Qwen2.5来提升内容质量。作者还基于Qwen2.5-Math和Qwen2.5-Coder来合成数学以及代码数据

训练包含三个stage：

上下文长度为4K tokens，训练数据为30T tokens, 目标是让模型掌握初步的语言能力和知识
上下文长度为4K tokens，训练数据为5T tokens,这部分数据主要是knowledge intensive的数据，比如STEM, coding和reasoning等
上下文长度扩展到32K tokens，训练数据主要是高质量长上下文的数据

Post-training

Post-training包含四个阶段，如下图所示

Stage 1, Long CoT code start: 作者基于math, coding, logical reasoning, STEM等domain的Long CoT数据来微调模型，让模型拥有初步的推理能力，这和Kimi-VL是类似的
Stage 2, reasoning-based RL: 作者使用rule-based rewards来提供奖励，然后使用RL来训练模型，经一部提高模型的exploration和exploitation能力
Stage 3, thinking mode fusion: 作者混合了一部分instruction-following和Long CoT数据来提升模型的non-thinking能力，这样可以让模型在两种思考模式之间切换
Stage 4,general RL：作者使用RL在20多个general-domain任务上进一步提高模型的通用能力，包括instruction following, format following以及agent capability等

Future work

作者希望在未来能够在模型架构和训练方式上进行提升，包括：scaling data, increasing model size, extending context length, broadening modalities, advancing RL with environmental feedback for long-horizon reasoning.

结论

与Gemini2.5 pro，Kimi-VL等reaosning model不同，qwen3可以在快思考和慢思考之间进行转换。感觉未来有两个趋势，一个是如何在快思考和慢思考之间进行切换，切换的逻辑是什么？第二个就是qwen3以及qwen2.5-vl都在强调的agent能力，也就是我们不仅仅是在做一个LLM，而是逐步延伸到了agent这个层面。

参考文献

blog

Notes on VAPO

Thu, 17 Apr 2025 09:41:51 +0800

Abstract

字节Seed团队提出了 Value-based Augmented Proximal Policy Optimization (VAPO), 用于提高 reasoning model 的表现。 VAPO 通过集成 DAPO 和 VC-PPO 的优点，进一步提高了 value-based 方法的表现。

Introduction

现有的RL 训练方法可以分为 value-free 和 value-based 两大类。其中 value-free 方法不需要使用 value model, 比如 GRPO 和 GRPO 的变体 DAPO, 这类方法通过多次采样，然后使用 leave-one-out estimate 来代替 value model. 这类方法的优点是不需要训练value model, 但是缺点是在复杂的任务中表现不是很稳定。

另一方面，value-based 方法需要训练一个 value model, 比如 VC-PPO, 这类方法的优点是：

提供更细粒度的奖励信号
提供lower-varaince value estimation, 从而提高训练的稳定性
拥有更好的泛化能力

但是，value-based 方法在训练过程中存在一些问题：

训练一个low-bias 的 value model 比较困难，尤其是在long trajectory 上，因为 bias 会随着 trajectory 的长度增加而增加
在heterogeneous sequence lengths during training 中表现不佳，对于短文本和长文本，我们需要考虑 bias-variance 的trade-off
在sparse reward signal 中表现不佳

为了解决这些问题，字节Seed团队提出了 VAPO, 一个基于value-based 的 RL 训练方法，VAPO 通过结合 DAPO 和 VC-PPO 的优点，进一步提高了 value-based 方法的表现。

Preliminary

Preliminary包括 token-level MDP, RLHF, PPO 三个部分，这部分请参考 VC-PPO. 这里不做重复。

VAPO

作者针对value-based 方法在训练过程中存在的三个问题，在VAPO中分别进行了解决。

Mitigating Value Model Bias over Long Sequences

在 VC-PPO 中，作者提到了 value model 和 reward model 的不一致性，这种不一致性会导致 bias, 尤其是在long sequences上。在VAPO中，作者就直接使用了 VC-PPO的做法，包括value pretraining 和 decoupled GAE 来解决这个问题。

Managing Heterogeneous Sequence Lengths during Training

针对heterogeneous sequence的问题，作者提出了 Length-Adaptive GAE. 在VC-PPO中， $\lambda_{\mathrm{policy}}$ 被设置为 $0.95$. 但是当 sequence 非常长时， TD-error会迅速下降，导致GAE被一部分TD-error所主导，从而不利于模型的训练。

为了解决这个问题，作者将 $\lambda_{\mathrm{policy}}$ 与 sequence 的长度 $\ell$ 联系起来，具体来说，两者的关系如下：

$$ \sum_{t=0}^{\infty}\lambda_{\mathrm{policy}}^t = \frac{1}{1-\lambda_{\mathrm{policy}}} := \alpha\ell $$

其中 $\alpha$ 是一个超参数，用来控制 bias-variance 的trade-off. 给定 $\ell$, $\lambda_{\mathrm{policy}}$ 可以被计算为：

$$ \lambda_{\mathrm{policy}} = 1 - \frac{1}{\alpha\ell} $$

同时，为了平衡短文本和长文本的贡献，基于 DAPO, 作者构建了 token-level policy gradient loss，其具体形式如下：

$$ \mathcal{L}_{\mathrm{PPO}}(\theta) = \frac{1}{\sum_{t=1}^G|o_i|}\sum_{i=1}^G\sum_{t=1}^{|o_i|}\min\left(r_{i,t}(\theta)\hat{A}_{i,t},\mathrm{clip}\left(r_{i,t}(\theta), 1-\epsilon, 1+\epsilon\right)\hat{A}_{i,t}\right) $$

Dealing with Sparsity of Reward Signal in Verifier-based Tasks

与DAPO一致，为了解决reward signal的稀疏性问题，作者提出了Clip-Higher, 来让更小概率的输出也能获得较大的更新概率。其更新公式如下：

Clip-Higher的介绍见DAPO

然后，作者还将next-token prediction loss 和 PPO loss 结合起来，来降低奖励的稀疏程度。

$$ \mathcal{L}_{\mathrm{NTP}}(\theta) = -\frac{1}{N}\sum_{o_i\in\mathcal{T}}\sum_{t=1}^{|o_i|}\log \pi_{\theta}(a_{i,t}|s_{i,t}) $$

其中 $\mathcal{T}$ 是正确答案的集合。最终的loss为：

$$ \mathcal{L}_{\mathrm{VAPO}}(\theta) = \mathcal{L}_{\mathrm{PPO}}(\theta) +\mu \mathcal{L}_{\mathrm{NTP}}(\theta) $$

其中 $\mu$ 是一个超参数，用来平衡PPO loss和NTP loss。

Experiments

模型使用Qwen-32B来进行训练, 大部分细节与VC-PPO和DAPO一致，这里不再赘述。最后与DAPO以及R1的对比结果如下：

Ablation study

针对本文使用的模块，作者进行了消融实验，结果如下：

从实验结果可以看到：

value pretraining 和 decoupled GAE 可以显著提高模型的表现
clip-higer 可以提升模型的探索能力
length-adaptive GAE 可以平衡模型在短文本和长文本上的表现

Training Dynamics

与DAPO类似，作者也分析了VAPO的训练动态，结果如下：

从上面三张图可以看到：

VAPO相比于DAPO来说，其训练更加稳定
从response length来看，VAPO的response length更长，说明VAPO的length scaling更强
从reward score来看，VAPO的reward score更高，说明VAPO的reward signal提供的指导信息更多
从generation entropy来看，在训练末期，VAPO的generation entropy更低，说明VAPO的生成多样性要更低一些，但这也说明了VAPO的生成更加稳定。作者认为这个时候稳定性更重要。

Conclusion

作者在DAPO和VC-PPO的基础上，提出了VAPO，一个基于value-based 的 RL 训练方法，VAPO 通过结合 DAPO 和 VC-PPO 的优点，进一步提高了 value-based 方法的表现。后续，作者又提出了Seed-thiking-1.5的技术报告。可以说，这一系列论文的连贯性是非常高的。

Reference

Notes on VC-PPO

Mon, 14 Apr 2025 17:36:15 +0800

Abstract

字节Seed团队提出了 Value-Calibrated PPO (VC-PPO), 用于解决PPO的value initialization bias 以及 reward signal decay 问题. 具体来讲：

VC-PPO增加了 value pretraining 来解决 value initialization bias 的问题
VC-PPO分离了 actor 和 critic 的GAE的计算，避免了 reward signal decay 的问题

Introduction

已有的reasoning model的训练方法，主要包括两个个stage:

SFT: 这个阶段主要使用了一些标注好的long CoT数据，初步激活模型的reasoning能力(参考KImi-VL)
RL: 这个阶段使用收集到的数据使用RL算法进行训练，任务包括math, code, logic reasoning等

已有PPO算法在处理Long CoT任务时，存在的问题在 DAPO 中已经介绍过了，GRPO的解决方式为使用leave-one-out estimate来替换value model. 但是GRPO相比于PPO能够提供token级别的奖励来说，只能提供response level的奖励，因此限制了模型的性能。

Preliminary

Preliminary包括MDP, RLHF, PPO三个部分，RLHF和PPO我们在 DAPO 中已经介绍过了，这里不再赘述。

Token-level MDP

给定prompt $x$ 和 response $y$, 我们可以将 $x$ 和 $y$ 分解为 token 序列，比如 $y = y_1, y_2, \cdots, y_n$, 其中 $y_i\in\mathcal{A}$, $\mathcal{A}$ 是我们的词表。

我们将 token-level MDP定义为：

$$ \mathcal{M} = \langle \mathcal{S}, \mathcal{A}, P, r, d_0, \omega \rangle $$

其中：

$\mathcal{S}$ 是状态空间，表示当前的token序列， $t$时刻的状态可以表示为 $s_t = (x, y_1, y_2, \cdots, y_t)$
$\mathcal{A}$ 是动作空间，表示下一个token， $t$时刻的动作可以表示为 $a_t = y_{t+1}\in\mathcal{A}$
$P$ 是状态转移概率，表示在 $t$时刻，从状态 $s_t$ 转移到状态 $s_{t+1}$ 的概率
$r$ 是奖励函数，表示在 $t$时刻，从状态 $s_t$ 采取动作 $a_t$ 转移到状态 $s_{t+1}$ 的奖励
$d_0$ 是初始状态分布，表示初始状态 $s_0$ 的概率分布
$\omega$ 是终止状态分布，表示终止状态 $s_n$ 的概率分布，通常表示 <eos> token

方法

首先作者分析了一下为什么 PPO 在处理 long CoT 任务时效果不佳。

PPO的传统设置为：

将GAE的参数 $\lambda$ 设置为 0.95
使用一个 reward model 来初始化 value model

一方面，作者认为，$\lambda=0.95$ 是为了减少模型在 Mujoco 以及 Atari 等环境中的variance，但是对于Long CoT任务，这个设置会导致模型缺乏足够的探索能力。

另一方面，作者认为 reward model 和 value model 虽然都是提供关于 response 的信息，但是他们之间还是存在一些差距的。作者给出了使用PPO来进行 long CoT 任务相关的实验结果

可以看到，随着训练的进行，模型的context length以及在AIME benchmark上的表现都出现了下降。

在本文中，作者主要是在 verifiable tasks 上进行实验，答案的正确性与 response length 长度关系不大，因此, response length 可以反应模型的训练情况。作者绘制了训练过程中， value 和 advantage 与 token position 的关系图

从上图可以看到，value 和 advantage 更倾向与给最初的 token 更高的 bias, 作者认为出现这个情况的原因是 value model 和 reward model 的目标并不匹配。 reward model 的目标是给 $\omega$ 也就是 <eos> token reward, 对于一开始的 token, reward model会给出比较低的奖励；而 value model 的目标是给整个 response 一个奖励，因此，value model 会倾向于给最初的 token 更高的奖励。我们对 GAE 进行改写得到：

$$ \hat{A}_t = \sum_{i=t}^{T-t-1} \lambda^{i} \left( r_{t+i} + V(s_{t+i+1}) - V(s_{t+i}) \right) $$

从上式可以看到，一开始 $r_{t+i}$ 的值会比较小，而 $V(s_{t+i})$ 的变化比较大，因此，一开始的 token 的 advantage 会比较大，这个 bias 会持续影响整个 trajectory。

为了解决这个问题，作者提出了 value-pretraining, 也就是对value model 进行离线与训练，直到其收敛到一个具体的 policy 上。具体训练步骤为：

基于一个policy，如 $\pi_{\mathrm{sft}}$ 进行采样，然后更新value model ($\lambda=1.0$)
基于收集到的数据训练 value model, 直到 value loss 或者 explain variance 收敛

接下来，在训练 value model 时，我们还需要考虑 variance reduction 的问题。作者首先改写了 GAE 的公式：

$$ \hat{A}_t = \begin{cases} \sum_{i=0}^{T-t-1} \lambda^{i} \left( r_{t+i} + V(s_{t+i+1}) - V(s_{t+i}) \right)+V(s_t) & \text{if } \lambda < 1.0 \\ \sum_{i=0}^{T-t-1} r_{t+i} & \text{if } \lambda=1.0 \end{cases} $$

可以看到，当 $\lambda < 1.0$ 并且 $T-t-1$ 比较大时，<eos> token 的reward就非常接近于0了，作者通过实验验证了这一点。

可以看到，当我们降低 $\lambda$ 时，reward signal 的稀疏性会显著增加，从而提高了模型的训练难度。

但是，我们又不能不使用 variance reduction, 因为这会导致训练的不稳定性。作者从 TD error 的角度分析了 variance:

$$ \begin{aligned} \mathrm{Var}[A_{t}^{\lambda}] &= \mathrm{Var}\left[\sum_{i=0}^{T-t-1} \lambda^{i}\delta_{t+i}\right] \\ &= \sum_{i=1}^{T-t-1} \lambda^{2i} \mathrm{Var}[\delta_{t+i}] + 2\sum_{i=1}^{T-t-1}\sum_{j=0}^{i-1} \lambda^{i+j} \mathrm{Cov}[\delta_{t+i}, \delta_{t+j}] \end{aligned} $$

因为 $\lambda\in[0,1]$, 因此未来的 TD error 会衰减的更快，因此就降低了 advantage 的 variance.

那么如何在降低variance的同时，又不至于使得 reward signal 过于稀疏呢？作者的解决方案是分离GAE的计算，也就是将 GAE 的计算分为两部分：

$$ G_{t:t+h} = \begin{cases} \sum_{i=0}^{h-1} r_{t+i} + \bar{V}(s_{t+h}) & \text{if } t+h基于这个公式，我们可以写出policy gradient的公式：

$$ \begin{aligned} \nabla_{\theta} J(\theta) &= \mathbb{E}_{t}[\nabla_{\theta} \log \pi_{\theta}(a_t|s_t)A_t] \\ &= \mathbb{E}_{t}\left[\nabla_{\theta} \log \pi_{\theta}(a_t|s_t) \sum_{i=0}^{T-t-1} \lambda^{i}\left( r_{t+i} + \bar{V}(s_{t+i+1}) - \bar{V}(s_{t+i}) \right)\right] \\ &= \mathbb{E}_{t}\left[\nabla_{\theta} \log \pi_{\theta}(a_t|s_t) \left((1-\lambda) \sum_{i=1}^{T-t-1} \lambda^{i-1}G_{t:t+i}+\lambda^{T-t-1}G_{t:T}- \bar{V}(s_{t})\right)\right] \\ &= \mathbb{E}_{t}\left[\nabla_{\theta} \log \pi_{\theta}(a_t|s_t) \left((1-\lambda) \sum_{i=1}^{T-t-1} \lambda^{i-1}G_{t:t+i}+\lambda^{T-t-1}G_{t:T}\right)\right] \end{aligned} $$

通过这种方式，我们就可以避免 value function 对 policy gradient 的影响，因而我们可以对 value model 和 policy model 使用不同的 $\lambda$ 进行训练。

最终，我们就可以得到 VC-PPO 的算法：

实验

setup

作者在AIME, GPQA 以及Codeforces三个数据集上进行评测
作者首先进行了code-start，作者构建了一批样本然后要求模型在 <thinking> 和 </thinking> 之间生成推理过程，然后使用Verifier来针对答案部分提供奖励，正确则奖励为1，错误则奖励为-1
RL的Baseline使用的是PPO
value pretraining 时，作者将 GAE的 $\lambda$ 设置为 1.0，其他参数与PPO一致
对于decoupled GAE，作者使用 $\lambda_{\text{critic}}=1.0$, $\lambda_{\text{actor}}=0.95$

实验结果

实验结果如下图所示

作者还分析了以下模型在AIME数据集上随着训练步数增加准确率的变化情况

Ablation Study

作者首先探究了 value pretraining 以及 decoupled GAE 对于模型性能的影响

从上图可以看到，直接使用PPO并不能提升模型的表现，而使用value pretraining 以及 decoupled GAE 能够显著提升模型的表现。

作者接下来探究了不同的value pretraining steps对模型的影响，结果如下图

从上表可以看到，value pretraining 的训练步数并不是越多越好，随着训练步数的增加，模型可能会出现过拟合的现象。

最后作者还分析了以下 $\lambda_{\text{actor}}$ 对于模型性能的影响，结果如下图

可以看到， $\lambda_{\text{actor}} =1.0$ 的效果是最差的，但是 $\lambda_{\text{actor}}$ 也不是越小越好，实验发现当 $\lambda_{\text{actor}} \in [0.95, 1.0)$ 时结果比较好

Findings

作者还提供了一些发现。

作者认为，在LLM中进行RL的训练与传统的RL训练不同，我们不再是从一个随机policy开始，而是从一个SFT之后的policy开始，因此，这就会引入 prior，我们需要将 value model 与 policy model 进行对齐，才能使得训练更加稳定。
作者认为，value pretraining 可以王value model 中注入先验知识，作者通过实验发现，value pretraining 的过程可以分为两个阶段，第一个阶段是random alignment，这个和传统的RL训练类似，第二个阶段是knowledge injection，这个阶段，value model 开始学习如何给重要的token 更高的权重。
作者发现， value model 倾向于更大的 $\lambda$, 因此结果会导致更小的 bias 和更大的 variance. 而 policy model 倾向于更小的 $\lambda$. 这种差异启发我们需要使用一些基于policy gradient 的目标来训练 value model.

结论

本文提出了VC-PPO，一个通过使用value-pretraining 以及 decoupled GAE 来解决PPO 的 value initialization bias 以及 reward signal decay 问题的算法。

参考文献

VC-PPO
DAPO

Notes on DAPO

Wed, 09 Apr 2025 21:40:33 +0800

Introduction

字节Seed团队和清华合作提出了DAPO，一种全开源的，基于PPO的强化学习方法，一用于提升LLM的reasoning能力。

Preliminary

Preliminary包括PPO，GRPO还有KL divergence

PPO

PPO的训练目标为：

$$ \mathcal{J}_{\mathrm{PPO}}(\theta) = \mathbb{E}_{(q,a)\sim\mathcal{D},o_{\leq t}\sim \pi_{\theta_{old}}(\cdot\mid q)}\left[ \min\left(r_t(\theta)\hat{A}_t,\mathrm{clip}\left(r_t(\theta), 1-\epsilon, 1+\epsilon\right)\hat{A}_t\right) \right] $$

其中

$$ r_t(\theta) = \frac{\pi_{\theta}(o_t\mid q, o_{< t})}{\pi_{\theta_{old}}(o_t\mid q, o_{< t})} $$

$(q,a)$ 是从数据集 $\mathcal{D}$ 采样的QA pair，$\epsilon>0$ 是一个超参数，$\hat{A}_t$ 是 $t$时刻的优势估计 (advantage estimator). 给定 value function $V$ 以及 reward function $R$, $\hat{A}_t$ 通过计算GAE得到：

$$ \hat{A}_t^{\mathrm{GAE}(\gamma, \lambda)}=\sum_{k=0}^{\infty}(\gamma\lambda)^k\delta_{t+k} $$

其中

$$ \delta_k = R_k + \gamma V(s_{k+1})-V(s_k),\quad 0\leq \gamma,\lambda\leq 1 $$

GRPO

相比于PPO，GRPO不依赖于value function, 因此不需要使用reward model. GRPO通过一组输出来估计value $V(s)$, 然后进行更新。具体来说，给定 QA pair $(q,a)$, 我们从$\pi_{\theta_{old}}$中采样$G$个输出 $\{o_i\}_{i=1}^G$, 接下来我们基于reward $\{R_i\}_{i=1}^G$ 使用如下表达式来估计group-level reward:

$$ \hat{A}_{i,t} = \frac{r_i - \mathrm{mean}(\{R_i\}_{i=1}^G)}{\mathrm{std}(\{R_i\}_{i=1}^G)} $$

最后，GRPO的训练目标与PPO类似，只不过将 $\hat{A}_t$ 替换为 $\hat{A}_{i,t}$, 然后在分组上进行了归一化：

其中，

$$ r_{i,t}(\theta) = \frac{\pi_{\theta}(o_{i,t}\mid q, o_{i,< t})}{\pi_{\theta_{old}}(o_{i,t}\mid q, o_{i,< t})} $$

KL divergence

在传统的RLHF框架中，我们再PPO的基础上，增加了一个KL divergence正则项，用于约束新旧策略的差异。具体来说，给定旧策略 $\pi_{\theta_{old}}$ 和新策略 $\pi_{\theta}$，我们实际上优化的损失函数为

$$ \mathcal{J}_{\mathrm{RLHF}}(\theta) = \mathcal{J}_{\mathrm{PPO}}(\theta) - \beta\mathrm{KL}\left(\pi_{\theta_{old}}(\cdot\mid q)\|\pi_{\theta}(\cdot\mid q)\right) $$

其中，$\beta$ 是一个超参数，用于平衡PPO损失和KL divergence损失。上面的PPO损失函数也可以改为GRPO损失函数。

作者在本文中认为，reasoning model和RLHF的训练目标是不一样的，RLHF加上KL divergence正则项，会使得模型在推理时过于保守，偏向于explotition而reasoning model则需要exploration。因此，作者在本文中去掉了这一项。

Rule-based reward modeling

作者基于final accuracy来作为outcome reward, 避免模型训练出现reward hacking问题。reward function 如下：

$$ R(\hat{y},y) = \begin{cases} 1, & \text{if is\_equivalent}(\hat{y},y) \\ -1, & \text{otherwise} \end{cases} $$

DAPO

DAPO基于GRPO改进，其优化的目标函数为：

$$ \mathcal{J}_{\mathrm{DAPO}}(\theta) = \mathbb{E}_{(q,a)\sim\mathcal{D},\{o_i\}_{i=1}^G\sim \pi_{\theta_{old}}(\cdot\mid q)}\left[ \frac{1}{\sum_{t=1}^G|o_i|}\sum_{i=1}^G\sum_{t=1}^{|o_i|}\min\left(r_{i,t}(\theta)\hat{A}_{i,t},\mathrm{clip}\left(r_{i,t}(\theta), 1-\epsilon_{low}, 1+\epsilon_{high}\right)\hat{A}_{i,t}\right) \right] s.t. \quad 0< \vert \{o_i\mid \text{is\_equivalent}(o_i,a)\} \vert < G $$

其中，$\alpha$ 是一个超参数，用于约束输出数量, $r_{i,t}(\theta)$ 和 $\hat{A}_{i,t}$ 的定义与GRPO相同。

接下来就是DAPO算法的几个关键点：

Clip-Higher

作者认为，PPO和GRPO存在entropy collapse问题，也就是policy的entropy会迅速的下降，导致最终每个group的输出基本都差不多，模型很难探索到新的知识。因此，作者提出了Clip-Higher策略，也就是说，让模型能够充分探索。

作者举了一个例子，当 $\epsilon=0.2$ 时，如果一个group的输出为 $[0.01, 0.9]$, 那么在clip之后，最大的更新幅度为 $[0.01, 0.9]\times 0.2=[0.002, 0.18]$. 这对于概率比较小的输出来说，很难帮助模型提升。因此，作者修改了其阈值。进一步地，作者分离了clip的阈值，分别使用 $\epsilon_{low}$ 和 $\epsilon_{high}$ 来约束更新幅度。作者通过实验验证了这个例子，结果如下图所示：

最后，加入Clip-Higher策略的DAPO的训练目标为，与GRPO的训练目标相比，我们使用 $\epsilon_{low}$ 和 $\epsilon_{high}$ 来约束更新幅度。

在实际训练中，作者使用了一个比较大的 $\epsilon_{high}$ 来保证概率较小的token也能有较大的概率被采样到。 Clip-Higher的实验结果如下图所示：

Dynamic Sampling

作者认为，在GRPO中，如果一个group的输出全都是对的，那么其advantage会趋近于0，结果导致 policy 不会得到更新，这样就降低了采样效率。

为了解决这个问题，坐着提出了over-sample以及filtering策略，来过滤accracy等于1或者0的prompts,也就是DAPO中的约束项：

$$ s.t. \quad 0< \vert \{o_i\mid \text{is\_equivalent}(o_i,a)\} \vert < G $$

通过增加这个约束项，我们可以保证每个group的输出不会趋同，从保证模型能够得到更新，以提高采样效率。实验结果如下图所示：

Token-level policy gradient loss

GRPO采用了一个sample-level的loss计算方式，也就是GRPO首先在每个sample中计算每个token的损失并进行平均，然后在不同的sample中在此进行平均。在这种方式下，每个sample对最终的loss贡献是相同的。这样就导致两个问题：

对于高质量的long answer，通过在token层面进行平均，因为answe比较长，因此整体loss会比较低，也就会降低这个sample的贡献，导致模型很难学习到长答案。
一些非常长的sample，其往往包含gibberish或者repetitive tokens，这些sample对模型训练是有害的

因此，为了解决这个问题，作者提出了Token-level policy gradient loss，也就是我们改变了求和方式

$$ \frac{1}{G}\sum_{i=1}^G\frac{1}{|o_i|}\sum_{t=1}^{|o_i|}\left(\cdot\right)\to \frac{1}{\sum_{t=1}^G|o_i|}\sum_{i=1}^G\sum_{t=1}^{|o_i|}\left(\cdot\right) $$

通过这种方式，我们让长回答的loss贡献更大，从而提高模型学习长回答的能力。另一方面，现在每个token对整体loss的贡献是相同的，一些关键token的loss也会被放大，从而提高模型的表现。

实验结果如下图所示：

Overlong reward shaping

与Kimi-k1.5类似，DAPO也增加了length penalty，用于惩罚过长的回答。作者首先使用了一个overlong filtering技巧，来mask掉truncated samples的loss，作者发现通过这个技巧，可以提升模型的训练稳定性以及表现，实验结果如下图所示

作者还提出了soft overlong punishment，用于reshape truncated samples的reward，表达式如下：

$$ R_{length}(y) = \begin{cases} 0, & \text{if } |y|\leq L_{\max}-L_{cache} \\ \frac{(L_{\max}-L_{cache})-|y|}{L_{cache}}, & \text{if } L_{\max}-L_{cache}<|y|\leq L_{\max} \\ -1, & \text{if } |y|>L_{\max} \end{cases} $$

算法

DAPO的算法流程如下：

实验

数据集

作者构建了DAPO-Math-17K数据集用于DAPO的训练，该数据集从AoPS得到，包含了17K prompts，每个prompt都对应一个整数作为答案。

训练细节

作者以GRPO作为baseline， $G=16$, $L_{\max}=16384$, $L_{cache}=4096$, $\epsilon_{low}=0.2$, $\epsilon_{high}=0.28$

评估时，使用AIME作为benchmark，以 avg@32作为指标，temperature设置为1.0, topp设置为0.7。

实验结果

DAPO与GRPO的对比如下图所示：

Ablation study

作者探究了每一个部分对最终表现的影响，结果如下：

traing dynamics

作者还探究了mean response length, reward score, generation entropy以及mean probability随训练轮数的变化，其中：

mean response length: 在一定程度上反应了模型训练的稳定性和表现
reward score: 反应模型的表现，作者认为，给定一个reliable reward signal， LLM可以很好的fit到训练数据集上，但是作者发现最终的reward与val score相关性比较大，很可能是因为模型过拟合了
generation entropy & mean probability: 代表了模型的探索能力，通过实验结果可以看到，DAPO初期的探索能力比较强，但是随着训练的进行，探索能力下降，利用能力增强。

结果如下：

结论

作者基于GRPO提出了DAPO，一种全开源的，基于PPO的强化学习方法，用于提升LLM的reasoning能力。作者首先分析了GRPO损失函数存在的不足，然后进行了针对性改进，包括Clip-Higher, Dynamic Sampling, Token-level policy gradient loss以及Overlong reward shaping。作者通过实验验证了DAPO的有效性，并探究了DAPO的训练动态。

Reference

DAPO: An Open-Source LLM Reinforcement Learning System at Scale

Notes on Gemma3

Sat, 15 Mar 2025 11:15:29 +0800

作者提出了Gemma3系列大模型，包括1B, 4B, 12B, 27B四个size。4B, 12B, 27B三个size均支持多模态，128K的上下文长度以及140种语言。

方法

数据处理

数据格式

1
2
3
4
5
6
7
8


[BOS]<start_of_turn>user
Who are you?<end_of_turn>
<start_of_turn>model
My name is Gemma!<end_of_turn>
<start_of_turn>user
What is 2+2?<end_of_turn>
<start_of_turn>model
2+2=4.<end_of_turn>

pretrain和SFT的区别在于,pretrain时模型输出以<eos>结束,SFT时模型输出以<end_of_turn>结束.

图片处理

输入的图片都会被resize到896x896，如果图片精度过大或者不是正方形，则会通过Pan & Scan技巧裁剪为多个子图，然后每个子图分别进行resize。核心处理代码为：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14


crop_size_w = int(math.ceil(width / num_crops_w))
crop_size_h = int(math.ceil(height / num_crops_h))

# Don't apply PaS if crop size is too small.
if min(crop_size_w, crop_size_h) < pan_and_scan_min_crop_size:
 return []

crop_positions_w = [crop_size_w * i for i in range(num_crops_w)]
crop_positions_h = [crop_size_h * i for i in range(num_crops_h)]

image_crops = [
 image[pos_h : pos_h + crop_size_h, pos_w : pos_w + crop_size_w]
 for pos_h, pos_w in itertools.product(crop_positions_h, crop_positions_w)
]

模型架构

模型包括4个size，分别是1B, 4B, 12B, 27B。1B的模型是单模态的，4B, 12B, 27B的模型是多模态的。对于多模态模型来说：

Vision encoder: Siglip-400M
Projection layer: linear layer
LLM: Gemma 3

模型参数

Model	Vision Encoder	Embedding Parameters	Non-embedding Parameters	context length	multilingual	training data
1B	0	302M	698M	32K	English	2T tokens
4B	417M	675M	3,209M	128K	140+ languages	4T tokens
12B	417M	1,012M	10,759M	128K	140+ languages	12T tokens
27B	417M	1,416M	25,600M	128K	140+ languages	14T tokens

Attention layers

为了提高效率，作者将部分layer的self-attention替换为sliding window attention。论文里是将6 layers为一组，替换一组中最后一层为sliding window attention，其余层为self attention。判断某layer是否为sliding window attention的代码为：

1
2


# config.sliding_window_pattern = 6
self.is_sliding = bool((layer_idx + 1) % config.sliding_window_pattern)

Long context

作者将global self-attention的RoPE的base frequency从10K提升到了1M，对于sliding window attention，RoPE的base frequency保持10k不变。

预训练

数据在模型参数里进行了汇总
tokenizer使用的是Gemini2.0的tokenizer，基于SentencePiece，vocab大小为262K
使用知识蒸馏的方法进行训练，每个token按照教师模型的概率采样256个logits，然后使用cross-entropy loss进行训练
4B, 12B, 27B的模型先在32K的context length上进行训练，然后在128K的context length上进行训练

Quantization Aware training

作者提供了quantized版本的模型，模型基于预训练好的模型使用QAT方法进行SFT 5000 steps左右。最后是各个模型的内存占用情况

	Raw (GB)		Quantized (GB)
Model	bf16	Int4	Int4(blocks=32)	SFP8
1B	2.0	0.5	0.7	1.0
+KV	2.9	1.4	1.6	1.9
4B	8.0	2.6	2.9	4.4
+KV	12.7	7.3	7.6	9.1
12B	24.0	6.6	7.1	12.4
+KV	38.9	21.5	22.0	27.3
27B	54.0	14.1	15.3	27.4
+KV	72.7	32.8	34.0	46.1

后训练

数据过滤：筛选掉包含PII，不安全的或者有毒的输出等。留下上下文依赖高，幻觉小的数据

训练包括升级版的知识蒸馏和基于RL的finetuning，其中RL的reward来自于weight averaged reward models, code execution feedback, ground-truth rewards

实验

表现

Gemma3的表现如下图所示

与PaliGemma 2的表现对比

消融实验

Local attention layers

sliding_window_pattern对模型的表现影响不大
sliding window size对模型的表现也不是很大，如下图
使用sliding window attention可以降低KV cache的内存占用

Long context ablation

结论为long context会降低模型的性能

知识蒸馏

训练token个数比较少的时候，使用小的教师模型效果更好；训练token个数比较多的时候，使用大的教师模型效果更好。

Pan & Scan

使用图片原始的aspect ratio训练的模型效果更好。

memorization

memorization指模型输出的文本与训练数据中文本的重复率。结果发现Gemma3的memorization要更低一些。

结论

sliding window attention在Qwen2.5-VL里已经验证过有效,这里在Gemma3上同样验证了有效性.
模型架构与PaliGemma系列基本一致，只是attention改变了，然后LLM从Gemma 2升级到了Gemma 3。

参考文献

Notes on QwQ-32B

Sat, 08 Mar 2025 09:46:16 +0800

Introduction

通义团队在3月6号发布了QwQ-32B，一个基于RL的，32B参数的reasoning model，QwQ-32B的表现可以与DeepSeek-R1相比

模型架构

QwQ-32B基于Qwen2.5-32B开发。主要通过基于outcome-based rewards的RL进行训练。训练过程包括两个stage

Stage 1：本阶段在math和coding domain上进行RL的训练，作者使用了一个verifier来保证最终数学问题结果的正确性，使用了一个code execution server来保证最终生成代码的正确性。
Stage 2：本阶段在通用领域上进行RL的训练。模型基于general reward model和一些rule-based verifier进行训练。应该是类似 DeepSeek-R1 的规则。

实验结果

实验结果如下图所示

Reference

blog
demo

compression is intelligence

Thu, 06 Mar 2025 17:57:51 +0800

我们知道，基于decoder-only transformer的LLM的训练目标是最小化next-token-prediction loss，即给定sequence $x=(x_1,\dots, x_n)\in D$，我们的目标为求解以下优化问题

$$ \min_{\theta} -\sum_{x\in D}\log P_{\theta}(x_i|x_1,\dots,x_{i-1}) $$

这里 $\theta$ 就是我们的模型参数，$D$ 是我们的训练数据集。

无数模型通过实际效果告诉我们，这个优化目标可以很好地训练出具有良好泛化能力的大语言模型。但是，我们的问题是，为什么这个优化目标可以训练出智能的模型？本文将从压缩即智能的角度来理解这个问题。

压缩即智能

一个例子

我们首先来看一个简单的例子。给定如下三个0-1字符串：

1
2
3


01010101010101010101
01001000100001000001
01101000101010100101

我们该如何描述这三个字符串的规律？显然，第一个字符串最简单，它是01字符串重复得到的结果；第二个字符串稍微复杂一些，它在每个1之前插入重复次数的0；第三个字符串则最复杂，它是我随手写的一个字符串，基本没有任何规律，因此，我们只能直接存储这个字符串。

这个例子告诉我们，一个字符串的规律越简单，我们越容易描述它，因此，我们越容易压缩它。实际上，大语言模型做的也是类似的事情。它们的核心思想是，压缩即智能。

结论

本文中，我们从压缩即智能的角度来理解大语言模型的原理。我们发现，大语言模型的next-token-prediction其实就是压缩。我们通过压缩让大语言模型学习到了语言中的规律，从而让模型具有了智能。

参考文献

Notes on Phi-4

Mon, 16 Dec 2024 17:33:52 +0800

Introduction

微软在12月12号发布了Phi-4的技术报告,主要关注了数据的构建，特别是。Phi-4是一个14B的大语言模型，主要在STEM相关的QA任务以及推理相关的任务上表现比较好。Phi-4主要在三个方面进行了改进：

在pre-training和mid-training阶段使用了合成数据进行训练

Data

Pre-training

Architecture

Post-training

DPO

SFT

References

Blog
Arxiv

ROUGE (Recall-Oriented Understudy)

Thu, 09 May 2024 17:35:20 +0800

ROUGE stands for Recall-Oriented Understudy for Gisting Evaluation. It includes several automatic evaluation methods that measure the similarity between summaries.

Preliminaries

ROUGE-N: N-gram co-occurrence statistics

Definitions

Given any string $y=y_1y_2\cdots y_K$, where $y_i,i\in{1,\dots,K}$ are characters and an integer $n\geq1$, we define the set of $n$-gram to be

$$ G_n(y) = \{ y_1\cdots y_n, y_2\cdots y_{n+1}, \dots, y_{K-n+1}\cdots y_K\} $$

NOTE that this is a set with unique elements, for example, $G_2(abab)=\{ab, ba\}$.

Given any two strings $s$ and $y$, we define substring count $C(s,y)$ to tbe the number of appearances of $s$ as a substring of $y$. For example, $C(ab, abcbab)=2$ since $ab$ appear in $abcbab$ twice in position $1$ and $5$.

Base Version

ROUGE-N is an n-gram recall between a candidate summary $\hat{y}$ and a set of reference summaries $S=\{y_1,\dots,y_n\}$. ROUGE-N is defined as follows:

$$ \text{ROUGE-N}(\hat{y}, S) = \frac{\sum_{i=1}^n\sum_{s\in G_N(y_i)}C(s,\hat{y})}{\sum_{i=1}^n\sum_{s\in G_N(y_i)}C(s, y_i)} $$

Features of ROUGE-N:

the denominator increases as we add more references, since there might exists multiple good summaries.
A candidate summary that contains words shared by more references is favored by the ROUGE-N measure.

Multiple references

When there are multiple references for a candidate summary, it is suggests to use the following formula:

$$ \text{ROUGE-N}(\hat{y}, S) = \arg\max_{i=1,\dots,n}\text{ROUGE-N}(\hat{y}, \{y_i\}) $$

the above formula is favored for the following reasons:

There is no single “best” reference summary, the multiple reference formula allows ROUGE-N to take into account of all the possible reference summaries and provide a more accurate measure of the quality of the generated summary.
The multiple reference formula is more robust. If a reference summary contains a typo or a grammatical error, this can affect the ROUGE-N score.
The multiple reference formula can provide a more comprehensive evaluation of the generated summary, since it can allow ROUGE-N to evaluate the generated summary against a wider range pf possible reference summaries.

ROUGE-L

A sequence $Z=[z_1,\dots,z_m]$ is a subsequence of another sequence $X=[x_1,\dots,x_n]$ if there exists a strict increasing sequence $[i_1,\dots,i_k]$ of indices of $X$ such that for all $j=1,\dots,k$, we have $x_{i_j}=z_j$.

Given two sequences $X$ and $Y$, the longest common subsequences (LCS) of $X$ and $Y$ is a common subsequences with maximum length.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21


def longest_common_subsequence(x: str, y: str) -> int:
 """compute the length of LCS of x and y

 Args:
 x (str): a string of length m
 y (str): a string of length n

 Return:
 int: the length of LCS of x and y
 """
 m, n = len(x), len(y)
 dp = [[0 for _ in range(m + 1)] for _ in range(n + 1)]

 for i in range(1, m + 1):
 for j in range(1, n + 1):
 if x[i] == y[j]:
 dp[i][j] = dp[i - 1][j - 1] + 1
 else:
 dp[i][j] = max(dp[i - 1][j], dp[i][j - 1])

 return dp[m][n]

Sentence-level LCS

The intuition of sentence-level LCS is that the longer the LCS of two summary sentences is, the more similar the two summaries are.

Given two summaries $X$ of length $m$ and $Y$ of length $n$, assuming $X$ is a reference summary sentence and $Y$ is a candidate summary sentence, the LCS-based recall, precision and F-measure are defined as follows:

$$ R_{LCS} = \frac{LCS(X, Y)}{m}, P_{LCS} = \frac{LCS(X, Y)}{n}, R_{LCS} = \frac{(1+\beta^2)R_{LCS}P_{LCS}}{R_{LCS}+\beta^2P_{LCS}} $$

the above formula is called ROUGE-L. $\beta$ is a hyperparameter.

Features of ROUGE-L are listed as follows:

It doesn’t require consecutive matches but in-sequence matches, which is more reasonable than n-grams.
It automatically includes longest in-sequence common n-grams, there fore no predefined n-gram length is necessary.
It’s value is less tan or equal to the minimum of unigram F-measure of $X$ and $Y$.
The disadvantage of ROUGE-L is that it only counts the main in-sequences words, therefore, other alternative LCSes and shorter sequences are not reflected in the final score.

Summary-level LCS

We can apply sentence-level LCS-based F-measure score to summary level. Given a reference summary of $u$ sentences $\{r_1,\dots,r_u\}$ containing a total of $m$ words and a candidate summary of $v$ sentences $\{c_1,\dots,c_v\}$ containing a total of $n$ words, the summary-level LCS-based recall, precision and F-measure are defined as follows:

$$ R_{LCS} = \frac{\sum_{i=1}^u\max_{j}LCS(r_i, c_j)}{m}, P_{LCS} = \frac{\sum_{i=1}^v \max_{j}LCS(r_i, c_j)}{n}, R_{LCS} = \frac{(1+\beta^2)R_{LCS}P_{LCS}}{R_{LCS}+\beta^2P_{LCS}} $$

ROUGE-W

The basic LCS has a problem that it doesn’t differentiate LCSes of different spatial relations within their embedding sequences.

To improve the basic LSC method, we can simply remember the length of consecutive matches encountered so fat to a regular two dimensional dynamic program table computing LCS. we call this weighted LCS (WLCS) and use $k$ to indicate the length of the current consecutive matches ending at words $x_i$ and $y_j$.

 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


def weighted_longest_common_subsequence(x: str, y: str, f: Callable):
 """use dynamic programming to compute WLCS with a weighted function f

 Args:
 x (str): a candidate summary containing m words
 y (str): a reference summary containing n words
 f (Callable): a function satisfies f(x+y) > f(x) + f(y)

 Return:
 the WLCS score

 Reference:
 https://aclanthology.org/W04-1013
 """
 x, y = x.split(), y.split()
 m, n = len(x), len(y)
 c = [[0 for _ in range(m + 1)] for _ in range(n + 1)]
 w = [[0 for _ in range(m + 1)] for _ in range(n + 1)]
 for i in range(1, m + 1):
 for j in range(1, n + 1):
 if x[i] == y[j]:
 # the length of consecutive matches at (i - 1, j - 1)
 k = w[i - 1][j - 1]
 c[i][j] = c[i - 1][j - 1] + f(k + 1) - f(k)
 # remember the length of consecutive matches at (i - 1, j - 1)
 w[i][j] = k + 1
 else:
 if c[i - 1][j] > c[i][j - 1]:
 c[i][j] = c[i - 1][j]
 w[i][j] = 0 # no match at (i, j)
 elif c[i][j] = c[i][j - 1]:
 w[i][j] = 0 # no match at (i, j)

 return c[m][n]

where c[i][j] stores the WLCS score ending at word x[i] of x and y[i] of y. w stores the length of consecutive matches at c[i][j]. f is a function of consecutive matches at c[i][j].

Recall, precision, F-score based on WLCS can be computed as follows:

$$ R_{WLCS} = f^{-1}\left(\frac{WLCS(X, Y)}{f(m)}\right), P_{WLCS} = f^{-1}\left(\frac{WLCS(X, Y)}{f(n)}\right), R_{LCS} = \frac{(1+\beta^2)R_{WLCS}P_{WLCS}}{R_{WLCS}+\beta^2P_{WLCS}} $$

where $f$ is the inverse function of $f$. We call the WLCS-based F-measure as ROUGE-W. Usually, a function $f$ that has a close form inverse is preferred.

ROUGE-S

Skip-bigram is any pair of words in their sentence order, allowing for arbitrary gaps. Skip-bigram co-occurrence statistics measure the overlap of skip-bigrams between a candidate translation and a set of reference translations. A sentence with $n$ words will have $\binom{n}{2}=n(n-1)/2$ skip-bigrams.

Recall, precision, F-score based on skip-bigram can be computed as follows:

$$ R_{\mathrm{SKIP2}} = \frac{\mathrm{SKIP2}(X, Y)}{m}, P_{\mathrm{SKIP2}} = \frac{\mathrm{SKIP2}(X, Y)}{n}, R_{\mathrm{SKIP2}} = \frac{(1+\beta^2)R_{\mathrm{SKIP2}}P_{\mathrm{SKIP2}}}{R_{\mathrm{SKIP2}}+\beta^2P_{\mathrm{SKIP2}}} $$

where $\mathrm{SKIP2}(X, Y)$ is the number of skip-bigram matches between $X$ and $Y$. The F-score is called ROUGE-S.

ROUGE-SU

One problem of ROUGE-S is that it doesn’t given any credit to a candidate sentence if the sentence doesn’t have any word pair co-occurring with its references.

To fix this problem, we extend ROUGE-S with the addition of unigram as counting unit. The extended version is called ROUGE-SU. We can also obtain ROUGE-SU from ROUGE-S by adding a begin-of-sentence marker at the beginning of candidate and reference sentences.

Reference

Formal Algorithms for Transformer

Thu, 02 May 2024 13:13:12 +0800

This post is a notes on understanding how transformer works in an algorithm perspective.

Introduction

Transformer is a neural architecture that is used for neural language processing. Transformer receives an embedding matrix, which represents a sentence as input, and outputs a matrix of the same size as the embedding matrix, then the output can be used for downstream tasks.

Notation

We denote $V=[N_V]:=\{1,\dots,N_V\}$ as the vocabulary of tokens or words or characters.
We denote $\bm{x}=x[1...n]:=x[1]...x[n]\in V^n$ be a sequence of tokens, for example, a sentence or a paragraph.
Given a matrix $M\in\mathbb{R}^{m\times n}$, $M[i,:]\in\mathbb{R}^n$ is the $i$-th row of $M$, $M[:, j]\in\mathbb{R}^m$ is the $j$-th column of $M$.

Tokenization

Tokenization determines how the text are represented. Given a piece of text, for example, "I have an apple", we seek to find a proper way to represent this sentence.

Character level tokenization. In this setting, $V$ is the English alphabet plus punctuation. This tends to yield very long sequences (depends on the character contained in the raw text).
Word level tokenization. In this setting, $V$ is the set of all English words plus punctuation. Word level tokenization is straightforward, but it tends to required a very large vocabulary and cannot handle new vocabulary at test time.
Subword tokenization. This is the most common way used in nowadays, $V$ is the set containing the commonly used word segments like “ing”, “est”. This can be computed via Byte-Pair Encoding (BPE) algorithm.
We suppose the length of the input text is $L$, if the length input text exceeds $L$, we chunk it.

After tokenization, each element in the vocabulary is assigned to a unique index $i\in\{1,\dots,N_V-3\}$, and a number of special tokens are added to the vocabulary. For example:

mask_token$=N_V-2$, used in masked language modeling
bos_token$=N_V-1$ and eos_token$=N_V$, these two tokens are used to represent the beginning and the end of the sequence.

Finally, a piece of raw text is represented as a sequence of indices, often called token IDs corresponding to its subwords, preceded by bos_token and followed by eos_token.

Embedding

The embedding layer is used to represent each token as a vector that contains richer semantic information. The embedding contains two parts:

token embedding, where each token is embedded into a vector space
positional embedding, where embeds the position information of the tokens.

Token embedding

Given a sequence of token ID, we now need to represent each token as a vector in $\mathbb{R}^d$.

The simplest way is to use one-hot embedding, where each token $i$ is represented a vector $[0,\dots,1,\dots,0]\in\mathbb{R}^{N_V}$ whose elements are all $0$ excepts that $i$-th position is equal to $1$. However, the problem is that the vocabulary size $N_V$ is two large.

To solve this problem, we can train a learnable embedding model, of which parameter is a matrix $W_e\in\mathbb{R}^{d\times N_V}$, its $i$-th row corresponds to vector representation of the token $i$:

$$ \bm{e} = W_{e}[:, i]\in\mathbb{R}^d $$

Position embedding

There is a problem in token embedding, that is, it doesn’t contain consider the order of tokens. In latter, we show that the self-attention mechanism is equivariant to a permutation matrix $X\Pi$ of data $X$, where $\Pi$ is a permutation matrix, that is,

$$ \mathrm{Sa}(X\Pi) = \mathrm{Sa}(X)\Pi $$

the above equation indicates that the self-attention layer learn no position information at all!

To solve this problem, we add a positional embedding to token embedding. There are two kinds of embeddings:

Absolute positional embeddings. In this setting, a matrix $W_P\in\mathbb{R}^{d\times N}$ is learned or design to indicate the position of tokens. Mathematically, we have

$$ \bm{e}_p = W_p[:, \mathrm{index}(i)]\in\mathbb{R}^{d} $$

where $\mathrm{index}(i)$ is the index of token $i$ in the input sequence. 2. Relative positional embeddings. We leave this in latter notes. Compared to absolute positional embeddings, relative positional embeddings uses offset information, which performs well when the input sequence is too long.

The final embedding of a token $i$ is given by

$$ \bm{e} = W_e[:, i] + W_p[:, \mathrm{index}(i)]\in\mathbb{R}^{d} $$

Attention

The idea of attention mechanism is: Given a sequence of token, to predict the current token, which token should I pay attention to? For example, I opened the door with my ___, we may answer key, password or fingerprint etc. This is because we notice that we opened the door, so to predict the next token, we should make use of the information. What attention mechanism does is to quantify this process and make them parallel and learnable.

Single query attention

We first consider a simple example. Given the embedding of the current token $\bm{e}\in\mathbb{R}^d$ and the list of context tokens $[\bm{e}_1,\dots,\bm{e}_N]\in\mathbb{R}^{d\times N}$, the attention is given as follows:

compute query vector: $\bm{q}=W_q\bm{e}+b_q\in\mathbb{R}^{d}$
compute key vectors: for $i=1,\dots,L$, $\bm{k}_i=W_k\bm{e}_i+b_k\in\mathbb{R}^{d}$
compute value vectors: for $i=1,\dots,L$, $\bm{v}_i=W_v\bm{e}_i+b_v\in\mathbb{R}^{d}$
compute attention weights: let $\bm{s}=[\bm{q}^T\bm{k}_1,\dots,\bm{q}^T\bm{k}_L]\in\mathbb{R}^{N}$, then:

$$ \bm\alpha = \mathrm{softmax}\left(\frac{\bm{s}}{\sqrt{d}}\right)\in\mathbb{R}^{N}$$

5. compute vector representation of the token and context combined:

$$ \bm{v}'= \sum_{i=1}^N\alpha_i\bm{v}_i\in\mathbb{R}^{d} $$

where $W_q,W_k,W_v\in\mathbb{R}^{d\times d}$, $b_q,b_k,b_v\in\mathbb{R}$.

General attention

To extend the single query attention to general form, we consider the embedding matrix $X\in\mathbb{R}^{D\times N}$, the context matrix $Z\in\mathbb{R}^{d\times C}$ and a mask matrix $M\in\mathbb{R}^{D\times D}$, then the attention is computed as follows:

compute query matrix:

$$ Q=W_qX+\bm{b}_q\in\mathbb{R}^{D\times N}$$

2. compute key matrix:

$$ K=W_kZ+\bm{b}_k\in\mathbb{R}^{D\times C}$$

3. compute value vectors:

$$ V=W_vZ+\bm{b}_v\in\mathbb{R}^{D\times C}$$

4. compute attention weights:

$$\mathrm{Sa}(X) = \mathrm{softmax}\left(M\odot \frac{K^TQ}{\sqrt{D}}\right) \in\mathbb{R}^{C\times N} $$

where $\odot$ is the element-wise product. 5. output the updated representations of tokens in $X$, folding the information from tokens in $Z$

$$ \tilde{V} = V\odot \mathrm{Sa}(X)\in\mathbb{R}^{D\times N} $$

There are two kinds of mask matrices depending on which attention we are using:

Bidirectional attention, in this case, $M=\bm{1}\bm{1}^T\in\mathbb{R}^{C\times N}$.
Undirectional attention, in this case, $M[i,j] = \bm{1}_{i\leq j}$, where $\bm{1}_{i\leq j}$ is the indicator function.

Multi-head Attention

The previous describes the operation of a single head. In practice, transformers run multiple attention heads in parallel and combine their outputs, this is called multi-head attention. The idea behind multi-head attention can be summarized as follows:

In high dimensional space, two vectors are usually far from each other, with multiple attention heads, we can reduce the dimension of the representation.
With multiple attention heads, each heads may focus on specific semantics of the representation. For example, one head focuses on positiveness, and one another head focuses on noun/verb semantics.

For simplicity, we denote the single-head attention as $\mathrm{attention}(X, Z\mid M)$, suppose we have $H$ heads, then we compute $\tilde{V}_i$ for each heads:

$$ \tilde{V}_i = \mathrm{attention}(X, Z\mid M)\in\mathbb{R}^{D\times N},i=1,\dots,H $$

then attention representation are concatenated together:

$$ V = [\tilde{V}_1^T, \dots, \tilde{V}_H^T]^T\in\mathbb{R}^{HD\times N} $$

combined via an output matrix $W_o\in\mathbb{R}^{D\times HD}$:

$$ \tilde{V} = W_oV+\bm{b}_o\in\mathbb{R}^{D\times N} $$

We denote the multi head attention as $\mathrm{MhSa}(X, Z\mid M)$.

Transformer layer

After computing the multi head attention, we can now construct a transformer layer, which can also be stacked as convolution neural networks. A transformer layer can be constructed by the following operations:

Multi head attention (residual), $X\gets X + \mathrm{MhSa}(X, Z\mid M)$.
Layer norm, $X \gets \mathrm{LayerNorm}(X)$
Multi layer perception $\bm{x}_i\gets \bm{x}_i + \mathrm{mlp}(\bm{x}_i), i=1,\dots,N$
Layer norm, $X \gets \mathrm{LayerNorm}(X)$

where $\mathrm{LayerNorm}$ is the layer norm operation. $\mathrm{mlp}$ is a multi layer perception, usually it consists of one hidden layer of size $4D$, that is, then umber of neurons in three layers are $D, 4D, D$.

Usually, a large language model consists of multiple transformer layers.

Unembedding

The unembedding learns to convert a vector representation of a token and its context $\bm{e}$ into a distribution over the vocabulary elements.

$$ \bm{p} = \mathrm{softmax}(W_u\bm{e})\in \Delta(V)\subseteq \mathbb{R}^d $$

where $\Delta(V)$ is a simplex over the set $V$.

Reference

BLEU (Bilingual Evaluation Understudy)

Thu, 25 Apr 2024 22:46:53 +0800

BLEU (Bilingual Evaluation Understudy) is a widely used metric that evaluates the quality of the translated text with respect to the reference translations.

Introduction

The formula of BLEU is defined as follows:

$$ \mathrm{BLEU}_ {w_n}(\hat{S}, S) = \mathrm{BP} \cdot \exp\left(\sum_{n=1}^Nw_n\log p_n(\hat{S}, S) \right) $$

where

$\mathrm{BP}$ represents the Brevity Penalty to penalize the translations that are shorter than the reference translations.
$p_n$ represents the modified $n$-gram precision score
$w_n$ represents the weight for $p_n$, it satisfies $0\leq w_n\leq1$ and $\sum_{n=1}^Nw_n=1$.

Now we interprets three parts in detail.

Interpretation

Definitions

Given any string $y=y_1y_2\cdots y_K$, where $y_i,i\in{1,\dots,K}$ are characters and an integer $n\geq1$, we define the set of $n$-gram to be

$$ G_n(y) = \{ y_1\cdots y_n, y_2\cdots y_{n+1}, \dots, y_{K-n+1}\cdots y_K\} $$

NOTE that this is a set with unique elements, for example, $G_2(abab)=\{ab, ba\}$.

Modified precision score

We start from the one candidate translation $\hat{y}$ and one reference translation $y$. The modified $n$-gram is defined as

$$ p_n(\hat{y}, y) = \frac{\sum_{s\in G_n(\hat{y})}\min (C(s,\hat{y}), C(s,y))}{\sum_{s\in G_n(\hat{y})}C(s,\hat{y})} $$

The quantity measures how many $n$-grams of the reference translation $y$ appears in the candidate translation $\hat{y}$. In case that $\hat{y}$ is too short, we take a minimum between $C(s,\hat{y})$ and $ C(s,y)$. Then we normalize to make $p_n(\hat{y}, y)$ comparable among multiple translation pairs.

Now, suppose we have a candidate translation corpus, $\hat{S}=\{\hat{y}^1,\dots,\hat{y}^M\}$, and for each candidate translation $\hat{y}^i$, we have a reference translation corpus (there are multiple translations can represent the same meaning) $S_i=\{y^{i,1},\dots,y^{i,N_i}\}$. We define $S=\{S_1,\dots,S_M\}$, then our modified $n$-gram precision is defined as

$$ p_n(\hat{S}, S) = \frac{\sum_{i=1}^M\sum_{s\in G_n(\hat{y})}\min (C(s,\hat{y}), \max_{y\in S_i}C(s,y))}{\sum_{i=1}^M\sum_{s\in G_n(\hat{y})}C(s,\hat{y})} $$

note that we have replaced $\min (C(s,\hat{y}), C(s,y))$ with $\min (C(s,\hat{y}), \max_{y\in S_i}C(s,y))$ since there are multiple reference translation, we use the most similar one. So this is to say: “There are multiple answer, go to choose the best one and compute the score.”

BP (Brevity Penalty)

Candidate translations longer than their references are already penalized by the modified $n$-gram precision measure. Now, to penalize those translations shorter than the reference translations, we need add an penalty term. This is when brevity penalty comes out:

$$ \mathrm{BP}=\begin{cases}1&\text{ if }c > r\\ \exp(1-\frac{r}{c})&\text{ if }c \leq r \end{cases} $$

where

$c$ is the number of words or tokens of the candidate corpus. That is,

$$ c = \sum_{i=1}^M\mathrm{length}(\hat{y}^i) $$

$r$ is the number of words or tokens of the effective reference corpus length, where the effective reference is defined as the reference translation whose length is as close to the corresponding candidate translation as possible. That is

$$ r = \sum_{i=1}^M\mathrm{length}(y^{i,j}),\text{ where } y^{i,j}=\arg\min_{y\in S_i}|\mathrm{length}(\hat{y}^i)-\mathrm{length}(y)| $$

with this penalty term, we wish the model to output the translations with the same length as the reference translations.

Weight

The weight measures the importance of different precision score, in the original paper, the uniform weights are adopted, that is

$$ w_i = \frac{1}{N}, \ \text{ for } i=1,\dots,N $$

Final definition

The final definition of the BLEU is given by

$$ \mathrm{BLEU}_ {w}(\hat{S}, S) = \mathrm{BP} \cdot \exp\left(\sum_{n=1}^\infty w_n\log p_n(\hat{S}, S) \right) $$

usually, the upper-bound of the above summation can be reduced to $\max_{i=1,\dots,M}\mathrm{length}(\hat{y}^i)$.

Analysis

Disadvantages of BLEU:

BLEU compares overlap in tokens from the predictions and references, instead of comparing meaning. This can lead to discrepancies between BLEU scores and human ratings.
BLEU scores are not comparable across different datasets, nor are they comparable across different languages.
BLEU scores can vary greatly depending on which parameters are used to generate the scores, especially when different tokenization and normalization techniques are used.
BLEU ignores synonym or similar expression, which causes refuses of reasonable translation.
BLEU is affected by common words.

Python Implementation

 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


import math
from typing import Set, List


def compute_n_gram_set(s: str, n: int) -> Set[str]:
 return set(s[i : i + n] for i in range(len(s) - n))


def modified_n_gram_precision(S_hat: List[str], S: List[List[str]], n: int) -> float:
 if len(S_hat) < n:
 return 0.0
 numerator: int = 0
 denominator: int = 0

 for index, y_hat in enumerate(S_hat):
 n_gram_set = compute_n_gram_set(y_hat, n)
 if not n_gram_set:
 continue
 # print(n_gram_set)
 for n_gram in n_gram_set:
 candidate_substr_count = y_hat.count(n_gram)
 best_ref_substr_count, best_ref_len = max(
 [(y.count(n_gram), len(y)) for y in S[index]], key=lambda x: x[0]
 )
 numerator += min(candidate_substr_count, best_ref_substr_count)
 denominator += candidate_substr_count

 return numerator / denominator


def brevity_penalty(S_hat: List[str], S: List[List[str]]) -> float:
 r: int = 0
 c: int = 0
 for index, y_hat in enumerate(S_hat):
 c += len(y_hat)
 best_match_ref = min(S[index], key=lambda x: abs(len(x) - len(y_hat)))
 r += len(best_match_ref)
 return 1.0 if c > r else math.exp(1 - r / c)


def compute_bleu_score(S_hat: List[str], S: List[List[str]]) -> float:
 assert len(S_hat) == len(S)

 # take N as a sufficiently large integer
 # N = max(len(y_hat) for y_hat in S_hat)
 N = sum(len(y_hat) for y_hat in S_hat)

 bp = brevity_penalty(S_hat, S)
 precisions = [modified_n_gram_precision(S_hat, S, n) for n in range(1, N + 1)]

 # bleu_score = bp * exp(p_n)
 return bp * math.exp(
 sum(math.log(precision) for precision in precisions if precision != 0)
 )

Reference

Notes on Llama3

Mon, 22 Apr 2024 16:22:19 +0800

Meta released Llama3 at April 18, which is evaluated on several benchmarks and achieves the SOTA on open-sourced LLMs.

Introduction

Instruct model performance

The performance of Llama3 8B compared with Gemma and Mistral:

Model	Llama3 8B	Gemma 7B - It	Mistral &B Instruct
MMLU (5 shot)	68.4	53.3	58.4
GPQA (0 shot)	34.2	21.4	26.3
HumanEval (0 shot)	62.2	30.5	36.6
GSM-8K (8 shot, CoT)	79.6	30.6	39.9
MATH (4 shot, CoT)	30.0	12.2	11.0

performance of Llama3 70B compared with Gemini Pro 1.5 and Claude Sonnet:

Model	Llama3 70B	Gemini Pro 1.5 (Published)	Claude 3 Sonnet (Published)
MMLU (5 shot)	82.0	81.9	79.0
GPQA (0 shot)	39.5	41.5 (CoT)	38.5 (CoT)
HumanEval (0 shot)	81.7	71.9	73.0
GSM-8K (8 shot, CoT)	93.0	91.7 (11 shot)	92.3 (0 shot)
MATH (4 shot, CoT)	50.4	58.5 (Minerva prompt)	40.5

Pre-trained model performance

The performance of Llama3 8B compared with Gemma and Mistral:

Model	Llama3 8B	Gemma 7B (Published)	Gemma 7B (Measured)	Mistral 7B (Published)	Mistral 7B (Measured)
MMLU (5 shot)	66.6	64.3	64.4	62.5	63.9
AGIEval English (3-5 shot)	45.9	41.7	44.9	-	44.0
BIG-Bench Hard (3 shot, CoT)	61.1	55.1	59.0	-	56.0
ARC-Challenge (25 shot)	78.6	53.2(0 shot)	79.1	78.1	78.7
DROP (3 shot, F1)	58.4	-	56.3	-	54.4

performance of Llama3 70B compared with Gemini Pro 1.5 and Claude Sonnet:

Model	Llama3 70B	Gemini Pro 1.0 (Published)	Mixtral 8 $\times$ 22B (Measured)
MMLU (5-shot)	79.5	71.8	77.7
AGIEval English (3-5 shot)	63.0	-	61.2
BIG-Bench Hard (3 shot, CoT)	81.3	75.0	79.2
ARC-Challenge (25 shot)	93.0	-	90.7
DROP (3 shot, F1)	79.7	74.1 (variable shot)	77.6

Model Architecture

Several improvements are made on Llama3 compared to llama2:

Llama3 uses a tokenizer with a vocabulary of 128K tokens.
Llama3 adopts grouped query attention (GQA) across both the 8B and 70B sizes.
Llama3 uses to context window of size 8192 tokens

Traning

Llama3 uses 15T tokens for pre-training. Compares to Llama2, it is seven times larger and includes four times more code.

5% data of the training dataset are non-English to support multi-lingual use cases.

Data processing includes:

Heuristic filters
NSFW filters
Semantic deduplication approaches
Text classifiers to predict data quality. Llama2 is used to generate training data for the text classifiers.

Data mixing strategy is explored to improve the performance of Llama3.

Scaling up pretraining

Llama3 developed a series of scaling laws for downstream benchmark evaluations.

Scaling laws help:

Select an optimal data mix and to make informed decisions on how to best use training compute.
Scaling laws allow Llama3 to predict the performance of the largest models on key tasks without training the models.

The authors finds our that the performance of the model continues to improve log-linearly as the training tokens increase. It is seen that Larger models can match the performance of these smaller models with less training compute, but smaller models are generally preferred because they are much more efficient during inference.

The authors combine three types of parallelization:

Data parallelization
Model parallelization
Pipeline parallelization

Instruction fine-tuning

The fine-tuning of Llama3 contains:

Supervised fine-tuning
Rejection sampling
Proximal Policy Optimization
Direct Preference Optimization

Learning from perference rankings via PPO and DPO also greatly improved the performance of LLma3 on reasoning and coding tasks. Since perference ranking helps the model to select answer when it is in a dilemma.

Reference

Notes on RAG

Sun, 14 Apr 2024 12:38:04 +0800

Problems of LLM

Out of date knowledge: the model cannot gain knowledge after training
Humiliation: the model may generate nonsense output
Specific domain: the generalized model is difficult to adapt to specific domain
Enthetic problems: the model may encounter

Fine-tuning

Fine-tuning is used to improve performance of foundation model on specific tasks with the help with some supervised data

Fine-tuning methods can be classified into:

Based on range of updated parameters:
- Full Model fine-tuning: update the parameters of the whole model
- Partial fine-tuning: freeze the top layer; freeze the bottom layer
Based on special technology:
- Adapter tuning
- LoRA
- Continual Learning fine-tuning
Based on input:
- Instruction tuning
Based on objective
- Multi-task fine-tuning

Problems of fine-tuning:

Requires task-specific labeled data, may cause overfitting and catastrophic forgetting。
The generalization ability is limited, and fine-tuning are required when adapting to new tasks
The performance may be destroyed after fine-tuning, for example, safety.

RAG

RAG consists of three major processes of retrieval, augmentation, and generation. The framework of RAG in LLM can be described as follows:

Retrieval

Retriever type

Retrieval methods can be generally categorized into two types: sparse and dense, based on the information encoding methods.

sparse retrieval usually relies on inverted index matching along with the raw data input, for example TF-IDF and BM25. The limitation of sparse retrieval in RAG is
1. its no-training nature, which makes the retrieval performance heavily rely on the quality of database construction and query generation.
2. Moreover, such fixed term-based methods only support similarity retrieval, while cannot be adapted for other retrieval considerations demanding in LLM applications, such as the diversity
dense retrieval, on the contrary, embeds the query and documents into continuous vector space with certain criteria, for example, semantic similarity. Examples include BERT, Dense Passage Retriever (DPR), etc.

Retrieval Granularity

Retrieval granularity denotes the retrieval unit in which the corpus is indexed, e.g., document, passage, token, or other levels like entity.

Chunk retrieval.
Token retrieval.
Entity retrieval.

Pre-retrieval and Post-retrieval Enhancement

Pre-retrieval and post retrieval strategies can be added to improve the quality of the retriever.

Pre-retrieval methods include:

Query rewrite. This method aims to close the gaps between the input text and the needed knowledge in retrieval, to reformulate the original question into a more conducive version to retrieve.
Query augmentation. This method aims to combine the original query and the preliminary generated outputs as a new query, which is further used to retrieve relevant information from the external database

Post-retrieval enhancement denotes the procedure to process the extracted top-k documents from the retriever before feeding them to the generator for the sake of better alignment between the retrieval and generation stages.

Database

Wikipedia
Domain specific database
search engine

Generation

Parameter-Accessible Generators (White-box). Allow parameter optimization.
Parameter-Inaccessible Generators (Black-box). Focus more on retrieval and augmentation processes, trying to enhance the generator by augmenting the input with better knowledge, guidances or examples for the generation.

Augmentation

Input layer integration
Output layer integration
Intermediate layer integration

Retrieval Frequency

If it is necessary to retrieve? Self-RAG

retrieval frequency:

One-time.
Every-n-token
Every token

RAG training

Training Free
Independent training
Sequential training
Joint training

Advance RAG

Module RAG

Applications

NLP applications
- QA systems: REALM
- Chatbot:
- Fact Verification: self-RAG
Downstream tasks:
- Recommendations
- Software engineering
Domain-specific Applications
- AI for science
- Finance: ChatDOC

Limitations of RAG

Long Context Window

Advantages of Long Context Window:

Improve the understanding and relativity: long context window allows model to refer to more context information when generating answers.
Handling complex tasks: long context window makes handling complex tasks such as writing a long article, coding
Improve users’ experience: the user expects the model remember the chat history and use them to interact with the user.

Disadvantages of long context window:

Only uses context once, Requires refeeding the data to use long context window.
Cost expensive due to input price.
Time expensive due to limit of tokens per second.
Needle in HayStack experiment show that there are problems with long context window.

Advantages of RAG:

Privacy projection.
Allow chunking the texts and retrieve the related information more accurately
Adaptive to the size of data.
Accepts multiple type of data source (multimodality).
Only uses a small part of the total data, which is cheaper compared with long context window .

problems of RAG

The quality of retrieval
- The retrieved text cannot be aligned with the queried text.
- The queried text are not retrieved all.
- Redundancy or out-dated data may cause inaccuracy.
the quality of response generation
- Model Humiliation
- Irrelevance
- Organize the output to make it reasonable
- Depends on the external information

Futures:

Trustworthy RA-LLMs
Multi-lingual RA-LLMs
Multi-modal RA-LLMs
Quality of External Knowledge

Other technologies

Query transformations
Sentence window retrieval
Fusion retrieval/ hybrid search
multi-document agents

Reference

A Survey on RAG Meets LLMs: Towards Retrieval-Augmented Large Language Models

What's next for AI agentic workflows

Sun, 14 Apr 2024 12:38:04 +0800

In AI Ascent, Andrew Ng shared his opinion about the agentic workflow, specifically, his talk focused on designs and future works on agentic workflows.

What is Agent

LLM based agents can be classified into two categories: non-agentic and agentic.

Non-agentic workflow requires LLM do something in a sequence of steps, usually this requires the “zero-shot” learning ability of LLM.
Agentic workflow requires LLM do something step by step, and revise its action by thinking or reflection.

Andrew then posts an image shows the agentic workflow can improve the performance of LLM on HumanEval benchmark.

Agentic Design Reasoning Design Pattern

Agentic design pattern can be divided into following kinds. Where the first two design patterns are robust and the latter two are emerging.

Reflection

Self-Refine: Iterative Refinement with Self-Feedback
Reflexion: Language Agents with Verbal Reinforcement Learning

Tool Use

Gorilla: Large Language Model Connected with Massive APIs
MM-REACT: Prompting ChatGPT for Multimodal Reasoning and Action

Planning

Chain-of-Thought Prompting Elicits Reasoning in Large Language Models
HuggingGPT: Solving AI Tasks with ChatGPT and its Friends in Hugging Face

Multi-agent collaboration

Communicative Agents for Software Development
AutoGen: Enabling Next-Gen LLM Applications via Multi-Agent Conversation

Future Perspective

The set of tasks that AI can do expand dramatically because of agentic workflows.
We have to get used to delegating tasks to AI agent and wait patiently for response.
Fast token generation is important. Generating more tokens even from a a lower quality LLM can achieve a good result.

Reference

What’s next for AI agentic workflows ft. Andrew Ng of AI Fund

变量	含义
\(P\)	number of parameters
\(L\)	layers
\(V\)	vocabulary size
\(d\)	hidden size
\(d_{\text{ff}}\)	FFN hidden size
\(s\)	sequence length
\(b\)	batch size
\(h\)	number of attention heads
\(d_h\)	attention head dimension

Components	Training	Inference	Optimization
weights	\(2P\)	\(2P\)	quantization
optimizer states	\(4P\)	0	ZeRO, offloading
gradients	\(2P\)	0	ZeRO
activations	\(\sim 4Lhs^2b\)	\(\sim 0.4P\)	ckpt, offloading, flash attention
KV cache	0	\(4sLhd_h\)	KV optimization, attention
TOTAL	\(8P + 4Lhs^2b\)	\(2.4P + 4sLhd_h\)

LLM on Jiang Yi(姜祎)'s Homepage

Notes on Kimi-k2.5

Introduction

Method

Architecture

Data

Pre-training

Post-training

Agent Swarm

Infra

Experiments

Conclusion

References

Notes on Qwen3-Next

Introduction

Method

Architecture

Hybrid Attention

MoE

Normalization and Training

Experiments

Efficiency

Performance

Conclusion

References

megatron-lm

Introduction

Method

FFN

Attention

Embedding

Experiments

Implementation

Conclusion

References

Notes on Gated Attention

Introduction

Related Work

Method

Experiments

Analysis

Conclusion

References

Appendix

LLM Memory Computation

Motivation

Background

Transformer Architecture

Notation

Assumptions

Training Memory Analysis

Training Memory Components

Optimizer States

Activation

Activation — Attention

Activation — FFN & LayerNorm

Activation — Output

Activation — Total

Total Training Memory

Experiments

Case Study

Inference Memory Analysis

Inference Components

KV Cache Mechanism

KV Cache Memory

Total Inference Memory

Dynamic vs. Static

Optimization

Overview

Mixed Precision Training

ZeRO

Model Parallelism

Activation Checkpointing

Flash Attention

KV Cache Optimization

Weight Quantization

Activation Offloading

Real Systems

Training Frameworks

Inference Frameworks