Never Saddle for Reparameterized Steepest Descent as Mirror Flow
Authors: Tom Jacobs
Chao Zhou
Rebekka Burkholz
Venue: ICLR 2026 Poster (2026)
Keywords: Implicit bias
mirror flow
sign gradient descent
Adam
AdamW
steepest descent
reparameterization
diagonal linear networks
Primary area: optimization
TL;DR: The connection between reparameterizations and steepest mirror flows shows that the geometry of steepest descent is directly shaped, affecting feature learning by enabling saddlepoint escape, promoting sparsity, and stabilizing invariances.

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.

Published as a conference paper at ICLR 2026
NEVER SADDLE FOR REPARAMETERIZED STEEPEST DESCENT AS MIRROR FLOW
Tom Jacobs Chao Zhou Rebekka Burkholz
CISPA Helmholtz Center for Information Security, Saarbrücken, Germany (tom. jacobs, burkhol @
ABSTRACT
How does the choice of optimization algorithm shape 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 a unifying theoretical framework. This framework reveals how optimization ge- ometry governs learning dynamics, implicit bias, and sparsity and it provides two explanations for why Adam and AdamW often outperform SGD fine-tuning. Focusing diagonal linear networks and deep diagonal linear reparameteriza- tions (a simplified proxy for attention), we show that steeper descent facilitates both saddle- point escape and feature learning. In contrast, gradient descent re- quires unrealistically large learning rates to escape saddles, an uncommon regime in fine-tuning. Empirically, we confirm that saddle-point escape central chal- lenge 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.
INTRODUCTION
Optimization central driver of modern machine learning. First-order methods are particularly com- mon in deep learning, where models heavily over- parameterized and trained on highly non-convex ob- jectives populated with many saddle points and mul- tiple global minima. In this regime, the choice of op- timizer not merely about convergence speed (Pas-) canu et al. 2025): different algorithms can converge to different solutions with markedly different prop- erties like generalization, sparsity, and robustness (Woodworth et al. 2020 AAror al. 2019 Jacobs
10
10"
x x
10
q 1, 1 (SignGF)
10-10
q 1.5
10
2 (GF)
10
10'
10"
& Burkholz 2025 Tsilivis et 2024
Time
To understand the solutions that are preferred due to interplay between overparameterrizatoon and the optimization algorithm, a geometric lens has proven especially useful. well known that overparame- terization under gradient flow (GF) can induce mirror
Figure 1: For a deep diagonal linear net- work initialized close to saddle point, sign gradient flow (SignGF) converges faster than gradient flow (GF).
flows, changing the effective geometry in which optimization proceeds (Li et 2022 This per- spective clarifies how symmetries and balance constraints are preserved, how implicit regulariza- tion emerges, and how specific design choices like large learning rates, stochasticity, momentum, and explicit regularization can shape learned solutions (Marcotte et al. 2023 Kunin et al.) 2024
Gunasekar et al.) 2017 Woodworth al. 2020 Pesme et al.) 202 Even et al. 2023 Jacobs & Burkholz 2025 Jacobs al. 2025b Papazov al. 2024 Wang & Klabjan 2024
Tarzanagh al,
2023). Yet, most theories still center on gradient descent/flow, while modern practice in fine-tuning often operates in setting where plain (Stochastic) Gradient Descent (SGD) with small learning rates underperforms. In contrast, Adam (Kingma & Ba 2017) or AdamW (Loshchilov & Hutter 2017) variants routinely deliver more robust and stronger results.

---

Published as a conference paper at ICLR 2026
2.0
Hurd Escape Saddle
10
Blow
Metrie Exponent
1.5
1.0
10
Feature L earning
10"
0.0
10
0
2
8
10
Depth (L)
Time t
1 (SignGF)
q = 1.5 q q= 2 (GF)
Figure 2: Illustration of different steepest mirror flows (with varied q). On the left side, the metric exponent is shown dependent the associated depth. A high metric exponent increases the diffi- culty to escape zero and the instability of the flow. The right side illustrates saddle escape by plotting the solutions of the ODE's corresponding the metric exponents, dr, x dt, with xo = 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.
Why do modern adaptive methods work so well in fine-tuning, and what solutions do they.faoor? We approach this question by analyzing overparameterization and steepest descent methods via their resulting steepest mirror flows. Concretely,
study optimizer family indexed by q [1,2] that interpolates between GF (q = 2) and SignGF (q = 1). where the latter closely related to sign-based methods exhibiting Adam-like behavior. Working in this broader geometric setting is technically more challenging than for gradient flow, as we lose the inner product structure, making the optimization process operate in a Banach instead of a Hilbert space.
Following Nam et al. 2025)'s call for simple, analytically tractable models that nevertheless reflect focus deep diagonal reparameterizations (a simplified diagonal proxy forthe dynamics of KQ between the K key and Q query matrices in attention at depth (L 2)) and diagonal linear networks. Within this setting, we derive new balance equations that characterize the induced mirror flows and the metric exponent governing dynamics as a function of depth (see Figure This reveals significant qualitative difference for varying q: steeper descent (smaller g. approaching SignGF) facilitates saddle-point escape and feature learning, while GF (larger q) typi- cally requires unrealistically large learning rates to escape saddles (Pesme & Flammarion 2023 Du et al. 2017), which is uncommon in fine-tuning. Here feature learning reters to induce sparsity in the learned representation. Moreover, we show that decoupled weight decay (AdamW) controls a different balance equation from GF, which stabilizes feature learning without driving the dynamics into high-exponent (m > 1) regimes that impede saddle escape. The high exponent regime corre- sponds initial (exponential) slow down of convergence and finite time blow up corresponding global instability. These findings are in line with empirical observations.
common phenomena,
scenario for which the implicit bias is known classification on separable data. Recently, in this setting, max-margin characterizations have been derived for steepest descent (Tsilivis et al. 2025) and Adam (Zhang et al. 2024), establishing an -margin for both SignGF and Adam. However, this margin does not see the full geometry induced by overparameterizzaion,, as our analysis shows. For diagonal deep networks, the Loe -margin would be independent of depth, whereas our findings reveal that the margin actually depends critically on depth through the geometry that controls feature learning by the metric exponent (see Figure where a larger metric exponent leads sparser representation.
We validate these predictions for linear regression and separable binary classification, demonstrat- ing ground-truth recovery and the predicted saddle-escape behavior. Fine-tuning experiments on standard vision tasks and LLM adaptation further corroborate the generality our insights. Empiri- cally, we find that Adam-like steepest flows escape saddles faster and achieve stable feature learning at small learning rates. Decoupled vs. coupled weight decay exhibits the anticipated sparsity and stability trade-offs, aligning with our balance-equation analysis.
2

---

Published as a conference paper at ICLR 2026
Contributions.
Steepest mirror flows for family of reparameterized steepest flow dynamics. We develop a framework connecting reparameterizations to steepest mirror flows for family of steepest descent methods in separable settings, combining steepest descent and mirror geometry.
Qualitative gap between GF and SignGF. For deep diagonal reparameterizations. we prove that steeper descent (lower simultaneously escapes saddles faster and supports feature learning for deeper networks, whereas GF requires time rescaling / large learning rates to achieve comparable escape.
Decoupled weight decay for stability and sparsity. We show that Adam AdamW-style decou- pled weight decay enforces distinct balance equations from GF, yielding more stable feature learning and needs higher depth for sparsity.
Empirical validation in fine-tuning, We corroborate our theory for diagonal linear models by fine-tuning vision models and LLMs, highlighting (i) faster saddle escape with Adam- like flows and (ii) the predicted differences between coupled vs. decoupled weight decay for sparse, reparameterized training.
RELATED WORK
Specific reparameterizations trained with gradient flow in- duce a mirror flow Li al. (2022). This finding has been used to describe the implicit regularization induced by overparameterizztioo (Azulay et al. 2021 Vaškevičius et al. 2019; Zhao et al. 2022 Li ett al. 2021 Gunasekar et al.) 2017 Woodworth et al.) 2020), explaining, why highly overparame- terized neural networks can generalize well despite the risk of overfitting. Even the effect of large learning rates, stochastic noise, explicit regularization, and momentum can be covered by the theory (Pesme et al. [2021] [ven et al. 2023 Jacobs & Burkholz 2025 Jacobs et al. 2025b) Papazov et al.) 2024). Generalizing these results that apply gradient flows, extend the mirror flow analysis to steepest flows. This includes sign gradient descent, which has similar implicit bias as Adam (see Appendix[A]. As a highlight, we characterize the mirror flow stability with respect to the depth and type of descent algorithm. From technical point of view, our derivations overcome the challenge that, unlike gradient flows that operate in Hilbert spaces, steepest descent algorithms live in Banach spaces. (Banach spaces have less mathematically convenient structure, norms but not necessarily scalar products are defined.)
Mirror flow and reparameterizaatoos
Application of reparameterization to sparsity Recent work has used the implicit bias of repa- rameterizations to induce sparsity. (Jacobs & Burkholz]] 2025 Gadhikar et al. 2025 Jacobs et al.) 2025a) employ the mirror flow framework for gradient flows guide the (re- parameterized) train- ing dynamics, which are controlled by explicit regularization (Jacobs et al., 20256). The analysis centered around vision benchmarks where stochastic gradient descent with momentum is usually preferred over Adam. Kolb et al. (2025): Ziyin & Wang (2022) also exploit that reparameterized loss functions with Iæ-regularizzaiin are equivalent to a differently regularized optimization problem in the original parameters. Combining deep pointwise reparameterizations with weight decay, Kolb ctt al. (2025) observe that higher depth leads to extreme sparsity and performance degradation. For sign gradient descent we show that decoupled weight decay, in contrast, actually needs higher depth to induce sparsity. This reveals a key difference between coupled and decoupled weight decay.
Steepest descent and saddles Recent studies have revisited steepest descent as unifying lens for understanding optimization in modern machine learning. Fan et al. (2023) and Tsilivis et al. (2025) analyze the implicit regularization induced by different steepest descent algorithms classification settings with separable data, showing that the iterates approach particular max margin solution. Building on this line of work, Large et al. (2024) and and Bernstein & Newhouse (2025) highlight how modular duality provides a basis for steepest descent based algorithm design. A similar max margin implicit bias characterization has been provided for adaptive algorithms, including Adam (Zhang et al. 2024). For AdamW, the effect of decoupled weight decay implicit bias can be expressed as a bound on the Lx norm for general objective functions (Li et al. 2025). The convergence of sign gradient descent, optimizer with implicit bias similar to Adam, has also been studied, connecting
3

---

