通过渐进式加深网络结构，稳定大模型训练中的梯度波动。
Taming Curvature: Architecture Warm-Up for Stable Transformer Training
- 用快速在线估计器追踪预条件海森值最大特征值，实现百亿参数级实时曲率监测。
- 发现训练不稳与预条件曲率突增同步，且曲率随网络深度增加而上升。
- 提出架构预热策略，可有效降低训练震荡，适合大规模Transformer部署。
百亿参数Transformer的训练常因瞬时损失尖峰和发散而失败，造成大量算力浪费。尽管近期提出的稳定边缘（EoS）理论可通过预条件曲率控制优化稳定性，但其在大规模训练中未普及，主要因曲率估计复杂。为此，我们提出一种基于海森向量积的暖启动幂迭代变体，实现快速在线估计最大预条件海森特征值（即曲率）。理论分析与实证表明，该方法可在百亿参数规模下实现每轮迭代的曲率追踪，且精度更高。利用该工具，我们发现训练不稳与预条件曲率激增高度相关，且曲率随网络深度递增。受此启发，提出架构预热：逐步增加网络深度，以精细控制预条件海森矩阵，从而稳定训练。大规模Transformer实验验证，该方法在不减慢收敛速度的前提下，显著优于现有最先进的稳定化技术。
原文摘要 · Abstract (English)
Training billion-parameter Transformers is often brittle, with transient loss spikes and divergence that waste compute. Even though the recently developed Edge of Stability (EoS) theory provides a powerful tool to understand and control the stability of optimization methods via the (preconditioned) curvature, these curvature-controlling methods are not popular in large-scale Transformer training due to the complexity of curvature estimation. To this end, we first introduce a fast online estimator of the largest (preconditioned) Hessian eigenvalue (i.e., curvature) based on a warm-started variant for power iteration with Hessian-vector products. We show theoretically, and verify empirically, that the proposed method makes per-iteration curvature tracking feasible at billion parameter scale while being more accurate. Using this tool, we find that training instabilities coincide with surges in preconditioned curvature and that curvature grows with depth. Motivated by these observations, we propose architecture warm-up: progressively growing network depth to carefully control the preconditioned Hessian and stabilize training. Experiments on large Transformers validate that our approach enables efficient curvature tracking and reduces instabilities compared to existing state-of-the-art stabilization techniques without slowing down convergence.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。