密集输出头和稀疏路由器中的Z-Loss反向几何

# 密集输出头与稀疏路由器中的Z损失反向几何
来源：https://arxiv.org/html/2609.16179
###### 摘要
Z损失已被广泛应用于语言模型输出头和稀疏混合专家（MoE）路由器的对数概率上。Z损失通过约束这些输出头和路由器的softmax对数归一化器，从而限制了大幅对数概率偏移，减少了有限精度舍入暴露，并避免了训练损失发散。这些用例出现在现代Transformer设置中，涉及大词汇表softmax头、Top-k路由、融合损失和混合精度优化器的交互。Z损失通常仅被理解为对对数归一化器的标量惩罚。本文从反向传播的角度分析Z损失，重点关注Z损失惩罚产生的梯度。我们将注入在Z损失分支反向传播对数概率边界的对数空间梯度称为“反向源”；因此，反向源的效果取决于梯度传输所依赖的架构和实现方式。我们为Z损失开发了一个反向传输视图，将源的标量幅度和softmax形状与传输因子分离。这些因子包括公共偏移坐标、绑定嵌入路径、输出到隐藏层增益、融合源一致性、面向优化器的更新以及Top-k路由器缩减比例。这些诊断表明，几乎相同的前向Z损失值可以与不同的对数空间Z损失梯度共存，并且在经过架构和优化器传输后，产生不同的参数更新。传输诊断还解释了原始对数概率Z损失为何能在不改变输出到隐藏层增益的情况下减少标量尾部，以及激活路由缩减如何改变有效的路由器系数。在GPT-2和Pythia系列模型在WikiText-103和FineWeb-Edu上的评估中，考虑架构感知的变体在降低反向几何尾部的同时，在低系数区域保持了可比的验证困惑度。实证套件包括继续预训练、混合精度和融合内核审计、Adam状态压力测试以及端到端MoE训练。在压力环境下，相同的变体揭示了在验证质量与更新尾部行为和混合精度裕量之间的明确权衡。总之，这些结果提供了一个传输感知的框架，用于衡量、报告和干预Z损失的对数梯度源、其传输路径以及由此产生的更新尾部。
## 引言
现代Transformer训练通常包含旨在约束最终softmax和路由器对数归一化器的辅助损失项。在密集输出头中，其动机是限制输出对数概率或公共偏移的偏移，并保持有限精度数值裕量。在稀疏路由器中，动机还包括减少对舍入敏感的路由和训练损失发散。一个显著的例子是对数归一化器惩罚，通常称为Z损失，它已被广泛应用于现代Transformer训练。Mesh TensorFlow和Pathways语言模型（PaLM）将Z损失应用于最终softmax对数概率（Shazeer等人2018（https://arxiv.org/html/2609.16179#bib.bib3）；Chowdhery等人2023（https://arxiv.org/html/2609.16179#bib.bib7））。PaLM报告称，保持softmax对数归一化器接近目标可增加训练稳定性（Chowdhery等人2023（https://arxiv.org/html/2609.16179#bib.bib7））。先前的工作因此使用对数控制方法作为训练稳定性措施。路由器Z损失保持门控对数概率较小，减少低精度舍入误差，并防止训练损失发散（Zoph等人2022（https://arxiv.org/html/2609.16179#bib.bib6）），而输出嵌入中心化移除了与输出对数概率发散相关的公共偏移自由度（Stollenwerk等人2026（https://arxiv.org/html/2609.16179#bib.bib9））。尽管Z损失被广泛使用，但先前的工作通常仅将其呈现为对对数归一化器的标量惩罚。本文将Z损失表述为依赖于架构的反向传输，并利用这一梯度视角来分析和诊断Z损失如何从根本上影响训练行为。我们开发了涉及标量源幅度、源形状、输出到隐藏层增益、路径耦合和路由器有效规模等常见反向传输现象的诊断方法。利用softmax的标准平移不变性，我们推导了Z损失和保持交叉熵（CE）的输出头中心化的公共偏移敏感性。我们还推导了绑定嵌入梯度分解和Top-k路由器规模扭曲。我们的实证评估使用预训练的Transformer语言模型（Vaswani等人2017（https://arxiv.org/html/2609.16179#bib.bib1））以及来自WikiText-103和FineWeb-Edu的文本。模型集包括来自生成式预训练Transformer 2（GPT-2）系列的四个模型：GPT-2、DistilGPT-2、GPT-2 Medium和GPT-2 Large。该框架在相同的密集继续预训练、有效未剪裁的高系数压力扫描、静态混合精度端点、低秩输出到隐藏层增益谱审计、面向优化器的更新审计以及端到端Top-k混合专家（MoE）语言模型训练中进行了测试。实现研究还审计了用Triton实现的按行融合CE+Z损失图形处理单元（GPU）内核，Triton是一种用于编写平铺GPU内核的语言和编译器（Tillet等人2019（https://arxiv.org/html/2609.16179#bib.bib26））。
##### 贡献。我们的主要贡献是一个反向传输框架，该框架将Z损失对数概率梯度视为一个源，其面向优化器的效果取决于架构和实现，而不仅仅是标量惩罚。我们推导了公共偏移、绑定路径、输出到隐藏层增益、融合源一致性和路由器缩减规模的诊断方法，以及特定坐标的探针和干预：中心化、分解、增益感知加权、源一致计算和规模匹配。匹配的预训练诊断、继续预训练、融合损失审计、压力端点和MoE训练表明这些坐标可能表现不同。低系数设置可以在保持验证困惑度（PPL）的同时，压力设置揭示了在验证质量、更新尾部和混合精度裕量之间的模型和环境相关的权衡。因此，由此产生的审计协议基于传输诊断和验证质量共同选择系数和干预措施，而不是规定一个普遍最优的Z损失变体。图1（https://arxiv.org/html/2609.16179#Sx1.F1）总结了该框架。这种表述也确定了与通常的配方级解释的对比。在那种解释中，Z损失通常被视为一个标量辅助项，其作用是使log Z变小，从而提高数值稳定性。然而，降低标量对数归一化器尾部并不等同于降低传输后的更新尾部。原始对数概率Z损失可以使报告的Z损失曲线看起来成功，同时输出到隐藏层增益、绑定路径耦合、自适应矩估计（Adam）状态压力或静态损失缩放裕量基本保持不变或更差。相反，中心化、分解、增益感知加权和路由器规模匹配并不是降低一个标量的可互换方法；这些方法是对同一反向源的不同坐标的干预。前向标量$L_Z \coloneqq \lambda (\log Z - c)^2$，Z损失对数概率梯度，即反向源$\delta_z^Z = 2\lambda (\log Z - c) p$。实现重用log-sum-exp统计$\Delta_{\rm src}, \Delta_{\theta}$，公共偏移$z = \tilde{z} + \mu \mathbf{1}$，移除$2\lambda \mu p$。输出和绑定路径$\delta_u^Z \coloneqq W_U^\top \delta_z^Z$，绑定$E_{\rm out} + E_{\rm in}$。路由器规模Top-k缩减$\lambda_R \alpha$和$N_{\rm tl} \alpha$。面向优化器的更新对数归一化器和增益尾部、Adam状态和半精度裕量，传输坐标和诊断。
图1：Z损失的反向传输视图。标量惩罚创建了对数概率梯度$\delta_z^Z$，称为反向源。架构和实现决定了该向量是被保留、重新缩放、放大还是耦合到共享参数中，然后传输源才到达面向优化器的更新。
## 反向传输视图
本节通过将Z损失对数概率梯度视为注入到反向传播中的源来开发反向传输视图。我们首先指定主要语言建模目标和辅助Z损失，然后将源构造与架构传输分离，识别修改反向源的坐标和路径，并利用这些结果来激发后续的诊断和实验。所有证明和更详细的理论推导见附录。
### 语言建模目标和Z损失设置
Z损失被添加到标准语言建模损失中。对于自回归标记序列$x_{1:T}$，主要目标是下一个标记预测的CE，$\mathcal{L}_{\rm CE} = -\frac{1}{T-1} \sum_{t=1}^{T-1} \log p_\theta(x_{t+1} \mid x_{\leq t})$。应用于已部署对数概率的对数归一化器，Z损失项作为辅助正则化器添加：$\mathcal{L}_{\rm total} = \mathcal{L}_{\rm CE} + \mathcal{L}_{Z}$。该设置也界定了我们稳定性主张的范围：Z损失直接约束该对数归一化器，而不是输出头精度或表示几何。较小的对数概率偏移可能减少有限精度舍入暴露，但任何由此产生的输出嵌入几何、数值裕量或面向优化器的更新的变化都是下游效应，必须单独评估。这里研究的惩罚具有形式$\mathcal{L}_{Z} \coloneqq \lambda (\log Z - c)^2$，$Z \coloneqq \sum_{i=1}^V \exp z_i$，$\lambda \geq 0$，(1)其中$z \in \mathbb{R}^V$是对数概率轴的对数概率向量，$c$是对数归一化器目标，$\lambda$是系数。目标$c$是$\log Z$的零惩罚值，$V$表示相关softmax轴的大小。先前最终softmax和路由器的公式通常使用$\lambda (\log Z)^2$，对应于$c=0$（Shazeer等人2018（https://arxiv.org/html/2609.16179#bib.bib3）；Chowdhery等人2023（https://arxiv.org/html/2609.16179#bib.bib7）；Zoph等人2022（https://arxiv.org/html/2609.16179#bib.bib6））。相反，我们研究公式（1）中的目标平移泛化，并使用$c = \log V$，除非另有说明；$c$是固定的，而不是调整的，而系数扫描则改变$\lambda$。这种选择为零对数概率均匀向量产生零违反。对于原始对数概率，设置$c = \log V$选择公共偏移目标，而对于中心化对数概率，相同的设置定义了剩余相对对数概率坐标的均匀参考。因此，我们表格中的标准Z损失行将公式（1）应用于$c = \log V$的原始对数概率；传统公式通常使用$c = 0$。公式（1）的Z损失的这种标量定义指定了前向损失值。然而，仅标量定义无法确定由此产生的基于梯度的参数更新。更新还取决于Z损失对数概率梯度（我们将在下一小节中称为反向源）如何通过架构和训练实现进行传输。因此，标量审计可以认证添加到目标中的数量，而错过实际注入训练中的向量。这一区分在Z损失被使用的有限精度和训练发散环境中至关重要。报告的惩罚在输出头量表、绑定嵌入、融合低精度损失、自适应优化器和Top-k路由器缩减之间可能保持不变，而这些因素改变了对数概率梯度到达参数的位置和强度。因此，完整的分析必须识别该梯度以及梯度作用的传输路径，此外还要衡量Z损失值。
### Z损失对数概率梯度和传输
设$z \in \mathbb{R}^V$为输出对数概率向量，$p \coloneqq \operatorname{softmax}(z)$，$Z \coloneqq \sum_i e^{z_i}$。对公式（1）求导得到$\delta_z^Z \coloneqq \frac{\partial \mathcal{L}_Z}{\partial z} = 2\lambda (\log Z - c) p$。(2)在全文中，“源”是该Z损失对数概率梯度的简称，等价于对数空间伴随$\delta_z^Z$。对数空间伴随是对数概率边界上的一个余向量。分量$\delta_{z,i}^Z \coloneqq \partial \mathcal{L}_Z / \partial z_i$表示对对数概率$z_i$的局部敏感性。术语“源”反映了该伴随作为起始反向信号的作用，该信号由反向模式微分的Z损失分支传输到隐藏状态和参数。该伴随既不是数据源也不是参数梯度。在密集输出头中，每个标记的源对每个词汇表条目有一个条目；在MoE路由器中，每个标记-层路由决策的源对每个专家有一个条目。除非明确限定为总CE+Z源，否则源仅指辅助Z损失贡献；CE对数概率梯度是单独的向量$p - y$。源有两个可分离的因子。第一个是带符号的标量系数$2\lambda (\log Z - c)$，其绝对值$2\lambda |\log Z - c|$我们称为标量源幅度。第二个是softmax形状向量$p$，我们称为源形状。因此，完整的欧几里得源幅度为$\left\lVert \delta_z^Z \right\rVert_2 = 2\lambda |\log Z - c| \left\lVert p \right\rVert_2$。对于固定的$\lambda > 0$，该恒等式也可以写成$\left\lVert \delta_z^Z \right\rVert_2 = 2\sqrt{\lambda \mathcal{L}_Z} \left\lVert p \right\rVert_2$，因此精确的每个实例源随标量惩罚而消失。然而，仅标量本身使源符号和softmax形状未指定，并且不能确定通过对数概率雅可比矩阵的放大。在讨论增益时，我们通过集中度$\left\lVert p \right\rVert_2$和归一化方向$p / \left\lVert p \right\rVert_2$来查看源形状。下游架构然后决定了这个有形状的源如何被传输。对于任何具有对数概率雅可比矩阵$J_\theta \coloneqq \frac{\partial z}{\partial \operatorname{vec}(\theta)}$的参数块$\theta$，Z损失梯度为$\operatorname{vec}(\nabla_\theta \mathcal{L}_Z) = J_\theta^\top \delta_z^Z$。(3)因此，反向术语具有直接的运算意义。完整的因果管道可以简洁地写为$\widehat{\delta}$