Taming Curvature: Architecture Warm-up for Stable Transformer Training
Authors: Sameera Ramasinghe
Thalaiyasingam Ajanthan
Hadi Mohaghegh Dolatabadi
Chamin P Hewa Koneputugodage
Gil Avraham
Violetta Shevchenko
Yan Zuo
Karol Pajak
Alexander Long
Venue: ICLR 2026 Poster (2026)
Keywords: Curvature
transformers
Primary area: foundation or frontier models, including LLMs
TL;DR: We propose a method to analyze and enforce the stability of training transformers.

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.

Published as a conference paper at ICLR 2026
TAMING CURVATURE: ARCHITECTURE WARM- UP FOR STABLE TRANSFORMER TRAINING
Sameera Ramasinghe
Ajanthan Thalaiyasingam
Hadi Mohaghegh Dolatabadi
Chamin Hewa Koneputugodage
Gil Avraham
Violetta Shevchenko
Yan Zuo
Karol Pajak
Alexander Long
Pluralis Research
ABSTRACT
Training billion-paramed Transformers is often brittle, with transient loss spikes and divergence that waste compute. Even though the recently developed Edge of Stability (EoS) theory provides powerful tool to understand and control the sta- bility of optimization methods via the (preconditioned) curvature, these curvature- controlling methods 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-pan parameter scale while being more accurate. Using this tool, 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 net- work 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-att stabilization techniques without slowing down convergence.
INTRODUCTION
Scaling up Transformers has driven remarkable progress across domains, from large language mod- els that power conversational systems to diffusion-based models for image generation VVaswan al. 2017: Kaplan et al. 2020 Brown al., 2020 Ouyang al. 2022 Ho al. 2020: Rom- bach et al.) (022). Yet, despite these gains, large models frequently exhibit training instabilities, l.e large loss spikes and even divergence, especially scale (Chowdhery et 2022 Dehghani al., 2023 Zhang et al. 2022 Molybog et al. 2023 OLMo et al.) 2024). As Dallion-parameter training becomes norm, improving training stability paramount: transient instabilities, i.e., loss spikes, gradient blow-ups, or full divergence, can consume vast compute budgets and wall-clock time. As models and datasets scale, stabilizing optimization reduces monetary and environmental costs while improving reproducibility and throughput, enabling dependable progress.
Stabilization at scale often relies on empirical controls for attention and optimization: soft-capping QC-clip to bound dot-product magnitudes of queries and keys (Henry et al. 2020 Dehghani et al. [2023] Team et al. 2025). and learning-rat rate or batch warmup to temper early steps (Gilmer et al. 2021 Dubey et al. (024).. In parallel, the Edge of Stability (EoS) literature shows that gradient methods gravitate toward regions where the product of step size and curvature approaches the stability boundary from classical quadratic optimization theory: for full-batch Gradient Descent (GD), training spends long phases with Amax (ff) x 2 (Cohen al. 2021 Wang al., 2022 where Amax (H) and 7 denotes the largest eigenvalue
the logits (Gemma Team 2024). Q -o-oomaalizatinn
"Corresponding author: samcera@pluraliis

---

Published as a conference paper at ICLR 2026
of the Hessian H and the step size, respectively while for preconditionecdaddppiy methods the relevant quantity is the preconditioned curvature (Cohen et al. 2022 Damian et al., 2023). Thus the stability threshold is inversely proportional to the step size and directly governed by the largest eigenvalue of the (preconditioned) Hessian. Although many of the previous works stabilizing Transformers can be interpreted attempts keep 1) Amaax (F) below this boundary (Zhai et al., 2023 Wortsman et al. 2023b Gilmer et al. [021] Shazeer & Stern 2018). verifying such claims has been difficult in practice, because estimating the curvature online for Dallion-parameter Transformers remains memory and compute-intensive
To this end, we first introduce an efficient method to estimate the curvature online using warm- started power-iteration with Hessian- Products (HVP) tailored for large models. Our key insight is that the top eigenvector of the (preconditioned) Hessian is slow-moving and warm-starting with the previous step's ' eigenvector significantly 1) reduces the iteration count and 2) improves ac- curacy. In particular, we require less five HVPs per step ($ 5), an order of magnitude lower than existing methods (Granziol 2025), while seamlessly extending to the time-varying preconditioned matrix for adaptive methods. We provide theoretical bounds for the change in the eigenvector and the resultant iteration saving. This makes online curvature tracking feasible for billion-parameter Transformers. We then use this approach to confirm that loss spikes in large-scale Transformers correlate with spikes in preconditioned curvature and show that the latter increases with the network depth.
Combining these insights, introduce an architecture warm-up strategy for stable training. The idea is to ensure the (preconditioned) curvature follows the trend of the stability threshold such that the stability criterion is satisfied throughout training. Precisely, we restrict the model have small curvature during the initial learning rate warm-up phase, and gradually relax this restriction (i.e., increase the curvature) when we start decaying the learning rate, noting that the stability threshold is inversely proportional to the learning rate. To control the curvature, we adopt a holistic approach of controlling the number of (effective) Transformer layers (i.e., depth), rather than making fine- grained modifications to each layer. Specifically, we freeze some Transformer layers to identity at initialization, and gradually unfreeze these layers as per a predefined schedule, ensuring a smooth in- crease in curvature. This architecture warm-up approach can be readily integrated to existing training recipes and standard architectures as it does not require dynamic computation graph surgery, outper- forms existing stabilization techniques, and expands the range of stable learning rates without any performance penalty. We provide extensive experiments demonstrating accurate curvature tracking and consistent stability gains across large transformer settings compared existing methods.
2 PRELIMINARIES
Below, we briefly review the literature on Edge of Stability (EoS) (Cohen et al. 2021 2022), and power iteration to compute the largest eigenvalue of the Hessian using Hessian- Vector Products (HVP) (Martens] (010). upon which we build our work. We refer the interested reader to the respec- tive papers for more details.
2.1 EDGE OF STABILITY
For a quadratic objective C()) L0T AO +00 +c, c, gradient descent with step size vj is stable only if n < 2/ Amax (A). Locally, neural network training admits the quadratic approximation:
C(0 ~ (()) + VL( TA + LAT H()A
(1)
so the Hessian (0)) plays the role of A, and Auaax (0)) determines the maximum stable step size: violating 7 < 2/Amax(H1 causes oscillation or divergence along the sharpest direction. Em- pirically, full- GD often operates near the Edge of Stability (EoS) where 17 )max(() = [Cohen et al.] 2021). Adaptive methods (e.g., Adam (Kingma & Ba 2014)) show an analo- gous behavior with the time-varying preconditioned curvature Xunase (P. 1/2 HP, 1/2) (Cohen al. 2022 where P.1 denotes the update preconditioning. For Adam, the preconditioner takes the
PPeecondiiineed Hessian should be considered for optimizers that use preconditioned updates, and adaptive methods are shown operate optimizer dependent stability thresholds (Cohen et al. 2022.
2

---

Published as a conference paper at ICLR 2026
form: Pt = diag( + 2) where 2t-1 = Bger + (1 82)gr. Note that the stability crite- and for Adam with 31 = 0.9, adaptive EoS is determined to be (P. 1/-11/2111 P HPP 1/2) ~ 38.
rion is optimizer-dependeer ?) Amax
This enables a powerful tool to understand and control the stability of optimization methods by controlling the learning rate and the preconditioned curvature. However, this theory has only been verified on small-scale models (5 25M parameters), mainly due to the memory complexity of com- puting the Hessian, the time complexity associated with estimating iterative approximation. Below, we first discuss a well-establiheed power-itcration method to compute the curvature without materializing the Hessian explicitly, and later introduce approach that reduces its iteration com- plexity by an order of magnitude, making online curvature tracking feasible and more accurate at callion-parametes scale.
2.2 COMPUTING THE LARGEST HESSIAN EIGENVALUE VIA HVP-BASED POWER ITERATION
Given A E Rdd and the loss f : TRd IR, we can estimate the top eigenpair (0max(11 Umax (0)),, where H()) v²)(0) without explicitly forming (0). For any vector
=(0)u= ve((() d g(0 Eu)
Thus, H())u, i.e., the Hessian Hessian-Vector-eee Vector-prouuct (HVP), is directional derivative of the gradient (which quantifies how much the curvature would change in the direction of v) and can be obtained without materializing full H (Martens, 2010). This costs roughly two backprop passes and uses O(1) extra memory beyond the retained graph. The efficient computation of HVP allows us to compute the Amax (H) and Umax (H))) with power iteration. Let (0) be symmetric. Suppose
its eigenvalues satisfy: A 2 A2 2 ... 2 Xd:
p=
E 10. Starting from a unit vector X1
()),, define the normalized power iteration:
((++)
(1+1) Hy)t)
((+2)
((+1) (1t1)) (y
[=(++1)1|22
With above iterations, (t) + Umax and (()) Amax as - 0o. The only primitive u H> Hu, i.e., an HVP. This in practice is expensive for large H. since power iteration requires multiple steps to converge from a random initialization, each requiring two backward passes. This is more prominent in high dimensions where the initial alignment with the leading eigenvector is O(1/vd) in expectation, where d is the dimension of the parameter vector.
3 METHODOLOGG
In this section, we first elaborate on our key insight that the (preconditioned) curvature of the Hessian is slow-moving with theoretical and empirical justifications, enabling us to develop warm-started power iteration variant to compute the curvature efficiently with less than five HVP steps. This allows verify that the loss spikes in large-scale transformers are also result of spikes in curvature, verifying the EoS theory at scale. Based on this, we later develop our architecture warm- up strategy that restricts the curvature of the model the early learning rate warm-up phase and increases the curvature as the learning rate starts to decay, ensuring the stability criterion of adaptive EoS is satisfied throughout training.
3.1 ONLINE CURVATURE TRACKING WITH WARM-STARTED POWER ITERATION
We show that under practical assumptions, the largest eigenvalue of the Hessian of neural networks evolves slowly. Specifically, under Lipchitz continuity of the Hessian and nonvanishing spectral gap 7 we can precisely bound the change in the leading eigenvector between successive steps:
LH ||0x+1
sin Sin.(11 v1.))
0xll
Thus, when step sizes and gradients shrink during training, U1,4 is an increasingly accurate initializer for Un 11 This result is formally presented below.
3

---

Published as a conference paper at ICLR 2026
Theorem 1. Let 18y h2o be parameter sequence and H₃; ((x))- Assume that there exists
LH < oc with ([) [ & LH|le 8' for all 8, 8' Equivalently, | f(0) I LH and Along the parameter path considered, ((6) == X1 (H) x2(H) 2 7 0. Then, with 11 unit top
eigenvector of H and ek == (U1, v. we have
H+++1
Hall
LA ||x++1 - 0&11
sin S
S
(5)
Further, with stochastic gradient descent with step size The such that Ok+1 = 02 - 17 with Elgg Ou] f(0s) and Ellgg 11? S G2, we have
E sinEk 0x] S (LH/7) The Elg s (La/7)mcc
(6)
Leveraging the above result, we propose a simple yet novel modification: warm-starting power iteration across training steps. At iteration k. suppose have obtained an estimate of the top eigenvector 1i.k of the Hessian H₂ = V f(xx). ). At the next training iteration, instead of reinitializ- ing power iteration from random vector, initialize from V1,4 and run power iteration H Below, we quantify the gain in iteration count due to warm-start.
Let H3 (0k) have eigenvalues λ1.k 2 λ2, 2.. and unit eigenvectors vik-- Set 0K+1 = x2,4+1/ / E [0, 1). Define the successive misalignment Sk == LVI,E41, UI,)). Run normalized power iteration on Hx+1:
Theorem 2.
(1+1) (x+18(1)
y( = -1,k
[1+++1001]]
and let a2 = ((y v1. Then:
2t 0 S A1 - (OT Hxtiy! S (A1:41-22a1 11441-22a0 228+1) sin ar I Pk-1 A1 Ph++1 tan? Ek.
(8)
Also, and baad be the number of iterations needed for warm start and random initialization to achieve convergence, respectively. Then, with high probability, we have,
2 log d - log Lu 1|0x+1 0xll)
rrudd
(9)
log 1/ l/Pt+1)
Hence warm-starting is strictly advantageous whenever Lu 10 - 0& < d'/2
3.1.1 EXTENSION TO THE PRECONDITIONED HESSIAN
When the optimization algorithm incorporates momentum and adaptive learning-rate scaling (i.e, preconditioning). stability depends on the preconditioned curvature. Let us consider the Adam update: 0.+1 0, 17 P. m4+1 where m, is the momentum, updated
exponential moving average m = 31 m, + (1 32) 9: with = (L(),) and the preconditioner is the square-root of the second moment of the gradients. The effective curvature, therefore, is the spectrum of
G. == = P (0.) P.
(10)
not of (0,) itself. Because the preconditioner Pt changes slowly when 83 1 and (1.)) is
Lipschitz smooth, G, evolves smoothly along the training trajectory. As a result, our warm-starting extends verbatim to the preconditioned Hessian. That is, the previous warm-start analysis for (0)
carries over by replacing with G, Define the top eigenvector U1,t of Ge and the eigengap 7fth \G AG,). From Theorem we directly obtain:
|G++1 Gell
sin Aui,4i, S
(1I)
78
So under Lipschitz Hessian ([H (0.+1) ((0) S L | lell) and slowly varying P. (c.g., ||PTT//2 P. p-11/2 Lp]16t+1 - 0ell with Lp x (1 B2), the top eigendirection of G, is slow- moving. Consequently, power iteration on Ge+1 warm-started from ,n, enjoys the same benefits as in the non-preconditioned case. This allows us compute HVP for G, without new primitives.
4

---

Published as a conference paper at ICLR 2026
Warm-started tracking of the preconditioned Hessian. At step t, form the diagonal pre- conditioner Pt = diag(y + E) from the optimizer state. We estimate the top eigenpair of G,= P. 1/2 (O.)) by power iteration with warm start from the previous step. Concretely: (i initialize y with the transported eigenvector yl) normalize P/ u/,.1/ u1.-1)) (or simply
y (0) 11 transport is omitted); (ii) for T = ......o.ooom one preconditioned HVP power step
= (/ == (0,) (HVP), P- y(++1) >/||=|12;
(12)
uTu and stop when the change in d falls below a tolerance. The output (Ae, u1,t = y(²")) is reused at step t+1. Each iteration costs single HVP plus cheap elementwise scalings by p:t1/2, with warm starts, r" is typically < 5, enabling efficient, step-by-step tracking of the effective curvature that governs stability under momentum/Adam (e.g., see Fig. 3)
(iii) compute the Rayleigh estimate d, = (y()))
3.2 ARCHITECTURE WARM-UP
Now, we are ready to introduce our approach to control the curvature. Specifically, we first review that deeper networks have the potential to increase the curvature, and provide an approach seam- lessly increase the number of trainable layers without introducing any function or gradient disconti- muities. Recall that, the EoS, the learning rate and curvature are inversely proportional; therefore, we keep the network shallow in the early learning rate warm-up phase, and gradually increase (ef- fective) depth when we start decaying the learning rate, ensuring the stability criterion is satisfied throughout training. The proposed method can be readily integrated into existing training recipes and assumes a standard transformer architecture without requiring hardware-level operations.
3.2. 1
DEEPER NETWORKS INCREASE CURVATURE AND SHRINK STABILITY THRESHOLD
&2 .... Φ2 be L-block residual Transformer with dx(x)) x BB(x). For the input Jacobian,
Let Ge
L
[J>90(>)||
S
II (B.( S exp Σ [BBr
(13)
8=1
For losses C(0) () [g(g (1)) y)] with Amax (V28)= S Le the Gauss-Newton bound yields
Amaxd (VBC(0)) S Le (sup |aga (x)
(14)
and Jago inherits the multiplicative growth in Eq. 13 through backprop. Hence Amax (and, under Positive Semi-Definite (PSD) diagonal preconditioner P., Amax (G.)) increases with depth L, shrinking the first-order stability margin. See App. B for more formal results and discussions.
Remarks. Although we can derive an upperbound as in Eq. 14 it not sufficient to conclude that curvature must strictly increase with depth. Empirically, the curvature might be lower than the bound due to architecture components such as residual connections, and normalization (Sagun et al.) 2016 Li al., 2018 Ghorbani et al., 2019: Yao et al. 2020), due the specific parameterization of the model (Dinh et al. 2017). Furthermore, since the preconditioned geometry is the relevant one for adaptive optimizers, efficient online tracking of (curvature) top eigenvalues, especially the preconditioned analogue Amax (G.), ). is more informative than pointwise analyses at initialization or local optima (Yao al. 20200 (Saunn et al., 2016 Ghorbani et al. 2019). Our method enables such tracking up to multi- -il-iilloooo-ameeterr parameter Transformers, where we observe (in Sec. 4.3) that (i) preconditioned curvature spikes predominantly during learning-rate warm-up and (ii) increases with
depth.
3.2.2 PROGRESSIVE DEPTH VIA CONSTRAINING BLOCK WEIGHTS TO ZERO
Recall that at the stability threshold, the (preconditioned) curvature Amax (G₂) scales inversely with the learning rate ny larger 7, the smaller the admissible Amax (G.) must be to remain stable. Con- sequently, while ramps up during warmup (tightening the threshold), we keep the network's ef- fective depth low to limit Xmax (G.). As training proceeds-ale after the peak learning rate during
5

---

Published as a conference paper at ICLR 2026
leamiing-rate decay, when the stability threshold relaxes we progressively enable additional depth, activating the full model only once the stability margin permits it. To this end, we wish to add depth keeping controlling curvature and avoiding function discontinuities. Note that the transformer block function can be expanded as follows: Let the input to the 1 layer be X E R where b, n, and d are the batch size, sequence length, and embedding dimension, respectively. Then the function admits recursive structure as
X
X midden W X P2
W CONCH P +X
(15)
where hidden and concat are the hidden layer of the feedforward network and the concatinated output of the attention heads, respectively. W P1 and W are linear projection weights. natural idea is to hold only the projection matrices at zero (i.e., W P: = w p2 = 0) so the block computes the identity with X(+1 = X,, and can later be "unlocked." However, this can still induce a discontinuous jump in the network's 's Lipschitz constant at unlock: even a small change in W P: or W
p2
Ps immediately injects the pre-existing (randomly initialized) attention/FFN paths into the residual stream. Formally, the input-output Jacobian of the 1-th block satisfies the first-order bound
J| $ I S
[ Pi |
X concatl + [1 Xindeenl
(16)
If all other weights retain their (nonzero) random initialization while the projections are zeroed, then Xomeaa ll and Ix Xnitdel can already be large at the moment of unlock: consequently, even tiny updates to w p1 or W ps can abruptly raise |xx:,yll. inflating sup,, l| (x) || of the bound Eq.[4 [4] and thus increasing the curvature. This sudden elevation can push the model over the stability threshold, often manifesting as loss spikes or instabilities.
To remedy this, and to guarantee contimuity of both the function and its first derivative at un- lock, we set all the weights (except RMSNorm weights and exclude these parameters from the optimizer while the block is locked. Under this constraint, .oncat X!+1= = XX and Dx I, i.e., the block is an exact identity with no increase in the Jacobian. When the block is unlocked, all paths start from zero, and the Jacobian perturbation grows smoothly as the newly trainable weights move away from zero. This prevents the instantaneous jump in effec- tive Jacobian and thus avoids the associated curvature spike. In practice, we keep all block weights at zero and frozen within the optimizer until curvature criterion met; then we start training them with zero initialization. This architecture warm-up keeps the product 1) Amax (Gr) within the stability envelope while depth increases, yielding smoother loss and more reliable training.
and Xaidden 0, so
Does architecture warm up compromise representation capacity? We provide intuitions from two perspectives: (i) Spectral bias / F-principle: deep nets fit low-frequency structure first, with higher frequencies learned later (Rahaman et al., 2019 Xu et al. 2019); a shallow stack suf- fices early, so temporarily limiting depth does not bottleneck what the model actually learns. (ii) Functioon-ppcc//VTKK view: early training operates in near-linear, low-curvature regime where learning aligns with dominant, low- low-complyxtyy components (Jacot et al. 2018);; additional layers can be enabled later to increase expressivity without restricting the attainable solution class. Our convergence results, and prior work on progressive, function- peretrinprresrvingg growth, corroborate that deferred depth does not harm final performance (Chen et al. 2016 Wei et al. 2016 Gong et al.,
2019 Chen et al. 2020).
4 EXPERIMENTS
We evaluate decoder-only Transformers (Llama 3- 3-ytyle (Dubey et al. 2024 on three large-scale corpora: Fineweb (Penedo et al., 2024). DCLM (Li et al., 2024) and Olmo-Mix (Allen Institute for AI. 2024). We reserve a held-out set for validation. Unless otherwise noted, we train with context length 1024, embedding dimension 2048, 32 heads, global batch size 1024, weight decay 0.01, AdamW optimizer with standard parameters, and a linear warmup over 2000 steps followed by linear decay. Tokenization is GPT-2 (vocabulary size 30,000). We use five power iterations with warm-start for online curvature tracking. Our experiments include models scaling up to 3- billion parameters. Architecture warm-up schedule: Unless noted otherwise, we keep the model at half depth until learning- -ate warmup completes, then unlock the remaining layers in four groups,
Excluding RMSNorm weights critical since they show inferior convergence from zero initialization.
6

---

Published as a conference paper at ICLR 2026
spaced by 500 training iterations, to reach full depth. Although this schedule, we observed that performance is not tightly coupled to this spacing (see App. E).
4.1 SLOWLY MOVING TOP EIGENVECTORS
Theorem 1 predicts that the leading Hessian eigendirection moves slowly along the optimization path. To verify this em- pirically (without any estimator bias), we compute the exact Hessian for a 4-layer Transformer and measure the princi- pal angle between successive top eigenvectors across training steps. We choose shallow network as computing the exact Hessian of large model is computationally prohibitive. Fig. [ shows that these angles are typically 0.1 rad (<s 5.7%) (ex- cept for few spikes), confirming the "slow drift" property.
0.25
0.20
0.15
0.10
0.05
0.00
200
400
500 eoo 1000 1200 1900
Training step
Figure 1: Leading eigenvector is slow-moving, For a 4-layer Trans- former, we compute the exact top Hessian eigenvector at each train step and plot the principal angle to the next step, Consistent with The- orem T angles are typically 0.1 radians, with only rare outliers.
4.2 EFFECTIVENESS
OF WARM-START POWER ITERATION
We compare our warm-started estimator to a cold-start (ran- dom) power method 4-layer model, using the relative error Aexel as the metric. Figure reports error versus iteration count. We ran each experiment five times and report the error bars. Warm start achieves high-accuracy estimates even with Re 5 iterations, while the cold start error is higher with even 2 20 iterations. Further, the variance of the cold-start error constantly high across the number of iterations. As shown in Fig. 2 this is due to the occasional large errors certain steps (sensitivity to the initializer) even with high iteration count. This confirms that reusing the previous step's top direction substantially reduces the HVP budget and stabilizes power convergence.
4.3 DEPTH, EFFECTIVE CURVATURE, AND STABILITY
We examine the impact of depth curvature and stability by training 8, 16, and 32-layer models on the Fine dataset with peak learn- ing rate 8 X 10-3 Figure 3 plots Amas (H) and Amax (G.) over training. The 8-layer net- work maintains low, stable curvature and smooth loss trajectory. As depth increases, both Amax (F) and (G,) exhibit higher levels and variability, and the training loss becomes prone to spikes and divergence. These re- sults support the hypothesis that deeper stacks go beyond the stability margin with increased curvature. This observation motivates our ar- chitecture warm-up: start shallow (low curva- ture), when during the learning rate warm-up (where the stability threshold increasingly be- comes lower) and progressively unlock addi- tional layers when the stability margin higher.
Moaa rolel
Moga toolll Moga Enoree
de 50 Yorking
(a) 5 HVP iterations.
(b) 20 HVP iterations.
(c) Error across differrent HVP iterations.
Effectiveness of warm-start HVP. (a.b) Online tracking error (vs. exact curvature) using 5 and 20 power-iteration steps, respectively, over the course of training. (c) Error vs. num- ber of power-iteration steps with error bars over random seeds (both model init and HVP probes).
Figure 2:
4.4 STABILITY
OF ARCIITECTURE WARM-UP
As discussed in Sec. 3.1.1 under the stability threshold, Amax (Gi) scales as O(1/Amaxx(C (Gt)) against 77 in first-order methods. Thus, larger 7 demands smaller effective curvature. We show that architecture warm-up, which suppresses
The warm-started estimator converges faster (even with 5 HVPs) and attains lower error and variance than cold-start baselines.
7

---

Published as a conference paper at ICLR 2026
byen lass
100
1.0
10-
0.1
0.
0.0
2300 2000
5030
6000
1000
3008
4000
6000
3000
2000
4000
6932
Traiiing
Tosinng
Traniig
(a) Training loss
(b) Amax (H)
(c) Amax > (G.)
Figure 3: Curvature vs. depth. We train transformers of increasing depth/ssze- layers (640M), 16 layers (IB), and 32 layers (3B)-wi)
learning rate x 10-3 and track training loss, Amax HH), and (Ga) Consistent with Sec. 3.2.1 curvature rises with depth: the 16 and 32-layer models exhibit pronounced spikes in Amax (F) and Am,)(G,) during learning-rate warmup, crossing the stability boundary and diverging. By contrast, the 8-layer model maintains lower curvature and converges stably (Amax (H) and )G,) are normalized).
curvature early and lets it grow in controlled manner, substantially widens the range of sta- ble learning rates: across peak-n sweeps, mod- els trained with architecture warm-up maintain bounded Amax (G.) and avoid loss spikes, whereas vanilla networks (no architecture warm-up) exhibit rapid curvature escalation and unstable conver- gence under the same settings (Fig.
Guscita
Duss
8
(a) 1 3 x 10
(b)
(c) 1 8 10
(d) (Gr)) 8 x 10
Figure 4: Convergence under varying stability thresholds. As the stability boundary scales as ((1/n) for fixed curvature, we sweep the peak learning rate 17 (i.e.. tightening/relaxing the thresh- old) and compare architecture warm-up to an unaltered baseline. As r) increases (smaller 1/n mar- gin), the baseline becomes increasingly unstable, whereas architecture warm-up maintains stable convergence across the range. Models are trained on Feh.
4.5 COMPARISON AGAINST OTHER STABILIZATION METHODS
We compare architecture warmup to three state-of-the-ae art stabilization baselines, QK-Norm (Henry et al. 2020), OK OK-Clip (Team et al. 2025). and Softcap (Gemma Team (024)). and an unmodified baseline (Fig. 5 with a peak L.R of 8 x 10-3 To this end, we use 16-layer, TB parameter models. Note that we intentionally use higher learning rate to observe the performance of methods under a smaller stability threshold. Architecture Warmup consistently reduces spike frequency and mag- nitude, avoids divergence, and exhibit faster convergence where other methods destabilize, yielding more reliable training across datasets. Interestingly, we found QK-CLIP to be quite unstable, of- ten diverging to NaN values mid training. Table [ shows validation perplexities. As QK-nomm was performing best on Fine we compare against it on a longer training run, up to Chinchilla (Hoffmann al. 2022) compute optimal (see App. D)
8

---

Published as a conference paper at ICLR 2026
Table Perplexity ()) across three datasets and five methods. indicates the diverged runs.
Dataset
Baseline
QK-Norm
QK-Clip
Softcap
Arch-Warmup
FineWeb
19.88+000003 61.07100001
51.41 +00002
25.02 + 0.001
DCLM
163.62 1:0.4
43.388.0.03
22.64 1: 0.002
OL.Mo-Mix
18.54 1: 0.001
Archo loas!
Arch-narmas lous))
Arth-nrres (oars)
10
0k-tomm
10
10
Qo-Marm
Satrag
Sattap
Settop
0-
6-
Training
Training
(a) Fin
(b) DCLM
(c) Olmo-Mix
Figure 5: Comparison against existing stabilization techniques. Across datasets, competing methods converge more slowly and exhibit frequent loss spikes, sometimes leading to divergence, whereas our method remains stable and consistently faster train.
4.6 CAN LEARNING-R ATE WARMUP BE REPLACED BY ARCHITECTURE WARM-UP?
Learning-rate warmup stabilizes early training by enlarg- ing the stability margin (x 0(1 ((/))) when curvature is
12
Baseline (w/o LR mu)
high; as 1) increases, this margin tightens. Architecture warm-up plays the complementary role by directly con- trolling curvature: keep depth (and thus Aumax (Gt) low initially, then increase depth as training progresses. We therefore ask whether LR warmup is necessary if curva- ture is gated by architecture. As shown in Fig. 66) models equipped with architecture warmup, achieves on par con- vergence while exceeding the stability of LR warmup.
Arch-wamup (w/l LR wu:
Baseline (w/ LR
2k 3k
5k
Training stap
Here, the baseline is trained with a standard, well-tuned learning-rate schedule with warm-up. Concretely, we use commonly adopted hyperparameters for LLaMA-style models: peak learning rate x 10-4, weight decay 0.1, 2000 warm-up steps, and cosine decay schedule.
Figure 6: Replacing LR warmup with Arch-warmup. A 16-layer (IB) Trans- former trained on FineWeb using ar- chitecture warm-up without learning- rate warmup attains on-par convergence while exhibiting more stable training.
variant obtained by taking this baseline configuration and removing only the LR warm-up: we set the number of warm- up steps to zero, keep the peak learning rate and decay schedule unchanged, keep all optimizer hyperparameters (including weight decay) fixed, and then enable architecture warm-up.
The siattrecture-aam-uppoooo
This suggests architecture warm-up has the potential to replace LR warmup in practice, or be com- bined with it for an even wider stable operating range.
RELATED WORKS
Transformer training stability. Stahilization strategies for transformers span attention- and optimizer-level interventions. On the attention side, seft-capping limits logit magnitudes to avoid softmax saturation (Gemma Team 2024), and QK -normalizationn bounds dot-product seales (Henry et al. 2020). On the optimization side, methods adjust or regularize updates (e.g., Adafactor and related stabilizers) (Shazeer & Stem, 2018: Wortsman et al. 2023a). These techniques target proxi- curvature- based diagnostics that we focus Edge of Stability. The Edge of Stability (EoS) describes the regime where training hovers near the stability boundary, with y) Xmax (H) 2 for full-batch GD (Cohen al. 2021 Wang et al. 2022
mate causes of loss spikes and complementary

---

Published as a conference paper at ICLR 2026
Follow-ups generalized the phenomenon to preconditiineet/adapii methods, replacing H by the preconditioned Hessian G. P.-1/2 HP.-1/2 and analyzed implicit self-stabilization dynamics (Cohen et al.) 2022 Damian al., 2023; Chen & Bruna 2023). Collectively, these works suggest that practical training often operates close to the spectral stability limit, motivating online control of the (preconditioned) top curvature rather than relying solely fixed hypertuning. Our work aligns with this direction, but extend the analysis from previously explored small scale models to billion parameter transformers proposing an efficient HVP estimation mechanism. Progressive growing of Transformers. Prior work increases Transformer capacity during training to cut cost while pre- serving accuracy: Progressively Stacking adds layers stagewise in BERT (Gong al. 2019 Chen et al., 2020); fuunciion prrriioiireeriingggg expansions (Net2Net, Network Morphism) enlarge depth/width without changing the realized function (Chen et al. 2016 Wei et al. 2016); and LiGO learns growth operators to expand pretrained Transformers with minimal regression (Li et al., 2023). Our focus is stability: we keep the full graph present from initialization and and ensure that depth unlocking pre- serves curvature. This ties growth to an explicit stability criterion, rather than fixed stage schedules or efficiency alone, and further avoids computational graph surgery.
CONCLUSION
We introduce scalable framework for curvature-aware training of large Transformers. First, we propose an online estimator for the top eigenvalue of the (preconditioned) Hessian that reuses the previous step's 's eigenvector we prove that the leading eigendirection is slow-moving, which yields rapid geometric convergence of warm-started power iteration. Then, propose architecture warm-up: function- preserring mechanism that progressively increases the effective depth according to a curvature budget, thereby controlling the growth of effective curvature. Empirically, this combination broadens the range of stable learning rates, reduces loss spikes, and improves training reliability.
warm start. Under standard smoothness and eigengap assumptions,
REFERENCES
Allen Institute for AI. Olmo-mix-1124 Stage-1 pretraining mixture for olmo 2. https:
///uuginggface. wttp///atttsssaaaaaaaaaa/a/aaaae ging Face).
mix- 1124 2024. Dataset card (Hug-
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, et al. Language models are few-shot learners. In Advances in Neural Infor- mation Processing Systems (NeurIPS), 2020.
Lei Chen and Joan Bruna. Beyond the edge of stability via two-step gradient updates. In Proceedings of the 40th International Conference on Machine Learning (ICML). volume 202 of Proceedings of Machine Learning Reseurch, pp. 2617-2665. PMLR, 2023. URL https: / /
mlr. /etesss/222// htt2222/ccee22b.... html
Shuo Chen al. Progressively stacking 2.0: A multi-stage layerwise training method for bert. arXiv preprint arXiv:2011. 2020.
Tianqi Chen, Ian Goodfellow, and Jonathon Shlens. Net2net: Accelerating leamning via knowledge transfer. arXiv preprint 1511:: 2016. URI. httt:ps://aa.iivoor//aab////////
05641
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm:
Scaling language modeling with pathways. arXiv preprint 2004v:2204.02311, 2022. URL
https://arxxt org/ abs/ 22004.4422 02311
Jeremy Cohen, Cameron Musco, and Christopher Musco. Gradient descent on neural networks
typically occurs at the edge of stability. In International Conference on Learning Representations (ICLR), 2021.
Jeremy M Cohen, Behrooz Ghorbani. Shankar Krishnan, Naman Agarwal. Sourabh Medapati, Michal Badura, Daniel Suo, David Cardoze, Zachary Nado, George E Dahl, et al. Adaptive gradient methods at the edge of stability. arXiv preprint arXiv: 14484, 2022.
10

---

Published as a conference paper at ICLR 2026
Alex Damian, Eshaan Nichani, and Jason D. Lee. Self-stabilizatioon: The implicit bias of gradient descent at the edge of stability. In International Conference on Learning Representations (ICLR), 2023. URL wpwpr//eeeeeeeeeeeeeeeeee net wwtp://aaad.dddd.aa.aa55...
Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, Rodolphe Jenatton, Lucas Beyer, Michael Tschannen, Anurag Arnb, Xiao Wang. Carlos Riquelme, Matthias Min- derer, Joan Puigcerver, Utku Evci, Manoj Kumar, Sjoerd van Steenkiste, Gamaleldin F. Elsayed, Aravindh Mahendran, Fisher Yu, Avital Oliver, Fantine Huot, Jasmijn Bastings, Mark Patrick Col- lier, Alexey A. Gritsenko, Vighnesh Birodkar, Cristina Vasconcelos, Yi Tay. Thomas Mensink, Alexander Kolesnikov, Filip Pavetic, Dustin Tran, Thomas Kipf. Mario Lučié, Xiaohua Zhai, Daniel Keysers, Jeremiah Harmsen, and Neil Houlsby. Scaling vision transformers to 22 bil- lion parameters. In Proceedings of the 40th International Conference on Machine Learn- ing (ICML), volume 202 of Proceedings of Machine Learning Research. PMLR, 2023, URL https::/proceeddi:
mmlr. www.mmww.....e..mmm.mm.mm. wwwpppressssseeee 2ttps/////2222/2eehaaaa222.2. dehghani2a. html
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. arXiv preprint arXiv: 2017. URL
org. /bs/
1703. 04933
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang. Angela Fan, et al. The llama 3 herd of models. arXiv e-prints, pp. arXiv-2407, 2024.
Gemma Team. Gemma 2: Improving open language models practical size. arXiv preprint arXiv:2408.000 00 2024. URL http:sss/arr.v..orgaa//////////0 tion logit soft-capping details.
Includes atten-
Amir Ghorbani, Vivek Krishnan, and Ying Xiao. An investigation into neural net hessians. In International Conference on Machine Learning (ICML). pp. 2232-2241. PMLR, 2019.
Justin Gilmer, Behrooz Ghorbani, Ankush Garg, Sneha Kudugunta, Behnam Neyshabur, David Car- doze, George Dahl, Zachary Nado, and Orhan Firat. A loss curvature perspective training instability in deep learning. arXiv preprint arXiv:2 10.04369, 2021.
Yichen Gong, Yelong He, Weizhu Li, et al. Efficient training of BERT by progressively stack- ing. In International Conference Machine Learning (ICML), 2019. URL https:// proceedings mlr.. wtt.:issm......r..m....mmmmm pttpsppp//ess//sss///// gongg//9/////gg../nna got.///.gggggaggggc
Diego Granziol. Hessformer: Hessians at foundation scale. arXiv:2501.XXXX 2025.
Alex Henry, Prudhvi Raj Dachapally, Shubham Pawar, and Yuxuan Chen. Query-key normalization for transformers. In Findings of the Association for Computational Linguistics: EMNLP 2020, pp.
2020. URL httpt::t///e// hattps://aclannhoology.. www.ollyyyyoggg2222 wrgg//20..0000... findings-entt -em
379.pdf
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Sally Henderson, Jacob Menick, Katie Millican, et al. Training compute-opting large language models. In Advances in Neural Information Processing Systems (NeurIPS), 2022. URL pattps://arxiv org/ 2ltp///.../2.2a.. 15556
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel:
Convergence and gener- alization in neural networks. In Advances in Neural Information Processing Systems (NeurIPS), volume 31, 2018 h
neurips. cc cc//paper/208//aahh .ttll
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, al. Scaling laws for neural language models. arXiv preprint arXiv:2001.08366 2020.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.680 1482.6980, 2014.
11

---

Published as a conference paper at ICLR 2026
Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss land- scape of neural nets. In Advances in Neural Information Processing Systems (NeurIPS), 2018.
org/ http//bb/11122 09913
Jeffrey Li, Alex Fang, Georgios Smymis, Maor Ivgi, Matt Jordan, Samir Gadre, Hritik Bansal, Etash Guha, Sedrick Keh, Kushal Arora, et al. Datacomp-Im: In search of the next generation of training sets for language models. arXiv preprint arXiv:2406.1179. 2024. URL https:
// aaaxiivorrgg/o.....alaa...i....a abs/ 2406.11794
Zhen Li, Xiuyu Chen, Ji Liu, and Zhangyang Wang. Learning to grow pretrained models for efficient transformer training. In International Conference Learning Representations (ICLR), 2023. wtw.eee.neeeeeeeenmeeeeeeeee Ldm URL hnttps://pppnttt /pppereeeeeeeeeeeeeeeeeeeee net forumm
James Martens. Deep learning via hessian-free optimization. In ICMI., 2010.
Igor Molybog, Peter Albert, Moya Chen, Zachary DeVito, David Esiobu, Naman Goyal, Punit Singh Koura, Diana Liskovich, Sharan Narang. Andrew Poulton, Ruan Silva, Binh Tang, Puxin Xu, Yuchen Zhang, Melanie Kambadur, Stephen Roller, and Susan Zhang. theory on adam instability in large-scale machine learning. arXiv preprint arXiv:2304.0987. 2023. URL
https://eeet http://ww..rr//////2344...0
Team OLMo, Pete Walsh, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Shane Arora, Akshita Bhagia, Yuling Gu, Shengyi Huang, Matt Jordan, et al. 2 olmo 2 furious. arXiv preprint wrXiv:2501.00655 2024.
Long Ouyang, Jeff Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pam Mishkin, Chong Zhang, et al. Training language models to follow instructions with human feedback. In Advances in Neural Information Processing Systems (NeurIPS). 2022.
Guilherme Penedo, Hynek Kydlfček, Loubna Ben Allal, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro von Werra, and Thomas Wolf. The fineweb datasets: De- canting the web for the finest text data at scale. In NeurlPS Datasets and Bench- URL https: nips pctppppp///2222//111111 Bariaatttssaan_BBeecaaaasssssss
marks Track, 2024.
Track.pdf See also 17X1:22406.17577.
Nasim Rahaman, Devansh Arpit, Aristide Baratin, Felix Draxler, Min Lin, Fred A. Hamprecht, Yoshua Bengio, and Aaron Courville. On the spectral bias of neural networks. In Proceedings of the 36th International Conference on Machine Learning (ICML), volume 97 of Proceedings of Machine Learning Research, pp. 5301-5310. PMLR, 2019. URL https:
h ptstps://prooeeedngg:
mlr. www.ppeesssvv..
/oass/v91.0 /atss//99//raaaan11 9a. html
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. Proceedings of the IEEE/CVF Con- ference on Computer Vision and Pattern Recognition (CVPR), 2022.
Levent Sagun, Léon Bottou, and Yann LeCun. Eigenvalues of the hessian in deep learning: Sin- gularity and beyond. arXiv preprint arXiv:1611.07476, 2016. URL httpt: http:///xxxiv.ovoo
abs/1611 07476
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pp. 4596-4600 PMLR, 2018
Kimi Team, Yifan Bai, Yiping Bao, Guanduo Chen, Jiahao Chen, Ningxin Chen, Ruijue Chen, Yanru Chen, Yuankun Chen, Yutian Chen, et al. Kimi k2: Open agentic intelligence. arXiv preprint arXiv:2507.2053: 2025.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Infor- mation Processing Systems (NeurIPS), 2017.
Zixuan Wang, Zhouzi Li, and Jian Li. Analyzing sharpness along gd trajectory: Progressive sharp- ening and edge of stability. Advances in Neural Information Processing Systems, 35:9983-9999 2022.
12

---

Published as a conference paper at ICLR 2026
Tao Wei, Changhu Wang, Yong Rui, and Changyou Chen. Network morphism. In ICML Workshop on Principles of Machine Learning, 2016.
Mitchell Wortsman, Tim Dettmers, Luke Zettlemoyer, Ari Morcos, Ali Farhadi, and Ludwig Schmidt. Stable and low-precision training for large-scale vision-language models. Advances in Neural Information Processing Systems, 36:10271-10298, 2023a.
Mitchell Wortsman, Peter Liu, Lechao Xiao, Katie Everett, Alex Alemi, Ben Adlam, John Co- Reyes, Izzeddin Gur, Abhishek Kumar, Roman Novak, et al. Small-scale proxies for large-scale transformer training instabilities. arXiv preprin! arXiv:2309. 14322, 2023b.
Zhi-Qin John Xu, Yaoyu Zhang, Tao Luo, Yanyang Xiao, and Zheng Ma. Frequency principle: Fourier analysis sheds light deep neural networks. arXiv preprint 1rXiv:1901.06523, 2019. URL
06523
Zhewei Yao, Amir Gholami, Kurt Keutzer, and Michael W. Mahoney, PyHessian: Neural networks through the lens the hessian. IEEE International Conference Big Data, pp. 581-590. IEEE, 2020.
Shuangfei Zhai, Tatiana Likhomanenko, Etai Littwin, Dan Busbridge, Jason Ramapuram, Yizhe Zhang, Jiatao Gu, and Joshua M Susskind. Stabilizing transformer training by preventing attention entropy collapse. In International Conference Machine Learning, pp. 40770 40803. PMLR, 2023.
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christo- pher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, Todor Mihaylov, Myle Ott, Sam Shleifer, Daniel Simig, Punit Singh Koura, Anjali Sridhar, Tianlu Wang, Luke Zettlemoyer, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2005.01668, 2022. URL https://arxivv org/ abs///b////// 221.
13

---

Published as a conference paper at ICLR 2026
7 APPENDIX
A PROOFS
A.1 PROOF FOR THEOREM 1
Proof. First, We recall a useful result from Davis- sin-O theorem.
Let Σ, E R2xP be symmetric, with eigenvalues x1 2 ... 2 Ap and ò 2 ... 2 Ap= and assume that
respectively. Fix 1 S
min Ar, X, Ast1) > 0,
where we define Ao = oc and Apt1 = -cc. Let = s - + 1, and let V (Ur, fr+1....,,u Us) Rexd and V (0r, 0rt1....... 0.) E TRpxd have orthonormal columns satisfying Σv, = u and Σ ú وو for j =r, 1, ......... Then
2 min( /222 12/ Ellops [ ElFF)
ssin [V,
min( \ Ar Xo (st1)
We set 1. Then we have, min( A Ar, d (o+1) λ1 -2. Further, V 0g and == U1, d 1. Substituting Σ = H&+1 and E Hg to above result, and with the Lipschitz condition, have
2 min( | Hallop: || - H&lle)
sin <
7
And we know, min( || Hallops || HxllF) = |1t++1 HAllF.
So have,
Hk+1 - H&ll
LH
sin I
10xx1 - 0xl1-
(17)
For gradient descent, ||0k+1 - 0kl My V f(x))
If f 02 77 with Elgn 0x]
(f(()))
and E 2 < G2, then IE sin Ex 0%. < (T.11 /7) E Elggh| S (In/r) m.G.
L
A.2 PROOF FOR THEOREM 2
Proof. For a symmetric matrix with simple top eigenvalue,
tan 0+11 II(T date Ayy
D2
tan Bg = p tan 0r,
Ay(o)
A1
hence by induction
tan 0 p² tan 0g.
Take = HK+1, U1 = U1.4+1 and y = U1.k Then
0u = 4y (1,041) = AV1,6, (1.+1) = Ek:
with Pk+1 A2,,41/A1,6111 10,
tan 02 < pf+n tan
Expand y in the H -eigenbasi:::
y H
Σ N1,41 (0,
i21
14

---

Published as a conference paper at ICLR 2026
Therefore,
A1,+11 1o HK+1Y=
Σ (A1,t1 (e.t1) (( S (A )2x+11
Σ
(y (1,++))
R22
$22
and since Ei (18+1) sin? Oc.
0 S X1 y H S (A1, x2,++1) sin² Oz.
Using equation A.2 and sin Og S tan O2 [0, Z/1.
11.++1 Ha+1y I (A1,041 - 22.+11 tan EN = =11- pk+1) A1,+11 pktt tan EK
21
which is equation
By Theorem
1
10x+1 - 0xll- tan 02 S pi+1 Lu Lu 11x+1 - 0xl 7
tan S sin Ek I
To ensure 0, 5 (0, 2), suffices that
Ph+1 L l10k+1 0xll S tan 8.
Since 0 < pb+1 < 1, taking logs yields
log( LH 10 0x log(tan 6)
2
(gg(1/bb+1)
Let xt Uni With high probability,
(z, 01,411)² = ((1/d)
tan =
Thus tan 0g S Sptt1 e @(vd) tan 5 implies
log d - log(tan
frand
log (1/pk+1)
and subtracting the warm-start bound gives equation
A.3 IMPLEMENTATION SKETCH OF WARM-START POWER ITERATION
Algorithm Warm-Start HVP Power Iteration for ()
1: input: parameters 0; optional warm vector Dwarm with ||ywarmlle I|2 1; optional previous esti-
mate Awarm: tolerance
2: initialize:
3: (0)) !/aam,
if provided
randn unit vector, otherwise
4: (early exit check) If Awarm provided, optionally compute residual (()) + [1(H(0)
Awarm (0)2 (1 HVP). If ( < return (Awerm (0))
5: for = 0, ..... do
6: ((+1)) <- (0) (() (Pearlmutter HVP)
7: guard: (++
0 (numerical underflow), reinitialize (*)) to fresh random unit vector
and continue.
8: y (++1) <- ((++1) (1++1++)) ||2 (power update)
((++)) +- (1)))) (0) ((+1) {one extra HVP or reuse cached)
9:
warm-start stabilization (optional):
11: (i) Momentum mix:
0 and Dwarm provided, set y(1) normalize (a + (1-a:) Dwarm) with small (1 a)) (e.g. 0.1).
if
12:
(ii) Cosine guard: if t 0 and ()) Mwarm)l < (e.g. 0-0.1), replace ()) +- Dwarm.
stopping test: If (H)(() ((++1) (y & [(++)) return (i(++1, ((+1)).
13:
14: end for
15

---

Published as a conference paper at ICLR 2026
Note that the proposed warm-start strategy is particularly well suited for large neural networks where HVPs are expensive and frequent monitoring of curvature is desirable. By amortizing eigenvector estimation across training steps, enable efficient online tracking of sharpness without compro- mising accuracy. To the best of our knowledge, this continuation-style use of power iteration has not been previously explored in the deep learning literature, despite being standard idea in numer- ical linear algebra. Our theoretical guarantees and empirical results demonstrate its clear superiority over random initialization for curvature estimation neural network training.
B DEPTH, LIPSCHITZ. GROWTH, AND CURVATURE
In this section, we discuss the connection between the depth and the stability margin with formal results.
Setup and notation. Let ge = 2 o ... 11 with residual blocks 22(x) = x + Be(x). Write Je(x) = I (Be(z), for the input Jacobian, and I, d2 (x( Norms operator (spectral) norms.
Lemma (Depth-Lipschitzz bound). For all x,
L II E=1
[+))
< [[ 3 exp Σ [0B
Proof. By the chain rule, (x) = Jr (ZL-)) ... (xr)). Submultiplcativity gives [>ago(>)]] < He [Je(ze-1)l |]. Next use || + All S + ||||| S exp(IIAll) obtain the exponential bound.
Lemma (Parameter sensitivity growth). Let (age(z) denote the Jacobian of ge w.r.t. parameters. Then [Jage (z)ll < C IT' ( for some constant C depending on block parametrization (e.g., linear maps and elementwise activations yield C 1 up to dimension factors). In particular,
Jago inherits the multiplicative growth in Lemmal
Proof. Differentiate the composition with respect to block parameters; each term contains a product of input Jacobians J& before and after the block where parameters appear. Bounding each product by IL Jell gives the stated incquality (constants collect plo-block linear maps).
Proposition (Curvature bound via Gauss-Newton). Assume & twice differentiable in = = go()) with Amax (V < Lt. Then
Amanx ViC())) < Le (sup [Iooo(())))))
x
Proof. For each (ar, v), the Geuss-Newton term is Jaga (x:) Jeg 3 Le Jego(1)) (x)) hence the stated bound after taking expectation and the maximum eigenvalue.
Corollary (Preconditioned curvature). Let Pe SPD diagonal with my 1 1 P. 3 Mel. Then for G. 1/2 PP=P-1/2HP /2HP,111
1 Amax(
M, Amax ()) < Amax (G.) S
mh
Hence the depth dependence of Amax (Gt) mirrors that of Amax (H) up constans (m+, M.).
Proof. For SPD Pt, Rayleigh quotients give Amax(Gr) = max u||=1 v T...TTPPP P.- max|julle; =1 u² Hu, where ||ul|? bounds.
P.u. Using mu]lu||2 < |ul Millull yields the
Corollary 2 (Stability margin scales as 1/Amax). For a quadratic model of the local dynamics, gradient descent is stable if < 2/ 2/>maxx (H); with momentun/Adam, the admissible region for (n, 3) scales as
Combining Lemmal) [ Lemma g and Proposition ] shows that increasing depth I, raises Xmax (H) (and Xmax (Gt)). thereby shrinking the stability margin.
16

---

Published as a conference paper at ICLR 2026
C SELF-STABILIZATION OF CURVATURE DURING TRAINING
C.I CURVATURE AS TRAINING PROGRESSES
We track the raw curvature mmax ( (0)). the effective (preconditioned) curvature >max(Ge), and the preconditioner inverse P/- for a 16-layer (IB) model. As shown in Fig. 7 after an initial transient the dynamics enter regime where Ama) oscillates within a narrow band, consistent with the -dge-if-stability picture. This self-stabilization supports increasing depth later in training, once curvature has settled. Notably, Pe' continues to grow time, indicating steadily strengthening preconditioning effect from the optimizer.
1.0
Biyers V
1.0
15 [18)
1.0
0.8
0.4
0.4
0.4
1x
24
16
9c
32
Taaniig step
Traniing
Training
(a) p-1
(b) Amax (H)
(C) Amax (G.)
Figure 7: Self stabilization of transformers. After an initial transient period, both the raw cur- vature and the preconditioned curvature settles down to a stable, lower band. This supports the effectiveness of increasing the depth later in the training. Interestingly. p-1 keeps growing as the training progresses (shown values are normalized).
C.2 SELF STABILIZATION OF CURVATURE DURING TRAINING
We observe self-stabilizzttion effect in Transformers (see Fig. 8): when curvature rises, the op- timizer's preconditioner counteracts In Adam-like methods with P. = diag( + e) and M++1 = 82 + (1 32)9r larger curvature typically coincides with larger gradients It. which increases 2 and thus P.. Since the effective eurvature is G, 1,-1/2 1/2 a larger Pt (smaller P. reduces Rayleigh quotients and damps Amax (Gt), yielding stabilizing feedback. See below for a formal discussion.
Self-stabilization via adaptive preconditioning (and its limits). Let :,
= ((.)) and consider an Adam-like update ft+1 = A: 17 M, m, with diagonal preconditioner M, diag( (vt + E), where
Ut+1 = 322 + (1 32) gi
= (L().).
The effective (preconditioned) curvature experienced by the optimizer is
G. = M, H, M, 1/2
Amax (G.) = max x. T M, 1/2 H,M, 1/2 |=|=1
When curvature or gradient energy surges, gt increases and (after EMA smoothing) M++1 1 con- sequently M and all Rayleigh quotients of Ge+1 decrease. A simple bound follows from the Rayleigh quotient and diagonal ordering:
Amaxe (1ax(11++1) Amax(Gr4) $ ||M Amax Amax(H+))) (H++1) min; (vut+1,5 + +
Thus, as Vt+1 grows, the preconditioner shrinks the effective curvature, exhibiting an implicit self- stabilization that nudges the product 17 Amax(G.) toward the stability band (the EoS).
However, this is not sufficient to ensure stable convergence. Despite this automatic balancing, there are three failure modes: (i) lag: i, reacts on time scale 1-%, steps, so sharp, step-scale spikes
in H, can push 7 Amax (G,) beyond the boundary before M, catches up: (ii) anisotropy: M, is diagonal, whereas H, can be highly anisotropic; coordinate-wis. preconditioner cannot instantly
17

---

Published as a conference paper at ICLR 2026
Hilt
An www Aw Aadl
Liril
terden
Lsitt
1344
Figure Self-stabilization in Transformers. Close-ups over particular training windows of second- and first-order statistics for 8-, 16-, and 32-layer models. Each row (left~rright) shows Amax (H). P-1, and (Gt). When the Hessian spikes, the preconditioner dips, reducing Amax G.) and partially stabilizing the effective curvature; however, this feedback is not always sufficient to keep Amax(G:)) within the stability threshold.
suppress a sharp direction that is dense combination of coordinates; (iii) coupling with momen- tum: with 31 > 0, large can overshoot even if M, starts to grow, transiently amplifying the update.
In summary, adaptive methods do provide reactive stabilizer, M, tends to drop curvature rises, reducing Amax(Ge), but this mechanism imperfect under fast spikes, strong anisotropy, or momentum coupling. Our architecture warm-up complements this by acting proactively Hμ itself: by keeping depth (and hence the operator norm of intermediate Jacobians) low early and unlocking blocks later in training, keep the system inside the stability envelope when the optimizer's s preconditioner has not yet adapted.
D VALIDATION AT COMPUTE OPTIMAL
Since QK-Norm was the strongest baseline our shorter FineWeb runs, we benchmark it at the Chinchilla compute-optimal setting (Hoffmann al. 2022): IB-parameter, 16-layer model trained for 25B tokens on Fibe This experiment tests whether early-phase gains from Arch- Bech-Wamupp persist compute-optimal budget. As Table 2 shows, Arch-Warc Warmup outperforms QK-Norm at this scale, indicating stronger convergence and stability that carry through to the compute-optimal regime
SENSITIVITY TO THE WARM UP SCHEDULE
By default, we keep the model at half depth until learning-rate warmup completes, then unlock the remaining layers in four groups spaced by 500 training iterations to reach full depth. Performance, however, is not sensitive to this spacing: Fig. 10 shows that even with zero spacing (i.e., unlocking all remaining layers at enabled blocks are zero-initialized, so their contribution to the Jacobian/eessian- and thus the total (preconditioned) curvature- groww gradually from the half-depth baseline. The model therefore
at peak learning rate), results similar. The reason is that newly
18

---

Published as a conference paper at ICLR 2026
Arch-warmup (spacing
Arch-warmu (spacing 250
Arch-warmu (spacing 0)
0k 1k
2k
3k
4k
5k
Training step
Figure 9: Convergence with different architecture warm-up schedules. We observe that our method not highly sensitive to the schedule.
Table Validation perplexity (PPL. on FineWeb at Chinchilla compute-optimal (IB, layers, 25B tokens). Lower is better.
Method
Val PPL +
QK-Norm
20.28
Softcap
32.44
Arch-Wir Warmup
18.35
never encounters sudden "full-depth" curvature jump at the unlock step; instead, capacity and curvature realized progressively those weights move away from zero.
Baseline
20
الله الل
warmup
Amax(G:)
15
0k
1k
2k
3k
4k
5k
Training step
Figure 10: Log-scale plot for the evolution of Xmnx (Gt) with peak learning rate of 8 10-3
19


Reviews:
[{"review_id": "0hY9eU34R6", "reviewer": "Reviewer_y4WB", "summary": "The paper makes two contributions. One is a to warm-up power iteration schemes to track the largest eigenvalue of the Hessian. The second is to propose architecture warm-up: unlocking more and more layers across training to account for natural increase in sharpness.\nThe warm-up scheme is quickly tested and we see that the largest eigenvalue of the Hessian evolves slow enough for the warmups to be efficient. The architecture warm-up appears to be an effective technique to improve stability during training. In particular the authors argue that it could replace the usual warm-up.", "questions": "- Why do the lambda max plots have a y-axis from 0 to 1? They could be larger than 1. And the plots should be in log-scale. If the plots were for example in log-scale or better centered, could we see to what value the preconditioned hessian converges (in the architecture warmup case)?\n- Most importantly, can the authors provide a plot of the largest eigenvalue of the preconditioned hessian with architecture warm-up? We cannot see its evolution in the plots for now (too small). It would be great to see that evolution with the exact times when the architecture is unlocked and to which scale each time. It would be great to see whether progressive sharpening increases/decreases with the increased depth. \n- The authors say line 366, \"progressively unlock additional layers only when the effective curvature remains below within the stability margin\". What does that mean? As long as the edge of stability is not reached, we could unlock the architecture. The sentence above does not reflect a trigger or a \"only when\".\n- As said in the weakness section, it would be great to have the detailed experimental setup for figure 6: how were hyperparameters tuned, how robust is this observation, could there actual be limitations to architecture warm-up etc...", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The proposed tool is imple but could be of use. We hope that the authors will do their best to open source the code in all relevant platforms.\n- The architecture warm-up is interesting. It is reminiscent of optimization techniques for deep learning that learned one layer at a time.\n- Experiments are done at a relatively large scale. \n- The improvements in stability are somewhat surprising. They also contribute to a very different approach to improve stability in training.", "weaknesses": "- The theoretical justifications are somewhat off: the spectrum is bounded above by something that depends on the number of layers. It is not bounded below. So these justifications are unclear.\n- The experiments are strangely presented (see questions below). Many baselines diverge. One culprit seems to be the numver of warmup steps (2000) that does not look standard for the baselines.\n- The authors claim that architecture warm-up could replace standard warm-up. Such a claim requires detailed experimental setup: how were all hyperparameters tuned (so learning rate, weight decay, warm-up steps).\n- The authors claim that spacing of unlocking does not matter for the architecture warm-up but there must be extreme cases where it does. Unlocking every step would not work for example. \n- The largest eigenvalue fo the Hessian generally drives dynamics mostly in the full batch case. The stochastic case may be governed by different mechanisms [1]. So even though architecture warm-up appears to work. It is unclear whether the authors have revealed why yet.\n- Aggressive warm-up or even learning rate tuners can also directly enforce a smaller curvature [2]. So, if curvature is the leading reason for improved training curves, there could be other ways to drive the network towards more stable regions.\n\n[1] Agarwala, A. and Pennington, J., 2024. High dimensional analysis reveals conservative sharpening and a stochastic edge of stability. arXiv preprint arXiv:2404.19261.\n[2] Roulet, V., Agarwala, A., Grill, J.B., Swirszcz, G., Blondel, M. and Pedregosa, F., 2024. Stepping on the edge: Curvature aware learning rate tuners. Advances in Neural Information Processing Systems, 37, pp.47708-47740.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Hd2LnxhCom", "reviewer": "Reviewer_4tpE", "summary": "This paper addresses training instabilities in large-scale Transformers through curvature control. The authors propose: (1) an efficient online estimator for the largest (preconditioned) Hessian eigenvalue using warm-started power iteration, reducing computational cost from >20 to <5 HVPs; (2) an \"architecture warm-up\" strategy that progressively increases network depth to maintain the stability criterion η·λmax(Gt) within safe bounds. Experiments on models up to 3B parameters demonstrate improved stability compared to QK-Norm, Softcap, and QK-Clip.", "questions": "1.\tComputational Cost: Can you provide concrete wall-clock time measurements and memory overhead for 1B and 3B models? What percentage overhead compared to standard training?\n2.\tGradient Clipping: How does your method compare against well-tuned gradient clipping? Can they be combined?\n3.\tInitial Depth Selection: How sensitive are results to starting depth (1/3 vs 1/2 vs 2/3)? Can you provide a principled selection rule?\n4.\tHVP Iterations: How does the required number of HVPs vary across training stages and model sizes? Should this be adjusted dynamically?\n5.\tRMSNorm Initialization: What specifically goes wrong with RMSNorm from zero initialization? Did you try alternatives?\n6.\tFailure Analysis: In Table 1, can you analyze the curvature behavior for methods that diverged (*)?\n7.\tCompute-Optimal: Can you provide compute-optimal comparisons with all baselines (not just QK-Norm)?\n8.\tLR Warm-up Replacement: Is Figure 6's result generalizable across model sizes and datasets, or specific to this setting?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.\tSolid Theory: Provides formal analysis connecting Edge of Stability theory with practical deep learning, with non-trivial theoretical bounds\n2.\tPractical and Simple: Easy integration into existing pipelines without architectural surgery or complex hyperparameter tuning\n3.\tComprehensive Experiments: Multiple datasets (FineWeb, DCLM, OLMo-Mix), various model sizes (640M-3B), and informative ablations\n4.\tEmpirical Validation: Figures 1-2 effectively validate theoretical predictions with proper error bars\n5.\tReproducible: Sufficient implementation details with publicly available datasets", "weaknesses": "1.\tLimited Scale: Maximum 3B parameters tested; no validation on models >10B where instabilities are most severe. The \"billion-parameter scale feasibility\" claim needs larger-scale demonstration.\n\n2.\tLack of Statistical Rigor: Most experiments show single runs without error bars or significance testing. Table 1 improvements need statistical validation.\n\n3.\tLimited Evaluation: Only validation perplexity reported; no downstream task performance to verify that stability improvements translate to better capabilities.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qEDdbBoCs0", "reviewer": "Reviewer_KvM5", "summary": "In this work the authors develop a technique called \"architecture warmup\" to control curvature. In short, they initialize many layers to 0, and slowly bring them online into training. Due to the residual connections the initial reduced network trains well, and the new layers are brought online continuously. The authors use a warmstarted power iteration to measure sharpness and show that this method initializes at lower sharpness, and ends the procedure at higher sharpness. They show that this can allow for training at larger initial learning rates, and that it can obtain similar training curves to traditional warmup.", "questions": "Do the minibatch curvature measurements well-reflect the full batch curvature statistics? Can evidence be provided using e.g. HVP aggregation over multiple batches?\n\nIs there any way to test whether or not networks trained with architecture warmup have benefits in the fine-tuning setting, where the curvature of the pretrained model may be more predictive of downstream success?\n\nSuggestion: for the warm started curvature measurement, it might be better to spend less time justifying the procedure, and more time describing the experiments.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The idea of architecture warmup is simple and, to my knowledge, novel; I admit that I am not familiar with the literature on growing networks so I will defer to other reviewers on this point. Overall the simplicity of the idea and its general compatibility with common training schemes is a big strength. The experiments seem sound to the best of my ability to understand them.\n\nThe idea of warmstarting the power iterations is also simple yet effective, and I appreciate the authors for running experiments demonstrating its utility as a measurement tool.", "weaknesses": "The main question is whether or not the proposed architecture warmup provides benefits that are better than learning rate warmup. Of particular interest is the case of the large transformer models where depth and width are often scaled simultaneously. Indeed the resulting learning curves are quite similar to transformers trained with warmup.\n\nThere is a broader point that simply allowing learning with higher learning rates is not in and of itself a win; for example if the gradient norms are smaller than the effect of a larger learning rate can be cancelled out. This suggests that a deeper investigation of the method is needed to understand why it might be useful.\n\nOne potential issue with the curvature measurements: typically experimentalists do power iteration on a small batch of data to compute an estimate of the large Hessian eigenvalues. However, this does not give an unbiased estimator of the full batch Hessian, and also may give a noisy estimator of the batch-Hessian maximum eigenvalue. This makes interpreting the curvature measurements in the paper more difficult.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)