Published as a conference paper at ICLR 2026
its behavior to Lipschitz smoothness and yielding looser convergence bounds than gradient descent [Balles et al. 2020), with comparable rates in settings with unbounded smoothness (Crawshaw et al.) 2022). As we show, overparameterization can lead to faster convergence for sign gradient flow than for standard gradient flow, which we attribute to better saddle point escape.
In finetuning, small learning rate is preferred to not alter the representation much too prevent catastrophic forgetting (Zhou et al. (025). This clashes with the fact that saddle point escape needs time rescaling in gradient flow dynamics (Pesme Flammarion 2023). Note that different mecha- nisms that have been shown and studied allowing for saddle point escape are large learning rate and noise perturbation (Jin et al. 2017; Fang al.] 2020, Roy al.) 2020). In contrast, our analysis reveals a different mechanism which only relies on the geometry of the dynamics. As we show in experiments (Figure 5a SGD with a small learning rate can not escape saddle points while Adam can.
Conservation and algebraic invariance The reason why reparameterizations can induce a mirror flow is that gradient flow satisfies symmetries that do not change during training (Marcotte et al. 2025 2024| 2023), i.c. so called balance equations. The scale and the relative scales of these invariances are important. Note that the relative scale also referred as X-balance (see Definition A slight initial imbalance can support feature learning, according to (Kunin et al. 2024). The gradient flow of deeper networks has also been studied under balanced invariance as dynamical
3.3
system Arora et al. (2019); Gadhikar & Burkholz (2024); Gadhikar et al. (2025); Boursier et al.) (2022). Even exact solutions have been derived for two layer networks using a Ricatti equation (Dominé et al.) 2024; Saxe et al. 2014, Xu & Ziyin| 2024). Less is known about steepest descent algorithms. We show that the relative scale for steepest descent optimizers can differ significantly, explaining, why sign gradient descent can train relatively faster than gradient descent.
BACKGROUND: REPARAMETERIZZTTI AND MIRROR FLOW
R). This can be accom- plished with gradient descent: I++1 - ( Tinit where > 0 is the learning rate. We study the resulting flow by taking the learning rate v - 0, resulting the differential equation: dx, [ dlt, Tinit-
Consider minimizing continuously differentiable objective f C (R",
Reparameterizations and mirror flow Training reparameterizztioos of x with gradient flow have been connected mirror flows (Li et al., 2022, Jacobs al. (0256). (See Appendix C for sum- mary). Concretely, consider the reparameterizztionn E C"MM. R"), assuming that M is smooth manifold. This corresponds to the gradient flow: dw; suitable conditions, this can be described by mirror flow:
wi [gg(() wp = Winit Under
dV, R(x:) = f
TD
Tinits
(1)
where R implicit bias (Sun al. 2022: Pesme et al. 2024 Gunasekar et al. 2018), i.e. the type of solution we converge to.
R" + R is Legendre function (see Definition 3.1 A mirror flow can control the
Definition 3.1. (Legendre Function, Definition 3.8 (Li et al. 2022 )) Let R Rd + R U {2o)
} be differentiable convex function. We say R is Legendre function when the following holds: 1) R is strictly convex on the interior of its domain int(dom R)). ) 2) For any sequence {x; {x:) going to the
boundary of dom/t, the gradient diverges, i.e. lim;- ||V 1.V..(((z)) (T,,)((;))l) (x;))
xc.
Example 3.2. Let the reparameterization g R R + I." be deep diagonal linear network g(m, w) = O w or equivalently g(m, w) = diag(m) diag(w). Assuming [Wa,maal < M4,inis the corresponding Legendre function is:
R(x) = 2
Σ 24 aresinh
- V2 + 2x3 - x; log
Mijinii + Wi.imit
Mijinit - Wi.init
sE(m
where A; = m w iinir This corresponds the hyperbolic entropy which interpolates between L Ln-ormm (A + 0) and L2- norm (A - oc) implicit bias (Woodworth et al. 2020). Moreover, R is also Bregman function B.9. which is property necessary for convergence.
4

---

Published as a conference paper at ICLR 2026
In Example 3.2 À controls the relative scale. This is connected to the preserved balance by gradient flow. Similar balance equations exist for products of matrices. The small scale is associated with sparsity and with this inducing feature learning. Furthermore, the reparameterization can be used as a proxy for the key K and query Q matrices in attention (Tarzanagh et al. 2023 Jacobs et al, 20256 Marcotte et al. 2025).
Definition 3.3. A product of parameters m E IR" and w E RR is called balancee iff m w² A1,, where we used the convention m² m2 i.e., element-wise multiplication and 1, the all one vector.
Marcotte et al. (2023) have shown that, if Definition 3.3 is satisfied, then balance is preserved under gradient flow for the more general matrix case. In other words, the parameters stay À X-balanced during training. This establishes connection between mirror flows and the balance equation.
Implicit bias and linear regression For mirror flows, the implicit bias for linear regression tasks can be characterized for general data sets. Let {2,, 9:) C R" x R be a dataset consisting of l samples with n features. The output of linear model with parameters The goal is to solve the regression predict the target vector Y (V1, 9/...... W²&) based input vector = (21, 22,..., 2b)) The next theorem establishes a mirror flow in this setting.
the i-th data is 2 x.
Theorem 3.4. (Theorem 3.9 (Li et al. 2022)) Given (Z,Y), suppose the objective f(x) is of the form f(x) ) = fZzx for some differentiable j IR" - R. Initialized at Ho = Rinin assume that the mirror flow Eq. ( converges Too lim; Tr which satisfies Zz Y, then
DR(Ise ro) min DR(Z, zo); where DR(2, == R(x) R(xo) (V.((zo): x zo). =E
DR is also known as the Bregman divergence (Definition B.8 with respect to R.
Theorem 3.4 associates the Bregman divergence DR with the limits of mirror flow. In Example 3.2 if R is the hyperbolic entropy (Eq. (2), a balancing constant À + 0 induces feature learning regime and controls the strength of the induced sparsity bias. conclusion, the reparameterization and À allow us to control the implicit bias.
Inducing sparsity with reparameterizations Reparameterizations have been used to induce spar- sity in deep learning architectures (Ziyin & Wang 2022, Kolb et al. [025 Jacobs & Burkholz 2025) by exploiting the equivalence between the following optimization problems:
min (m @ w) (11ml)2, + || Il,,)
and min fxx + 2o Ral|=l|zz. 2E3"
m, WR
Hence, their local minima correspond to each other, see (Theorem 2 in Ziyin & WWang 2022
4 THEORY: STEEPEST MIRROR FLOW AND DEEP REPARAMETTERU.ATIONS
To characterize the difference between modern optimizers Adam (= SignGF) and SGD (= GF), study reparameterized steepest flows as steepest mirror flow. Our analysis is especially relevant for the finetuning setting, where small learning rates are used.
Steepest flows We consider a class of algorithms that is based on steepest descent with respect to the Lp norm. These are captured by the unnormalized steepest flow:
dx, = sign (v=f(x:)) @ (=f(x))) dt,
Finit
(3)
where 9 satisfies + = 1. Most interesting to us are gradient flow (GF) p = 2 (q = 2) and sign gradient flow (SignGF) = oo (q= 1), which is proxy for Adam (see Appendix Aj. On technical note, we mention that the unnormalized flow is equivalent the normalized flow up to time rescaling (see Appendix(B). The solution to the studied ODE does not have to be unique but can be interpreted the Filippov sense (Filippov) 1988). In this setting, Gunasekar et al. (2018) argue that a similar implicit bias characterization in Theorem 3.4 is not possible, except for which corresponds to standard GF. Accordingly, this also not possible for reparameterizattons still study the induced dynamics to analyze the feasibility of
2,
trained by Eq. 33.. However, feature learning. Our main objective is to make qualitative statements about the dynamics such as saddle point escape, stability and the effect of decoupled weight decay.
5

---

Published as a conference paper at ICLR 2026
Steepest mirror flows Consider Legendre function R. (Definition 3.1 with respect to the Lp norm is given by:
A steepest mirror flow
dV R(z,) - sign( Dign(Ve.lte v.l(=))) @ NaI(s)))
Fo
Pinit-
(4)
For this class of flows, we can show convergence using the second order condition of coercivity as Definition 4.1 i.e. the inverse Hessian is bounded from below by a positive constant.
Definition 4.1. We call function R E C² (R"; ; R) inversely p-coercive iff there exists a constant j O, the coercivity constant, such that for all IR:
V R (xx)x 24 1
Theorem 4.2. Let R CR R) be separable function (Definition B.9 that inversely p--
coercive (Definition 4.1, Moreover, assume that the set for E Dom R min (())) is non-empty and
there exists constant > such that for all > (..((x)) S B for all [n]. Then the loss decays and satisfies:
( S (f(zo) - -(()))) / (NB²²9).
Assume that E C (R" R) is strongly convex. Then for the iterates of Eq. 4 converges such that we have lim Jt. = where x² is the unique minimizer of f with linear rate 2A.-2A.
Proof. The proof follows from tracking the evolution of the loss f and the observation that for strongly convex functions the sign only zero when the minimum is reached (see Theorem E.I
Theorem 4.2
highlights the dependence of the convergence rate on the coercivity constant. As we will show, the coercivity will effectively correspond how hard it is to escape the saddle point set.
Deep diagonal reparameterizations For the deep diagonal reparameterization given by x g(w) = T7 Wt; in Example 3.2] we can study the steepest flow with respect the Li norm with decoupled weight decay as AdamW (Loshchilov & Hutter 2017) with = 1. The flow described for each E [Z.] by:
dw,,t -iign wV f (g (wr,)))) @ (o a, 01O
iiiit
(5)
As additional result, we show that all separable steepest mirror flows have corresponding reparam- eterization in Appendix G
Deep diagonal parameterization have inherent saddle points characterized next by Theorem 4.3
Theorem 4.3. Assume that (0) * 0. Then, in addition the saddle points of f. the deep diagonal reparameterization æ = g(w) = IT introduces saddle points at:
$ $= {ww, wy,) Vije(()) ut; wy 0, wh 0 i,j and i + j}-
Proof. Apply the saddle point condition from Definition D.I (see Theorem D.2).
implies that small initializations are close to the set S. Our next derivation shows how steepest mirror flows can escape such saddle points. The escape rate depends the following balance equations, which are satisfied by the dynamics. Remark 4.4. The points of the set S would not be saddle points of the regularized dynamics with coupled or decoupled weight decay. However, larger indicating that escaping from near the set S would harder for GF (q = 2) than SignGF (q 1).
Theorem 4.3
will see, the metric would still be smaller for
Balance equations The balance equations of the next lemma are needed to derive a mirror flow.
Lemma 4.5. Consider steepest descent with respect Top and weight decay, with 1 1. Then, for a deep diagonal reparameterization, i.e., I = g(w) = II wi satisfies the following balance equation for almost everywhere:
-wu,al = (Iwi,ol - 115,01%) exp (-9 (dds) for all i,j E [Z].
(6)
Proof. follows from deriving the evolution the left hand side of Equation
(see Lemma
E.2

---

Published as a conference paper at ICLR 2026
Lemma 4.5 leads the following natural extension
of Definition 3.3
Definition 4.6. A product of parameters m RR and
I" is X Ly -alannee with 1 +
iff
p 9
0.2
|m| |w]² A1n
where L, E R is the all-one vector.
0.2
We illustrate Def. 4.6 Fig. Observe that for smaller q,
1 (SignGE)
move faster away from the origin in both parameters, providing intuition for the saddle escape. Note, there is analogue that holds for gen- eral deep reparameterizationss as recently shown by Marcotte et al. (2025) for q 1.
q=
1.5
-0.4
-0.4
-0.2
0.0
0.2
0.4
Remark 4.7. We focus on fixed value for all æ E R". However the analysis is pointwise, therefore, we can have different values for per parameter.
Figure 3: The balance equation for q 11, 2] and initialization 0. Observe that the (curved) path away from the initialization to point the curve mw = I with x = =0.1 (in the plot) is shorter for smaller q, indicating faster sad- dle escape.
0.1, =
Saddle escape and stability The next theorem shows that the invariances above induce steepest mirror flow when weight decay is turned off. This allows us quantify the coercivity constant and also the stability the dynamics. Furthermore, we can derive explicit expressions for the seperable Bregman functions by considering À = or ===2
Theorem 4.8. Initialize deep diagonal reparameterization such that it is X Lp-balanced for with respect the first parameter w1. Then, steepest descent satisfies separable Ly-mirror flow almost everywhere:
dV A.R.
-iign (V.f (x,)) (Z) dt,
To
Finit
where R...L:: RR" + R² is a seperable Bregman and Legendre function when L I completely characterized by the balances of Lemma 4.5 For 2, we explicitly get
V.R 2),))((1)
4 X
Proof. First, express the metric in terms of [ using the derived balances. Second, use the implicit function theorem to express |wi]" function of x and A. For L 2, we can do this analytically using the quadratic formula. To show RL,.L is Bregman we use the properties of function V R such being separable, bounded from below, asymptotic behavior near the boundary and being an even function. (Full proof see Theorem E.3)
Corollary 4.9. For A Lp balanced initialization, steepest descent has coercivity constant μ 1-1
Corollary 4.9 allows us to directly apply Theorem 4.2 for globally stable configurations such that ابمم < 1. Furthermore, at face value, Corollary 4.9 could indicate that all steepest descent methods have the same coercivity constant. However, the same initialization corresponds to very different values for different p.
Corollary 4.10 Initialize the reparameterization such that wi and w; 1, > 0. Then, training in Eq. (5) is X2 Lp balanced and X9(L-1)
Proof. Plug into Eq.
13 Lemma 4.5
Corollary 4.10 indicates
indicates that, for smaller and thus larger p, indeed have large cocrcivity stant and therefore can escape the saddle set faster. For small X, the coercivity constant dominates the escape rate, shown in Figure

---

Published as a conference paper at ICLR 2026
Remark 4.11. The case p == oc, L corresponds to the same mirror map structure smoothed sign gradient descent in WWang& & Klabjan 2024).
For deeper insights into the dynamics, we are also interested in the shape of the Bregman function and its metric exponent, defined next. This we can derive explicitly in case of X = 0.
(²)=1(x) e for a constant & E (IR+)"
Definition 4.12. m called metric exponent, if litm litm |||+00
Lemma 4.13. For and = O, we huve;
im - - =:
R (x) = - Σ
(1z10g(13)) - By - -x,1og((,,0))
sE(m]
if = - # 1:
Jry |ry12 1
1
(x) = L (L 1)9 JE(n]
q -
m 1, RL,.L L is Bregman function with metric exponent the domain Rsign(as,o) x Raign(Zun.)) If < 1, the domain is R Otherwise, RI,.L not a Bregman function.
Proof. 1) Derive the inverse metric in terms |oo|_ 2) Integrate the metrics twice and use that V, R(ro) = (See proof of Lemma E.4)
Theorem (..)and Lemma 4.13 reveal a key distinction between GF (- SGD) and SignGF (= Adam). For GF with balanced initializations function is not satisfied, but it is for SignGF. This distinction has implications for the stability of the dynamics. Accordingly, SignGF cannot escape beyond the boundaries of the Bregman function, making it globally stable which is captured by Corollary 4.14 Moreover, this corresponds to large metric exponent (m 1) as in Figure 2b.). As illustrated in the figure, the large metric exponent also leads to an initial (exponential) slow down of the convergence. Together this characterizes the stability of the dynamics. Furthermore, the gradient now may grow unbounded violating the assumptions in Theorem 4.2
higher depth, the smoothness condition of the Bregman
Corollary 4.14. If X 2 then for p = only 2 is a valid Bregman function. Furthermore, for 00, L 2 2 are all valid Bregman functions. For p < 2, there is no valid Bregman function.
Recall that À needs become very small for feature learning it has to approximate the Bregman functions in Lemma 131t to induce sparsity. This we can accomplish with weight decay as shown in Lemma 45
The effect of weight decay For gradient flow, the effect of explicit regularization can be integrated into time-varying mirror flow (Jacobs al., (025b). For steepest flows, can only study the Rie- mannian gradient flow, or, more specifically, the induced regularization on the manifold generated by the separable metric tensor V.R.. This informs how regularization is affected by the geometry.
Definition 4.15. For the regularizer h(x)) = Σiejn h,(2;) with each h; E C (IR, IR), the on manifold regularizer with respect to a separable T.p steepest mirror descent characterized by R is (x)((r) == -Eiei) S of (x.)),,,,
such that we have
R(r.) = sign (V.f (at)) @ V=/(E=)) di VzMreg(D)d)
F0 Tinit-
Theorem 4.16. Assume a) m ب # 2 or b) =qh-- 2. The manifold regularizer for decoupled weight decay with I.p steepest descent on the manifold for a reparameterization of depth
L with balanced initialization () = 0) is: a) Ie--149 Ese(m] 12:12 2- or b) Eie(a) log log((=s)).
Proof. Use V R from Corollary 4.13 and use (,h,(zz) Læ,, (See Theorem E.5 0
Example 4.17. For q = 2 (GF) and L 2, we recover |ı<]|r, as on manifold regularizer like (akesacobs)
& Burkholz 2025). ). For finite depth I., we get a - ||,, sparsity bias for q = - implying that for
q FS get L co.
8

---

