Oral on Jiang Yi(姜祎)'s Homepage

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 机制并不能提高模型的表现

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

Oral on Jiang Yi(姜祎)'s Homepage

Notes on Gated Attention

Introduction

Related Work

Method

Experiments

Analysis

Conclusion

References

Appendix

Notes on DPO

Introduction

Background

DPO

Theoretical Analysis of DPO

Experiments

Conclusion

References

Notes on olmoe

Introduction

Method

Pretraining

Post-training

Ablation Study

MoE Settings

MoE vs. Dense

Expert Granularity

Shared Experts

Expert Choice vs. Token Choice

Sparse Upcycling

Load Balancing Loss

Router Z-loss

General Pre-training Settings

Initialization

QK-Norm

AdamW Epsilon

Adaptation Settings

Load Balancing Precision

MoE Analysis

Router Saturation

Expert Co-activation

Domain Specialization

Vocabulary Specialization

Conclusion

References