Never Saddle for Reparameterized Steepest Descent as Mirror Flow

Tom Jacobs; Chao Zhou; Rebekka Burkholz

cs.LGMar 2, 2026

Never Saddle for Reparameterized Steepest Descent as Mirror Flow

Authors: Tom Jacobs, Chao Zhou, Rebekka Burkholz

Organizations: CISPA Helmholtz Center for Information Security, Saarbrücken, Germany

Abstract

How does the choice of optimization algorithm shape a model's ability to learn features? To address this question for steepest descent methods --including sign descent, which is closely related to Adam --we introduce steepest mirror flows as a unifying theoretical framework. This framework reveals how optimization geometry governs learning dynamics, implicit bias, and sparsity and it provides two explanations for why Adam and AdamW often outperform SGD in fine-tuning. Focusing on diagonal linear networks and deep diagonal linear reparameterizations (a simplified proxy for attention), we show that steeper descent facilitates both saddle-point escape and feature learning. In contrast, gradient descent requires unrealistically large learning rates to escape saddles, an uncommon regime in fine-tuning. Empirically, we confirm that saddle-point escape is a central challenge in fine-tuning. Furthermore, we demonstrate that decoupled weight decay, as in AdamW, stabilizes feature learning by enforcing novel balance equations. Together, these results highlight two mechanisms how steepest descent can aid modern optimization.

Figures & tables

Figure 1: For a deep diagonal linear network initialized close to a saddle point, sign gradient flow (SignGF) converges faster than gradient flow (GF).

Figure 2: Illustration of different steepest mirror flows (with varied

q

). On the left side, the metric exponent is shown dependent on the associated depth. A high metric exponent increases the difficulty to escape zero and the instability of the flow. The right side illustrates saddle escape by plotting the solutions of the ODE’s corresponding to the metric exponents,

dx_{t}=x_{t}^{q}dt

, with

x_{0}=0.1

(from the origin). Concluding, SignGF does not get stuck near saddles and still allows feature learning by entering the green strip in the plot on the left, effectively inducing sparsity.

Figure 3: The balance equation for

q\in\{1,1.5,2\}

and initialization

m=0.1,w=0

. Observe that the (curved) path away from the initialization to a point on the curve

mw=x

with

x=\pm 0.1

(in the plot) is shorter for smaller

q

, indicating faster saddle escape.

	Coupled	Decoupled
$q=1\text{ (SignGF)}$	$(L_{1},\text{log})$	$(L_{\frac{3}{2}},L_{1})$
$q=1.5$	$(L_{1},\text{log})$	$(L_{\frac{5}{4}},L_{\frac{1}{2}})$
$q=2\text{ (GF) }$	$(L_{1},\text{log})$	$(L_{1},\text{log})$

Table 1: Comparison of the effect of coupled or decoupled weight decay (

M_{\text{reg}}

) for two reparameterization depths, namely,

(L=2,L=\infty)

. Note that the infinite depth would lead to a non-convex logarithmic regularizer (log) in the coupled case, potentially leading to instability.

Figure 4: The

L_{\infty}-

margin for Adam with high and low depth

L

. The green region indicates the non-zero ground truth features. Higher depth leads to sparse ground truth recovery in line with Corollary 4.13 .

Panel of Figure 5 — Figure 5: Eigenvalue spectra in finetuning for an ImageNet pretrained ResNet-18 on CIFAR-10 (a) and weight sparsity in reparameterized training for a ResNet-50 on Imagenet (b).

Model	Finetune	SGD (lr $<<1$ )	SGD (lr $>0$ )	Adam (lr $<<1$ )
ResNet-18	CIFAR-10	$19.15\pm 2.82$	$93.60\pm 0.38$	$\bm{95.19\pm 0.21}$
ResNet-18	Flowers	$1.22\pm 0.53$	$62.13\pm 1.10$	$\bm{80.50\pm 1.38}$
ViT-large	CIFAR10	$73.27\pm 3.68$	$99.07\pm 0.35$	$\bm{99.28\pm 0.07}$
ViT-large	Flowers	$1.03\pm 0.82$	$98.94\pm 0.05$	$\bm{99.37\pm 0.08}$
Bert-base	MRPC	$43.87\pm 24.02$	$84.80\pm 1.00$	$\bm{85.95\pm 0.64}$