Published as a conference paper at ICLR 2026
Table 1: Comparison of the effect of cou- pled or decoupled weight decay (Mrcg) for two reparameterization depths, namely, (L = 2, L. = 20). Note that the infi- nite depth would lead to a non-convex log- arithmic regularizer (log) in the coupled case, potentially leading to instability.
assume balanced initialization (A = 0). However, with sufficient amounts of weight decay, we know À 0 "fast enough" during training
In Theorem 4.16
according to Lemma 4.5 Hence, our insights generally also apply to À 0.
Example 4.17 establishes for SignGF (g = 1) that we need + oo induce sparsity with explicit decoupled weight decay. This stands in stark contrast to coupled weight decay, which would induce extreme sparsity, as shown in Theorem by Kolb et al. (2025). Table ! provides an overview of the effect weight decay on the induced regularization Mreg for I. = and I, oo. Note that these results imply that the respective flow cannot correspond to a time-varying steepest mirror flow, except for q 2 (GF), which is covered by Jacobs et al. (2025b). This follows from
Coupled
Decoupled
9 (SignGF)
log)
LLg,11
9 -= 1.5
(Lg,
(Lg. Lg)
q= 2 (GF)
(I,, log)
(L1, log)
Corollary [.6] in the appendix, according to which the manifold regularizer Mreg would need to match weight decay, which is impossible for q # 2.
5 EXPERIMENTS
The purpose of our experiments is to substantiate our theoretical findings. First, we verify our theoretical predictions on deep diagonal linear networks. Next, we show how our predictions hold practical settings such as reparameterized sparse training and finetuning of vision and language models. Appendix H. we study the natural invariance extension of Definition 4.6 for matrices and ablate the matrix product formed by the Q query and K key matrices in attention (as men- tioned in Example (32) for family of LLama mod- els (Grattafiori al. 2024). In practice, gradient flow is implemented as gradient descent with small learning rate (i.e. 17 0.0001 in Fig. T and 17 - 0.01 in Fig.
L -magang
L-1
L 1-10
L_ mmargii
20
40
60
80
100
Features
Figure 4: The Loe -margin for Adam with high and low depth L. The green region indicates the non-zero ground truth fea- tures. Higher depth leads to sparse ground truth recovery in line with Corollary 4.13
Diagonal linear network In line with our theory, we consider a diagonal deep network x 14 for re- gression and binary classification with respect to the mean squared error or exponential loss, respectively. x* denotes the sparse ground truth. This setting corresponds to Theorem B. 10 and Theorem B.12 Our initialization follows Corollary 4.110 for small À close to the saddle point set S. For the experieen- tal details, see Appendix and
In Fig. D first illustrate Theorem 4.2 by reporting the overdetermined setting for linear regression with & = 300 > n samples, n = 100 features, and depth L = 3. With high probability this ensures the existence of unique minimum. that is, strong convexity. We observe that takes significant more time for gradient descent with small learning rate to escape the saddle point initialization and reach the global minimum. For higher depth, this effect is intensified, as can be seen in the abla- where we also consider coupled versus decoupled weight decay to demonstrate Lemma 4.5 and study the effect of less data and small batch size in detail.
tions in Appendix[[] [[
In the classification setting, we consider k = 80 samples and a sparse ground truth (see Appendix F). Fig. 4 shows how higher depth leads to sparse ground truth Loe -margin recovery. This is in line with Corollary 4.13 for SignGF (I> Adam), where higher depth corresponds to a higher sparsity inducing Legendre function. This geometric bias was not covered before by max-margin results, as illustrated in Theorem
F.T Moreover, margins of SignGF and GF are compared in Appendix F
Finetuning scenario Fig. 5(a) illustrates a mechanism by which Adam can outperform SGD in fine-tuning vision task, despite SGD typically achieving better performance in vision pre-training scenarios. The top 50 eigenvalues of the Hessian spectrum were calculated with software from (Gol-

---

Published as a conference paper at ICLR 2026
1.0
SGD, n 0.001
Adam wd
Normalized Eigenvalues
SGD, 0.8
Adam
0.5
Adam, = 0.001
10
0.0
0.5
10
-1.0
10"
20 Eigenvalues
75
100
Epochs
(a) Top 50 eigenvalues of Hessian at solution ob- tained by SGD and Adam after finetuming on CI- FAR10. SGD with small learning rate has difficulty escaping the saddle point contrast Adam.
(b) T1 norm the weights during training for Adam with coupled weight decay strength 1 - and AdamW with 1e - 1. The dashed lines correspond depth = 10 and solid lines
2.
Eigenvalue spectra in finetuning for an ImageNet pretrained ResNet-18 18 on CIFAR-10 (a) and weight sparsity reparameterized training for a ResNet-50 on Imagenet (b).
Figure 5:
mant et al. 2018) for a ResNet-18 pretrained on ImageNet (Deng et al. 2009) after fine-tuning on CIFAR-10 10 Krizhevsky 2009). They highlight how far the optimizer has moved away from the ini- tial saddle point. We observe that Adam exhibits fewer and weaker negative eigenvalues, indicating that it escapes saddle regions more effectively than SGD, while achieving higher performance. In Appendix K additional ablations are provided, including additional experiments on different archi- tectures (ViT-large, Bert-base) and datasets (Flowers (Nilsback & Zisserman (008), MRPC). The validation accuracy is reported in Table which shows that Adam outperforms SGD for both small and (tuned) large learning rates. The specific learning rates are given in Appendix K
Table 2: Validation accuracy for finetuning scenarios (95% confidence interval.).
Model
Finetune
SGD << 1)
SGD (Ir 0)
Adam (Ir <<<<<>>>> 1)
R8
CIFAR-I0
19.15 - 2.82
93.60 = 0.38
95.19 0.21
ResNet-18
Flowers
1.22 0.53
62.13 = 1.10
80.50 1.38
Vare
CIFARIO
73.27 3.68
99.07 = 0.35
99.28 0.07
ViT-large
Flowers
1.03 + 0.82
98.94 - 0.05
99.37 + 0.08
Bert-base
MRPC
43.87 + 24.02
84.80 = 1.00
85.95 = 0.64
Sparsification
Next, we analyze how decoupled weight decay alters the sparsity bias in a repa- rameterized ResNet-50 trained on Imagenet. As shown in Figure[5]b) AdamW exhibits sparsity- inducing effect only for very deep reparameterizations and sufficiently large weight decay, aligning with Table T The effects of weight decay strength and reparameterization depth are reported in Appendix and the validation accuracy in Table
DISCUSSION
We have studied training dynamics through a geometric lens that derives mirror flows for fam- ily of steepest-descent optimizers, moving beyond gradient flow into Banach space setting. This framework clarifies how optimizer geometry interacts with architectural choices (e.g., attention and reparameterizations).. While our analysis applies to deep diagonal reparameterizatioos, we corrob- orate its relevance more broadly via fine-tuning experiments on L.LM and vision tasks. The theory yields concrete, testable predictions that match practice: Compared to gradient flow GF (~ SGD), sign gradient flow SignGF (= Adam) escapes saddles faster, is more stable at small learning rates, and behaves differently under decoupled weight decay, as inducing sparsity with decoupled decay requires deeper reparameterizations. These insights translate into actionable levers for efficient fine- tuning: Select optimizer geometry to control saddle escape and tune depth to target sparsity. We view this as a step toward co-design of optimizers and architectures, and a foundation for extending our analysis to non-diagon -oa-diagonal models and discrete, stochastic training.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGEMENTS
The authors gratefully acknowledge the Gauss Centre for Supercomputing e.V. for funding this project by providing computing time on the GCS Supercomputer JUWELS at Jülich Supercom- puting Centre (JSC). We also gratefully acknowledge funding from the European Research Coun- cil (ERC) under the Horizon Europe Framework Programme (HORIZON) for proposal number 101 16395 SPARSE-ML.
REPRODUCIBILITY STATEMENT
For the theory, detailed proofs have been provided for the main statements in Appendix and used previously known statements have been provided in Appendix B and C Additional derived statements are provided in Appendices D and G For the experiments, the details are provided in
Appendices[F and and
LLM STATEMENT
To improve fluency of the text sentence level editing has been done using large language models.
REFERENCES
Felipe Alvarez, Jérôme Bolte, and Olivier Brahic. Hessian riemannian gradient flows in convex pro- gramming. SIAM Journal on Control and Optimization, 43(2):477-501, January 2004. ISSN 1095-7138.
doi: 10. 10.11377:00660012924499999
URL http: doi. org/ 10. 1137/
$036301290241991
Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep ma- trix factorization. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Cur- ran Associates, 2019.
http:c//ppppppeeee////p/pp files/
Shahar Azulay, Edward Moroshko, Mor Shpigel Nacson, Blake E Woodworth, Nathan Srebro, Amir Globerson, and Daniel Soudry. On the implicit bias of initialization shape: Beyond infinitesimal mirror descent. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 468 477. PMLR, 18-24 Jul 2021. URL https: http:::/pproceediing azu azulay21a. html
mwr. wwt.pprss//v33.9
Lukas Balles, Fabian Pedregosa, and Nicolas Le Roux. The geometry of sign gradient de- scent. ArXiv, abs/2002.0805 080066 2020. URL https: http://ppp..enattttccchhool.rr
CorpusID:2111711
Jeremy Bernstein and Laker Newhouse. Modular duality in deep learning. In Forty-second Interna- fional Conference on Machine Learning, 2025. htll:/htpss//ppeeeee.ee..eee///
id=hErdfffse FTstu
Jeremy Bernstein, Jiawei Zhao, Kamyar Azizzadenesheli, and Anima Anandkumar. signSGD with majority vote is communication efficient and fault tolerant. In International Conference on Learn- ing Representations, 2019. URL wtt.hhttss://peneeeee
net httt://frmmmiiem set. JACY7
Etienne Boursier, Loucas PILLAUD-VIVIEN and Nicolas Flammarion. Gradient flow dy- namics of shallow relu networks for square loss and orthogonal inputs. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neu- ral Information Processing Systems, volume 35, pp. 20105-20118 Curran Associates, Inc., 2022. URL httpss//ppoceeedns...errrirtto.e wwwpeeeddddnn...niiiissocosoe

---

Published as a conference paper at ICLR 2026
Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Yao Liu, Hieu Pham, Xu- anyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, and Quoe V. Le. Symbolic discovery of optimization algorithms, 2023.
owg./vivorrraaoo www.:bsgg///o..... http:/bbb/2220... 230. 06675
Frank H. Clarke. Optimization and Nonsmooth Analysis. Society for Industrial and Applied Math- ematies, 1990. doi: 10. 171113777778111199991110 URL ltttps://epubs... siam. org/doi/
Michael Crawshaw, Mingrui Liu, Francesco Orabona, Wei Zhang, and Zhenxun Zhuang. Robustness to unbounded smoothness of generalized signSGD. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022.
URL
net.:wwwwwwwweeee.eeeee..e wttnnee//oorrue Eorum id 801 2Ypp0]
J. Deng, W. Dong, R. Socher, L.J.. Li, K. Li, and L. Fei-Fei. ImageNet: A Large-Scale Hierarchical Image Database. In CVPRO9, 2009.
Clémentine Carla Juliette Dominé, Nicolas Anguita, Alexandra Maria Proca, Lukas Braun, Daniel Kunin, Pedro A. M. Mediano, and Andrew M Saxe. From lazy to rich: Exact learning dynamics in deep linear networks. In UniReps: 2nd Edition of the Workshop on Unifying Representations in Neural Models, 2024. URL (tthttts::/ /ttp://///ppeeeeieeee/eeemee net www.eeooomemiiomoooeooeeoeeeetee /tte:/orrrmmmmm11333111..1
Simon S. Du, Chi Jin, Jason D. Lee, Michael I. Jordan, Barnabás Póczos, and Aarti Singh. Gradient descent can take exponential time to escape saddle points. In Proceedings of the 31st International Conference on Neural Information Processing Systems, NIPS' pp. 1067-1077, Red Hook, NY, USA, 2017. Curran Associates Inc. ISBN 9781510860909
Mathieu Even, Scott Pesme, Suriya Gunasekar, and Nicolas Flammarion. (s)gd over di- agonal linear networks: Implicit bias, large stepsizes and edge of stability. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances Neu- ral Information Processing Systems, volume 36, pp. 29406-29448. Curran Associates, Inc., 2023. URL. htttps://proced// ngs. www.innn..nnnneeeeeeeeeeee neurips wtteriiss.cccpppeeeeeeeeeeeeoee http:pppppr_lllepppppereree File/ /ttp:///cc0000077777000000000010
Caraper--Coffeenccee.e. pdff
Chen Fan, Mark Schmidt, and Christos Thrampoulidis. Implicit bias of spectral descent and muon on multiclass separable data. In Hiigh-dimenssonal Learning Dynamics 2025, 2025. URL https:
//openreview www.:eet//oooe oott://ooom
Ying Fang, Zhaofei Yu, and Feng Chen. Noise helps optimization escape from saddle points in the synaptic plasticity. Frontiers in Neuroscience, 14:343, 2020.
A. F. Filippov. Differential Equations with Discontinuous Right-Hand Sides. Springer Dordrecht, 1988.
Advait Gadhikar and Rebekka Burkholz. Masks, signs, and learning rate rewinding. In Twelfth International Conference Learning Representations, 2024. /tRL/htttss////pereeeee
met/forune htt//oorumiiddie wid.:mid-dd0oovxxT00
Advait Gadhikar, Tom Jacobs, Chao Zhou, and Rebekka Burkholz. Sign-in to the lottery:
Repa-
rameterizing sparse training from scratch, 2025. URL
12801
Noah Golmant, Zhewei Yao, Amir Gholami, Michael Mahoney, and Joseph Gonzalez. pytorch- hessian-eigenthings:: efficient pytorch hessian eigendecomposition, October 2018. URL https:
Aaron Grattafiori, Abhimanyu Dubey, Abhinay Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, al. The llama 3 herd of models. arXiv preprint arXi::2407.21783, 2024
Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Srebro. Implicit regularization in matrix factorization. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Ad- vances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL UtR./tttess/pppreeeiineeeeesrsss Fife:/Saf1..iSSSeeceeee333eeeee1
12

---

Published as a conference paper at ICLR 2026
Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Characterizing implicit bias in terms of optimization geometry. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 1832-1841. PMLR, 10-15 Jul 2018. URL htttps://proceedings
mlr. porress/v8//e wwr..peess/v/b//uneege httpaa///wwb/gnassaaaaaaaa
html
Tom Jacobs and Rebekka Burkholz. Mask in the mirror: Implicit sparsification. In The Thirteenth International Conference Learning Representations, 2025. URL
net/
Tom Jacobs, Advait Gadhikar, Celia Rubio-Madright and Rebekka Burkholz. Ham: A hyperbolic step to regulate implicit bias, 2025a. URL
wwwpz//..rrr//ooooooooooo absa http:/lv...ggaa///550/e// 02630
Tom Jacobs, Chao Zhou, and Rehekka Burkholz. Mirror, mirror of the flow: How does regularization shape implicit bias? Forty-second International Conference on Machine Learning. 2025b.
URL https: httt:tttt///peeeeeee /ttps//i//peeeiiee...omom
net http:eviwww.nee//ourrreee Fow.:rmmmmdd...11SSS555-
Chi Jin, Rong Ge, Praneeth Netrapalli, Sham M. Kakade, and Michael I. Jordan. How to escape saddle points efficiently. In International Conference on Machine Learning, 2017. URL https:
//api www.emmnnnscchhhlaaaaaaacaamaaee htt.:oaaa..orrrooooopss::::::1:1
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017. URL
https::/arxiv.crrt wttp://www../////////o////////// 6980
Chris Kolb, Tobias Weber, Bernd Bischl, and David Rügamer. Deep weight factorization: Sparse learning through the lens of artificial symmetries. In The Thirteenth International Confer- ence on Learning Representations, 2025. URL ptttps://openrevieet net. http:/ww.ew/eeettt/e/ forum?id-
7
Alex Krizhevsky. Learning multiple layers of features from tiny images. 2009. URL https:
//api. www.nticscooolrr...c............ wrr.echooaaa.rrg/oooruus..o. 18268744
Daniel Kunin, Allan Ravent'os, Cl'ementine Carla Juliette Domin'e, Feng Chen, David Klindt, Andrew Saxe, and Surya Ganguli. Get rich quick: exact solutions reveal how unbalanced
initializations promote rapid feature learning. ArXiv, abs/2406.06158,, 2024. URL https:
seaant :ttpaaarrcrrroooposi:::0000
Tim Large, Yang Liu, Minyoung Huh, Hyojin Bahng, Phillip Isola, and Jeremy Bernstein. Scalable optimization in the modular norm. CoRR, abs/2405.14813, 2024. URL htttp:://di....
10. 1188550/aaX10 24011...440.... 2405 14813
Bingrui Li, Wei Huang, Andi Han, Zhanpeng Zhou, Taiji Suzuki, Jun Zhu, and Jianfei Chen. On the optimization and generalization of two-layer transformers with sign gradient descent. In The Thirteenth International Conference on Learning Representations, 2025. URL https:
//openzevies wttp://ttottooooo//ee/m
Jangyuan Li, Thanh V. Nguyen, Chinmay Hegde, and Raymond K. W. Wong. Implicit sparse regularization: The impact of depth and early stopping, 2021. URL https:///axxiii abs/2108. 05574
Zhiyuan Li, Tianhao Wang. Jason D. Lee, and Sanjeev Arora. Implicit bias of gradient descent on reparametrized models: On equivalence mirror descent. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Sys- fems, 2022. URL Intpts:/ pestts:///pperevvieet ht////ppervviiieeee net wtt.://orumm..ckk....... _
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Reprexentations, 2017. URL httph:tt/////. btttss://ap.. api. www.ssssaaattsccclrr...ror
Corpus ID: 535922235522
Sibylle Marcotte, Rémi Gribonval, and Gabriel Peyr'e. Abide by the law and follow the flow: Conservation laws for gradient flows. ArXiv, abs/2337.001144 2023. URL https: sattps://ags //pp
semant cschol wrtph//oarrorr//oortattooooo 6537
13

---

Published as a conference paper at ICLR 2026
Sibylle Marcotte, Rémi Gribonval, and Gabriel Peyré. Keep the momentum: Conservation laws beyond euclidean gradient flows, 2024. URL http:://axxivvvr.o / http://ab//2400005 12888
Sibylle Marcotte, Rémi Gribonval, and Gabriel Peyré. Transformatiie or conservative? oonpr//ommmmtttooooooocccettteve conservation laws for resnets and transformers, 2025. URL
Ionut- Vlad Modoranu, Mher Safaryan, Grigory Malinovsky, Eldar Kurtic. Thomas Robert, Peter Richtárik, and Dan Alistarh. Microadam: Accurate adaptive optimization with low space over- head and provable convergence. In The Thirty- Annual Conference Neural Information Processing Systems, 2024. URL
Yoonsoo Nam, Scok Hyeong Lee, Clémentine Carla Juliette Dominé, Yeachan Park, Charles Lon- don, Wonyl Choi, Niclas Alexander Göring, and Seungjai Lee. Position: Solve layerwise linear models first to understand neural dynamical phenomena (neural collapse, emergence, lazy/rich regime, and grokking). In Forty-second International Conference on Machine Learning Position Paper Track, 2025. URL/attps://ppereviie net.eeewwwwwweeeeeeeeeeee.eeeeee fote/ffrum?? fid-iriddddddoooo
Sree Harsha Nelaturu, Advait Gadhikar, and Rebekka Burkholz. TurboPrune: High-Speed Distributed Lottery Ticket Training. URL httpp:://gitini com/nelaturuharshal
TurboPrune
Maria-Elena Nilsback and Andrew Zisserman. Automated flower classification over a large number of classes. In Indian Conference on Computer Vision, Graphics and Image Processing, Dee 2008
Hristo Papazov, Scott Pesme, and Nicolas Flammarion. Leveraging continuous time to understand momentum when training diagonal linear networks. In AISTATS, pp. 3556-3564, 2024. URL
httpss:/proceedinggg www.innsssr......mmmm.m...mm.mae mlr wwwpsppeeeeeee... /etess/vv330 2ttp2//ppppaaaeeeee
Razvan Pascanu, Clare Lyle, Vondt-Vlad Modoranu, Naima Elosegui Borras, Dan Alistarh, Petar Velickovic, Sarath Chandar, Soham De, and James Martens. Optimizers qualitatively alter solu- tions and we should leverage this, 2025. URL
12224
Scott Pesme and Nicolas Flammarion. Saddle-to-saddle dynamics in diagonal linear networks. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https://
spenreview.net/ feceet/fooummffom (otoorummn--Lucgg11116
Scott Pesme, Loucas Pillaud- Pillaud-Vivien, and Nicolas Flammarion. Implicit bias of SGD for diagonal linear networks: a provable benefit of stochasticity. In A. Beygelzimer, Y. Dauphin, P. Liang, and Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL net. fttt//oorm?iddvve 7K
Scott Pesme, Radu-Alexandru Dragomir, and Nicolas Flammarion. Implicit bias of mirror flow on separable data. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. httLtttts///:pnnna
wttp//////ootee formm? www.:mrdddwwwwwssooro 1d=
Jesus Rios, Pierre Dognin, Ronny Luss, and Karthikeyan N. Ramamurthy. Sparsity may be all you need: Sparse random parameter adaptation, 2025.
Abhishek Roy, Krishnakumar Balasubramanian, Saeed Ghadimi, and Pragyya Mohapatra. Escap- ing saddle-point faster under interpolation-like conditions. In Advances in Neural Information Processing Systems, volume 33, pp. 12414-12425, 2020.
Andrew M. Saxe, James I.. McClelland, and Surya Ganguli. Exact solutions to the nonlinear dy- namics of learning in deep linear neural networks, 2014. URL MtRttt:/t///axtt org. abs/ 1312.67 6120
Haoyuan Sun, Kwangjun Ahn, Christos Thrampoulidis, and Navid Azizan. Mirror descent maxi- mizes generalized margin and can be implemented efficiently. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (cds.), Advances in Neural Information Processing Sys- fems, 2022. URL WtRLhttps:// /ttpss///ppeeeeeeeeeee net. fww.mrrrmmioooooomm dd -
Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transformers as support vector machines. In NeurlPS 2023 Workshop on Mathematics of Modern Machine Learning, 2023. URL htt.stt::: U wtwp:eeewnwt/ooooom//eeemmm
14

---

Published as a conference paper at ICLR 2026
Nikolaos Tsilivis, Natalie S. Frank, Nathan Stebro, and Julia Kempe. The price of im- plicit bias in adversarially robust generalization. In A. Globerson, L. Mackey, D. Bel- grave, A. Fan. U. Paquet, J. Tomezak, and C. Zhang (eds.), Advances in Neural In- formation Processing Systems, volume 37, pp. 58023-58057. Curran Associates, Inc., 2024. URL https::/prrceed ngs neur wttprriis.cccccpeeeeeeeeeeeee http:/paper_fles/ppppppeeeep /
Conference p pdf
Nikolaos Tsilivis, Gal Vardi, and Julia Kempe. Flavors of margin: Implicit bias of steepest de- scent in homogeneous neural networks. In The Thirteenth International Conference on Learning Representations, 2025. URL httpsttpss. :ettess//ei http:///op///eeeeeeeeeeeee net. ftt://orrmmm (et.rrmmmn-pppaaeiii G=BRpaPHD] 9r
Tomas Vaškevičius, Varun Kanade, and Patrick Rebeschini. Implicit regularization for optimal sparse recovery, 2019. URL BRLhhtts::// htttt:tt/tsxiiii// www.wi.www..a..ooooooo 05112
Shuyang Wang and Diego Klabjan. A mirror descent perspective of smoothed sign de- scent. In Conference on Uncertainty Artificial Intelligence, 2024. URL httpss//api.
hwtpattiiicccooaoooooooooooo..ee 2773422212
Blake Woodworth, Suriya Gunasekar, Jason D. Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes overparametrized models. In Jacob
Abernethy and Shivani Agarwal (eds.), Proceedings of Thirty Third Conference Learning The- ory, volume 125 Proceedings of Machine Learning Research, pp. 06635-3673.. PMLR, 09-12 Jul 2020. htLhhtts:/// ptattss:/pprcceednnsss mlr press http:////////////2/2222222222
html
Yizhou Xu and Liu Ziyin. Three mechanisms of feature learning in linear network. In International Conference on Learning Representations, 2024. URL
org/CorpusID: : 266999353
Chenyang Zhang, Difan Zou, and Yuan Cao. The implicit bias of adam on separable data. In The Thirty-eighth Annual Conference Neural Information Processing Systems, 2024. URL
Peng Zhao, Yun Yang, and Qiao-Chu He. High-dimensional linear regression via implicit regular- ization. Biometrika, 10944):0333-1046 February 2022. ISSN 1464-3510. doi: 10.1093/biomet/
asac010. /tt.://tt///....................
Chao Zhou, Tom Jacobs, Advait Gadhikar, and Rebekka Burkholz. Pay attention to small weights. In The Thirty-nimth Annual Conference on Neural Information Processing Systems, 2025. URL https: httt:::///penneviee...o hwt.:eeww.eet/toooo/ddodoooooooo
7MhC2)
Liu Ziyin and Zihao Wang. spred: Solving 11 penalty with sgd. In International Conference on Machine Learning, 2022. URL //api semanticscholes org/CorpuslI:
259075663
15

---

Published as a conference paper at ICLR 2026
A EQUIVALENCE BETWEEN SIGNGD AND ADAM
We recall the optimization algorithms Adam Kingma & Ba 2017) and SignGD here to highlight their connection. Moreover the equivalence SignGD with coupled and decoupled weight decay is mentioned. We can set e 0 and 31 82 = 0 in Algorithm [T. then we recover Algorithm 2 Similarly recover the equivalence of AdamW (Loshchilov & Hutter 2017) and SignGD with decoupled weight decay. Note that just setting = 0 already gives us sign like update as well. Note another related optimizer LION which sign gradient descent with momentum (Chen al.,
2023
Algorithm Adam with Coupled (a1) and Decoupled (a2) Weight Decay
1: Input: parameters To. learning rate 7,, decay rates B1, 12, e for stability, weight decay coeffi- cients (1, Q2
2: Initialize my - 0, dy + O, +- 0
3: while not converged do
4: 1- 1+1
5: Compute gradient:
9t x- Vol(re-1) + 4
6:
m4 f,m, - (1 B1)g,
7:
Dt - 820; + (1 32)9t
8:
m m1/(1 31)
9:
& + 0-1(1-13
Update rules:
11:
Coupled (Adam + a1):
th
2t. Tt-1 1)
Vbz + C
12:
Decoupled (AdamW W 022):
mp
Dt 81-1
11
02 +
13: end while
Algorithm SignGD with Coupled (a1) and Decoupled (a2) Weight Decay
1: Input: parameters Zo, learning rate 7,, weight decay coefficients 01,, Oz
2: t 2-00 0
3: while not converged do
4: /<(+]
5: Compute gradient (with coupled a1): ):
ge Vaf(x=-1) a13-11
6: Update rules:
7: Coupled (SignSGI on):
It-1 1J sisin(())
8:
Decoupled (SignSGD + az):
F4-1 - sigin (oognVa.foe (2=f(()11)) 70281-1
9: end while
16

---

Published as a conference paper at ICLR 2026
B CONVEX ANALYSIS, LINEAR REGRESSION, AND CLASSIFICATION
In this section we recall definitions from convex analysis and known results from the implicit bias literature.
Convexity an PI inequality For convergence to minimizer the objective function needs to sat- isfy some condition. Two common ones are convexity and the Pl.-inequality. Note that strong convexity implies both.
Definition B.1 (Convex Function). 0 E 1).
A function f R + IR is convex if for all x, E R" and
f(Or + (1 0)v) S (x) + (1-0) ()f
Definition B.2 Soolyak-Lojasiewiiee (PL) Condition). A differentiable function f satisfies the PL condition with parameter A > if
1
v.f(z)l 2 A((x)) J*) for all sr,
2
where f inf (1:).
Steepest descent
The family of steepest descent algorithms generalizes classical gradient descent arbitrary normed optimization geometries. We consider the same setting as in (Tsilivis et al,,
2025). Given a norm || with dual norm || ||+, the steepest descent update for loss f(z) is defined as
T++1 x2 7) Azz, where Az, arg
min (x.)ll
(u, v
|||||
When || 11-112, this reduces the familiar gradient descent method. More generally, the steepest flow in continuous time is given by
dr E arg
min (u,gt) 9t) E af(x)) ||ld
(8)
dt
where 0g(0:) denotes Clarke's subdifferential (Definition B.4 to allow for non-differentiable acti- vations such ReLU. For the Lp norm this reduces to:
dx, = sign( V f(x)) @ V (= 2-9 dt
TD
Tinits
where qg satisfies
1. Now define time rescaling - So || (r. 112 ds giving:
dx, sign v=f(x>)) V=](x=))
Tn Finit
This recovers the flow investigated in the main text.
Differential inclusion In order to study these flows we need to introduce what a Clarke subdiffer- ential is and differential inclusion. This is needed as the flow not be interpreted in the classic sense where there exists a unique solution. Instead we can use a set valued interpretation. Definition B.3 (Differential Inclusion). differential inclusion is a generalized ODE:
dxg F(x)).
t2 0,
where F: 1 = R" is set-valued. Definition B.4 (Clarke Subdifferentia).. For locally Lipschitz function f TR - R, the Clarke subdifferential at x: is
8 f(x) = convt lim vaf( as - f differentiable at 2 1k-00
Remark B.5. Gradient flows for nonsmooth convex functions be written as =tt) E (x)( (using the convex subdifferential), and more generally for Lipschitz functions using the Clarke sub- differential.
17

---

Published as a conference paper at ICLR 2026
Remark B.6 (Clarke subdifferential viewpoint on sign descent). Let g(u) ||u/|1- Its Clarke subd- ifferential is
8g(u) { RR" 8; (u:) u, + 0, 8; E if 0}.
ER
Hence, for any differentiable the sel-valued sign map satisfies
sign (vf(x)) 8 v=) I11.
Consequently, the sign gradient flow can be written as the differential inclusion
dæ
dt (a) ll,
which well-posed in the sense of Filippov for locally Lipschitz right-hand sides (Clarke 1990 Filippov 1988).
To avoid notation overload, will use the classical notation for steepest descent and write:
dx, sign v=f((x)) @ V =z(()=))--ddt
Zo Finit
Mirror flow A mirror flow can be defined in the classical sense:
R(=)((x))
(xx)dl
TD = Pinit-
where R is Legendre function (Defintion B.7]. The overparameterization in deep linear networks can be interpreted as mirror flow as we will see in Appendix C
Definition B.7. (Legendre function Definition 3.8 ((Li et al. 2022 Let R R.
R U {oo} be
differentiable convex function. We say R is Legendre function when the following holds:
R is strictly convex int(domR..
For any sequence { going to the boundary of domR, lim | [
For convergence of the itterates of the mirror flow as in Theorem 4.14 of (Li al. 2022 the function R also needs to be Bregman divergence and function, which we define Definitions B.8 and B.9
Definition B.8. Bregman divergence for generator function R points æ1, 32 E RomR:
R
R is defined for two
DR(1I, R(xi) - (V,R(22); w2)
-
Definition B.9. (Bregman function Definition 4.1 (Alvarez et al. 2004)) A function R is called
Bregman function if it satisfies the following properties:
domA is closed. strictly convex and continuous on domA. is C on int(dom/tt).
For any E domR and 7 E R, {y E domRt Dn(x, S 7 is bounded.
For any E dom/ and sequence { C int(dom/t) such that lim = x. it holds that lim;-+oc DR(1, 2x)
Implicit bias in linear regression We recall known result for the linear regression setup as also highlighted in Theorem B.10 We denote the data matrix with Z and outputs with Y. This includes, gradient flow, sign gradient flow and mirror flow. Note the mirror flow case covers the gradient flow case as it corresponds to R(x) := 1
Theorem B.10 (Implicit bias of gradient and mirror flow). (Gunasekar al. 2018) Let be Legendre function and initialize empty and that f R - T. is convex and or satisfies the Pl-inequality. Among interpolants, the mirror-flow limit (when it exists) minimizes Bregman divergence to Dmnii"
Vimit Assume that the set {o E domli Zu Y) is non-
r² argmin Dn(D,,(((.... Fina) such that Zr = Y.
Remark B.11. As shown in (Gunasekar et al. 2017) steepest descent algorithms do not nessecary allows similar characterization for linear regression in Theorem B.10
18

---

Published as a conference paper at ICLR 2026
Implicit bias for classification For steepest descent there is recent result seperable data for binary classification (Tsilivis et al. 2025). Similarly result for general mirror flow exists (Pesme et al., 2024), not steepest mirror flows. We focus the steepest descent result here as this includes our steepest descent reparameterization as well (it a homogeneous network). By exploiting invariances,
show that the margin has satisfy additional constraints for deep diagonal networks. Their analysis relies on the following assumptions, which are satisfied by many practical neural network architectures and our reparameterization:
1. Local Lipschitzness: For any 2; Rd. the mapping x + f(z;; x)) is locally Lipschitz.
2. L-Homogeneity: The network is homogeneous of degree L, i.e. f(;;o)) ct ((ix) for any c >>> 0.
3. Realizability: There exists tu > such that L(xw) < 1, ensuring that perfect training accuracy is eventually achieved.
We now recall the main result of the paper regarding the implicit bias of steepest descent.
Theorem B.12 (Convergence to KKT Points (Tsilivis et al. 2025, Theorem 3.4)). Under assump-
tions (3),3), consider steepest flow with respect to a norm on the exponential loss
L(x)=
(c(m)
Then, any limit point & of the normalized trajectory { TTxll J:20 lies in the direction of Kaanshh-KuuTTuukke (KKT) point of the margin maximization problem
{
min ZERP
|01 s.t. y.f(2x;)) 2 1, vi [m].
(10)
This theorem establishes that steepest descent algorithms implicitly bias the solution towards maxi- mizing a geometry- geometry-dependen margin.
19

---

Published as a conference paper at ICLR 2026
C REPARAMETERIZATIONS AS MIRROR FLOW
This section recaps the general results for reparameterizations and mirror flows and is based on Appendix A in (Jacobs et al.) 20256). For gradient flow we present the existing results for the mirror flow framework and time varying mirror flow framework. Consider an objective function f:R²
min f(x). mER
We can use the implicit bias framework to study the effect of overparameteriizatoon An overparame- terization can be accomplished by introducing function, y - IR", with M smooth manifold. For particular g. the reparamcterization of the loss function f leads to mirror flow. The general framework is given in (Li et al., 2022) and extended in (Jacobs et al.] 2025b) to study the implicit bias in terms of mirror flow. (Li al. 2022) provide sufficient condition for the reparameter- ization 9 such that induces a mirror flow Eq. ( The Legendre function R, see Definition B.7
controls the implicit bias and steers the trajectory of the dynamics.
In order to recover the convergence result in Theorem 4.14 in (Li et al. 2022) the function It also needs to be Bregman function, which is defined in Definition B.9 For reparameterization to induce a mirror flow with corresponding Legendre function we first have to give two definitions. Furthermore, we define ôg as the Jacobian the function g.
Definition C.1. (Regular Parmeterization Definition 3.4 (Li et 2022) Let M be smooth sub- manifold of R.. A regular parameterization M - R is a C" parameterization such that @G(w) of rank for all w E M.
For the second definition, we first need to define what a Lie bracket is.
Definition C.2. (Lie bracket Definition 3.4 (Li et al. 2022)) Let M be a smooth submanifold of R.. Given two C' vector fields X, Y M, define the Lie Bracket of X and Y as [X, Y(() (Y X (w) (x)))((w))
Definition C.3. (Commuting Parameterization Definition 4.1 eLi al. (022)) Let M be smooth submanifold of IRP A C2 parameterization
.M R2 is commuting in subset $ c M iff S. Moreover, we call g a
for any i (n], the Lie bracket [Vwgi Vudji (w) = for all commuting parameterization if it is commuting in the entire M.
One additional assumption is need ed on the flow of the solution. We define the solution of the gradient (descent) flow of function f M + R" initialized at x M
dx. V ((x)))
To x
(11)
as x4 (x) which is well defined if the solution exists. Using this we can make the following assumption.
Assumption C.4. (Assumption 3.5 Li et al. 2022) Let be a smooth submanifold of R and g: . TR be reparameterization. We assume that for any w M and [m), &g (w) is well- defined for (T) such that either lim,- Il (w) llz oc or T4 and similarly for T_. Also, we assume that for any w E M and i, E [n]. it holds that for (1, E R2 that de o, (w: is well-defined iff dit 63 (w)
Using these definitions state the known result for mirror flow.
Theorem C.S. (Theorem 4.9 (Li et al. 2022)) Let M be a smooth submanifold R.. and A + In be a commuting and regular parameterization satisfying Assumption C.4 For any initialization Wintt E M, consider the gradient flow for any objective f Ra - RR:
dw, = - J )d(
w6 = Whit-
Define x, = g(we) all 2 0, then the dynamics of is a mirror flow with respect to the Legendre function R. given by Lemma 4.8 in (Li et al.) 2022). i.e.
d( - (xt)
To ((wai))
Moreover, this R only depends on the initialization Wini and the reparameterizzation g, and is inde- pendent of the loss function Tup.
20

---

Published as a conference paper at ICLR 2026
Explicit regularization The above framework got extended recently in (Jacobs et al. 2025b) in- cluding explicit regularization. Consider the optimization problem:
min ((((w) ah(w). WEM
Then the dynamics becomes time varying mirror flow as described in Theorem C.6
Theorem C.6. Let (g, h): M + I++1 be regular and commuting reparameterization satisfying Assumption C.4 Then there exists fime-dependen Legendre function R. such that
Ray (24) - - (x=)d,,
To = g(Winit),
(12)
where a, = So as Moreover, Ra; only depends on the initialization Winu and the reparameter- ization g and regularization h, and is independent of the loss function f.
The deep diagonal linear reparameterizations do not satisfy time varying steep mirror flow as shown in Corollary E.6
21

---

Published as a conference paper at ICLR 2026
D DEEP DIAGONAL LINEAR REPARAMETERIZATION AND SADDLE POINTS
We characterize the saddle points induced by the deep diagonal linear reparameterization. For this first define what saddle points Definition D.I
Definition D.1. A saddle point E IR" of objective function f CA R) is characterized by:
(x) and (1)()) 2
L.e. it is critical point while the Hessian is not positive semidefinite.
Consider the product of parameters, Wi....; WT E R" as in the main text. Then the loss landscape of an objective function fxx with II w; has additional saddle points as characterized by the set S in Theorem D.2
Theorem D.2. Assume that (0) + 0. Then, in addition to the saddle points of J. the deep diagonal reparameterization x g(w) IT Wi introduces saddle points at:
= (w WL): V Vije(() Wt wy O, we + for + i and + j}
First we calculate the resulting gradient and Hessian using the chain rule:
(=)(x)=
Mi (E(L]
Vns(x)).
This implies that at least two wy = 0 to induce a critical point. Assume now that exactly two are indeed zero, then for the Hessian term depending V does contribute and get
(0,(()) v & H
where H2 is Hessian of ITE
= wi i.e. block matrices for every coordinate of x.
Every block matrix has two nonzero entries i.e. we have:
TT Me We if (k. = (f.j) or (j.i)
Ha.km
else
This matrix is indefinite with eigenvalues + We Since vzf( # 0 there is at least one negative eigen value.
Theorem D.2 highlights that if already one coordinate vector w; for E [4] is zero, the model is already close to a saddle point. This highlights that for the À b-balance, for small X, we are very close to saddle point.
22

---

Published as a conference paper at ICLR 2026
E
MAIN RESULTS: STEEP MIRROR FLOW AND INVARIANCE
We provide proof's here for the main results in the main text. The correspondence is:
Theorem
is
Theorem 4.2
Lemma
8 Lemma 4.5
Theorem
Theorem 4.8
Corollary
Corollary 4.13
Theorem 5 is Theorem 4.16
Theorem E.1. Let R C² R) be separable function (Definition B.9 that inversely p--
coercive (Definition
4.
Moreover, assume that the set {x: E Dom R min [(x) is non-empty and
there exists constant decays and satisfies:
> such that for all > (x,f( & B for all E [n]. Then the loss
I & dt S (f(zo)
(NB²-9).
Assume that f C (R", R. is strongly convex. Then for the iterates of Eq. 4 converges such that have lim x where x² is the unique minimizer of with linear rate uB9-2A.
Proof. The proof follows from tracking the evolution of the loss f and the observation that for strongly convex functions the sign is only zero when the minimum is reached.
First note the loss is decreasing:
(f) = - (V v ( sim
< - 11 V ( di
< 0
where we used that V2 R 1 is μr coercive and that it is separable. Rewriting the above equation gives us:
[.))(,,,, s (f(oo) fxxoo) // oc.
This resembles the classic sufficient descent lemma for L smooth functions. Moreover we have that:
V Vx,( di < fo B² V
implying that
[>=f(>)) dt - (f(#o) /) (μB²-9)
Note that if f is strongly convex then it satisfies the Pl.-inequality and we have:
(f())) = (x)) (x,) sien(Va.(az)))V<((()))))
S - (x:0)1
s - (
< --BO-2A (f(xx) f(x))) dt
where we use the bounded gradients and the fact that ya 2 B for for y RR. Then by Grönwall Lemma we have that:
((x) ((x)) S (f(x)) f(x")) exp (-1nB²-9A)
recovering linear convergence depending on μ and A. We now can use that for A-strongly convex functions we have for all æ E n" and the unique minimizer =
|| x= 112, < (f())) f(x"))
using this also have:
||,, x" 112, < A exp ( (ttuBa-21)
This concludes the proof.
23

---

Published as a conference paper at ICLR 2026
Lemma E.2. Consider steepest descent with respect to Lp and weight decay, with 1 + 1. Then, for a deep diagonal reparameterization, i.e., I = g(w) = II wi satisfies the following balance equation for 0 almost everywhere:
p
wi,r |w = (ww,,l - lw;,019) exp (-o a for all i,j [2].
(13)
Proof. This can be checked by deriving the flow of the left hand side:
d (lwa,el lw q sign (w.e))luc, ( dw,,t sign( W3,+) / 1 dww sign(w,, w/,, - sign (V w: (2re)) V we:/ (2;))1 dl
+ sign w3,+ - (V,,, (x+)) Vwy (xt)) (-
q a4 (1wr,1 (
It remains be shown that the first terms cancel out. We
the decompositions of signs and absolute values i.e. sign (ab) sign (a) sign (b) and lab = |a|lb]. Using this we get for all E [L]:
Sign(w..)) w 1 sign (V w., (mx)) aw f(x))(+=
sign( Wi Jign (M)))4)(()))))))))) |, 1 sign (v+((z)))=
sign( 1 sign (x:() (v)(.,)
which holds for all absolutely continuous solutions wy- Therefore, have that the evolution is given by:
([wareel |w,,,(9) = -q qq Oz (Iwa,r) (wyaa9) dt.
This is linear ODE of the form dž, = -q a,2,dt which has solution 2t = 2o exp (-9 a,
Plugging in 2t = [w,,e 1 |w yields the result. Note that this result has be interpreted in the Filippov sense i.e. for all absolutely continuous solutions this holds almost everywhere. M
Theorem E.3. Initialize a deep diagonal reparameterization such that it is X Lp balancede for with respect the first parameter w1. Then, steepest descent satisfies a separable Tipmirror flow almost everywhere:
dV dV RV, -sign (v (v.(t)))) de, To
Finits
(z) is seperable Bregman function completely characterized by the balances of where V. R.R.L.L (xx)
Lemma 4.5 For L = 2, we explicitly get
v 22 /4|x/1 + X2
Proof. First we derive an expression for the metric in terms of for i E [L]. We then use Lemma 4.5 to characterize 2 R- (xx). From the chain rule and decomposition of signs and absolute values
follows that:
dx,
Σ 11 1 sign (Va.(az) (E(T]]
Now using the invariance and balance assumption with respect to the first parameter w1 that holds a.e.:
|urrrl | 2 for all j E [4] \ (1),
express the inverse metric in terms of |22,a|9 and A:
R
(x) diag ([ww) + 2 + (L - 1)lun (/wi) + X) 2)2-2)
(14)
This is a continuous differentiable function in | 1 Moreover, we have that: lo| |w) (wwi)9 2
24

---

Published as a conference paper at ICLR 2026
By the implicit function theorem from calculus we know there exists continuous function w1 (z, A) for all x E I" and À > O. For this we need to have that there exists a unique positive solution to the polynomial equation of the form:
|x|* = z(z + A)m
where the left hand side is non-ne non-negative constant. We can show that the right hand side increasing for 2 O implying unique solution:
d dz (=(= + (((=+)) = (= 1) + (T. x))
for À >> 0. Thus there is unique solution. In case À 0 have that
|r|# +
Therefore in the case X > 0 we can guarantee using the implicit function theorem that we can express W1 in terms and x. Moreover, for À = an explicit expression is available. Plugging this into Eq. (14) yields the result.
For L = 2 we have that
|-|9 = |wn] 1 (Iw1 1 A).
This is quadratic equation in terms of |w1|9. We need to select the sole nonnegative solution, giving:
12 + 4
We can plug this into V (2R-1(x) giving
2/w1| +
/ X2 + |z/9
This concludes the first part.
remains to be shown that the implicit constructed mirror map is a separable Bregman function. We will use the connection between Legendre functions and Bregman functions to show this. We use that if the domain of a Legendre functions is R and its convex dual R has this as its domain well then R is Bregman function according Theorem 4.7 in (Alvarez et al. 2004). Therefore, we need to show R L is Legendre function and characterize the domains.
We first note that it separable by construction. This allows case. By construction, know that V. V. has domain R and range [Xa(t-11) oo). Therefore, V RLP.L has domain I and range (0, 1 (0(1-11. This holds for all [n]. This implies that R is strictly convex and (²(R",, (0, >-v((-1)))) proving the first condition of being Legendre function. For the essential smooth condition, we can use the asymptotic behavior near the boundary of the domain of V Rr,., This provides lower bound on [V.Rr,,.tl Concretely we use the triangle inequality and lower bound the growth of R.RLL.L
focus on the one dimensional
2
V V.RLp_L
V RL (t).
Rr. (1) (y)
|vl - dy
24 20
[-1
The right hand side only diverges if and only if ايوو < 1. Hence RI,.L a Legendre function. In order to show RI,.L is Bregman we use the following two observations. 1) The anti-derivative of an even function is odd 2) V. is an even function. It follows from 2) that also the reciprocal V RL even. Now we integrate and this implies that V, V,RL,.L is odd. Now using continuity and essential smoothness imply that the range of V,RL,.L is R. Therefore, the domain of the V,RL,., is R. This implies R Ri has domain R", Hence RI,, is Bregman function accordingly. 0
25

---

Published as a conference paper at ICLR 2026
Lemma E.4. For L 2 and À 0, candidates for the Legendre function are given by:
if = 1-1 1:
RI (x) L Σ
(zylog(Hz) - 2j Izgog((3,00)
jE(n]
ifm = - #
Jay 2- (1. -
= = L (L 1)9 jE(n]
13030013,00
/ ((t-t))
If mf 1, RI,.L L is Legendre function with metric exponent the domain (sig((21,0) X Rsigg(Z...o). If < the domain is R Otherwise, RE,.L not a Legendre function.
Proof. Plug in À = 0 and calculate wi (x. 0). This gives an explicit expression for the inverse metric:
V (x) -|
We now integrate the metric to get the Legendre function, keep notation clean omit the sum- ming over r; E [n] as the calculation is the same for all. Integrating the inverse twice and using that V, R((o)) = 0 gives: If q I-1 we have that
L
1" L/ul
V R(v)dvddu
dudu
1 / log(u) log(xo)
T.
T (zlog(()) x xllg((a)))
Moreover, if 1-1 #
we have that:
1
1" /
1 |o| 10 -² dudu
R(o)dduuu
L
/ u k//l" 8-
20 [z) 1-9
+
(L - (L-1)q- L du
(L
1)
L
|x --42
To |zo|| t z-
+ x (L 1)
(1 + 2)
((L 11 L)
|ar| -11+2
(1 - rzol (( c(+
L (L 1)4
This concludes the result. In order for RL,.L to be strictly convex need 9-او 1 the other conditions to be Legendre function such as essentially smooth are then also satisfied. The domains follow from the derived Legendre function cases. 0
Theorem E.S. Assume a) m = -اام # 2 or b) m = لااو 2. The manifold regularizer for decoupled weight decay with Lp steepest descent on the manifold for a reparameterization of depth L with balanced initialization (A = 0) is: a) T(2-9)+9 Eve(n) |2,| 1- or b) EiE(n) log((x;)).
Proof. The regularization rebalances the balance equation leading to the balance with À = 0. We can use Corollary to derive the metric. A key difference now is that the regularization is still on we have dynamics of the form:
dr; - -| 9 - (sign(V (xx)) (=)())))(+)
Londl,
Tp Pinit-
This can be equivalently written as:
dr, = - -|x, "
(sign(Va/(x:))) O V (x:()) + x. kx/ -/-0 dl,
55
Vinit-
26

---

Published as a conference paper at ICLR 2026
Similarly this can written as the mirror flow due to the equivalence of Riemannian GF and mirror flow:
dv R = (ign(Va)(zz)) @ (a/(rz))
mp/or 1 dt
Therefore, the on manifold regularization is the Mcg((1) (ar):
L12 Σie se 2- if 1-1 +
Mee (a) = Σ
|o:/ =. -x,d;
Eie(n] log( x if - 2.
iE(m]
This concludes the result.
Corollary F.6. Iff q = 2, weight decay is equal to the on manifold regularization May for for
Proof. Since À 0, the weight decay is given by
| li
|
2
sE|n]
We can match this with Mreg (r). For this we need to have:
L
+ L(2 9) + - 2 21 (1-1((11))))) L) = =2(1- L.)
2 L(2 - 9) +
which is true if and only if q = M
Corollary E.6 highlights that Theorem
C.6 can not be extended directly to steeper flows. This is due to the fact that the possible limiting regularization Mrg on the manifold mismatches with the weight decay i.e. À = 0, so in the end of training the time-varying mirror flow has to break down. Furthermore, the result Theorem CaoooremC.6jalrradyy breaks for as mentioned (Jacobs et al. 2025b).
27

---

Published as a conference paper at ICLR 2026
F IMPLICIT BIAS OF STEEP MIRROR DESCENT FOR BINARY SEPARABLE CLASSIFICATION
We present margin characterization for SignGF using recent result from (Tsilivis al. 2025) We observe that the margin should be independent of depth The margin now becomes dependent on maximum of | but this is an increasing function with the magnitudes as input thus the maximum would not change. In other words, the margin does not see what happens zero. However, our mirror flow analysis suggests that the movement speed of the parameters near initialization will influence the solution reached by slowing down movement near zero and accelerating further away. This helps with sparse ground truth recovery.
Theorem F.1. Consider a X-balanced deep diagonal linear networks trained in the linear separable classification setting as Theorem B.12 with sign descent then i, == || direction of a KKT point of margin maximization problem:
2t limit point lies the
min max | such that y, (2, za) 2 1 for all i e [ #ER" RE|d]
Proof.
follows from Theorem 1.12 that We 1
wa is in the direction of KKT point:
min Wiy E
|w,, .... Wr Li such that y (.) 2 1 for all i E [M]
where g(w) n wy. In addition, know the iterates
||w|| L.n
Combining this with
Lemma 4.5 it follows that for all i E [L]:
A 0
|
|| 1..
These additional constraints reduce the optimization problem to:
max [1x]] 2 such that y (x, for all E [k]
CR =y sE(d]
It is easy to show that 3; |T satisfies the KKT conditions above as well by using that in the
limit |wwlee max |x| and 1
have that:
St = lim 1 lim II 1,3,,
lim Rt =
lim
||
||wel
where the middle equality follows from the invariance relationship. This concludes the proof.
Experimental illustration We conduct experiment binary classification with exponential loss as described above. The main goal to illustrate the effect depth which would not have an influence according to Theorem F.T. However, our dynamics description would predict that higher depth will lead to a relative slow down near zero of the dynamics effectively creating a sparsity bias.
We generate sparse ground truth x² (1, 1, 0......... 0) R and k 80 data samples from random Gaussian such that Zi, ~ (0, 1) with i, E [100, 80). The labels are then determined by the classifier groundtruth i.e. Yj == sign(= = x Then we initialize at zero with = 0.1. We train for 10000 steps with learning rate 0.01. The optimizers used are SignGD, GD and Adam.
We report the final margin in Figure Observe that for higher depth the margin much sparser than for low depth. This highlights new implicit bias mechanism caused by depth, leading to feature learning. Note that for GD depth L = 10, did not converge, as expected. This explains the spiky nature of the Loe margin.
28

---

Published as a conference paper at ICLR 2026
La mmsggi
La moggia
mmagge
L-11
L-
L-11
20
60
80
100
20
60
80
100
20
CU
80
300
Festuure F
Feaitance Fn
(a) SignGD
(b) Adam.
(c) GD.
Figure 6: Resulting Lo 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.
G SEPARABLE MIRROR REPARAMETERZATION CONSTRUCTION
For completeness, we show how each separable steepest mirror flow can be seen reparameteri- zation of steepest gradient flow. This done by construction.
Theorem G.1. Consider a one dimensional steepest mirror flow with Legendre function R and is ju-coeriive. Then there exists a reparameterizattion g R - IR such that we have x = (
Proof. We can show this by construction in the one dimensional case.
A valid invertible reparameterization is (using ju-coercive::
| (02 (02((z)) dæ,
see this we can calculate the evolution of
dz, (02 dr. = - (000(((2) R(x)) sign (O.f(x))) (x:)())))" dt.
Now we use the implicit function theorem for the derivative of f with respect to 2:
0.f(x) = (0PR(=))) - 0xf(x).
Plugging this gives us:
d2, (8 [()) --1=) sign (3.f(( sign (0.f(x:)) (z)
Therefore I can be seen as the inverse of 2. Hence there exists steep gradient flow with respect to the reparameterizztion 2 -1 that corresponds a chosen mirror flow by construction. O Remark G.2. The proof in the one-dimensional case is quite simple as it is by construction. How- the proof in higher dimensions for standard mirror flow already relies on the Nash embedding theorem (Li et al. 2022) which not constructive.
29

---

Published as a conference paper at ICLR 2026
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 negative delta
Model
Q (%)
K (%)
Total (%)
Avg ÅLi
%, L
Avg Ar %, L2
LLaMA-3.1 8B
1.25
0.87
1.57
-624.13
100
-20.28
100
LLaMA-3.. 3.2MAA-3.220 3B
4.37
3.50
5.04
-1757.04
100
-101.71
100
LLaMA-3.2 IB
4.73
3.24
6.11
-891.76
100
-66.15
100
H INVARIANCE ISSUE FOR STEEPEST DESCENT FOR MATRIX INVARIANCES
The main hurdle for more general balance equation to hold is that the sign operator does not distribute over matrices. In other words for two matrices W1 and W2 we do have
sign (Wi W2) sign (Wi) sign (W2)
If this condition would hold plus the same condition with respect to the gradient then we would expect for a reparameterization g 1W1, W2) = W1 W2 trained with a sign gradient flow the following to hold during training:
Wi,ee L, [|2,,lle, (llWi,olle, (|W2,olle,) exp (-/ agds
This would then hold instead of the balance equation for gradient flow:
[i.elii, [/2,,li.... (llW (1W2,olli.g) exp -2 122 a,
which is known hold for gradient flow. To see this, compare for family of LLama models the base version with their tuned instruct version. Their tuning (partially) has been done with Adam W. Even tough, sign flips occur during training, effectively ruining the balance for wider reparameteri- zations. We empirically observe that for finetuning setting with small learning rate, less sign flips occur, making the insights from our example potentially relevant to larger scale finetuning. We track the direct generalization of the balance as in Definition 4.6 K key matices the attention mechanism:
for the matix product of the Q query and
L 11Qu112, 11,,
|1 pre pme 12,
pre 11,
In Table we observe that indeed the L1 balance is minimized more than the L2 balance which is an indication that our balance result might be able to generalize to the fine tuning setting where AdamW used. In addition, we observe for finetuning scenarios, that the signs of parameters change minimally. This we can capture by Definition H.T. which could lead to a bound on the invariance. However, this needs further assumptions on the nature of the gradients and how they evolve.
Definition H.1. Let g R=xm R + R be a reparameterizztion defined by 9W1, W2) = Wi W2. Then is called sign stable during training if for t 2 0,
sie,(We.r) sign(We.o)
for é E [2].
30

---

Published as a conference paper at ICLR 2026
1 ADDITIONAL EXPERIMENTS ON DIAGONAL DEEP LINEAR NETWORKS
For linear regression with mean squared eπror we set the groundtruth to (1, 1, 1, 1, 1.0,..... ................................ 0) E 18.00 and sample Z (0, 1) for [100], j E [k]. For our experiments we will train with steepest descent i.c. the discretization of Eq. ( and train with learning rate 17 1e 4 for 1e + 6 steps. For our experiments in the main text we will set wo = and w; = for E [L] (1), with X = 0.1. This ensures we start close to a saddle point as described in Appendix D Moreover, we vary the parameters q E [1, 1.,, 2J. [300, 80), L E [1. 2, 10], and study the effect of coupled and decoupled weight decay.
First we consider the underdetermined case with 80, to illustrate the different implicit biases at each depth L. In Figure 7 we that for high depth (L 10) sign gradient descent recovers the sparse ground truth and gradient flow can not escape the saddle, which is in line with our dynamical description. Moreover, for L = in line with the implicit bias of the hyperbolic entropy see Example 3.2
that gradient flow gets close the ground truth which is
Depth L=1
Depth L=2
Depth L=10
10
10
10
10-
10
10
10 Time (1)
10"
10
10
10
10'
10 Time (1)
10
Time (1)
q 1. (SignGF)
q 1.5
q= (GF)
1 we do net get close to the ground truth in all cases, for T, = 2 gradient flow gets close to the ground truth as in line with Theorem B0 and in for higher depth 10 the sign gradient flow (SignGF) converges close to ground truth which we would expect based on the dynamic reformulation.
Figure 7: Underdetermined linear regression (k 80), for depth
Next observe in Figure 8 and 10 that smaller batch size is beneficial for feature learning when the depth L plus steepest descent method q leads close to an L, bias. Furthermore, in Figure 9 with less data, the implicit bias argument does not prevail and we do not observe feature learning. This highlight that there is no guarantee for feature learning. However, it seems to be possible to remedy it with smaller batch size.
Moreover, we conduct an additional experiment for sign gradient descent with coupled and decou- pled weight decay of which the results reported in Table We use the same setting as described in the main text with & = 80 data samples and the same ground truth. We report the average L distance the theoretical predict balance value at the end of training which denote with Balance Distance, Observe that for coupled weight decay (a2) the distance increases while for decoupled weight decay (a1) we stay close to the theoretical predicted value. To add to this, high depth and decoupled regularization leads to recovering the ground truth the best.
The benefit of noise majority voting interpretation in (Bernstein et al. 2019) where it is used for convergence guarantees. If parameter needs be zero to reach the ground truth and starts at zero, the gradient is poten- tially small, however, it still has sign direction which might pull it away from the ground truth. Nevertheless, if we train with stochastic estimates we might equally moved either direction. This is captured by the following thought experiment, consider the gradient and stochastic gradient estimate:
The benefit of noise for feature learning could be seen from re-purposing the
-0.01 w.p.
V fz( 0.01 and g(x)
0.00 w.p. 1 2
31

---

Published as a conference paper at ICLR 2026
These estimators would have the same gradient expectation but the sign expectation is different i.e. we have
sign V and E [sign(g()))] =.
This indicates need stronger pull away from zero to actually move in the stochastic case. In other words, a larger majority of the gradients need vote for a certain direction.
Depth L-1
Depth L-2
Depth L-10
10
10"
10
10
10
10
10"
10
10'
Time (1)
Time (1)
Time (1)
1, (SignGF)
4.1.5
q= (GF)
Figure 8: Recovering the ground truth with small batch size for underdetermined regression with k
80.
Depth L-1
Depth L-2
Depth L-10
||x x I.
10a
10
10"
10"
10'
10
10"
10 Time (1
10'
Time (t)
Time ()
1. (SignGF)
q 1.5
2= (GF)
Figure Recovering the ground truth with full batch for underdetermined regression with 40.
Depth L-1
Depth L-2
Depth L-10
10
10"
10
10
10
10"
10 Time (1)
Time
Time (1
1, (SignGF)
1.5
q= (GF)
Figure 10: Recovering the ground truth with small batch size 5 for underdetermined regression with 40.
k
32

---

Published as a conference paper at ICLR 2026
Table Effect of Regularization Strengths on theoretical balance according Lemma 4.5 and the distance the ground truth. For the decoupled weight decay (0r) indeed stays close the theoretical predicted balance.
Depth L
Q1
02
Balance Dist.
Groundtruth Dist.
0
le-4
7.1
le-4
0
0
7.0
0
le-3
0
7.1
1e-3
0
5.8
0
1e-2
0
7.1
le-2
0
1.0
0
0
7.1
2
0
le-4
5.3e-4
4.8e-1
le-4
7.3e-5
4.7e-1
0
le-3
5.le-3
4.7e-1
le-3
1.8c-3
4.1e-1
0
le-2
3.5e-2
4.7e-1
le-2
6.7e-4
4.8e-1
0
1.le-4
4.8e-1
10
0
le-4
1.2e-1
2.5e-4
le-4
0
1.5e-4
1.4c-4
0
1e-3
3.9e-1
2.9c-4
le-3
0
2.9e-4
4.9e-5
0
le-2
7.8e-1
3.0e-3
le-2
0
1.6e-3
7.9e-6
0
1.5e-4
2.1e-4
J SPARSITY EXPERIMENT
In this section provide additional experiments for the reparameterized sparsity bias. More- over, we provide additional experimental details Table 5 The tunable parameters are depth
L E {2, 4, 10} and weight decay strength {le 1e 4). In the case for coupled weight decay we are effectively optimizing:
min ER
S(I(2 +
Σ
|| 122,
sE[L]
or equivalently
min f(x) La Σ zER"
|| 12/2
RE[L]
Theorem 1 in (Kolb et al. 2025). The code used is based Turboprune (Nelaturu et al.). The initialization of the depth reparameterization based on (Gadhikar et al. (025) and for deeper the balancing equation inform initialization i.e. and w: = for i f 1. This is closely related to the closed form formula for initialization of depth 2:
reparameterizations
w1
V + 2
V 2
=
v22
and
V2
where / /
/ x
with 7
}.
We can this from Taylor approximation around
d² and then 1/0 = v2 (1 x + t putting this
I = 0. Then we have ~ V2 x - together give:
mg + 2 and w/p T.
So when x² is negligible it matches proposed initialization for deeper reparameterization.
33

---

Published as a conference paper at ICLR 2026
In Figure and show the L1 nomm during training for Adam with coupled weight decay and Adam W. Moreover, we compare them directly in Figure [3] Observe that for coupled weight decay see that for both little and strong weight decay, the sparsity bias becomes more when the depth increases. In contrast, with less weight decay, Adam for higher depth, the I1-norm increases more. This is in line with the prediction for SignGF, which has the stationarity condition || < 4. Therefore, the parameter x can move more freely and the geometry has less effect. see higher sparsity bias
However when the weight decay is increased observe the opposite: for deeper reparameterization. Furthermore, we report the corresponding validation accuracies in T Observe the significant accuracy drops for Adam with coupled weight decay for increasing the regularization, an indication for extreme sparsity.
We conduct the same experiment for a ResNet-50 on Imagenet (Deng et al., 2009). We report for depth I. 2, 10 the 7, norm during training for both Adam with coupled weight decay and AdamW Figures a3d and 14 Validation aceuracy values are reported Table [ We observe the same behavior as for ResNet-20 CIFAR-IO, coupled weight decay leads to sparsity faster and with that drop in generalization performance.
Table 5: Training details for all experiments presented on sparse reparameterizations.
Dataset
Model
LR
Epochs
Batch Size
Optim
Schedule
CIFAR-10
ResNet-20
0.001
150
512
Adam, Adam
Triangular
Imagenet
ResNet-50
0.001
100
1024
Adam, AdamW
Triangular
10a
10a
10a
50
100
100
Epochs
Epochs
L=10
Figure 11: Adam with coupled weight decay trained with various depth reparameterizations for ResNet-20 on CIFAR-10. On the left is high regularization le and on the right less regular- ization le - 4.
34

---

Published as a conference paper at ICLR 2026
10a
10¹
10a
50
100
150
100
150
Epuchs
Epochs
L-10
Figure 12: AdamW (decoupled weight decay) trained with various depth reparameterizations for ResNet-20 on CIFAR-10. On the left is high regularization le 1 and on the right is less regular- ization le 4.
Table 6: Test Accuracy (%) + 95% CI for AdamW and Ad across depths and weight decays training a ResNet-20 on CIFAR-10.
Optimizer
Depth
Weight Decay
Accuracy + CI
AdamW
le-1
89.75 1 0.20
wd
2
1e-1 1
64.36 + 2.70
AdamW
1e-4
89.29 0.28
Ad
2
1e-4
88.27 = 0.08
AdamW
le-1 1
89.73 + 0.18
Adam+wd
le-1 1
58.23 + 4.98
AdamW
4
1e-4
89.38 + 0.35
Ad
le-4
86.55 1 0.25
AdamW
10
le-1 1
89.33 = 0.23
Ad
10
le-1
43.13 + 3.73
AdamW
10
le-4
89.49 1 0.06
Ad
10
1e-4
81.99 + 0.05
10
10
10a
10
20
20
60
Epochs
Epochs
L=10
Figure 13: Adam with coupled weight decay trained with various depth reparameterizations for ResNet-50 on Imagenet. On the left high regularization le-1 and on the right is less regularization 1c 4.
35

---

Published as a conference paper at ICLR 2026
6000
Adam + wd
4000
=
AdamW
2000
50
100
150
Epochs
Figure 15: L1 norm of the weights during training for Adam with coupled weight decay strength
1e and A with le 1. The dashed lines correspond to depth L 10 and solid lines to
L 2. The training setup is ResNer-20 CIFAR-10
10
5 10
10a
10
20
20
60
100
Epochs
Epochs
L=10
Figure 14: AdamW (decoupled weight decay) trained with various depth reparamcterizations for ResNet-50 on Imagenet. On the left high regularization 1e-1 and on the right is less regularization 1e 4.
Table Test Accuracy (%) + 95% CI for AdamW and Adam+wd across depths and weight decays training a Resnet 50 on Imagenet.
Optimizer
Depth
Weight Decay
Accuracy + CI
AdamW
le-1
76.23 1 0.07
Adam+wd
2
le-1
1.95 - 0.48
AdamW
le
73.32 0.11
Adam+wd
2
1e-4
73.35 + 0.05
AdamW
10
le-
62.20 + 0.25
Adam+wd
10
le-1
0.58 + 0.06
AdamW
10
1c-4
73.19 0.04
Adam+wd
10
1c-4
9.78 = 0.94
36

---

Published as a conference paper at ICLR 2026
K SADDLE ESCAPE FOR FINETUNING
In this section we present the saddle escape experiment for finetuning, We finetune a ResNet-18 that was pretrained on ImageNet on CIFAR-I0 and Flowers. To do this, we have replace the classifier layer with new randomly initialized one. We finetune the model with two different optimizers: SGD and Adam. Both cases are run for 15 epochs with the best learning rate selected after a sweep for both Adam and SGD. The learning rates are selected from preliminary sweep for
Adam 1j E {8e 4, le 3, 2e - 3, 3e - 3) and SGD 1j E (le 2, 5e - 2, le 1, 2e - 1, 3e 1, 4e - 1, 5e 1, 6e 1, Te 1, 8e 1, 9e 1}. We also run the best learning rate for Adam for SGD to illustrate our main point of the saddle point escape. Note that for vision tasks, SGD usually outperforms Adam. However, in finetuning observe the opposite. We track the top-50 largest eigenvalues during finetuning. For the experiment presented in the main text, we show the final eigenvalue distribution for the corresponding best validation accuracy.
In Ban the validation accuracy for both the CIFAR-10 10 and Flowers finetuning scenario are reported. Observe that Adam outperforms SGD in both cases. In addition, traveled by all parameters (including the classification layer) in terms of the I, and L2 norm, Adam has much larger [. norm indicating more uniform movement of the parameters. In other words, the adaptiveness of Adam allows all parameters to move more, which is as expected. In Figures 17 18 [9 and 20 we report the top 50 eigenvalues for each seed, not normalized and similar for the Flowers finetuning in Figures 21 22 23 and 24 We observe that the difference between the seeds is quite large. We believe that this is due to the randomly initialized classification layer. Furthermore, we report the normalized eigenvalues for each best seed also for Flowers finetuning in Figure [6 We observe less negative eigen values for Adam. Note that here we used standard SGD and Adam, not using parameter efficient versions such in (Zhou et al. 2025 Rios et al. 2025).
report the distance
that is.
Modoranu et al.
2024
Table 8: Validation accuracy and parameter distance traveled in terms of 1₁ and L2 norm for fine- tuning ResNet18 CIFAR-10.
Metric
SGD (17 = 0.001)
SGD (1) 0.8)
Adam (1 0.001)
Val Acc
19.15 - 2.82
93.60 = 0.38
95.19 = 0.21
L1
424911.48 + 34308.92
477750.60 + 10343.88
693101 .67 + 13509.59
L
29640.98 + 985.56
28409.58 + 219.53
27833.50 + 494.50
Table 9: Validation accuracy and parameter distance traveled in terms of Lı and L2 norm for fine- tuning ResNet18 on Flowers.
Metric
SGD (n = 0.002)
SGD (1 0.4)
Adam (1) = 0.002)
Val Ace
1.22 + 0.53
62.13 - 1.10
80.50 = 1.38
L1
206325.76 + 1327.51
173882.76 + 2967.69
618592.47 + 3445.18
L2
10124.38 76.52
7015.04 + 226.30
11432.11 + 350.38
K.1 ADDITIONAL VISION FINETUNING EXPERIMENTS
We now present finetuning experiments using large-scale transformer architecture, ViT-Large. We finetune ViT-Large pretrained on ImageNet on CIFAR-10 for 30 epochs and on Flowers for 15 epochs. As is standard in finetuning, the original classifier head is replaced with a newly initialized one. We evaluate two optimizers SGD and Adam with learning rates selected via a sweep: nj e (9e-5, -, 1e - 4, 1e - 4, 5e - 4} for Adam and 7 E {le- - 3, 5e 3, 1e 2, 5e 2, le - 1} for SGD. Additionally, we run SGD with the best Adam learning rate to further illustrate our observations on saddle escape. All experiments use batch size 1 128, weight decay O, cosine annealing learning rate scheduling, and label smoothing of 0. Because of the large model size and limited compute, we track only the top-25 eigenvalues. Table. 10 and report the validation accuracy and Flowers, along with the L1 and L2 parameter distance traveled (including the classifier layer). Adam consistently achieves higher validation accuracy than SGD both tasks. As in our earlier experiments, Adam induces a larger L, parameter shift, reflecting its more uniform adaptive updates.
CIFAR-10
Figure. 25
26
27
28
29
30 show the eigenvalue spectra across seeds and tasks. We additionally
37

---

Published as a conference paper at ICLR 2026
Normalizer Iaaaaiizdd Eigenvalnes
0.5
0.0
SGD, 0.002
0..5
SGD, 0.4
Adam, 0 0.002
-1.0
40
Eigenvalues
Figure 16: Normalized top-50 eigenvalues for ResNet-18 finetuned on Flowers.
Seed 123 - LR =0.001
Seed 123 - SGD LR=0.8
Seed 123 -AdamW LR-0.001
Vilue
-100
200
0
20
40
20
40
20 Bigenvalues
40
bigenvalues
Bigenvalues
Epoch
Epoch
Epoch
Figure 17: The eigen value evolution for seed 123 on CIFAR-IO.
Scod 456 LR =0.001
Seed A56-SGD LR-0.8
Sced 456 -AdamW LR=0.001
Value
1000
20
40
Eigenvalues
Eigenvalues
Eigenvalues
Epoch
Epoch
Figure 18: The eigen value evolution for seed 456 CIFAR-10.
Seod 789 - LR =0.001
Seed 789 SGD LR=0.8
Sced 789 AAdamW LR=0.001
Value
-200
20
40
20
40
Figenvalues
Eigenvalues
Eigenvalues
Epoch
Epoch 7
Epoch 15
Figure 19: The eigen value evolution for seed 789 on CIFAR-10.
38

---

Published as a conference paper at ICLR 2026
Seed 1000 LR --R00.000 =0.001
Seed 1000
SGD LR=0.8
Seed 1000-A400 -ddamW LR=0.001
100
Vilue
0
20
40
20
40
20 Bigenvalucs
401
Bigenvalues
Bigenvalues
Epoch
Epoch 7
Epoch
Figure 20: The eigen value evolution for seed 1000 on CIFAR-10 10.
Seed 123 - LR -0,002
Seed 123- SGD LR -.0.4
Seed 123 -AdamW LR-0.00 0.002
50
Value
-50-
-100
0
20
40
40
20
40
Eigenvalues
Eigenvalues
Eigenvaluce
Epoch
Epoch 7
Epoch
Figure 21: The eigen value evolution for seed 123 on Flowers.
Seed 456 - LR =0.002
Seed 456 SGD LR-0.4
Seed 456 AAdam LR-0.003
100
Value
-100
200
0
20
40
20
40
20
40
Bigenvalues
Eigenvalues
Bigenvalues
Epoch
Epoch 7
ppoch
Figure 22: The eigen value evolution for seed 456 on Flowers.
Seed 789 - LR -0.002
Seed - SGD LR-0.4
Seed 789 Addm LR-0.002
50
Vilue
-100
0
20
40
0
20
40
20
40
Bigetvalues
Bigenvalues
Bigenvalucs
Epoch
Epoch
Epoch
Figure 23: The eigen value evolution for seed 789 on Flowers.
39

---

Published as a conference paper at ICLR 2026
Seed 1000 -.R =0.002
Seed 1000 SSG LR-0.4
Seed 1000 Adam W 1R-00.002
50
Value
-50
-100
20
40
Eigenvalues
Bigenvalues
Bigenvalues
Epoch
Epoch 7
Epoch
Figure 24: The eigen value evolution for seed 1000 on Flowers.
provide unnormalized and normalized spectra in Figure and Figure 32 for different tasks. In the unnormalized CIFAR-10 spectra (Figure. B1aa. SGD with learning rate of le 4 produces substantially larger eigenvalues than the other configurations, obscuring the trends for Adam and SGD with le 2. Removing this outlier (Figure. (1bb) reveals that Adam exhibits fewer negative cigenvalues. The same behavior holds for finetuning ViT-Large on Flowers.
Table 10: Validation accuracy and parameter distance traveled in terms of L1 and L2 norm for finetuning ViT-Large on CIFAR-10.
Metrie
SGD (7) 0.0001)
SGD (7) 0.01)
Adam (7) 0.0001)
Val Acc
73.27 + 3.68
99.07 + 0.35
99.28 0.07
L1
460.47 219.86
24617.83 14406.4
453934.900 + 22278.43
Lg
0.48 + 0.059
6.25 4.29
39.47 = 1.01
Table 11: Validation accuracy and parameter distance traveled in terms of Lı and L2 norm for finetuning ViT-Large on Flowers.
Metric
SGD (1) = 0.0001)
SGD (1) = 0.0T)
Adam (1) = 0.0001)
Val Ace
1.03 0.82
98.94#0.05
99.37 + 0.08
LA
25.71 + 30.48
4655.83 + 576.49
108583.62 + 2078.48
L2
0.04 0.02
1.50 + 0.12
8.35 + 0.16
K.2 ADDITIONAL LANGUAGE FINETUNING EXPERIMENTS
In addition to our experiments on vision tasks, we conduct parallel study on language models. Specifically, we fine-tune pretrained BERT-base model on the MRPC task from the GLUE bench- mark, following the setup Zhou et al. (025)). The model is fine-tuned for 5 epochs using both SGD and Adam. Learning rates selected via sweep: 7) {5 x 10-5, 7 10 5 9 x 10-53 for Adam, and E {10
x 10 10-1) for SGD. We additionally evaluate SGD using the best learning rate obtained for Adam. As before, track the top-50 eigenvalues through- out training. Table 12 reports the validation accuracy along with the parameter displacement measured in L1 and La norms. Figures 33 34 and [35 show the evolution of eigenvalues across different random seeds. the model achieving the best validation performance. The conclusions mirror those observed in our vision experiments.
x 10 2, 10
the unnormalized and normalized eigenvalue spectra for
40

---

Published as a conference paper at ICLR 2026
Seed 7
=0.0001
Seed 5CG LR=0.01
Seed 7 AdamW LR 0.000
500
4-
Value (x10404)
250
2-
220
-2
5500
-4
10 Eigenvalues
10
20
Eigenvalues
Eigenvalues
Epoch 0
Epoch 15
Epoch 30
Figure 25: The eigen value evolution for seed on finetuning ViT-Large on CIFAR-10.
=0.0001
Seed 42 SGD LR=0.01
Seed 42 AdamW LR=0.0001
0.2
Value (x1e+04)
2
-0.2
-0.5
-0.4
10 Eigenvalues
Eigenvalues
Eigenvalues
Epoch 0
Epoch 7
Eooch 15
Figure 26: The eigen value evolution for seed 42 on finetuning ViT-Large on CIFAR-10.
1234
SGD =0.0001
1234 SGD LR=0.01
Seed 1234 AdamW LR=0.000
5.0
Value (x10444)
50
0.0
-25
-50-
Eigenvalues
Eige nvalueu
Eigenvalues
Epoch 0
Epoch
Epoch 15
Figure 27: The eigen value evolution for seed 1234 on finetuning ViT-Large on CIFAR-10.
Seed SGD =0.0001
Seed LR=0.01
Seed
AdamW LR=0.000
400
Value (x1e+04)
200
2
2
0-
200
-400
10 Eigenvalues
Eigenvalues
Eigenvalues
Epoch 0
Epoch
Eooch 15
Figure 28: The eigen value evolution for seed on finetuning ViT-Large on Flowers.
41

---

Published as a conference paper at ICLR 2026
Seed SGD =0.0001
Seed SGD LR=0.01
Seed AdamW LR=0.0001
0.75
20
1.0
Value (x1e404)
0.50
0.5
10-
-0.5
0.25
-20
-1.0-
-0.50
10
10
20
Elgenvalues
Eigenvalues
Eigenvalues
Epoch 0
Epoch
Epoch 15
Figure 29:
The eigen value evolution for seed 42 on finetuning Vir-Large on Flowers.
Seed 1234 - SGD LR =0.0001
Seed 1234 SGD LR=0.01
Seed 1234 AdamW LR=0.000
200
0.4
Value (x1e+04)
0.2-
0-
0.0
-100
0.2
-10
-200-
0.4
10
20
10
20
10
20
Eigenvalues
Eigenivalues
Eigenvalues
Epoch
Epach
Figure 30: The eigen value evolution for seed 1234 on finetuning ViT-Large Flowers.
SGD, D.CL
Eiyneetures (xTe-04)
50
Normaloed Eigerrokurs
25
Adam, a
Elpenatuees
Adams n 3.0001
0.0
30
0.5
73
Elpenwatue Index
Eigenvake
(b) Unnormalised, removing SGD
(a) Unnormalised.
0.0001.
(c) Normalised
Figure 31: Top 25 eigenvalues Hessian at solution obtained by SGD and Adam after finetuning ViT-Large on CIFAR10.
20
SGD, SGD. n 8.02
10
Adam. p= 0.5001
Noemabee Eipervation
1.0
SCD,
SGD, 0.0000
Adam. q 0.0001
Adam, 0.0001
Eggewatars
-30
0.5
Indox
-10
10 15 Eganvalic
10
15
25
Elgonvalues
(b) Unnormalised, removing SGD
(a) Unnormalised.
7 0.0001.
(c) Normalised
Figure 32: Top 25 eigenvalues of Hessian at solution obtained by SGD and Adam after finetuning ViT-Large on Flowers.
42

---

Published as a conference paper at ICLR 2026
Seed 7 SGD LR =0.1
Seed 7 SGD LR=7e-5
Seed 7 AdamW LR=7e-S
Valuue (x1e+04)
20
40
0
20
40
20
40
Elgenvalues
Eigenvalues
Elgenvalues
Epoch
Epoch 2.30
Epoch
Figure 33: The eigen value evolution for seed finetuning Bert-base MRPC.
Seed 42 SGD LR =0.1
Seed 42 SGD LR=7e-5
Seed 42 AdamW LR=7e-5
0.4
0.4
Value (x1e+04)
0.2
0.2
0
-0.2
-0.2
-5
-0.4
-0.4
20
40
20
40
20
40
Eigenvalues
Eigenvalues
Eigenvalues
Epach
Epoch 2:0
Epoch
Figure 34:
The eigen value evolution for seed 42 finetuning Bert-base MRPC.
Seed 1234 SGD LR =0.1
Seed 1234 SGD 7
Seed 1234 AdamW 5R=7e-5
0.50
0.50
Value (xle+04)
0.0
0.00
0.00
-0.2
-0.25
-0.5
-0.50
-0.50
20
40
20
40
20 Eigenvalues
Eigenvalues
Eigenvalues
Epoch 0
Esoch 230
Epoch 575
Figure 35: The eigen value evolution for seed 1234 on finetuning Bert-base on MRPC.
Egenvalues (x1e+04)
0.2
SGD, n- 7e -5 Adam. n 7
SGD, /- 0.1
Normalizer lzed Eigenvalues
1.
SGD, nー 7e -
SGD, n 0.1
0.5
Adam, n 7
0.0
0.0
-0.2
-0.5
-0.4
-10
20
30
40
40
50
Eigenvalue Index
Eigenvalues Index
(a) Unnormalised.
(b) Normalised
Figure 36: Top 50 cigenvalues of Hessian at solution obtained by SGD and Adam after finetuning Bert-base on MRPC.
43

---

Published as a conference paper at ICLR 2026
Table 12: Validation accuracy and parameter distance traveled in terms L₁ and L2 norm for finetuning Bert-base on MRPC.
Metric
SGD (n 7e 5)
SGD (17 0.1)
Adam (1) 7e - 5)
Val Ace
43.87 + 24.02
84.80 +: 1.00
85.95 + 0.64
L1
5002.44 0.0
6066.93 = 34.33
31079.54 754.26
Lg
0.73 0.00
1.26 + 0.01
5.57 + 0.24
44


Reviews:
[{"review_id": "19kme4BkUP", "reviewer": "Reviewer_f8Av", "summary": "This paper studies continuous-time steepest descent methods, specifically algorithms taking the form:\n$$\\begin{align} d x_t = - \\mathrm{sign} \\left( \\nabla_x f(x_t) \\right) \\odot \\left| \\nabla_x f(x_t) \\right|^{q-1}, \\end{align}$$\nwhere $q \\in [1,2]$.\n\nDifferent values of $q$ result in different trajectories. For example:\n For $q=2$, the algorithm becomes gradient flow, which is the continuous-time approximation of gradient descent. For $q=1$, it becomes SignGF (Sign Gradient Flow), which serves as a good proxy for studying optimizers like Adam. The architecture on which the paper focuses is deep diagonal reparameterization, defined as $x = g(w) = \\prod_{i=1}^{L} w_i$, where $x$ is represented as the product of $L$ scalars.\n\nThe authors demonstrate that under $\\lambda$-balanced initializations, these steepest flows can be re-parameterized as steepest mirror flows. Using this reparameterization, the paper analyzes how quickly the flow escapes saddle points and investigates the effect of both coupled and decoupled weight decay.", "questions": "1) How is the manifold regularizer derived ? Is it due to the fact that steepest descent with de coupled weight decay can be written as \n\n$$ d( \\nabla_x R(x_t) ) = - \\mathrm {sign} \\left( \\nabla_x f(x_t) \\right) \\odot \\left[   \\nabla_x f(x_t) \\right]^{q-1} dt  - \\nabla M_{reg}(x) dt  $$\n\nit would be nice to detail this as the manifold regularizer appears a bit abrupt. \n\n2) The saddle points defined in Theorem 4.3 are not saddle for with coupled or decoupled weight decay so the feature learning results only work without the weight decay. This point needs to be clearly mentioned in the manuscript. \n\n3) From the abstract, how is the deep diagonal reparameterizations a proxy for attention ?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper provides a seperation result for signGD with coupled and decoupled weight decay and show that they have different regularization properties which is interesting.", "weaknesses": "a) The paper focuses on deep diagonal reparameterizations which is a product of one-dimensional variables with a particular initialization shape. The setting is restrictive to generalize the results of the paper.", "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": "Pk1qk3Apd2", "reviewer": "Reviewer_iSvT", "summary": "The paper proposes “steepest mirror flows” to explain why Adam/AdamW beat SGD in fine‑tuning via faster saddle escape and different implicit regularization.\nSteeper (sign‑like) geometry helps saddle escape and that decoupled weight decay (AdamW) stabilizes feature learning, with theory for deep diagonal reparameterizations and small fine‑tuning case studies. Most formal results hold only in separable/diagonal settings; the transformer link is indirect; and empirical support for fine‑tuning claims is narrow.", "questions": "-", "limitations": "", "rating": null, "confidence": 5, "soundness": 1, "presentation": 3, "contribution": 1, "strengths": "The message that adam escapes saddles better than SGD is believable and may be cool. However, it is not well established. The rest of the claims are not substantiated.", "weaknesses": "#### How is this about fine-tuning transformers?\nThere is not even a linear diagonal attention there, no one ever claimed that a diagonal network is a good model for a transformer, because it is not. How do you argue this? \n\n#### Mirror flow study is incremental and does not adequately support the thesis.\nWhile the diagonal‑network analysis is neat, it is very similar to existent ones and does not bring any real novelty to the community. \n*I believe, it is extremely incremental.*\nIt is way too limited to show that your sign-mirror-descent escapes saddles to claim that *adam* escapes saddles *in transformers*.\nThere is a too big of a jump from the mathematical argument to the goal. Moreover, I think that such a result is a perturbation of existent ones.\n\n#### Order-2 saddles\nSo what? We know they are present in neural networks, actually, arguably also of higher order. A long line of works that you do not cite addresses this issue much more in general. That same line of work empirically and sometimes theoretically notice that they are not a problem in practice. You should discuss this line of research. On top of it saddles are not an issue in linear networks generally cause the standard initialization is with high probability outside of the area with saddles.\n\n#### Title mismatch\n\nEven the title is an oversell. It missmatch with the paper, this is not a paper about generally reparameterizing steepest descent as mirror flow. This is a paper about adam escaping saddles which lack novelty and does not support their claims. \n\n#### Experiments\nFor cifar10 actually SGD generalizes better than Adam with CNNs, and this is also a classical result. I really do not understand the whole point of the paper and how these are supporting experiments.\n\n#### Conclusions\nEven though experiments are present, I thus believe the central claims are not adequately supported and the research methodology is not sound.", "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": "qAkkYpBxIR", "reviewer": "Reviewer_e8j7", "summary": "This work introduces steepest mirror flows as a unifying geometric framework to study how optimization algorithms influence reparametrization. By analyzing diagonal linear networks and deep diagonal reparameterizations, the authors show that steeper descent methods, such as sign-based variants, escape saddle points more efficiently than gradient descent. Empirical results from linear regression, classification, and fine-tuning experiments confirm the theoretical predictions of faster saddle escape, stable learning, and distinct sparsity dynamics between coupled and decoupled weight decay.", "questions": "What is meant by: iterates x_t converge? I assume it is meant lim_t->\\infinity x_t", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This is an interesting study with insightful findings. Combining mirror descent with reparameterization is a great idea.", "weaknesses": "The statement of the first contribution is misleading. This is not the first work studying GF and reparametrization. It is probably meant with respect to the family. \n\nFor this type of work, assumptions are always problematic. They are very simplistic (diagonal networks) but it is very hard to potentially show more general results. \n\nIn experiments, rather than studying the networks considered in the analysis, they should see if the results hold when the assumptions are not fulfilled (for example, other networks).", "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": "JyZhHJuMj3", "reviewer": "Reviewer_NKAm", "summary": "This paper studied and compared the dynamical behaviors of different types of steepest descent algorithms. The authors focused on deep linear diagonal models and showed that a steepest flow in the weight matrices with respect to $L_p$ norm induces a \"steepest mirror flow\" in the end-to-end matrix. They analyzed how the choice of $p$ influences the dynamics, convergence rate and the effects of weight decay. Experiments on both linear and practical models are provided for the comparison between different algorithms.", "questions": "I find some parts of the presentation unclear. Please see the questions below. \n\n**Notations** \n\n* In Example 3.2, is the expression \"a deep diagonal linear network $g(m,w)=m \\odot w$\" actually referring to a shallow network, as there are only two weight matrices $\\operatorname{Diag}(w), \\operatorname{Diag}(m)$? \n\n* In Definition 3.3 and 4.5, does $\\mathbb{I}_n$ denote the all-one vector? This symbol is commonly used for the identity matrix, so clarification would be helpful. \n\n* In Corollary 4.8, the notation $w_i=\\lambda>0$ is ambiguous. Does it mean the vector $w_i$ has all entries equal to $\\lambda$? \n\n**Regarding Theorem 4.6**\n\n* $\\lambda−L_p-$balancedness in Definition 4.5 is defined for shallow network with two weight matrices. Then what does \"$\\lambda−L_p-$balanced with respect to the first parameter $w_1$\" mean in Theorem 4.6, which is stated for deep networks? \n\n* In Lemma 4.4 and Theorem 4.6, in what sense does \"almost everywhere\" mean (e.g.,\"steepest descent satisfies a separable $L_p$-mirror flow almost everywhere\")? Does it mean the result may fail only for a measure-zero set of initializations?\n\n* Corollary 4.12 discussed when the function $R_{L_p,\\ L}$ is a Legendre function. However, in Theorem 4.6, there seems no restriction on $p$ or $L$. Does Theorem 4.6 indicate that such an $R_{L_p,\\ L}$ always exists, even in cases where it is not a Legendre function?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The paper extends the study of reparametrization in gradient flow to steepest flow. This provides a useful platform for studying how reparametrization and optimization geometry interact to shape the training dynamics. I find this extension meaningful and valuable. Several established results are quite interesting: \n\n* Lemma 4.4: The balance equation for steepest flow in training deep diagonal linear network is new and interesting. \n\n* Theorem 4.6: The equivalence between steepest flow with reparametrization and \"steepest mirror flow\" is quite interesting. To my knowledge, previous work only showed the equivalence between gradient flow with reparametrization and mirror flow. \n\n* Corollary 4.8: It sheds light on how different choices of optimization geometry and reparametrization (network depth) could affect the convergence rate.", "weaknesses": "The main weakness is that the main claims are not well supported by the presented theory. The title begins with \"Never saddle\" and the authors emphasized both in the abstract and in Contributions that \"we prove that steeper descent (lower $q$) simultaneously escapes saddles faster and supports feature learning\", and that \"decoupled weight decay ... yielding more stable feature learning\". While these statements are strong and intriguing, I do not find them adequately supported by the theoretical results. \n\n**Regarding saddle escaping**\n\nAfter Lemma 4.4 and in Figure 3, the authors claimed that \"the (curved) path away from zero is shorter for smaller $q$, indicating faster saddle escape\". This reasoning is not logically sound: First, the notion of a \"shorter path\" is ambiguous (e.g., between which endpoints?). Second, the geometry of the invariant manifold alone does not determine how fast the algorithm moves along the manifold. \n\nBesides the above, the only theoretical evidence for the claim that \"lower $q$ escapes saddles faster\" is based on two results: (i) Corollary 4.8: under initialization \"$w_1 = 0$ and $w_i = \\lambda > 0$\", a smaller $q$ yields a larger coercivity constant $\\mu$; and (ii) Theorem 4.2: when $R$ is a separable Bregman function, a larger $\\mu$ yields a faster linear convergence rate. However, I do not find these two results sufficient for the claim, for the following reasons: \n\n* It is not established that the $R$ used in Corollary 4.8 (or Theorem 4.6) is a separable Bregman function. Thus it is unclear whether Theorem 4.2 applies; \n\n* The considered initialization \"$w_1 = 0$ and $w_i = \\lambda > 0$\" lies in a measure-zero set; \n\n* Only a subset of saddles (those near the considered initializations) are analyzed, whereas, as noted by the authors, other saddles exist elsewhere in the parameter space. \n\nTherefore, even if Theorem 4.2 is applicable, the results only indicate that: for certain specific saddles and for a measure-zero set of initializations in their neighborhoods, a smaller $q$ leads to a faster escape. I thus think it is overstated to use the title \"never saddle\", or to claim that smaller $q$ results in faster saddle escaping, which sounds like a general statement, just based on these results. \n\n\n\n**Regarding feature learning** \n\nThe authors claimed that \"smaller $q$ supports feature learning\". However, in the theory section (Section 4), I could not find a clear discussion on how the choice of $q$ affects feature learning. The following two aspects might be relevant, but it is unclear how they support the claim: \n\n* In Corollary 4.11 and 4.12, the authors discussed whether the function $R_{L_p, L}$ is a Legendre function under different $q$, network depth $L$, and metric exponent $m$. Then in the left panel of Figure 2, the authors indicated that a metric exponent slightly smaller than $1$ correspond to a feature learning regime (the green band in the figure). However, it is not explained why this is true and how does metric exponent relate to feature learning. For example, why does $m=0.95$ lead to feature learning, while $m=0.5$ or $m=1.5$ does not?\n\n* In Theorem 4.14, the authors discussed the on manifold regularization induced by weight decay under different $q$. However, as shown in Table 1, when $q=2, L=2$, the decoupled weight decay induces an $L_1$ bias; whereas when $q=1, L=2$, it induces an $L_{3/2}$ bias, which is less sparse than $L_1$. As sparsity biases have been linked to feature learning in some settings, this observation actually suggests that $q=2$ may facilitate feature learning instead of a smaller $q$.  \n\nThe authors also claimed that \"decoupled weight decay, as in AdamW, stabilizes feature learning by enforcing novel balance equation\". The balance equation in Lemma 4.4 indicates that the weight decay encourages the weights to become more balanced during training. But it is not clear to me how this balancing translates into a more \"stable\" feature learning. \n\nI suggest that the authors clarify in the paper what they mean by \"feature learning\" and \"stable feature learning\", and then compare the algorithms with different values of $q$ as well as with and without weight decay.\n\n\n**Initial recommendation**\n\nOverall, while the study of how optimization geometry and reparametrization affect the dynamics and the proposed framework are very interesting, I find the main claims of this paper, in particular those on saddle escaping and feature learning, insufficiently supported. Therefore, my initial recommendation is rejection.", "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)