[{"content":"我们从DeepSeek V4的技术报告出发，以此为蓝本来查漏补缺，并进一步建立相关知识储备。万丈高楼，始于平地。\nSpeculative Decoding与MTP 在进入核心的模型创新点之前，我们首先回顾一下Speculative Decoding以及从DeepSeek V3一并沿用至V4的MTP策略。需要明确的是，MTP是一种 训练目标或模型结构，speculative decoding 是一种推理加速方法。在DeepSeek的模型中，MTP本质上是为了更好地支持speculative decoding，同时还能提升模型能力。\n对于一个标准的autoregressive LLM，在给定了$x_{\\lt t}$的上下文之后，一次forward只生成一个token，训练目标为$p(x_t|x_{\\lt t})$。这个策略所存在的问题是显然的：在每一次inference的时候，每一个token都会经历一个完整的forward过程。而为了提升效率，在训练时就引入MTP作为训练目标，也即每一个step都预测多个token：在DeepSeek-V3的架构中，除了主模型之外还引入了MTP modules专门用于预测主模型之后的token。\n而在推理时，一般情况下的speculative decoding可以理解为先用便宜模型或模块草稿式地猜几个 token，再让大模型一次性检查这些token是否正确。注意：speculative decoding的操作可以保证其token的概率分布和朴素策略一致，推导如下\n假设目标模型，也就是原本的大模型，其条件分布为$p(x_t \\mid x_{\\lt t})$；而 draft model 或 MTP module 给出的近似分布为$q(x_t \\mid x_{\\lt t})$，在普通 decoding 中，每一步都直接从$p$中采样：$x_t \\sim p(\\cdot \\mid x_{\\lt t})$；因此，生成$K$个 token 需要进行$K$次目标模型 forward：$x_t, x_{t+1}, \\ldots, x_{t+K-1}$. 而在 speculative decoding 中，我们首先用较便宜的draft model一次性生成一段候选 token：$\\hat{x}_t, \\hat{x}_{t+1}, \\ldots, \\hat{x}_{t+K-1} \\sim q$，随后，目标模型$p$并不是逐个重新生成这些token，而是对整段draft token进行一次并行验证。对于第$i$个draft token $\\hat{x}_{t+i}$，目标模型会计算$p(\\hat{x}_{t+i} \\mid x_{\\lt t}, \\hat{x}_t, \\ldots, \\hat{x}_{t+i-1}),$并将其与 draft model 的概率$q(\\hat{x}_{t+i} \\mid x_{\\lt t}, \\hat{x}_t, \\ldots, \\hat{x}_{t+i-1})$进行比较。接受该 token 的概率通常写作：\n$$ \\alpha_i = \\min\\left( 1, \\frac{ p(\\hat{x}_{t+i} \\mid x_{\\lt t}, \\hat{x}_{t:t+i-1}) }{ q(\\hat{x}_{t+i} \\mid x_{\\lt t}, \\hat{x}_{t:t+i-1}) } \\right). $$ 直觉上，如果draft model给出的token在目标模型看来也足够合理，即 $p/q$ 较大，那么该token就会被接受；如果draft model过于自信地提出了一个目标模型并不认可的token，那么该token就更可能被拒绝。一个关键点是，标准speculative decoding并不是简单地“猜对就用，猜错就丢”，而是通过接受-拒绝采样保证最终生成分布仍然与直接从目标模型$p$解码一致。为了看清这一点，可以先考虑单个token的情况。draft model先采样：$\\hat{x} \\sim q(\\cdot)$，若 $\\hat{x}$ 被接受，则输出 $\\hat{x}$。此时某个 token $x$ 被接受并输出的概率为： $$ q(x) \\cdot \\min\\left(1, \\frac{p(x)}{q(x)}\\right) = \\min(p(x), q(x)). $$ 如果draft token被拒绝，则不能直接重新从$p$中采样，否则会破坏分布一致性。正确做法是从residual distribution中采样： $$ p_{\\text{res}}(x) = \\frac{(p(x)-q(x))_+}{ \\sum_{x'} (p(x')-q(x'))_+ }, $$ 其中$(a)_+ = \\max(a, 0)$；因此，一个token $x$最终被输出的总概率为： $$ \\underbrace{\\min(p(x), q(x))}_{\\text{accepted from draft}} + \\underbrace{(p(x)-q(x))_+}_{\\text{sampled from residual}} = p(x). $$ 这说明，只要接受-拒绝和residual sampling的过程设计正确，speculative decoding的最终输出分布就和直接从目标模型$p$中采样保持一致。它改变的是计算方式，而不是目标分布。 回到DeepSeek的MTP设计，其作用正是在这里体现出来。MTP module相当于为主模型提供了一个轻量级、与主模型高度对齐的internal drafter。由于MTP module在训练阶段就被要求预测未来token，并且与主模型共享部分表示或输出空间，它给出的draft token往往比一个外部小模型更加贴近主模型分布。\nDeepSeek MoE 接下来是MoE的部分，我们一步步回顾一下目前已经成为基模主流的MoE设计的基本原理以及经典的优化与改进策略。我们知道，一个标准的Transformer层包含了attention和FFN的部分，而FFN往往占据了大量的参数和计算，因此MoE就希望不让每一个token都经过一个巨大的FFN，而是准备多个expert FFN，并让router为每一个token都选择少数几个expert。具体而言，一个标准的FFN可以写成\n$$ \\text{FFN}(h)=W_2\\sigma(W_1h) $$ 而MoE层把一个FFN替换成多个expert：$E_1,\\dots, E_N$，并通过router产生每一个expert的权重$g_i(h)=\\text{Softmax}(W_rh)$，实际计算时只选择top-$K$的结果，记为$\\mathcal{T}(h)$（注意最终的hidden feature还需要加上residual connection的结果，在这里我们只考虑FFN转换成MoE的逻辑）： $$ \\text{MoE}(h)=\\sum_{i\\in \\mathcal{T}(h)}g_i(h)E_i(h) $$ 但MoE往往会遇到以下类型的问题，严重影响训练的效率以及性能： expert collapse：router总是偏向少数expert，导致部分expert过载，部分expert几乎不被使用。\nexpert redundancy：多个expert 学到相似知识，没有形成足够清晰的分工。\nrouting instability：训练早期router不稳定，导致token-expert分配波动较大。\ncommunication overhead：在分布式训练中，不同token被送到不同expert，通常需要all-to-all communication，这会带来额外系统开销。\n所以MoE的核心难点不是“把 FFN 拆成很多expert”这么简单，而是如何让expert既被均衡使用，又真正形成specialization。DeepSeek MoE针对这个需求提出了以下两点改进方案：\nFine-Grained Expert Segmentation：本质上就是把较大的expert拆成较小的expert，例如原来是8个专家选2个，那么拆分之后可能变成64个专家用16个，这自然也就意味着模型可以为不同的token构造更加细粒度的expert组合，从而增强specialization Shared Expert Isolation: 虽然上述操作能让不同的expert学到更加细粒度的知识，但有些通用知识实际上所有的expert都会需要，如果不加处理的话就可能会导致不同expert有部分参数需要处理类似的知识。因此，DeepSeekMoE更进一步引入了$K_s$个shared expert（我们记这一集合为$\\mathcal{S}$），不经过router选择而对所有token都激活 最终得到的MoE层可以写成\n$$ \\text{DeepSeekMoE}(h)=\\sum_{i\\in \\mathcal{S}}E_i^{\\text{shared}}(h)+\\sum_{j\\in \\mathcal{T}(h)}g_j(h)E_j^{\\text{routed}}(h) $$ 在负载均衡的设计逻辑上，DeepSeekMoE引入了Expert-Level balance loss和Device-Level balance loss来防止MoE坍缩到特定的几个expert或者设备上。而对于DeepSeek V4，除了基础细节修改之外基本沿用了DeepSeekMoE的设计思想与架构。 Manifold-Constrained Hyper-Connection 接下来就是mHC这一个核心的贡献，我们从Transformer中使用的最基本的residual connection出发，一步步推导mHC的来历和原理。对于一个普通的Transformer层，可以抽象地写成\n$$ x_{l+1}=x_l+F_l(x_l) $$ 残差连接之所以在深层网络中非常有效，一个重要原因在于它保留了一条identity path，使得浅层特征可以较为直接地传递到深层，同时也为反向传播提供了一条直接的梯度通路，从而缓解梯度消失或梯度爆炸等训练不稳定问题。 这也可以帮助理解为什么当前许多大模型训练更倾向于采用 Pre-Norm 结构，即\n$$ x_{l+1}=x_l+F_l(\\text{Norm}(x_l)) $$ 如果采用Post-Norm，那么我们实际上有 $$ x_{l+1}=\\text{Norm}(x_l+F_l(x_l)) $$ 这实际上就破坏了$x_l$到$x_{l+1}$的无损传播路径，归一化参数更新时的梯度变化有可能导致梯度消失或爆炸 我们把Identity Mapping逐层展开，得到的结果如下所示，可以发现最初的信号$x_0$能够一直传递到深层。\n$$ x_L=x_0+\\sum_{l=0}^{L-1}F_l(x_l) $$ 然而，普通residual的结构其实很死板，包含了两个固定的假设：残差项的权重永远是1，且层输出$F_l(x_l)$也直接加入进去，也就是说，它无法学习形如$x_{l+1}=a_lx_l+b_lF_l(x_l)$这样的结构，更不能学习$l+1$层对于$l-2,\\dots, l-k$层的直接依赖关系。在Hyper-Connection的论文中，也指出Pre-Norm在传统的residual connection的设定下很容易导致深层相似的问题。为了解决这个挑战，HC把一条residual stream扩展为了$n$条。需要明确的是，residual stream依然是对应于每一个layer，相当于每一个layer同时维护多条残差流；并不是$x_{l+1}$就直接依赖$x_{l}, \\dots, x_{l-n+1}$，具体而言，我们设多条residual stream为 为了解决这个挑战，HC把一条 residual stream 扩展为了$n$条。具体来说，我们记第$l$层开始前的多条residual streams为\n$$ X_l= \\begin{bmatrix} x_l^{(1)}\\\\ x_l^{(2)}\\\\ \\vdots\\\\ x_l^{(n)} \\end{bmatrix} \\in \\mathbb{R}^{n\\times d}, $$ 其中$x_l^{(i)}\\in\\mathbb{R}^d$表示第$i$条 residual stream。普通Transformer block $F_l$本身仍然只接收一个$d$-维输入，因此HC首先需要从这$n$条stream 中读出当前层真正要处理的输入。我们记这一读入权重为$\\alpha_l\\in\\mathbb{R}^{n}$，于是当前层的输入可以写成$\\alpha_l^\\top X_l$，同时记 $$ y_l=F_l(\\alpha_l^\\top X_l)=\\sum_{i=1}^{n}\\alpha_{l,i}x_l^{(i)}. $$ 接下来，HC 还需要决定如何把当前层输出$y_l$写回到$n$条 residual streams 中。我们记写回权重为$\\beta_l\\in\\mathbb{R}^{n}$，并引入一个 residual stream 之间的混合矩阵$R_l\\in\\mathbb{R}^{n\\times n}$，于是HC的整体更新可以写成 $$ X_{l+1}=R_lX_l+\\beta_l y_l. $$ 这里需要注意，$\\beta_l y_l$表示把同一个层输出$y_l$按照不同权重写入不同的stream，即 $$ \\beta_l y_l = \\begin{bmatrix} \\beta_{l,1}y_l\\\\ \\beta_{l,2}y_l\\\\ \\vdots\\\\ \\beta_{l,n}y_l \\end{bmatrix}. $$ 将上式按stream展开，可以得到 $$ x_{l+1}^{(i)} = \\sum_{j=1}^{n}R_{l,ij}x_l^{(j)} + \\beta_{l,i} F_l \\left( \\sum_{j=1}^{n}\\alpha_{l,j}x_l^{(j)} \\right). $$ 从这个形式可以看出，HC 相比普通的残差连接多出了三类可学习的连接关系：第一，$\\alpha_l$决定当前层从哪些residual stream中读取信息；第二，$\\beta_l$决定当前层输出被写回到哪些residual streams；第三，$R_l$ 决定不同residual stream之间如何混合和传递。 另一种理解HC的方式是回忆LSTM的结构：LSTM中$h_{t}$其实同时依赖当前的输入$x_{t-1}$以及$h_{t-1}$；原始的$h_{t-1}$只是一个$d$维的向量，而HC的思想相当于$h_{t-1}$变成了一个$d\\times n$的矩阵，其中$n$为residual条数。普通的residual connection自然就可以看成是HC的一个特例\n接下来我们可以证明，HC在提供了相比于普通残差连接更加丰富通路的同时，却也重新引入了梯度消失和梯度爆炸的问题：如果我们HC的更新函数进行展开，会得到\n$$ X_L=R_{L-1}R_{L-2}\\dots R_0X_0 + \\sum_{l=0}^{L-1}(R_{L-1}\\dots R_{l+1})\\beta_lF_l(\\alpha_l^\\top X_l). $$ 我们不难发现，对比普通残差连接的展开式，HC的展开式中能发现早起层的信息不再只能沿着一条固定的identity path向后传播，而是可以通过$R_l$在不同residual stream之间不断混合，并在后续层中被不同的$\\alpha_l$重新读取。但如果$R_l$是一个完全不受约束的学习矩阵，那么多层复合矩阵$R_{L-1}R_{L-2}\\dots R_0$就可能在深层网络中产生不稳定的放大与衰减，普通residual connection中恒等映射所保证的稳定性就会被破坏。 因此，DeepSeek-V4中所引入的mHC正是为了解决这一问题，它通过约束residual stream mixing matrix中$R_l$的取值，使其满足$R_l\\in\\mathcal{B}_n$，其中$\\mathcal{B}_n$被称为Birkhoff Polytope，为所有$n\\times n$双随机矩阵的集合：\n$$ \\mathcal{B}_n=\\{R\\in\\mathbb{R}_+^{n\\times n}|R\\mathbf{1}=\\mathbf{1},\\quad R^\\top\\mathbf{1}=\\mathbf{1}\\} $$ 这种情况下，我们能够证明$R_{L-1}R_{L-2}\\dots R_0\\sim \\mathbf{I}$，类似residual connection保留一条稳定的通路；在这里，我们依然可以利用简单的线性代数来推导一下： 假设我们只考虑skip path：\n$$ X_{l+1}^{\\text{skip}}=R_lX_l $$ 我们定义所有residual streams的平均表示为 $$ \\tilde{x_l}=\\frac{1}{n}\\mathbf{1}^\\top X_l. $$ 那么我们有 $$ \\tilde{x}_{l+1}=\\frac{1}{n}\\mathbf{1}^\\top X_{l+1}^{\\text{skip}}=\\frac{1}{n}\\mathbf{1}^\\top R_lX_l $$ 由于$R_l\\in\\mathcal{B}_n$，因而我们有 $$ \\tilde{x}_{l+1}=\\frac{1}{n}\\mathbf{1}^\\top R_lX_l=\\frac{1}{n}\\mathbf{1}^\\top X_l=\\tilde{x}_l $$ 这说明，在mHC中，residual stream mixing并不会改变多条streams的平均表示，只是对于信息的重新分配。进一步地，我们可以证明对于任意深度的复合residual mapping，$A_{0\\rightarrow L}=R_{L-1}R_{L-2}\\dots R_0$，依然有$A_{0\\rightarrow L}\\in\\mathbb{B}_n$ 在实际操作时，mHC通常不会直接学习一个已经满足双随机约束的矩阵，而是先学习一个uncontrained matrix，然后再进行投影和归一化操作来进行近似，使其符合双随机矩阵的性质。\nCompressed Sparse Attention (CSA) 接下来我们分析DeepSeek-V4中另一个核心结构：Hybrid Attention with CSA and HCA。对于标准的attention，当序列长度为$n$时，每一个query token都需要和前面所有KV tokens做attention，因此整体复杂度随序列长度近似呈二次增长。为了解决这一问题，DeepSeek-V4 引入了Compressed Sparse Attention以及Heavily Compressed Attention，也即CSA和HCA。在本章节中，我们主要讨论CSA的思路：CSA同时结合了KV压缩和稀疏选择这两种思想：它首先把每$m$个token的KV cache压缩成一个compressed KV entry，然后再通过DeepSeek Sparse Attention的方式为每个query选择top-$k$ 个compressed KV entries参与核心attention计算。\n我们先来看KV压缩的部分，这一块可以通俗理解为$m$个token的信息融合为一个token的pooling过程，但不同于最naive的average pooling，这里我们希望模型自己去判断不同token之间的重要性以及融合的策略。由于这一块之前接触的不是很多，因此我们先从high-level的角度对比一下传统的KV Cache以及CSA的思想\n在标准attention里，我们通常有$K=[k_1,\\dots, k_n]$，$V=[v_1,\\dots, v_n]$，其中$n$为token总数，同时我们有\n$$ o_t=\\text{softmax}(q_tK^\\top)V $$ 因此在标准的KV Cache中，会分别缓存$K_{1:t}$以及$V_{1:t}$来重复利用；但CSA作为一个KV Cache压缩的策略，并不是分别生成$K^{\\text{comp}}$以及$V^{\\text{comp}}$，而是生成一个$C^{\\text{comp}}\\in\\mathbb{R}^{\\frac{n}{m}\\times c}$，其中每一行$C_i^{\\text{comp}}\\in\\mathbb{R}^c$表示第$i$个压缩之后的KV entry。如果我们想要和标准的KV Cache对齐，我们可以粗略理解为$K_i^{\\text{comp}}=C_i^{\\text{comp}}$，$V_i^{\\text{comp}}=C_i^{\\text{comp}}$，也即$C^{\\text{comp}}$同时作为压缩之后的key matrix和value matrix。 总结而言，对于传统的KV Cache，我们是考虑每一层的$K$以及$V$；但在CSA中，我们是直接利用$H$来计算$C^{\\text{comp}}$。我们将在之后的推导中，逐步展开这一套混合注意力机制的完整流程\n设其输入的hidden states为$H\\in\\mathbb{R}^{n\\times d}$，其中$n$为序列长度，$d$为hidden size。CSA不直接为每一个token都保存完整的key和value，而是先通过线性映射得到两组候选KV Entries：\n$$ C^a=HW^{aKV}, \\quad C^b=HW^{bKV} $$ 其中$C^a,C^b\\in \\mathbb{R}^{n\\times c}$，这里的$c$就可以理解为compressed KV entry的维度，实际上先做了一层降维的操作。与此同时，模型还会生成两组compression weight logits： $$ Z^a=HW^{aZ}, \\quad Z^b=HW^{bZ} $$ 直观上而言，$C^a, C^b$提供要被压缩的内容；而$Z^a,Z^b$则决定每一个token在压缩时应该占有多大的权重。进一步地，为了让这种压缩操作不仅依赖token本身的hidden state，还能感知token在当前压缩窗口中的相对位置，CSA还引入了两组learnable positional bias，记为$B^a, B^b\\in \\mathbb{R}^{m\\times c}$，其中$m$为压缩窗口的大小。也即对于一个长度为$m$的token block，这个learnable positional bias可以让模型区分block内部不同位置的token。 接下来我们看$i$个压缩之后的KV entry是如何得到的，它对应的主窗口为$\\{mi,mi+1,\\dots, m(i+1)-1\\}$。一个朴素的想法可能是直接把这$m$个token的候选KV entries聚合成一个向量，但CSA的设计逻辑稍微复杂一些：它同时使用两组候选entries $C^a$和$C^b$，并在生成第$i$个compressed entry时，使用$C_{mi:m(i+1)-1}^a$以及$C_{m(i-1):mi-1}^b$，分别负责当前block以及前一个block。对应的compression logits也就变成了$Z_{mi:m(i+1)-1}^a+B^a$以及$Z_{m(i-1):mi-1}^b+B^b$。将这两者拼接起来，可以得到一个大小为$2m\\times c$的压缩打分矩阵，然后CSA对这个矩阵在token维度上做softmax操作，就得到了对应的压缩权重\n$$ \\left[ S^a_{mi:m(i+1)-1}; S^b_{m(i-1):mi-1} \\right] = \\mathrm{Softmax}_{\\mathrm{token}} \\left( \\left[ Z^a_{mi:m(i+1)-1}+B^a; Z^b_{m(i-1):mi-1}+B^b \\right] \\right). $$ 需要特别注意的是这个地方softmax函数的含义，并不是给每个token一个单独的标量权重，而是对于每一个feature channel上都单独生成一组token-level的权重，也就意味着不同的feature channel可以关注不同的token。在得到了压缩的权重之后，第$i$​个压缩后的KV entry的结果就可以通过加权求和来得到： $$ C_i^{\\mathrm{comp}} = \\sum_{j=mi}^{m(i+1)-1} S_j^a\\odot C_j^a + \\sum_{j=m(i-1)}^{mi-1} S_j^b\\odot C_j^b. $$ 最终得到的$C_i^{\\text{comp}}\\in \\mathbb{R}^c$。此时我们也就能理解，$C^a,C^b$相当于提供了两套不同的压缩视角，分别从当前block以及前一个block中获取信息，从而减轻block boundary所带来的信息割裂问题。最后需要注意的是，标准意义上的KV Cache是一个纯inference优化的策略，但是CSA的KV压缩是架构上的创新，在训练时也会涉及。 要特别注意的一点是，对于CSA而言，attention中KV侧的组织方式已经和原始的Transformer有了明显的区别。在标准的Transformer中，每一个token都会对应一组token-level的$k_j$以及$v_j$，当前query通常需要在所有可见历史token的$K,V$上做attention。而现在这里我们得到的$C^{\\text{Comp}}$可以理解为基于CSA层的输入$H$所得到的一个memory representation，而不是主干网络中继续向后传播的hidden feature，它的作用是作为后续attention计算中被query读取的压缩KV memory。\n接下来，我们将进一步考虑attention的具体计算过程。假设原始序列长度为$n$，CSA经过压缩之后得到的compressed KV entries的数量大约是$n/m$，在超长上下文的context之下依然可能非常大。如果每一个query在计算attention的过程中依然需要dense attend（也即query和每一个compressed entry都要交互），那么attention的计算量依然很高。因此，CSA在得到了$C^{\\text{comp}}$之后还会引入一个轻量级的sparse selection模块，也即论文中的Lightning Indexer，用来为每一个query token动态选择最相关的top-$k$个compressed KV entries\n那么如何设计这个indexer的逻辑呢？我们首先需要明确这个indexer的输入：原始给到attention层输入自然是$H=[h_1,\\dots, h_n]$，但是既然此处的目标是选择对应的query token和哪些Compressed Token来计算attention，因此这个indexer的输入应该也是Compressed之后的KV Entry。但是，这里并没有直接使用$C_i^{\\text{comp}}$，而是使用刚刚介绍的一整套压缩策略来构造出另一组Compressed Token $K^{I,\\text{Comp}}\\in\\mathbb{R}^{n/m \\times c_I}$， 其中$c_I\\ne c$为indexer的head dimension，一般$c_I\\lt c$（Index中每一个token的特征维度会更小，这一设计也相对符合直觉）。也就是说，$K^{I,\\mathrm{Comp}}$ 只用于计算index score，负责判断“哪些 compressed block 值得被选中”；而真正进入后续 core attention、作为key/value被读取的，仍然是对应位置上的$C^{\\mathrm{Comp}}$\n在得到了用于indexing的KV indexer $K^{I,\\text{Comp}}$之后，接下来就需要考虑query token的计算了。这里同样要注意区分的是，相比于原始Transformer中的$Q=W_QX$的基本计算，在这里用于选择到底哪些KV Cache比较重要的query并不是原始值，而也是专门的index query，计算方式如下：\n$$ c_t^Q = h_t\\cdot W^{DQ} $$ $$ \\left[q_{t,1}^I,q_{t,2}^I,\\dots, q_{t,n_h^I}^I\\right]=q_t^I=c_t^Q\\cdot W^{IUQ} $$ 其中$h_t\\in \\mathbb{R}^d$是query token所对应的hidden state，$c_t^Q\\in \\mathbb{R}^{d_c}$是query压缩之后得到的结果，$d_c\\lt d$为query压缩维度；$n_h^I$则定义了用于索引的query一共有几个head；$W^{DQ}\\in\\mathbb{R}^{d\\times d_c}$以及$W_{IUQ}\\in\\mathbb{R}^{d_c\\times c^In_h^I}$为索引query的投影矩阵，注意到最终投影得到的维度是$c^In_h^I$，也即$n_h^I$个head，每一个head所对应的维度$c^I$与前面计算得到的$K^{I,\\text{Comp}}$是对齐的。这样，我们也就可以计算query token $t$以及第$s$个compressed block之间的index score $I_{t,s}\\in\\mathbb{R}$，其实就是$n_h^I$个head所对应的$q_{t,r}^I$分别与前面专门用于index的第$s$个索引$K_s^{I,\\text{Comp}}$计算点积（$s\\lt\\text{Floor}(\\frac{t}{m})$，相当于选取的是query token的$t$所对应的前一个完整的block），外部再加一层可学习的权重即可，具体而言有：\n$$ [w_{t,1}^I;w_{t,2}^I;\\dots;w_{t,n_h^I}^I]:=w_t^I=h_t\\cdot W^w $$ $$ I_{t,s}=\\sum_{r=1}^{n_h^I}w_{t,r}^I\\cdot \\text{ReLU}\\left(q_{t,r}^I\\cdot K_s^{I,\\text{Comp}}\\right). $$ 其中$W^w\\in\\mathbb{R}^{d\\times n_h^I}$为一个可学矩阵，$w_{t,h}^I\\in\\mathbb{R}$也即表示第$h$个索引头的权重。在上述推导中，需要注意$q_{t,r}^I\\in \\mathbb{R}^{c_I}$，$K_s^{I,\\text{Comp}}\\in \\mathbb{R}^{c_I}$；在得到了index query $t$与compressed block $s$之间的分数之后，对于第$t$个token而言就只会保留top-$k$个compressed block来用于计算：\n$$ C_t^{\\text{SprsComp}}=\\left\\{C_s^{\\text{Comp}}\\Big|I_{t,s}\\in \\text{Top-k}(I_{t,:})\\right\\}. $$ 这样，我们也就完成了对于每一个token各自需要关注的KV Entry的选择。CSA之后计算attention的方式就类似于MQA的思想（用共同的一组$K,V$来处理不同head的query），并更进一步地设置$K=V=C_t^{\\text{SprsComp}}$。对于一个query token $t$，现在我们才从latent representation $c_t^Q$中计算最终传统意义上的Query值 $$ [q_{t,1};q_{t,2};\\dots;q_{t,n_h}]:=q_t=c_t^Q\\cdot W_{UQ} $$ 其中$n_h$为query head的个数，$W^{UQ}\\in\\mathbb{R}^{d_c\\times cn_h}$为上采样矩阵。我们注意到$c_t^Q$在这里同时用于query $q_t$以及索引query $q_t^I$的计算。最后，输出的attention就定义为 $$ o_{t,i}=\\text{CoreAttn}(\\mathtt{query}=q_{t,i},\\mathtt{key}=C_t^{\\text{SprsComp}},\\mathtt{value}=C_t^{\\text{SprsComp}}) $$ 其中$o_{t,i}$表示第$i$个head在token $t$处的输出值；在MQA的最后，会涉及到多个head合并的output projection的操作。注意到$o_t\\in\\mathbb{R}^{cn_h}$，而希望重新投影回$d$，因此在投影的时候额外引入了一个trick：先把$n_h$个head的输入分成$g$个group，然后对于每一个group $o_{t,r}^G\\in \\mathbb{R}^{c\\frac{n_h}{g}}$，我们将其投影到$d_g$维度的中间层$o_{t,r}^{G'}\\in \\mathbb{R}^{d_g}$，其中$d_g\\lt c\\frac{n_h}{g}$，最终再把中间结果进一步投影到$\\hat{o}_t\\in \\mathbb{R}^d$。 Heavily Compressed Attention (HCA) 在这一章中，我们将讨论HCA：HCA的思想与CSA基本类似，但是使用更加激进的压缩比，同时删除了CSA中使用的前后block组合压缩以及sparse selection的思想。假设输入的hidden states为$H\\in\\mathbb{R}^{n\\times d}$，HCA先计算出待压缩的原始KV entry $C\\in\\mathbb{R}^{n\\times c}$以及对应的weight logit $Z\\in \\mathbb{R}^{n\\times c}$：\n$$ C = H\\cdot W^{KV} $$ $$ Z = H\\cdot W^Z $$ 其中$W^{KV}, W^Z\\in\\mathbb{R}^{d\\times c}$为可训练参数。类似CSA的思路，$C$中每$m'\\gg m$个token都会根据压缩权重以及新引入的learnable positional bias $B\\in\\mathbb{R}^{m'\\times c}$被压缩到一个token之中，得到$C^{\\text{Comp}}\\in \\mathbb{R}^{\\frac{n}{m'}\\times c}$，其中每一个压缩之后项$C_i^{\\text{Comp}}\\in \\mathbb{R}^c$如下计算：\n$$ S_{m'i:m'(i+1)-1}=\\text{Softmax}_{\\text{token}}(Z_{m'i:m'(i+1)-1}+B), $$ $$ C_i^{\\text{Comp}}=\\sum_{j=m'i}^{m'(i+1)-1}S_j\\odot C_j. $$ 通过这一操作，HCA就把KV表征的序列长度压缩到到原来的$\\frac{1}{m'}$；之后，HCA也类似CSA来进行KV共享的MQA以及output projection。在这里，由于没有了sparse selection的环节，因此query token会和所有满足因果约束的历史的compressed KV entries来计算attention（而不是CSA中的其中一个子集）。具体而言，对于一个给定的token $t$，我们有\n$$ c_t^Q = h_t\\cdot W^{DQ} $$ $$ [q_{t,1};q_{t,2};\\dots;q_{t,n_h}]:=q_t=c_t^Q\\cdot W^{UQ} $$ 其中$h_t\\in\\mathbb{R}^d$为token t的hidden state；$n_h$表示query head的数量；$W^{DQ}\\in\\mathbb{R}^{d\\times d_c}$以及$W^{UQ}\\in\\mathbb{R}^{d_c\\times cn_h}$为可学的投影矩阵。最终，我们的便可以得到输出：\n$$ o_{t,i}=\\text{CoreAttn}\\left(\\mathtt{query}=q_{t,i},\\mathtt{key}=C^{\\text{Comp}},\\mathtt{value}=C^{\\text{Comp}}\\right), $$ 最后projection也先采用分组降维再投影的方式，以减少计算量。 总结 在第一部分中，我们总结了DeepSeek-V4在模型结构层面的核心设计逻辑；而在之后的分析中，我们将进一步深入DeekSeek-V4的更多设计细节，例如Muon优化器的使用，推理框架以及预训练和后训练的框架。\n","permalink":"https://jubsteven.github.io/blog/posts/deepseek-report-1/","summary":"我们从DeepSeek V4的技术报告出发，以此为蓝本来查漏补缺，并进一步建立相关知识储备。万丈高楼，始于平地。","title":"DeepSeek-V4技术报告（一）"},{"content":"在前段时间基于VeRL的Search-R1仓库进行了一些Agentic Search相关的探索，也算是通过一些实践来积累了一点Agentic RL的基本常识。之前朋友推荐了几个帖子来讨论了pg_loss在训练开始时为0的分析，自己阅读之后发现过去对于Agentic RL的理论推导还有些生疏，借此机会重新梳理一下。\n单轮与多轮SFT 首先，我们先来回顾一下传统SFT的参数更新策略，我们从最基本的单轮对话开始，假设输入的prompt为$p:=[p_1, p_2,\\dots, p_m]$，对应的标准答案为$\\mathbf a^*:=[a_1^*,\\dots, a_n^*]$，而模型$\\theta$针对问题$p$输出的分布为$\\pi_\\theta(\\cdot \\mid p)$，此时SFT的目标就是最大化条件似然： $$ \\mathcal{L}_{\\text{SFT}}=-\\sum_{t=1}^n\\log \\pi_\\theta(a_t^*|p,a_{\\lt t}^*). $$ 注意在训练的时候模型并不会针对$p$直接输出一整串完整答案$\\mathbf{a}$，而是采用teacher forcing的方式逐token来进行监督，通过展开上述$\\mathcal{L}_{\\text{SFT}}$的公式即可说明。在这里，每一个输出的token都有对应的label予以监督，因此也不存在credit assignment的问题。在目前大多数实现中，对于一个样本会将其问题和回答拼接成一条回复$[p_1,\\dots, p_m, a_1^*,\\dots, a_n^*]$，然后用$p$预测$a_1$，用$(p,a_1^*)$预测$a_2$等等，在一个forward过程中得到$a_1^*,\\dots, a_n^*$每一个位置的logit，并将loss累加。\n上面我们讨论了单轮对话的情况，如果将其扩展为多轮，那么 SFT 的目标其实仍然没有本质变化：仍然是在teacher forcing下，对所有需要监督的assistant token做条件似然最大化；只是这里的条件前缀不再只是单个prompt $p$，而是截至当前轮为止的整个对话历史。设一个多轮对话一共有$K$轮，我们将第$i$轮用户的输入记为$u_i$，模型在这一轮的回复记为$\\mathbf a^{(i)*}$，那么整段对话可以写成一条轨迹$\\tau=(u_1,\\mathbf a^{(1)*},\\dots, u_K, \\mathbf a^{(K)*})$，而其中每一个assistant的回复都可以进一步写成一个token序列$\\mathbf a^{(i)*}=[a^{(i)*}_1,\\dots, a^{(i)*}_{n_i}]$。当模型在生成第$i$轮的回复时，它所条件化的上下文是此前的全部历史： $$ h_i:=[u_1,\\mathbf a^{(1)*},u_2,\\mathbf a^{(2)*},\\dots, u_i]. $$ 因此，第$i$轮第$t$个token的条件分布就可以写为$\\pi_\\theta(a_t^{(i)*}|h_i,a_{\\lt t}^{(i)*})$，目标函数可以理解为单轮SFT公式的直接推广： $$ \\mathcal{L}_{\\text{SFT}}=-\\sum_{i=1}^K\\sum_{t=1}^{n_i}\\log \\pi_\\theta(a^{(i)*}_t|h_i,a^{(i)*}_{\\lt t}). $$ 在训练的时候和单轮一样，不是按轮一轮轮独立生成再分别backward，而是通常把整段对话直接拼成一个长序列，一次forward得到所有位置的logits。需要注意的是在多轮SFT中，后续轮次虽然仍然条件化于此前完整对话历史，但该历史由数据集中的gold trajectory给定，而非由当前模型rollout得到；因此，轮次之间不存在由模型早先动作诱导出的状态转移依赖，也就不涉RL意义下对前序离散动作的credit assignment。\n单轮与多轮RL 在完成了对于SFT的基本理解之后，我们进入到RL的部分；由于过去一段时间的大多数实验使用了GRPO，因此后续分析也以此为例。首先，我们依然从单轮的情况开始考虑：依然给定输入prompt为$p=[p_1,p_2,\\dots,p_m]$，此时模型不会再进行teacher forcing，而是直接按照当前策略$\\pi_\\theta$自行采样得到完整回复。记一次采样得到的回复为$\\mathbf a:=[a_1,\\dots,a_n]\\sim \\pi_\\theta(\\cdot \\mid p)$，在RL的状态下模型在生成第$t$个token时所条件化的前缀$(p,a_{\\lt t})$本身就是由模型之前的采样结果构成的，而非SFT中的golden trajectory。对于单轮任务，环境通常只会在模型生成完整回复$\\mathbf a$之后给一个标量奖励$r(\\mathbf a,p)$。由此，我们希望优化的目标写为 $$ J(\\theta)=\\mathbb{E}_{\\mathbf a\\sim \\pi_\\theta(\\cdot\\mid p)}[r(\\mathbf a,p)]. $$ 进一步地，由于整条回复的生成概率可以按token进行分解，我们有 $$ \\pi_\\theta(\\mathbf{a}|p)=\\prod_{t=1}^n \\pi_\\theta(a_t|p,a_{\\lt t}). $$ 因此就有了 $$ \\log \\pi_\\theta(\\mathbf{a}|p)=\\sum_{t=1}^n\\log\\pi_\\theta(a_t|p,a_{\\lt t}). $$ 如果我们使用REINFORCE这个最经典的policy gradient的方法，可以将目标函数的梯度写为 $$ \\nabla_\\theta J(\\theta)=\\mathbb{E}_{\\mathbf{a}\\sim \\pi_\\theta(\\cdot|p)}[r(\\mathbf{a},p)\\nabla_\\theta \\log\\pi_\\theta(\\mathbf{a}|p)]=\\mathbb{E}_{\\mathbf{a}\\sim \\pi_\\theta(\\cdot|p)}\\left[r(\\mathbf{a},p)\\sum_{t=1}^n\\nabla_\\theta \\log\\pi_\\theta(a_t|p,a_{\\lt t})\\right]. $$ 注意到在这种情况下，虽然奖励只在整条回复完成之后才给出，但它最终会共同作用在整条回复的所有token上，作为一个序列级的reward signal，这与传统RL并无差异。当然，在具体实现时为了降低方差，我们会使用Advantage而非原始奖励作为信号。而GRPO则是抛弃了PPO中的value model，使用旧策略$\\pi_{\\theta_{\\text{old}}}$采样出一组回复$\\mathcal{A}=\\{\\mathbf{a}^{(1)},\\dots, \\mathbf{a}^{(G)}\\}$，并分别计算它们对应的奖励$r^{(g)}:=r(\\mathbf{a}^{(g)},p)$，并计算group-relative advantage $$ A^{(g)}=\\frac{r^{(g)}-\\text{mean}\\left(\\{r^{(j)}\\}_{j=1}^G\\right)}{\\text{std}\\left(\\{r^{(j)}\\}_{j=1}^G\\right)} $$ 最终的优化目标写为PPO形式的 $$ \\mathcal{L}_{\\text{GRPO}}=-\\frac{1}{G}\\sum_{g=1}^G\\frac{1}{n_g}\\sum_{t=1}^{n_g}\\left(\\min\\left(\\frac{\\pi_\\theta(a_t^{(g)}|p,a_{\\lt t}^{(g)})}{\\pi_{\\theta_{\\text{old}}}(a_t^{(g)}|p,a_{\\lt t}^{(g)})}A^{(g)}, \\text{clip}\\left(\\frac{\\pi_\\theta(a_t^{(g)}|p,a_{\\lt t}^{(g)})}{\\pi_{\\theta_{\\text{old}}}(a_t^{(g)}|p,a_{\\lt t}^{(g)})}, 1-\\epsilon ,1+\\epsilon\\right)A^{(g)}\\right)\\right) $$ 其中$n_g$表示第$g$个样本的回复长度，$\\epsilon$为clipping的系数。这个公式我们自然是见过很多次了，但是在这里希望进一步补充的是它在backward时的具体含义。对于一条固定采样得到的回复$\\mathbf{a}^{(g)}\\in\\mathcal{A}$，训练时并不是沿着rollout的过程逐步回传，而是会将整条sampled sequence固定下来，重新拼接成$[p_1,\\dots, p_m, a_1^{(g)},\\dots, a_n^{(g)}]$做一次并行forward，并在每一个assistant token所对应的位置上重新计算logprob。因而，对于单个位置$t$来说，其局部loss本质上依然是一个交叉熵项，只是target不再是SFT中的gold token，而是当前策略采样出的token带上advantage项作为权重。下面我们可以做一个简单的推导：如果我们设重要性采样$\\rho_t^{(g)}$，假设不考虑clipping激活的场景，则这条sequence第$t$个token的loss可以近似为 $$ \\ell_t^{(g)}=-\\frac{1}{Gn_g}\\rho_t^{(g)}A^{(g)}. $$ 对应位置的梯度便可以表示为 $$ \\begin{aligned} \\nabla_\\theta \\ell_t^{(g)}\u0026=-\\frac{1}{Gn_g}A^{(g)}\\nabla_\\theta\\rho_t^{(g)}\\\\ \u0026=-\\frac{1}{Gn_g}A^{(g)}\\nabla_\\theta\\left[\\frac{\\pi_\\theta(a_t^{(g)}|p,a_{\\lt t}^{(g)})}{\\pi_{\\theta_{\\text{old}}}(a_t^{(g)}|p,a_{\\lt t}^{(g)})}\\right]\\\\ \u0026=-\\frac{1}{Gn_g}A^{(g)}\\rho_t^{(g)}\\nabla_\\theta \\log \\pi_\\theta(a_t^{(g)}|p,a_{\\lt t}^{(g)}).\\\\ \\end{aligned} $$ 注意最后一步使用了RL中最常用的一个变换函数 $$ \\nabla_\\theta f(\\theta)=f(\\theta)\\nabla \\log f(\\theta) $$ 这说明对于单个token而言，它所接收到的梯度方向，本质上依然是提高或者压低当前token的logprob的方向。这个形式和普通的交叉熵函数是完全一致的，只是前面乘上了一个$\\rho_t^{(g)}A^{(g)}$的权重项。因此，从backward的角度看，AR policy model的训练逻辑其实并不神秘：虽然 reward 是在整条回复结束后才给出的，但在真正更新参数时，它会被转化为每一个token位置上的一个加权交叉熵梯度；模型仍然是在每个位置输出一个词表分布、计算对应logprob，再将这些位置上的梯度累加到同一套共享参数上。与SFT的根本区别并不在于backward的计算图结构，而在于这里的监督信号不再来自token-level的 gold label，而是来自sequence-level reward所诱导出的advantage。\n接下来我们也将其推广到多轮对话的场景，假设一轮完整的轨迹一共有$K$轮，第$i$轮的用户输入记为$u_i$，模型在该轮生成的回复记为$\\mathbf{a}^{(i)}=[a_1^{(i)},\\dots, a_{n_i}^{(i)}]$，则整条轨迹可以写为 $$ \\tau=(u_1,\\mathbf{a}^{(1)},u_2,\\mathbf{a}^{(2)},\\dots, u_K, \\mathbf{a}^{(K)}). $$ 和多轮SFT不同的是，这里的历史并不是数据集中给定的gold trajectory，而是由当前的policy逐轮rollout的sampled trajectory。我们设模型生成第$i$轮回复时，它所条件化的上下文为此前全部的交互历史： $$ h_i:=[u_1,\\mathbf{a}^{(1)},u_2,\\mathbf{a}^{(2)},\\dots, u_i]. $$ 由此，整条多轮轨迹的概率可以分解为如下形式，其中$n_i$为第$i$轮assistant回复的token数量： $$ \\pi_\\theta(\\tau)=\\prod_{i=1}^K\\pi_\\theta(\\mathbf{a}^{(i)}|h_i)=\\prod_{i=1}^K\\prod_{t=1}^{n_i}\\pi_\\theta(a_t^{(i)}|h_i, a_{\\lt t}^{(i)}) $$ 如果环境在整条轨迹结束之后给出一个最终奖励$r(\\tau)$，那么对应的优化目标也就可以写为 $$ J(\\theta)=\\mathbb{E}_{\\tau\\sim \\pi_\\theta}[r(\\tau)] $$ 如果使用REINFORCE算法，其梯度形式为 $$ \\nabla_\\theta J(\\theta)=\\mathbb{E}_{\\tau\\sim \\pi_\\theta}[r(\\tau)\\nabla_\\theta\\log\\pi(\\tau)]=\\mathbb{E}_{\\tau\\sim \\pi_\\theta}\\left[r(\\tau)\\sum_{i=1}^K\\sum_{t=1}^{n_i}\\nabla_\\theta \\log\\pi_\\theta(a_t^{(i)}|h_i,a_{\\lt t}^{(i)})\\right] $$ 这个式子和单轮情形在形式上几乎完全一致，只不过求和对象从单条回复中的各个 token，扩展成了整条多轮轨迹中所有assistant token。它所表达的含义也非常直接：如果最终奖励只在整条轨迹结束后给出，那么这个sequence-level reward就会共同作用在整条轨迹上的所有动作token 上。\n在具体实现中，我们同样不会沿着rollout过程逐轮逐token地单独回传，而是会将整条sampled trajectory固定下来，拼接为一个长序列后统一做一次forward，并在所有assistant token对应的位置上重新计算logprob。于是，从backward的角度看，多轮RL和单轮RL并没有本质区别：对于轨迹中任意一个位置 $(i,t)$，其局部 loss 仍然可以理解为“该位置 sampled token 的对数概率乘上一个由整条轨迹表现决定的权重”；差别只在于，这里的前缀$h_i$本身也已经包含了模型在更早轮次作出的决策，因此后续奖励会间接归因到这些更早的动作上。\n对于GRPO而言，我们依然是对同一个初始输入采样一组完整轨迹$\\{\\tau^{(1)},\\dots, \\tau^{(G)}\\}$，并计算每一条轨迹对应的奖励$r^{(g)}:=r(\\tau^{(g)})$以及group-relative advantage $$ A^{(g)}=\\frac{r^{(g)}-\\text{mean}(\\{r^{(j)}\\}_{j=1}^G)}{\\text{std}(\\{r^{(j)}\\}_{j=1}^G)}. $$ 而多轮的GRPO的优化目标就可以写成 $$ \\mathcal{L}_{\\text{GRPO}}=-\\frac{1}{G}\\sum_{g=1}^G\\frac{1}{N_g}\\sum_{i=1}^{K_g}\\sum_{t=1}^{n_i^{(g)}}\\min\\left(\\rho_{i,t}^{(g)}A^{(g)},\\text{clip}(\\rho_{i,t}^{(g)},1-\\epsilon,1+\\epsilon)A^{(g)}\\right). $$ 其中$N_g=\\sum_{i=1}^{K_g}n_i^{(g)}$表示第$g$条轨迹中的assistant token总数，$K_g$为总轮数，$n_i^{(g)}$则表示该采样轨迹在第$i$轮的token数，而重要性采样 $$ \\rho_{i,t}^{(g)}=\\frac{\\pi_\\theta(a_t^{(g,i)}|h_i^{(g)},a_{\\lt t}^{(g,i)})}{\\pi_\\text{old}(a_t^{(g,i)}|h_i^{(g)},a_{\\lt t}^{(g,i)})}. $$ 因此，从形式上看，多轮 GRPO 相比单轮并没有引入新的目标结构，只是把单条回复替换成了完整轨迹，把回复长度$n_g$替换成了轨迹中所有assistant token的总数$N_g$。但从建模含义上看，两者有一个根本差异：在单轮场景下，某个token只会影响当前这条回复本身；而在多轮场景下，某一轮生成出的token还会进入后续轮次的上下文，进而改变之后整条rollout的演化路径。而从credit assignment的角度考虑，每条轨迹 $\\tau^{(g)}$ 中所有assistant token共享同一个scalar advantage $A^{(g)}$，包括第1轮的思考token、第3轮的工具调用token、以及最终summary的每个token——它们都被等权重地上移或下压。这意味着：如果一条轨迹最终得到高奖励，模型根本无法区分\u0026quot;第2轮搜索策略好\u0026quot;还是\u0026quot;第5轮总结写得好\u0026quot;，两者都得到相同的正向强化。\n最后，在梯度计算上，多轮情形依然可以类比单轮形式进行理解：对每一个assistant token而言，其局部梯度本质上仍然对应于该sampled token的log-prob，再乘上一个由整条轨迹最终表现决定的advantage-related coefficient。RL下的多轮对话，往往是应用在Agentic场景之下，此时的user message转化为环境或者工具调用给出的observation。\n相关分析 Question: 解释一下batch、micro_batch、mini_batch这三者之间的关系，并分析它们取值变化所可能带来的后果？\n在VeRL的训练框架下，一次完整的GRPO训练并不是直接rollout一批数据，然后立刻对整批数据同时更新一次。我们会进一步拆成三个层次来理解，即batch、mini_batch与micro_batch。它们分别对应着数据采样的范围、一次optimizer step所使用的数据量，以及为了适配显存而进行的梯度累积粒度。\n首先，在每一个training step开始时，会先从数据集中取出一个batch的prompt。若记batch_size为$B$，GRPO中的采样数量为$G$，则真正进入GRPO更新阶段的trajectory总数自然为$N_{\\text{rollout}}=B\\times G$。对于这$N_{\\text{rollout}}$条序列，会首先基于当前的旧策略$\\pi_{\\text{old}}$计算并冻结它们的old_logprobs。也就是说，这一整批的rollout数据共同对应同一个固定的旧策略$\\pi_{\\text{old}}$，后续的所有GRPO更新都会围绕这批冻结数据予以展开。\n接下来，在真正做参数更新时，这$N_{\\text{rollout}}$条序列并不会一次性全部送入optimizer，而是会进一步切分成若干个mini_batch，记每一个mini_batch中包含的序列数为$M$，则mini_batch的总数则为$N_{\\text{rollout}}/M$，也对应$N_{\\text{rollout}}/M$次参数更新。需要注意的是，这$N_{\\text{rollout}}/M$次参数更新虽然使用的都是同一轮rollout得到的数据，但是随着每次mini_batch更新之后参数$\\theta$发生变化，当前策略$\\pi_\\theta$会逐渐偏离$\\pi_{\\theta_{\\text{old}}}$，这也正是重要性采样会逐渐偏离1并需要使用clipping进行约束的原因。\n然而，即使一个mini_batch中只有$M$条序列，在大模型训练的过程中往往也无法一次性完成forward与backward，因此还需要进一步切分成micro_batch。在VeRL的实现中，micro_batch的相关参数定义为actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu，即每一张卡每一次梯度累积的样本数量。从数学上来看，micro_batch只是把一个mini_batch拆开来适配显存大小，并不改变最终的优化目标，梯度多次累积与一次性forward/backward在不考虑工程误差的情况下，数学上是等价的。\n综合上述分析，一个典型的batch优化pipeline如下所示\nRollout阶段 (π_θ_old固定): 256 prompts × 5 samples → 1280条序列 → 计算 old_logprobs, advantages (全部冻结) PPO更新阶段 (π_θ被更新): for each minibatch of 128: # 1280/128 = 10次optimizer step for each microbatch of 32: # 128/32 = 4次梯度累积 forward: 计算 new_logprobs # 用当前π_θ backward: 累积梯度 optimizer.step() # π_θ更新, old_logprobs不变 下一轮: π_θ_old ← π_θ (同步), 重新rollout 在取值变化所带来的后果上，这三者的影响也并不相同。增大batch会使每轮rollout包含更多prompt与trajectory，从而让reward和advantage的统计更稳定、方差更小，但代价是rollout成本线性上升、单轮迭代变慢；反之，batch过小则数据量不足，reward统计波动更大，训练更容易不稳定。增大mini_batch会让每次参数更新基于更多trajectory，梯度估计更平滑、噪声更小，但会提高单次step的显存与算力压力，并减少固定rollout数据上的optimizer step次数；减小mini_batch则会让step更频繁，但单步噪声更大，同时同一批旧数据会被重复用于更多次更新，使当前策略更快偏离$\\pi_{\\theta_{\\text{old}}}$，importance ratio更容易偏离1，训练稳定性也会变差。相比之下，micro_batch更多是工程层面的显存切分参数：增大它可以减少梯度累积次数、提升吞吐，但受显存约束；减小它虽然不改变理论上的梯度期望，却会带来更多累积step、额外通信与更低的硬件利用率，从而拖慢训练。\nQuestion: 在实际Agentic RL的训练过程中，pg_loss在一开始可能为0，试解释这种现象？\n在VeRL的训练设定之下，pg_loss表示policy gradient loss，也即上述分析中$\\mathcal{L}_{\\text{GRPO}}$，不包括大多数时候会添加的KL散度约束。我们可以将这个问题分成两个维度来分析：\n区分pg_loss为什么一开始时可能为0 为什么它等于0时并不一定意味着这一轮没有梯度，模型不能更新 在进入具体公式之前，首先需要结合前述batch、mini_batch、micro_batch的流程来理解：对于某一轮fresh rollout得到的数据，所有trajectory的old_logprobs都是由同一个旧策略$\\pi_{\\theta_{\\text{old}}}$ 计算并冻结的；而在随后的PPO/GRPO更新中，当前参数$\\theta$会在这些固定数据上重新计算new_logprobs并进行优化。也就是说，所谓的“训练一开始”更准确地说，并不是指整个训练任务的最开头，而是指针对某一批fresh rollout数据，刚开始做第一个optimizer step的时刻。\n我们先忽略clipping激活的细节，使用单轮对话的GRPO作为案例，那么有 $$ \\mathcal{L}_{\\text{pg}}(\\theta)=-\\frac{1}{G}\\sum_{g=1}^G\\frac{1}{n_g}\\sum_{t=1}^{n_g}\\rho_t^{(g)}A^{(g)}. $$ 其中$\\rho^{(g)}$为重要性采样的系数，在这里重新补充一遍 $$ \\rho_t^{(g)}=\\frac{\\pi_\\theta(a_t^{(g)}|p,a_{\\lt t}^{(g)})}{\\pi_{\\theta_{\\text{old}}}(a_t^{(g)}|p,a_{\\lt t}^{(g)})}. $$ 在优化开始时，对于这一批刚刚rollout完成并冻结了old_logprob的数据，在处理第一个mini_batch时，当前策略实际上尚未发生更新，因此有$\\theta=\\theta_{\\text{old}}$，此时$\\rho_t^{(g)}=1$；将其代入上式，自然有 $$ \\mathcal{L}_{\\text{pg}}(\\theta_{\\text{old}})=-\\frac{1}{G}\\sum_{g=1}^G\\frac{1}{n_g}\\sum_{t=1}^{n_g}A^{(g)}=-\\frac{1}{G}\\sum_{g=1}^GA^{(g)}. $$ 而由于GRPO的advantage天然就满足组内均值为0，因此$\\sum_{g=1}^GA^{(g)}=0$，也就得到了$\\mathcal{L}_{\\text{pg}}(\\theta_{\\text{old}})=0$。因此总结而言，pg_loss在“一开始”为0，首先可能只是一个非常直接的数值结果：对于某一批fresh rollout数据，在第一个minibatch的第一个optimizer step开始之前，当前策略与旧策略完全一致，故importance ratio恒为1；而GRPO的advantage又经过了零均值归一化，因此最终的policy loss在数值上会恰好抵消为0。\n不过，这里最容易产生误解的地方在于：loss取值为0，并不等价于梯度为0，继续对上式求导，我们有 $$ \\nabla_\\theta\\mathcal{L}_{\\text{pg}}(\\theta)=-\\frac{1}{G}\\sum_{g=1}^G\\frac{1}{n_g}\\sum_{t=1}^{n_g}A^{(g)}\\nabla_\\theta\\rho_t^{(g)} $$ 依然使用上文提到的经典变换公式：$\\nabla f(x)=f(x)\\nabla \\log f(x)$，我们有 $$ \\nabla_\\theta \\rho_t^{(g)}=\\rho_t^{(g)}\\nabla_\\theta \\log \\pi_\\theta(a_t^{(g)}|p,a_{\\lt t}^{(g)}) $$ 因此当$\\theta=\\theta_{\\text{old}}$也即$\\rho_t^{(g)}=1$时，我们有 $$ \\nabla_\\theta\\mathcal{L}_\\text{pg}(\\theta)\\Big|_{\\theta=\\theta_{\\text{old}}}=-\\frac{1}{G}\\sum_{g=1}^G\\frac{1}{n_g}\\sum_{t=1}^{n_g}A^{(g)}\\nabla_\\theta\\log\\pi_\\theta(a_t^{(g)}|p,a_{\\lt t}^{(g)}). $$ 也就是说，这样这些样本的$A^{(g)}$并不全为0，那么即使此时pg_loss的标量值恰好为0，梯度通常仍然不为0。而当模型完成了一个mini_batch转入新的mini_batch之后，由于policy已经被更新过一次，因此重要性采样的值往往就不再为1，pg_loss也就一般不为0了。\n回到一开始的问题，在我们进行的实验中，并没有观察到step = 1时pg_loss为0的情况，这是因为VeRL在wandb记录数据时是按照step level来进行平均的，而我们mini_batch的大小并不等于batch的大小，即使在mini_batch的第一步出现了pg_loss为0的情况，后续非0的结果也会使得最终记录的均值偏离0。事实上，pg_loss在数值上是否为0，与是否处于全局训练的step = 1并没有直接关系，而是取决于当前更新是否仍然处在某一批fresh rollout数据的第一次optimizer step之前。而如果按照GRPO原本的one-iteration-per-step完全on policy的配置，也即$M=B$，在ppo_epoch = 1（也即mini batch只把batch中的每一个样本过一遍）的情况下得到的结论是：pg_loss始终为0，但这并不影响最终的梯度。\n","permalink":"https://jubsteven.github.io/blog/posts/rl-loss/","summary":"在前段时间基于VeRL的Search-R1仓库进行了一些Agentic Search相关的探索，也算是通过一些实践来积累了一点Agentic RL的基本常识。之前朋友推荐了几个帖子来讨论了\u003ccode\u003epg_loss\u003c/code\u003e在训练开始时为0的分析，自己阅读之后发现过去对于Agentic RL的理论推导还有些生疏，借此机会重新梳理一下。","title":"Loss in Agentic RL"}]