Table 2: Validation accuracy for finetuning scenarios (

95\%

confidence interval.).

Appendix figures & tables43 assets

Supplementary material from the paper’s appendix.

Appendix

g_{t}\leftarrow\nabla_{x}f(x_{t-1})+\alpha_{1}x_{t-1}

Appendix

Algorithm 1 Adam with Coupled (

\alpha_{1}

) and Decoupled (

\alpha_{2}

) Weight Decay

g_{t}\leftarrow\nabla_{x}f(x_{t-1})+\alpha_{1}x_{t-1}

Appendix

Algorithm 2 SignGD with Coupled (

\alpha_{1}

) and Decoupled (

\alpha_{2}

) Weight Decay

Panel of Figure 6 — Figure 6: Resulting $L_{\infty}$ margins for optimizers SignGD, Adam and GD, where the green strip indicates the contributing ground truth features. Observe the similarity between Adam and SignGD for all depth.

Model	Q (%)	K (%)	Total (%)	Avg $\Delta_{L1}$	$<\%,L_{1}$	Avg $\Delta_{L2}$	$<\%,L_{2}$
LLaMA-3.1 8B	1.25	0.87	1.57	-624.13	100	-20.28	100
LLaMA-3.2 3B	4.37	3.50	5.04	-1757.04	100	-101.71	100
LLaMA-3.2 1B	4.73	3.24	6.11	-891.76	100	-66.15	100

Appendix

Table 3: Parameter sign flips per group type, overall, and average L1/L2 differences for LLaMA models. We also indicate with

<\%

the percentage of the layers that have a negative delta

Figure 7: Underdetermined linear regression (

k=80

), for depth

L=1

we do net get close to the ground truth in all cases, for

L=2

gradient flow gets close to the ground truth as in line with Theorem B.10 and in for higher depth

L=10

the sign gradient flow (SignGF) converges close to ground truth which we would expect based on the dynamic reformulation.

Figure 8: Recovering the ground truth with small batch size

5

for underdetermined regression with

k=80

Figure 9: Recovering the ground truth with full batch for underdetermined regression with

k=40

Figure 10: Recovering the ground truth with small batch size

5

for underdetermined regression with

k=40

Depth $L$	$\alpha_{1}$	$\alpha_{2}$	Balance Dist.	Groundtruth Dist.
1	0	$1\text{e-}4$	$0$	$7.1$
	$1\text{e-}4$	0	$0$	$7.0$
	0	$1\text{e-}3$	$0$	$7.1$
	$1\text{e-}3$	0	$0$	$5.8$
	0	$1\text{e-}2$	$0$	$7.1$
	$1\text{e-}2$	0	$0$	$1.0$

Appendix

Table 4: Effect of Regularization Strengths on theoretical balance according to Lemma 4.5 and the distance the ground truth. For the decoupled weight decay (

\alpha_{1}

) indeed stays close the theoretical predicted balance.

Dataset	Model	LR	Epochs	Batch Size	Optim	Schedule
CIFAR-10	ResNet-20	$0.001$	$150$	$512$	Adam, AdamW	Triangular
Imagenet	ResNet-50	$0.001$	$100$	$1024$	Adam, AdamW	Triangular

Appendix

Table 5: Training details for all experiments presented on sparse reparameterizations.

Figure 11: Adam with coupled weight decay trained with various depth reparameterizations for ResNet-20 on CIFAR-10. On the left is high regularization

1e-1

and on the right is less regularization

1e-4

Figure 12: AdamW (decoupled weight decay) trained with various depth reparameterizations for ResNet-20 on CIFAR-10. On the left is high regularization

1e-1

and on the right is less regularization

1e-4

