Taming Curvature: Architecture Warm-Up for Stable Transformer Training

2026-06-15

AI Summary

This work addresses the instability commonly encountered in large-scale Transformer training, where abrupt loss spikes and divergence lead to wasted computational resources. The authors propose an efficient online curvature estimation algorithm that combines power iteration with warm starts and Hessian-vector products to dynamically track the dominant eigenvalue of the preconditioned Hessian in real time. Furthermore, they introduce a novel architecture warm-up mechanism based on progressively increasing network depth, which actively controls curvature guided by the Edge of Stability theory. Experiments on billion-parameter Transformers demonstrate that the proposed approach significantly enhances training stability, effectively mitigates divergence, and maintains the original convergence rate without degradation.

Abstract

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.

Problem

• Transformer training
• training instability
• curvature estimation
• large-scale optimization
• Hessian eigenvalue

Innovation

• curvature estimation
• architecture warm-up
• preconditioned Hessian
• stable training
• Transformer optimization