近零词元嵌入为何会让门控序列模型训练悄然停摆

导语

一个几乎为零的词元向量，可能成为门控序列模型中的“隐形故障点”。论文《Near-Zero Token Embeddings Silently Stall Training in Gated Sequence Models》指出，某些罕见词或模板词在预训练中长期得不到输入梯度；如果输入嵌入与输出嵌入不绑定，AdamW 的权重衰减仍会持续压低这些行的数值，最终形成近零嵌入。

核心要点

• 在调查的291个开源权重语言模型仓库、280张不同嵌入表中，未绑定的门控与混合模型检查点里，有20/38包含范数低于10^-6且仍可被访问的输入嵌入；权重绑定模型则为0%。
• 这一现象并非对所有架构影响相同。论文测试的Transformer配置中，非初始位置或BOS词元能够避免该故障；而在Mamba、Gated DeltaNet和xLSTM等门控序列模型中，近零词元可能让输入嵌入梯度放大超过30个数量级，最高达到35个数量级。
• 机制上，接近零的输入会让SiLU门或线性查询等“当前词元读出”因子趋近于零，历史状态的读出受到抑制。连续层中的残差RMS可能低于RMSNorm的数值下限平方根，Pre-RMSNorm于是放大反向梯度。
• 当梯度条目超过约1.84×10^19时，PyTorch、Hugging Face Trainer和FSDP2中的FP32全局范数裁剪可能溢出为无穷大，继而把当前批次梯度清零。论文在Nemotron-3-Nano-4B的Telugu持续预训练测试中观察到，约20%的步骤因此丢失更新，触发问题的嵌入行也会继续冻结。

意义与影响

这项工作提醒开发者，训练稳定性检查不能只关注损失曲线、激活范围或整体梯度范数，还应定期扫描可达词元的输入嵌入范数，尤其是未绑定嵌入和门控架构。对于已经发布的检查点，作者提出GroupClip，将输入嵌入的裁剪从其他参数中分离，并使用64位或重缩放范数，实验中的重缩放实现消除了更新丢失，并能在一步内修复触发的词元。对于预训练阶段，GateFloor在SiLU门上加入δ，取值范围为0.01至0.1，使零输入时门保持正值；在最高3.7亿参数的实验中，它将近零梯度降低了28个数量级，同时保持与基线相近的困惑度。

其更广泛的启示是：一个极少出现的词元，也可能通过架构中的门控、归一化和数值裁剪链路影响整个训练过程。嵌入表的边界状态因此应被纳入模型发布和持续训练的质量检查，而不能被视为无关紧要的异常值。

来源：Hugging Face Daily Papers