Optimizer	Depth	Weight Decay	Accuracy $\pm$ CI
AdamW	2	$1\mathrm{e}{-1}$	$89.75\pm 0.20$
Adam+wd	2	$1\mathrm{e}{-1}$	$64.36\pm 2.70$
AdamW	2	$1\mathrm{e}{-4}$	$89.29\pm 0.28$
Adam+wd	2	$1\mathrm{e}{-4}$	$88.27\pm 0.08$
AdamW	4	$1\mathrm{e}{-1}$	$89.73\pm 0.18$
Adam+wd	4	$1\mathrm{e}{-1}$	$58.23\pm 4.98$

Appendix

Table 6: Test Accuracy (%)

\pm

95% CI for AdamW and Adam+wd across depths and weight decays training a ResNet-20 on CIFAR-10.

Figure 13: Adam with coupled weight decay trained with various depth reparameterizations for ResNet-50 on Imagenet. On the left is high regularization

1e-1

and on the right is less regularization

1e-4

Figure 14: AdamW (decoupled weight decay) trained with various depth reparameterizations for ResNet-50 on Imagenet. On the left is high regularization

1e-1

and on the right is less regularization

1e-4

Figure 15:

L_{1}

norm of the weights during training for Adam with coupled weight decay strength

1e-4

and AdamW with

1e-1

. The dashed lines correspond to depth

L=10

and solid lines to

L=2

. The training setup is ResNet-20 on CIFAR-10

Optimizer	Depth	Weight Decay	Accuracy $\pm$ CI
AdamW	2	$1\mathrm{e}{-1}$	$76.23\pm 0.07$
Adam+wd	2	$1\mathrm{e}{-1}$	$1.95\pm 0.48$
AdamW	2	$1\mathrm{e}{-4}$	$73.32\pm 0.11$
Adam+wd	2	$1\mathrm{e}{-4}$	$73.35\pm 0.05$
AdamW	10	$1\mathrm{e}{-1}$	$62.20\pm 0.25$
Adam+wd	10	$1\mathrm{e}{-1}$	$0.58\pm 0.06$

Appendix

Table 7: Test Accuracy (%)

\pm

95% CI for AdamW and Adam+wd across depths and weight decays training a Resnet 50 on Imagenet.

Metric	SGD ( $\eta=0.001$ )	SGD ( $\eta=0.8$ )	Adam ( $\eta=0.001$ )
Val Acc	$19.15\pm 2.82$	$93.60\pm 0.38$	$95.19\pm 0.21$
$L_{1}$	$424911.48\pm 34308.92$	$477750.60\pm 10343.88$	$693101.67\pm 13509.59$
$L_{2}$	$29640.98\pm 985.56$	$28409.58\pm 219.53$	$27833.50\pm 494.50$

Appendix

Table 8: Validation accuracy and parameter distance traveled in terms of

L_{1}

and

L_{2}

norm for finetuning ResNet18 on CIFAR-10.

Metric	SGD ( $\eta=0.002$ )	SGD ( $\eta=0.4$ )	Adam ( $\eta=0.002$ )
Val Acc	$1.22\pm 0.53$	$62.13\pm 1.10$	$80.50\pm 1.38$
$L_{1}$	$206325.76\pm 1327.51$	$173882.76\pm 2967.69$	$618592.47\pm 3445.18$
$L_{2}$	$10124.38\pm 76.52$	$7015.04\pm 226.30$	$11432.11\pm 350.38$

Appendix

Table 9: Validation accuracy and parameter distance traveled in terms of

L_{1}

and

L_{2}

norm for finetuning ResNet18 on Flowers.

Figure 16: Normalized top-50 eigenvalues for a ResNet-18 finetuned on Flowers.

Figure 17: The eigen value evolution for seed

123

on CIFAR-10.

Figure 18: The eigen value evolution for seed

456

on CIFAR-10.

Figure 19: The eigen value evolution for seed

789

on CIFAR-10.

Figure 20: The eigen value evolution for seed

1000

on CIFAR-10.

Figure 21: The eigen value evolution for seed

123

on Flowers.

