论文

死方向调节器:深度网络的规范等变预处理

Dead-Direction Conditioners: Gauge-Equivariant Preconditioning for Deep Networks

模型训练预训练

摘要

深度网络的损失对其参数的连续对称性是不变的:logits 移位、ReLU 重新缩放、LayerNorm 缩放、每头注意力旋转。 Adam 的每坐标预处理器沿着每个对称轨道漂移,这将轨迹拉离优化所在的对称商,并模糊了商使可读的奇异学习率。我们构建了 DDC,一种死向调节器,它将基本优化器提升为 $G$ 等变优化器:它在 $G$ 不变度量的轨道分解中调节优化器的状态,因此轨迹在商 $\barθ= θ/G$ 上保持预先调节的梯度流。该结构带有四种架构规范(交叉熵移位、ReLU 和 SwiGLU 重新缩放、LayerNorm 和 RMSNorm 缩放、以及与 RoPE 匹配的每头 $O(d_{\rm head})$ 注意力旋转),在 Adam 基上证明完全等变,并通过规范等变正交器与 Muon 基组合。尊重对称性会改变优化器达到的最小值以及在那里留下的可测量值。在训练超过拟合点的语言模型上,DDCAdam 可以抵抗 AdamW 陷入的过度训练崩溃,将验证训练损失差距保持在 0.67 比 5.88,并读取 65 个逐层可观察单元中的 32 个中的死方向率,其中 AdamW 在 7 中读取它。从头开始训练的视觉 Transformer 在压缩备用数据的同时达到了较低的验证损失(1.71 比 2.12)。匹配的 AdamW 的前馈能力完好无损。在 Mu 子基座上,旋转规精确组成,DDCMuon 在深度 24 处摸到了 11 个种子中的 10 个,而普通 Mu 子永远无法到达。网络的规范对称性内置于优化器中,可锐化其找到的最小值,并将该最小值的几何形状转变为轨迹可以测量的东西。