# Taming Curvature: Architecture Warm-Up for Stable Transformer Training

arXiv: https://arxiv.org/abs/2606.16768 · HTML: https://mlsanity.com/paper/2606.16768/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 39 of its 45 references that are in the map.
- First posted on arXiv 15 Jun 2026.

## Builds on (most cited first, 30 of 39)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [OPT: Open Pre-trained Transformer Language Models](https://mlsanity.com/paper/2205.01068/) (2022)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [Gemma 2: Improving Open Language Models at a Practical Size](https://mlsanity.com/paper/2408.00118/) (2024)
- [On the Spectral Bias of Neural Networks](https://mlsanity.com/paper/1806.08734/) (2018)
- [The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale](https://mlsanity.com/paper/2406.17557/) (2024)
- [Visualizing the Loss Landscape of Neural Nets](https://mlsanity.com/paper/1712.09913/) (2017)
- [Adafactor: Adaptive Learning Rates with Sublinear Memory Cost](https://mlsanity.com/paper/1804.04235/) (2018)
- [Scaling Vision Transformers to 22 Billion Parameters](https://mlsanity.com/paper/2302.05442/) (2023)
- [Sharp Minima Can Generalize For Deep Nets](https://mlsanity.com/paper/1703.04933/) (2017)
- [Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability](https://mlsanity.com/paper/2103.00065/) (2021)
- [DataComp-LM: In search of the next generation of training sets for language models](https://mlsanity.com/paper/2406.11794/) (2024)
- [Frequency Principle: Fourier Analysis Sheds Light on Deep Neural Networks](https://mlsanity.com/paper/1901.06523/) (2019)
- [2 OLMo 2 Furious](https://mlsanity.com/paper/2501.00656/) (2024)
- [PyHessian: Neural Networks Through the Lens of the Hessian](https://mlsanity.com/paper/1912.07145/) (2019)
- [Query-Key Normalization for Transformers](https://mlsanity.com/paper/2010.04245/) (2020)
- [Net2Net: Accelerating Learning via Knowledge Transfer](https://mlsanity.com/paper/1511.05641/) (2015)
- [Small-scale proxies for large-scale Transformer training instabilities](https://mlsanity.com/paper/2309.14322/) (2023)
- [Stabilizing Transformer Training by Preventing Attention Entropy Collapse](https://mlsanity.com/paper/2303.06296/) (2023)
- [Eigenvalues of the Hessian in Deep Learning: Singularity and Beyond](https://mlsanity.com/paper/1611.07476/) (2016)
- [Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability](https://mlsanity.com/paper/2209.15594/) (2022)
- [Adaptive Gradient Methods at the Edge of Stability](https://mlsanity.com/paper/2207.14484/) (2022)

## Abstract (arXiv)

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.