Figure 22: The eigen value evolution for seed

456

on Flowers.

Figure 23: The eigen value evolution for seed

789

on Flowers.

Figure 24: The eigen value evolution for seed

1000

on Flowers.

Metric	SGD ( $\eta=0.0001$ )	SGD ( $\eta=0.01$ )	Adam ( $\eta=0.0001$ )
Val Acc	$73.27\pm 3.68$	$99.07\pm 0.35$	$99.28\pm 0.07$
$L_{1}$	$460.47\pm 219.86$	$24617.83\pm 14406.4$	$453934.906\pm 22278.43$
$L_{2}$	$0.48\pm 0.059$	$6.25\pm 4.29$	$39.47\pm 1.01$

Appendix

Table 10: Validation accuracy and parameter distance traveled in terms of

L_{1}

and

L_{2}

norm for finetuning ViT-Large on CIFAR-10.

Metric	SGD ( $\eta=0.0001$ )	SGD ( $\eta=0.01$ )	Adam ( $\eta=0.0001$ )
Val Acc	$1.03\pm 0.82$	$98.94\pm 0.05$	$99.37\pm 0.08$
$L_{1}$	$25.71\pm 30.48$	$4655.83\pm 576.49$	$108583.62\pm 2078.48$
$L_{2}$	$0.04\pm 0.02$	$1.50\pm 0.12$	$8.35\pm 0.16$

Appendix

Table 11: Validation accuracy and parameter distance traveled in terms of

L_{1}

and

L_{2}

norm for finetuning ViT-Large on Flowers.

Figure 25: The eigen value evolution for seed

7

on finetuning ViT-Large on CIFAR-10.

Figure 26: The eigen value evolution for seed

42

on finetuning ViT-Large on CIFAR-10.

Figure 27: The eigen value evolution for seed

1234

on finetuning ViT-Large on CIFAR-10.

Figure 28: The eigen value evolution for seed

7

on finetuning ViT-Large on Flowers.

Figure 29: The eigen value evolution for seed

42

on finetuning ViT-Large on Flowers.

Figure 30: The eigen value evolution for seed

1234

on finetuning ViT-Large on Flowers.

Panel of Figure 31 — Figure 31: Top $25$ eigenvalues of Hessian at solution obtained by SGD and Adam after finetuning ViT-Large on CIFAR10.

Panel of Figure 32 — Figure 32: Top $25$ eigenvalues of Hessian at solution obtained by SGD and Adam after finetuning ViT-Large on Flowers.

Metric	SGD ( $\eta=7e-5$ )	SGD ( $\eta=0.1$ )	Adam ( $\eta=7e-5$ )
Val Acc	$43.87\pm 24.02$	$84.80\pm 1.00$	$85.95\pm 0.64$
$L_{1}$	$5002.44\pm 0.0$	$6066.93\pm 34.33$	$31079.54\pm 754.26$
$L_{2}$	$0.73\pm 0.00$	$1.26\pm 0.01$	$5.57\pm 0.24$

Appendix

Table 12: Validation accuracy and parameter distance traveled in terms of

L_{1}

and

L_{2}

norm for finetuning Bert-base on MRPC.

Figure 33: The eigen value evolution for seed

7

on finetuning Bert-base on MRPC.

Figure 34: The eigen value evolution for seed

42

on finetuning Bert-base on MRPC.

Figure 35: The eigen value evolution for seed

1234

on finetuning Bert-base on MRPC.

Panel of Figure 36 — Figure 36: Top $50$ eigenvalues of Hessian at solution obtained by SGD and Adam after finetuning Bert-base on MRPC.

Neural Network Training Dynamics

Neural network training dynamics studies how parameters and representations change during optimization and how those changes shape learning, stability, and generalization.

See all 499 papers

Feb 18, 2026

The Implicit Bias of Adam and Muon on Smooth Homogeneous Neural Networks

