等变网络中学习率不均导致性能下降,通过块级归一化可显著提升Adam效果。
Equivariance Breaks the Learning Rate
- 在等变层中按不可约表示块分组更新,避免学习率失衡。
- 块级归一化使Adam在各块内步长更一致,提升训练稳定性。
- 适合研究等变神经网络优化机制或想替代复杂优化器的开发者。
等变网络通常使用Adam优化,但近期研究表明矩阵结构优化器如Muon表现更优却未解释原因。我们发现根源在于等变线性层:每个不可约表示(irrep)块共享一个通道混合矩阵 $W_l$,其扩展形式为 $W_l imes I_{2l+1}$。单次前向传播中,$W_l$ 的梯度由 $2l+1$ 个外积贡献相加而成,秩至多为 $2l+1$。Adam对存储权重进行逐元素缩放,未考虑irrep边界,导致同一层内不同块的学习步长差异。为此,我们提出对每个块的更新分别归一化,不引入新超参数。该方法仅改变更新尺度,保持Adam的动量估计和块内方向不变。在受控的 $ ext{SO}(3)$-等变模型与 e3nn 原子间势能模型上验证,块归一化结合独立调节Adam动量系数可使Adam在所有数据集上媲美Muon,表明块级步长控制与动量累积是Muon优势的主要来源。
原文摘要 · Abstract (English)
Equivariant networks are commonly trained with Adam, yet recent work reports that matrix-structured optimizers such as Muon can perform better on these architectures without explaining why. We identify one source of this difference inside equivariant linear layers. Each irrep block learns a channel-mixing matrix $W_l$ shared across its $2l+1$ components, giving the expanded map $W_l \otimes I_{2l+1}$. For a single application of the layer, the gradient of $W_l$ sums $2l+1$ outer product contributions and has rank at most $2l+1$. Adam rescales stored weights individually without using the irrep boundaries, so one learning rate can produce different spectral step sizes across blocks within a layer. We address this mismatch by normalizing each block update separately, without introducing a new hyperparameter. This changes only the scale of the update, leaving Adam's moment estimates and its direction within each block unchanged. We evaluate the mechanism in a controlled $\mathrm{SO}(3)$-equivariant model with a matched dense control and in an e3nn interatomic potential model trained on rMD17 and MD22. The toy setup isolates a mismatch that grows with width while the dense control shows no corresponding growth. In the interatomic potential model, block normalization and tuning Adam's momentum coefficients independently improve performance, but neither alone matches Muon. Combined, they make Adam competitive with Muon on all datasets, indicating that blockwise step control and momentum accumulation account for much of Muon's advantage.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。