Taming Curvature: Architecture Warm-Up for Stable Transformer Training
For practitioners training large Transformers, this work provides a practical curvature-tracking method and a simple stabilization technique to mitigate training instabilities and loss spikes.
This paper introduces a fast online estimator for the largest preconditioned Hessian eigenvalue to track curvature during Transformer training, and proposes architecture warm-up (progressively growing network depth) to control curvature and stabilize training. Experiments show that the method reduces instabilities compared to existing stabilization techniques without slowing convergence.
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.