We study the implicit bias of momentum-based optimizers on smooth homogeneous models. We show that momentum steepest descent algorithms like Muon (spectral norm), MomentumGD (

\ell_2

norm), and Signum (

\ell_\infty

norm) are approximate steepest descent trajectories under a decaying learning rate schedule, proving that these algorithms have a bias towards KKT points of the corresponding margin maximization problem. We extend the analysis to Adam (without the stability constant), which maximizes the

\ell_\infty

margin, and to Muon-Signum and Muon-Adam, which maximize a hybrid norm. Our experiments corroborate the theory and show that the identity of the margin maximized depends on the choice of optimizer. Overall, our results extend earlier lines of work on steepest descent in homogeneous models and momentum-based optimizers in linear models.

Eitan Gronich, Gal Vardi

Department of Computer Science and Applied Mathematics, Weizmann Institute of Science, Rehovot, Israel

Gradient Descent Adaptive Gradient Methods Algorithmic Bias

Ask PDF

Sep 17, 2026

Stiefel-AdamW: Geometry-Aware AdamW for Linear Factorization Blocks

A pervasive structural pattern in modern deep learning is the linear factorization block: a submodule of the form

W = BA

in which two parameter matrices are multiplied directly, with no intervening nonlinearity. Such blocks appear in LoRA adapters, low-rank compressed layers, query-key products of self-attention, and share a common pathology: the factorization is non-unique, which can destabilize training and limit usable learning rates. Despite this, factorization blocks are typically optimized with standard Euclidean methods that ignore the underlying geometry. We introduce Stiefel-AdamW, a near drop-in replacement for AdamW for use wherever such blocks appear. By constraining one factor on the Stiefel manifold while leaving the other Euclidean, Stiefel-AdamW relaxes the full

\mathrm{GL}(\mathbb{R}^r)

gauge symmetry to a compact orthogonal symmetry, ruling out factor blow-up while retaining the coordinate-wise diagonal preconditioning that gives AdamW its practical strength. Moment estimation is performed in the ambient Euclidean space, with geometry entering only through a tangent-space projection and a manifold retraction. The implementation overhead over AdamW is minimal, and we show that the resulting optimizer inherits both the stability benefits of Riemannian methods and standard convergence guarantees. We validate Stiefel-AdamW on LoRA-style fine-tuning of GPT2, ViT, and Mistral 7B and on full pretraining of GPT2 on OpenWebText, showing consistent improvements over strong baselines at essentially no additional cost over AdamW.

Emanuele Zangrando, Marco Sutti, Francesco Tudisco

Gran Sasso Science Institute, L’Aquila, Italy · University of Edinburgh, Edinburgh, UK

Deep Learning Optimization Riemannian Optimization Matrix Factorization Stiefel Manifold Optimization

Ask PDF

May 7, 2026

A Rod Flow Model for Adam at the Edge of Stability

Cohen et al. (arXiv:2207.14484) observed that adaptive gradient methods such as Adam operate at the edge of stability. While there has been significant work on continuous-time modeling of gradient descent at the edge of stability, extending these models to momentum methods remains underdeveloped. In the gradient descent setting, Regis et al. (arXiv:2602.01480) introduced rod flow, which models consecutive iterates as an extended one-dimensional object -- a "rod." Here we extend rod flow to Adam by working in the joint phase space of parameters and first moment

(w, m)

and treating the second moment

ν

as a smooth auxiliary variable. We also develop rod flows for heavy ball momentum, Nesterov momentum, and scalar and per-component versions of RMSProp, Adam, and NAdam. For all eight optimizers, we empirically evaluate rod flow on representative machine learning architectures, where it tracks the discrete iterates through the edge-of-stability regime significantly more accurately than the corresponding stable flow.

Eric Regis, Sinho Chewi

Yale University

Edge of Stability Momentum Methods Adaptive Gradient Methods Gradient Descent Dynamics

Ask PDF

Abstract

Neural Network Training Dynamics

Explore similar work

The Implicit Bias of Adam and Muon on Smooth Homogeneous Neural Networks

Stiefel-AdamW: Geometry-Aware AdamW for Linear Factorization Blocks

A Rod Flow Model for Adam at the Edge of Stability