密集输出头和稀疏路由器中的 Z 损失后向几何

Z-loss 已广泛应用于语言模型输出头和稀疏专家混合路由器的logits。 Z 损失约束这些输出头和路由器的 softmax 对数归一化器，从而限制大 logit 偏移，减少有限精度舍入暴露，并避免训练损失发散。这些用例出现在现代 Transformer 设置中，其中大词汇量的 softmax 头、top-$k$ 路由、融合损失和混合精度优化器相互作用。 Z 损失通常仅被理解为对数归一化器的标量惩罚。本文从后向传递的角度分析 Z 损失，重点关注 Z 损失惩罚产生的梯度。 logit空间梯度，我们称之为后向源，被注入到反向传播的Z-loss分支的logit边界处；因此，后向源的效果取决于传输梯度的架构和实现。我们开发了 Z 损失的后向传输视图，将源的标量幅度和 softmax 形状与传输因子分开。这些因素包括公共移位坐标、绑定嵌入路径、输出到隐藏增益、融合损失源一致性、面向优化器的更新以及 top-$k$ 路由器缩减规模。这些诊断表明，几乎相同的前向 Z 损失值可以与不同的对数空间 Z 损失梯度共存，并且在架构和优化器传输之后，可以进行不同的参数更新。传输诊断还解释了为什么原始对数 Z 损失可以在不改变输出到隐藏增益的情况下减少标量尾部，以及为什么主动路由减少会改变有效路由器系数。在 WikiText-103 和 FineWeb-Edu 上对 GPT-2 和 Pythia 系列模型的评估中，架构感知变体减少了后向几何尾部，同时在低系数体系中保持了可比的验证困惑度。