Minor First, Major Last: A Depth-Induced Implicit Bias of Sharpness-Aware Minimization
Authors: Chaewon Moon
Dongkuk Si
Chulhee Yun
Venue: ICLR 2026 Poster (2026)
Keywords: sharpness-aware minimization
implicit bias
gradient flow
Primary area: optimization

Abstract: We study the implicit bias of sharpness-aware minimization (SAM) when training $L$-layer linear diagonal networks on linearly separable binary classification. For linear models ($L=1$), both $\ell_\infty$- and $\ell_2$-SAM recover the $\ell_2$ max-margin classifier, matching gradient descent (GD). However, for depth $L = 2$, the behavior changes drastically—even on a single-example dataset where we can analyze the dynamics. For $\ell_\infty$-SAM, the limit direction depends critically on initialization and can converge to $0$ or to any standard basis vector; this is in stark contrast to GD, whose limit aligns with the basis vector of the dominant coordinate in the data. For $\ell_2$-SAM, we uncover a phenomenon we call *sequential feature amplification*, in which the predictor initially relies on minor coordinates and gradually shifts to larger ones as training proceeds or initialization increases. Our theoretical analysis attributes this phenomenon to $\ell_2$-SAM’s gradient normalization factor applied in its perturbation, which amplifies minor coordinates early and allows major ones to dominate later, giving a concrete example where infinite-time implicit-bias analyses are insufficient. Synthetic and real-data experiments corroborate our findings.

Published as a conference paper at ICLR 2026
MINOR FIRST, MAJOR LAST: A DEPTH-INDUCED IM- PLICIT BIAS OF SHARPNESS-AWARE MINIMIZATION
Chaewon Moon
Dongkuk Si Mobilint, Inc. dongkuk@mob) lint.... com
Kim Jaechul Graduate School of AI, KAIST chaewon. wwwn@ aai stssstt.t. ac kr
Chulhee Yun
Kim Jaechul Graduate School of AI, KAIST chulhee, yun@kaist. ac.
ABSTRACT
We study the implicit bias of Sharpness- Paarre Aware Minimization (SAM) when train- ing L-layer linear diagonal networks on linearly separable binary classification. For linear models (L = 1), both 30-- and f2-SAM recover the lz max-margin classifier, matching gradient descent (GD). However, for depth L = 2, the be- havior changes drastically- even on a single-example dataset. For Lx -SAM, the limit direction depends critically initialization and can converge standard basis vector, stark contrast to GD, whose limit aligns with the basis vector of the dominant data coordinate. For 22-SAM we show that although its limit direction matches the L max-margin solution as in the case of GD, its finite- time dynamics exhibit phenomenon we call sequential feature amplification. in which the predictor initially relies on minor coordinates and gradually shifts to larger ones as training proceeds or initialization increases. Our theoretical analysis attributes this phenomenon to f2-AM'ss gradient normalization factor applied in perturbation, which amplifies minor coordinates early and allows major ones to dominate later, giving concrete example where infinite-time implicit-bia -iis analy- ses are insufficient. Synthetic and real-data experiments corroborate our findings.
to any
INTRODUCTION
Modern deep networks often generalize well despite extreme wer-paramettrizattoo.. One explana- tion emphasizes the geometry of the objective: models perform better when optimization settles in flatter regions of the landscape (Hochreiter & Schmidhuber, 1994; Keskar et al., 2016; Neyshabur et al., 2017; Jiang al., 2019). Motivated by this view, Foret et al. (2020) introduce Sharpness- Aware Minimization (SAM), which seeks parameters that minimize the worst-case loss within a small neighborhood. Following its empirical success (Chen al., 2021; Bahri al., 2021; Kaddour et al., 2022a), various theoretical works have analyzed SAM's implicit bias to understand its effec- tiveness (Andriushchenko & Flammarion, 2022: Behdin & Mazumder, 2023a; Zhou et al., 2025). However, these analyses primarily apply to scenarios with attainable finite minimizers (c.g.. squared loss), leaving open the case of losses whose infimum lies infinity (e.g.. logistic loss).
We consider the implicit bias of SAM when training L-layer linear diagonal networks on linearly separable classification datasets with logistic loss. We study two variants of SAM, foe SAMM and (2-SAM, named after the norm defining their local perturbation (See Section 2). For L. = 1 (linear models), gradient descent (GD) is known to converge in direction to the l max-margin classi- fier (Soudry al., 2018). For both l SAM and SAMS we show that they also align with the same limit direction. Thus, SAM does not change the implicit bias here, as shown in Figure la.
However, for 2-1ayer diagonal linear networks, we find that the trajectory of the linear coefficient vector (t)) under both E and SAM can differ substantially from the maximum m-mmagi implicit bias of GD (Gunasekar al., 2018b). In Figure 1b, consider toy separable dataset {(u, }1} with μ = 21, In this case, the f max-margin direction e₂ = (0. 1). the standard basis vector for the major component of μ.. As predicted, all GD trajectories and some SAM tra-

---

Published as a conference paper at ICLR 2026
1.2
1.2
1.-
1.0
0.8
0.8
0.6
5AM
0.6
0.2
0.2
0.0
0.0
0.5
1.0
1.5
0.0
0.5
1.0
1.5
(a) Depth (linear network)
(b) Depth (linear diagonal network)
Figure I: Trajectories of the predictor ((t) E R from identical initial conditions under discrete GD, Lo -AM and f2-SAM on {(u, +1)} with μ = (1, We used rj = 0.3 and p = 1 for SAM.
jectories show increasing alignment of ((t) with 62- However, for some initializations, that some trajectories of ((t) under &o -SAM and f2-AM instead converge to zero, or even align with e (1, 0) a seemingly paradoxical implicit bias favoring the minor feature rather than the major one. It is interesting that the addition of a single layer-from L 1 to L 2-introduces this peculiar behavior of SAM different from GD, even for the simple setting of linear diagonal networks trained with a single example.
observe
1.1 SUMMARY OF OUR CONTRIBUTIONS
We analyze the optimization trajectory and implicit bias of SAM and f2-SAM in L-layer linear diagonal networks trained linearly separable data with logistic loss. For theoretical analysis, analyze the evolution of the linear coefficient (()) of the linear diagonal network under continuous- time versions of SAM, foo -SAM flow and f2-S flow. We characterize their limit directions, obtained when training on general linearly separable data, and their pre-asymptotic behavior before aligning with the limit directions, analyzed on single-example dataset (1)}.
Depth 1 (linear). For linear models (L = 1), both SSM flow and f2MSAM flow have the same l max- margin implicit bias as GD linearly separable data; in the single-example setting, further show that the foo SAM coincides exactly with the GD trajectory.
Depth I., -SAM. For 2 2 and -SAM flow, we characterize the coordinate- wise trajectory of (t) determined by the relative scale of each coordinate initialization and the perturbation radius of foc -SMM (Theorem 3.2). For almost all initializations, (tt) diverges and its limit direc- tion is one of the standard basis vectors ci..... Cd or it converges a finite point (Corollary 3.5). Compared to GD, the limit direction of fo SAM becomes more sensitive to initialization.
Depth 2, f.SAM. For I. = and SAM flow, we first prove that the limit direction (if con- vergent to zero loss) is the E max-margin solution (Theorem 4.2); however, this infinite-time characterization does not explain our observation from Figure 1b. We empirically investigate the finite-time trajectory of (t) and identify the sequential feature amplification phenomenon, in which B(t) initially relies minor coordinates and gradually shifts larger ones f increases initialization scale grows. We provide a theoretical explanation of both time-wise (Theorem 4.4) and initialization-wiss (Theorem 4.5) aspects of the phenomenon. This example shows that fo- cusing only on the - xc limit can overlook aspects of the training dynamics. SAM provides a clear instance where a finite-time view is essential to understanding how its implicit bias emerges. We present synthetic and real-data experiments to corroborate our findings.
1.2 RELATED WORK
Implicit Bias of GD on Linear Diagonal Networks. Soudry et al. (2018) show that under linearly separable data with logistic loss, the weight of a linear model diverges while the direction converges to the & max-margn classifier. For linear diagonal networks, gradient descent biases toward sparse predictors (Gunasekar al., 2018b), with 2-layer models converging to max-margin direction un- der the assumption of directional convergence. This directional convergence has later been formally established for gradient flow (Ji & Telgarsky, 2020), supporting the validity of this assumption.
2

---

Published as a conference paper at ICLR 2026
Subsequent papers have studied linear diagonal networks in sparse regression, in which initializa- tion scale governs the implicit bias: large initialization favors ty-type bias, while small initialization favors L tyype sparsity (Woodworth al., 2020; Yun et al., 2020; Moroshko et al., 2020). Stochas- tic gradient descent (SGD)'s noise provides implicit regularization toward sparser solutions (Pesme al., 2021), amplified large learning rates (Even et al., 2023). Nacson et al. (2022) show that large GD step sizes push solutions out of the kernel regime, enabling sparse solutions. Beyond GD and SGD, recent works analyze implicit bias in diagonal linear networks through mirror-flow and related continuous- time formulations (Jacobs et al., 2025; Wang & Klabjan, 2024; Papazov et al,, 2024: Jacobs & Burkholz, 2024); we provide brief overview Appendix A.2.1. Prior work on small-initialization GD under squared loss in the same diagonal network setting shows incremental muldle-to-sadde learning dynamies, where coordinates become active in discrete stages as the pre- dictor moves between saddles (Berthier, 2023; Pesme & Flammarion, 2023). We provide detailed comparison between our setting and these saddle-to-sadee dynamics in Appendix A.2.2.
Sharpness-Aware Minimization. Motivated by the relationship between sharpness and general- ization (Hochreiter & Schmidhuber, 1994; Keskar et al., 2016; Jiang et al., 2019; Neyshabur et al., 2017), Foret et al. (2020) propose SAM. Extensive empirical work has shown the superior perfor- mance of SAM and its variants across tasks and architectures (Sun et al., 2024; Kwon et al., 2021; Li al., 2024b; Liu al., 2022: Yun & Yang, 2023; Bahri al., 2021; Zhuang al., 2022; Kaddour et al., 2022b). Complementing these empirical findings, theoretical work has analyzed SAM's opti- mization dynamics, generalization, and implicit bias (Li al., 2024a; Behdin & Mazumder, 2023b; Zhang et al., 2024; Agarwala & Dauphin, 2023; Wen et al., 2023; Long & Bartlett, 2024; Zhou et al., 2024; Springer al., 2024; Baek et al., 2024: Chen et al., 2023), including results in simplified set- tings such as diagonal linear networks on MSE loss (Andriushchenko & Flammarion, 2022; Clara et al., 2025). Recent work highlights convergence instabilities of SAM near local minima (Si Yun, 2023; Kim et al., 2023). We defer detailed discussion to Appendix A.2.3.
PRELIMINARIES
Notation. We write the i-th standard basis vector as e;. For N, let [n] {1. {1...... ... For a vector v E Rdd, we denote its coordinates by v (vi, va). For any block vector Z = (z(1) 2(0) E (Rd)=, we denote its ε-th block by 2 == (() E Rd. For a, b E Rd. a@b denotes the element-wise product; for a collection (a(9)/_, we write O a:
Model. We consider I.-layer linear diagonal networks, a simple family of homogeneous networks widely used for the study of implicit bias (See Section 1.2). Let (w ....),. w(l) E (IRd) be the parameter vector. For E Rdd, let the linear coefficient vector (()) and output f(x) be
( = O 2=1 ((()) E IRd, f(x) = (3(0), z).
Data and Loss. We consider the standard supervised learning setting where a binary classification dataset {(x,, y,) y)) is given. Let the logistic loss be f(u) log (1 + exp(-u)). Then the training loss function is defined as ((0) : / EN (vr 2x))) We write the gradient of with respect to 0 in block form, as V C = (V wi) L()). V L()).
Optimization Algorithms. In this paper, we mainly consider the implicit bias Sharpness-Aware Minimization (SAM, Foret et al. (2020)) and how depth causes it to deviate from the baseline algorithm, gradient descent (GD). At iteration a GD update reads ((t+1) == (t) ( where 77 > 0 is called the step size or learning rate.
On the other hand, SAM updates parameters by evaluating the gradient at a perturbed one: @(t) = (t) + (O(t)) O 1) == 0(t) v/ C()()),
where the perturbation (0)) is the approximate worst-case direction inside the ép-ball of per- turbation radius p 0: (0) arg max e (()).. We refer to θ as the ascent point. Since 0 = (w
-ii(1) ...., (t) and ()) = (e (1) (0), , Ep (0)) so that we can say (i(i) = ( + Ep (0). For = 2 and cx,
(L)) has block structure, we also write θ
(L)
the perturbation (0) has clean closed-form solutions:
V2 ( [V
Exc (0) psign
PCC))(VC(()), (VC())),
3

---

Published as a conference paper at ICLR 2026
and we consider the two variants, referred to f2-SAM when 2 and &0 -SAM when p For = oc, the maximizer is not unique when a coordinate of the gradient is zero. To make sure that the update is uniquely determined, adopt the convention sign(0) := 0, applied coordinate-wiset
Contimuous- time Flows. In the study of optimization algorithms, it is often useful to reduce the original discrete-time updates of optimizer to corresponding continuous-tine flow. Unless the step size is too large, continuous-timg flows offer a good approximation of the discrete-time optimizers, while allowing for clean and simplified analyses.
For GD, common continuous-time counterpart is gradient flow (GF): 0(r) =- - -V With gradient flow, the analysis GD trajectory boils down to solving an ordinary differential equation (ODE). Likewise, we define and study the flow counterparts of SAM, governed by the ODE
( - VV((()))
Depending on the choice of norm, we will use the terms fox -SAM flow and f2-SAM flow refer to the continuous-tiin time versions of SAM. Figure 6 in Appendix A. plots the trajectory of -SAM flow and f2-SAM flow under the same setup of Figure almost the same and the surprising implicit bias of SAM carries over to SAM flows. Hence, we aim to understand this unusual behavior of SAM by studying the corresponding SAM flows.
We observe that the trajectories stay
Rescaled Flows. As shown in Appendix A.3, for the special case of single-example dataset {(u, +1)}, the f,-SAM flow (p = 2, 00) of the i-th layer weight follows the same sputial trajectory as the following rescaled l, SAM flow:
(t) 44
(O Exii
19 + (0(1)))
(2)
obtained by taking out the loss derivative - -L' (((((((()) 11)) > 0 from the original 6,-SAM flow. Note that the original fp-SAM flow (1) and the rescaled flow (2) differ only by a reparameteriza- tion of time. Let 6 ari (forig) denote the original SAM flow and ((t) the rescaled flow. Then there
exists a strictly increasing map forig = r(t) such that (Tt)) @().. Applying the chain rule yields the relation
d0 dforgg dr ([(Oorig(((t)) dt
dr
dr dt e(((O(t))TT)
dt (((((()))
Since (u( f 0 as co, the rescaled flow accelerates time in the margin regime. Formally,
T(t) = S. ((((((s))
1 ds.
The rescaled flow makes the analysis easier due to the omitted term. Since our goal to gain better understanding of the spatial trajectory, study the rescaled SAM flows in our analysis.
Directional Convergence. Let B Tmax) + Rd be a trajectory with maximal existence time T Tiax (0, oc). We say that (Bt) converges in direction directionally converges if the limit .(t)) exists. this case, 8 is called the limit direction of 3. [[C
8 limm:
3 SAM WITH fox SE-PPUURUBATIONS
We begin with SAM.. For single-example data, its counterpart reseclelerttrreaaeeddd SAM the nice property that each coordinate evolves independently, enabling an exact characterization of the trajectory for any depth L.
3.1 DEPTH-1 NETWORKS
We start with the depth-1 case, in which the implicit bias of -SA coincides with that of GD. Theorem 3.1. For almost every dataset which is linearly separable, any perturbation radius and any initialization, consider the linear model f(x) = =w, x)) trained with logistic loss. Then, l SAM flow directionally converges in the la max-margin direction.
The proof is deferred to Appendix C.I. Since Theorem 3.1 holds for any p. it also recovers the implicit bias GF. While Theorem 3.1 characterizes the limit direction for almost all linearly sepa- rable datasets, Theorem C. shows that, for the single-example data, the foe SAM flow follows the
4

---

Published as a conference paper at ICLR 2026
same trajectory as GF. The yellow lines in Figure 6a depict the flows. As t + co, w(t) converges in direction to the le max max-margin direction μ4. Hence, when L = 1, GD and -SAM share the same bias toward the & max- max-marginn solution, independent of the initialization.
3.2 DEEPER NETWORKS (L 2 2)
To isolate the depth-induced implicit bias of SAM from effects of data-poin configuration and ob- tain a tractable characterization of the SAM dynamics, we analyze the minimalist separable dataset DN = {(H, +1] with feature vector μ E R satisfying < μu < ... < pd; without loss of generality, we assume this monotone ordering of μ4; s. The additional technical difficulties that arise the multi-point setting are deferred Appendix C.2. Appendix C.6, we empirically verify that the same behaviors persist under multi-point datasets and discrete SAM updates, suggesting that our insights extend beyond the single-point setting.
In contrast to the depth-1 1 case, for deeper (linear diagonal) networks, the implicit bias of loc -AAM differs from GD-even on this single-example dataset. For example, when L = 2, while GD always aligns with the major feature, -SAM can favor minor features depending on the initial condition. For L 2 3. we show that the implicit bias of S is more sensitive to initialization than GD, in the sense that wider range of initialization leads to solutions focusing on minor features. The next theorem characterizes the trajectory selected by the flow for different choices of initialization.
Theorem 3.2. For E [L], suppose w (0) a R.. Let w (t) follow the rescaled -SAM flow (2) with perturbation radius > on the dataset Du. Then, for the j-th coordinate of B(t):
Ifa; < p. then (t) converges to 0 if is even, or to p if odd.
Ifa; = p. then 31(t) = pp for all 2
for
Ifaz > p and L = 2, then (t) grows exponentially: 8y(t) = ((exp(22x)t).
If oj > P and L > 2, let == arg maxja;2ell 1g - -))_-2 and also let T =
minkeJ (((1.-2)j,,o,, bounded for all t << T.
If E J. then 3y(t)
- T; otherwise, 3,(t) stays
We provide the proof of Theorem 3.2 in Appendix C.3. The behavior of each coordinate B3 (t) completely determined by whether the initialization aj lies below, at, or above the threshold p. In each of these three regimes, (1) is monotone in L. Recall that (0) := psign(VL(()) For Dj,, the sign of the gradient (5) is determined coordinate-wise. Thus, the rescaled -SA flow (2) decouples across coordinates, and each (1) evolves independently, allowing us to state Theorem 3.2 for each separate trajectory of (t).
Remark 3.3 (Interpretation of the Finite-time Blow-up). For > 2, the rescaled -SAM flow (2) exhibits finite-time time blow-up: some coordinates satisfy (t) - oo as t - T. Interpreting this phenomenon in the original SAM time scale, the blow-up corresponds to infinite time in the original SAM flow, Indeed, as B(t) + CO, we have (B((t) 4) + +0, and therefore
1 de + +0 as t T(I) /'' ((((s))TT)
Thus, in the original SAM flow, only the coordinates in J diverge as the original time (() - co, while all other coordinates remain bounded.
Remark 3.4 (Interpretation of Exponential Growth). For L 2. each coordinate (t) with aj p grows exponentially as t - oo. Since r(t) + oo as + co, divergence occurs on the same infinite- time limit in both the rescaled and original f0 -SA flows. Nevertheless, because the dynamics are obtained after a time reparameterization, the exponential rate observed in the rescaled flow should not be directly interpreted as the actual divergence speed in the original SAM dynamics. Still, for fixed I, = 2, all coordinates share the same rescaled time, so their relative growth can be compared. Among the coordinates with a, > p. the one with the largest feature weight μ4 dominates asymp- totically and the lx SSAM flow therefore converges in that coordinate direction. We formalize these conclusions for general in the following corollary, characterizing the dominant direction.
Corollary 3.5. Under the assumptions of Theorem 3.2, let S [j aj: > p} and assume S + If there is a unique maximizing index jr == arg maxjes H1 (a; -p)= -)2-2, then the ESM flow converges in the ey' direction. In particular, when L = 2, we have j" := arg maxjes Pú.
5

---

Published as a conference paper at ICLR 2026
2
2
pz
8
1-
4 1
8: 8
&
%
Amp
P
(a) GD (L (b) SAM (L
(c) GD (L
(d) 1550-AM (L = 3)
Figure 2: Trajectories 3(t) from identical initializations under GF and -SAM flow with d = 2
and μ (1, For SAM, p
The proof is deferred to Appendix C.4. When I. = 2 and E R4+,, setting p = 0 in Corollary 3.5 yields S [d]. Hence, Corollary 3.5 recovers that the GF always aligns in the e, direction-the & max-margin direction-regardee
of the initialization.
Illustrative Example. Figure 2 shows the trajectories of (()) under GF and lo -AM flow with L == 2.3 and (1, 2). Figure 2a depicts the L = 2, GF case, where GF always aligns in the e2 direction. For = 2 and fx -SAM (Figure 2b), the plane (31, B2) is partitionned by the thresholds S), the SAM flow shows directional convergence in e2 (red/blue regions). In the yellow region, 2 & S and 1 S, the limit direction is e the "minor" feature. If all coordinates satisfy 0g < p. the flow converges to O (purple region), by Theorem 3.2.
3, a3 p². If a2 > P (so
For L 2 (Figures 2c and 2d), the blue regions get partitioned once more because large 0n leads to /1 -) > H2 H22(az - p).--, leading to directional convergence toward e1- Comparing the green regions in Figures 2c and 2d shows that the slope of the boundary between blue and green regions is steeper in -SM flow than that of GF. Considering that initializations in the yellow region also result in the limit direction e1. these together indicate that foo -SAM exhibits a greater sensitivity to initialization and stronger implicit bias toward minor features than GD.
SAM WITH 2 SS-PERTRRBATIONS SEQUENTIAL FEATURE AMPLIFICATION
We now turn to -AMS, which is the form most commonly used in practice.
4.1 ASYMPTOTIC BELAAVIOR ON D AND DEPTI--2 NETWORKS
For depth-1 models, 2a-S converges in the le max-marget direction regardless of initialization, matching the implicit bias of GD and -SAM. We prove the following theorem in Appendix D.I:
Theorem 4.1. For almost every dataset which is linearly separable, any perturbation radius p and any initialization, consider the linear model f(x) (w, x) trained with logistic loss. Then, S2M flow directionally converges in the l2 max-margin direction.
While Theorem 4.1 characterizes the limit direction for linearly separable datasets, Theorem D.1 shows that, for the single-example data, the loc--SS -SAM flow follows the same trajectory GF.
For depth-2 models, f2-SAM asymptotically converges in the ε max-margin direction as the loss converges zero, independently of the initialization scale. This parallels the well-known behavior of GD (Gunasekar et al., 2018b). We formalize this below, with the proof in Appendix D.3.
Theorem 4.2. For almost every dataset which is linearly separable, and any perturbation radius p. consider the linear diagonal network of depth 2, f(x) (w (2) w(()) x)) trained with logistic loss. Let (1) loss vanishes (1w wo()) (t)) + 0, (b) the predictor B(t) == w (1) O w converges in direction. Then the limit direction of (tt) is the f, max-margin direction.
follow the Sa-SAM flow with w()) (0) 20(2) Assume (a) the
Since Theorems 4.1 and 4.2 holds for any p, it also recovers the implicit bias of GF. We now revisit Figure 6, which is the flow counterpart of Figure 1, and compare the trajectories with the asymptotic directional convergence results above. First, the green lines in Figure 6a visualize the trajectories of 22-SAM flow for L = 1. and we can check that the trajectories coincide with GD's, as expected by theory. In the I 2 case (Figure 6b). the green SAMS flow curves include ones that (i) drift

---

Published as a conference paper at ICLR 2026
0.8
Regime
5
0.6
j=1
124
Initialization scale
0.6
and max 8
0.5
Reg
0.4
2
0.2
Regime
1
0.2
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
time
time
(a) Dominant index 1 = arg max, (t(() over a and t.
(b) Normalized 3,(1)/ for 0.4.
Figure 3: Rescaled 22MSAM flow on D with μ 6, E R and p
loward the origin, and those that (ii) initially align with e1, a direction orthogonal to the ε max- margin direction e2. Such behaviors are not explained by Theorem 4.2. Hence, account for what is observed in Figure 6b, we move on to analyze the dynamics of f2-SAM in finite time.
4.2 PRE-ASYMPTOTIC BEHAVTOR ON DEPTH-2 NETWORKS
We investigate the pre-asymptotic dynamics of f2-SAM on depth-2 linear diagonal networks and show that the trajectory exhibits behavior markedly different from its asymptotic limit. This trast highlights the need for finite-time analysis to understand how the implicit bias of SAM actu- ally emerges. In this section, we study the toy dataset D,, == ((,++)) with μ I& satisfying 0 < < < μa.- We further present experiments on multi-point datasets, discrete-time f2- SAM, and deeper models identified in the depth-2 single- -ogitttpontt& E2-SAM flow persist in these more realistic settings. Moreover, to capture the effect of the initialization scale with a single parameter, we adopt a coordinate-wise and layer-wise uniform initialization (1)) (0) (2) a1 throughout this subsection. We additionally report similar empirical results under random Gaussian initialization in Appendix E.2.
2 in Appendix D.8, which confirm that the qualitative behaviors
4.2. SEQUENTIAL FEATURE AMPLIFICATION
We begin by describing newly observed and surprising phenomenon of 22-SAM sequential fea- ture amplification. For certain initialization scales a and times t, f2-SAM first aligns with minor features; increases diate to major features. In contrast, GD selects the major feature regardless of a and t. We visualize this using rescaled f2-SAM flow in Figure 3a and show the GF and - flow counterparts in Figure 7. To quantify the phenomenon along the two axes-time and initialization scale a-al each t and a, we track the index j arg min, (t) and color the grid (t, a)) according j!. Regions where B is negligibly small are shown in gray, indicating convergence to 0. Based on the observations from Figure 3a, we partition the initialization scale a into three regimes.
increases, the dominant coordinate transitions from minor, interme-
(Regime Starting from any a in this range, the trajectory eventually collapses the origin as training proceeds; effectively no feature is expressed and the loss does not vanish.
(Regime 2) Time-wise sequential feature amplification emerges. With fixed a chosen from this regime and increasing there exists the period where the dominant coordinate index ji increases over time, transitioning from minor to major features. As shown in Figure 3b, jt sequentially changes from to 5 over time for 0.4.
(Regime 3) B aligns with the major feature from the outset and maintains this alignment throughout.
Beyond the time-wise phenomenon, Figure 3a also suggests that sequential feature amplification also happens in the a-axis. To see this, consider fixed slice of time t and navigate through the axis: for small a, the predictor 33 remains near the origin with no feature discovered. As a grows, the dominant coordinate at shifts sequentially-3,, becomes largest first, then 32, and so on. However, this is not a fair comparison between trajectories, because Figure 3a obtained from the rescaled flow; each trajectory (for each a) has different time scale.

---

Published as a conference paper at ICLR 2026
Nevertheless, we can compare between trajectories if base comparison trajectory maxima. Specifically, in Regime 2 (the sequential feature amplification phase), we define the trajectory-wise nost-amplified index, to understand how the initialization scale affects the "am- plification" of minor components. For each coordinate j. we track the ratio & (0/8)(t) over the entire trajectory, and define j*(a) = arg max, max, B,(t)/a(tt) imum relative amplification. In Figure 3a, for each value of a in Regime 2, we plot the time step that attains the maximum value of ()(()(t))patt in black dots; we can clearly observe that j (a) increases from the minor index 1 to second-most major index d - 1 in Regime 2. We call this phenomenon initialization-wsse sequential feature amplification.
the coordinate with the greatest max-
4.2.2 UNDERSTANDING THE EFFECT OF 2--SM
Before analyzing sequential feature amplification, we describe the rescaled &2M flow for depth-2 linear diagonal networks and offer an intuitive explanation of the sequential feature amplification
phenomenon. With initialization (0) = (0) 144, we have w (t) = (2) =; w()
for all t 2 0. Using this, derive in Appendix D.2 that the rescaled f2- SAM flow for w(t) reads
(t) 1 @ w p p where (t) /211p @ w(t)l
(3)
Compared to the p case, the extra term scales ja w(() coordinate-wise by P [ When (t) is large (e.g., under large initialization or after sufficient training), this factor is close to one and the dynamics becomes close to GF. When (t) is small (e.g., small initialization), the coordinate- wise scaling factor multiplies different scalars to different coordinates, some of which can even be negative and decrease the corresponding coordinates of w()). Notice that larger μ. _:___
1.
leads to smaller - 1-p (e( Thus, in the early stage of training, major features are suppressed while minor features are comparatively amplified, yielding the observed emphasis on minor features.
4.2.3 ANALYSIS OF TIME-WISE SEQUENTIAL FEATURE AMPLIFICATION
We next provide a theoretical account of the time-wise sequential feature amplification. At each time we analyze the instantaneous growth rate of each coordinate 3, (L), viewed a function of both t and the initialization scale a. This reveals how the growth behavior of different coordinates evolves across the training trajectory. In particular, we derive a coordinate-wän
growth rule of B, (t), in form analogous to Equation (3). The proof is provided in Appendix D.4.3, and an extension to the T--layer setting-where where an analogous growth rate can derived given Appendix D.5.
Lemma 4.3. The rescaled SAM flow (2) is (t) (t) with rt(() == 2/5r (1
AM
(e)
By Lemma 4.3, the rate (t) controls the instantaneous growth or decay of (y) For fixed t, (t) is concave quadratic in My. maximized at μ4 = (t(t) == =a(!)) Hence, indices with H4 closest mt(t) attain the largest (t); coordinates with feature strength μij nearest to me(() are amplified the most, while those farther away may even decay. Consequently, the trajectory of me(() dictates the catuureamplification dynamics, and it exhibits three regimes depending on the initialization scale. Recall that 0 < P1 <...< Hd.
2p
Theorem 444. There exists a unique an such that aig oo
an
< &
< =
p-a and the trajectory of m²(t) falls into one of the following three regimes. /Zllnlla
(Regime 1) If < on, then m²(t) strictly decreases for all t.2 O and there exists Tr such that for j E [d], (t( strictly decreases for all t 2 Ti.
(Regime 2) If af < a < a2 there exists T2 such that me(TT) increases for all L T2.
Pe- and me(t) strictly <
(Regime 3) Ifa a2 then m mt(t)= > Pa-1tHa and Ba(t) has the largest growth rate for all t 2 0.
The proof of Theorem 4.4 is provided in Appendix D.4.5. Theorem 4.4 identifies three regimes of the me)()) dynamies, each corresponding to qualitatively different pattern of feature amplification.
Regime 1. (t) decreases for all t 2 O, and reaches 4 at time T,. Once m²(t) < 4 every coordinate satisfies (1) by the form of (C), and thus (t) strictly decreases for all j (d].
8

---

Published as a conference paper at ICLR 2026
When (t) Patthaa1 the closest feature strength to (t) is pdd, so Ba(t) attains the 2 largest growth rate. This explains why the major feature remains dominant throughout this regime.
Regime 3.
Regime 2. When me (T( M the closest index je satisfies je < d. At this time, the largest growth rate is therefore achieved by the non-major coordinate B, (T2). Since me (t) strictly increases for all T2, the coordinate with the largest growth rate increases, exhibiting the time- wise sequential feature amplification observed empirically in Section 4.2.1. In Regime 2, there also exist instances where (t) initially decreases and later increases, leading non-monotonic sequential feature amplification phenomenon. We discuss this in Appendix A.5.
Regime also leaves a clear trace the training loss. f2-SAM exhibits an early plateau while mainly amplifies minor coordinates, and the loss drops quickly only after it shifts to major coor- dinates, whereas GD shows a steadier decrease without this minor-to-major transition. The corre- sponding loss curves and further explanation are given in Figure 4 and Appendix E.I.
0.6
0.6
Domimant
3 0.4
0.44
0.2
0.2
0.0
0.0
2000
4000
6000
2000
4000
6000
Epoch
Epach
Figure 4: Loss curves of GD (left) and f2-SAM (right) 2-layer diagonal network in Regime 2 (a = 0.35, / = (1,2,3,4, 5, 6), p 0.1). Colored regions mark the coordinate with highest growth.
4.2.4 ANALYSIS OF INITIALIZATION-NE
SEQUENTIAL FEATURE AMPLIFICATION
In the previous subsection, we examined which coordinate attains the maximal instantaneous growth rate. We now turn to the cumulative update over time and study iiitiaiiaationwiis sequential feature amplification. In Theorem 4.4, we characterize the range of a (Regime 2) in which sequential feature amplification can occur. Here, we quantify the strength of amplification within Regime 2 as a function of a. Since a coordinate (t) can diverge, we assess which feature is amplified-and and by how much- via the ratio of the j-th feature the major feature, (,(0)/aa(t)- For a given initialization scale a, we track and bound how large the amplification ratio (,(0//2((1) can be along the trajectory.
Integrating the rescaled f2-SAM flow (3) (derived in Appendix D.6.1) yields the coordinate ODE
(5(() = (0) exp (2pzt - (oppT(()) where T(t) = -6 mo(s) ds
for j E [4]. (4)
The behavior of ,3 in (4) is determined by I(t). Recall that n controls the behavior of f2-SAM
in Section 4.2.2 and is used to characterize the instantaneous growth rate in Section 4.2.3. Here, focus on cumulative updates over time, where the time integral (t) of 1/ng becomes decisive. By bounding [tt). we quantify how strongly each feature is amplified relative to the major feature.
Theorem 4.5. Let aip, a2 be defined in Theorem 4.4 and aj be the threshold from there. Suppose tre /2 az.. Then, for j E [d], there exists I, such that
01 <a S p
(,(T,) PaCT,T
> (a) exp 22 ((R3 1)
1)log log (I-0%/) + log
C((R)))
where Rg = (u/+ma)/ 2, R2 == (na- (12-05)/11 and C(R)= == Rlog R - (R 1) log( R 1).
The proof follows from lower bound on I(t), and is deferred to Appendix D.6.2. A numerical illustration of (a) for several choices of 11 is provided in Appendix D.7. Theorem 4.5 applies to the small-o portion of Regime 2. For each coordinate j, we select the time T, maximizing (,(1) Bale) over
the entire trajectory, and obtain a nontrivial lower bound LB, ) for this maximal amplification.
The theorem goes beyond the qualitative picture in Figure 3a, which only identifies which coordinate becomes dominant (the index j') Theorem 4.5 additionally quantifies how large this dominant co- ordinate must grow: shown in Appendix D.7, LB, (a) often exceeds 10, indicating that the minor to intermediate coordinates can take values more than ten times larger than the major coordinate.
Dependence For all in Regime 2, the ratio ag/a lies in (0, 1), both logarithmic terms in (a) are positive. Since Rg > 2, the first logarithmic term dominates the exponent, making (a) grow rapidly a 01- Thus smaller a in Regime 2 produces stronger amplification as
9

---

Published as a conference paper at ICLR 2026
shown in Appendix D.7. This is substantiated by Figure 3a: smaller in Regime 2 keeps the dy- namics aligned with minor-intermenn intermediate features for longer time t, leading to greater amplification.
Dependence on Feature Geometry. The coefficients R2 and R2 increase with the spectral gap Ma/Mi, so datasets with larger feature contrast amplify more strongly shown in Appendix D.7.
Since (a) varies across j. it is natural to ask which coordinate experiences the strongest ampli- fication. Proposition 4.6 identifies the maximizing index j (a), with the proof in Appendix D.6.3.
Proposition 4.6. Under the conditions of Theorem 4.5, define ja() == arg maxj (a) and set ag ag. Then, there exist thresholds og aj < am S Hitee for some m S d such that j (a) for (0]
V2l|a||a
Proposition 4.6 shows (a) monotonically increases sequentially from 1 to m on a E (ax, amnl- Namely, as the initialization scale grows, the index that maximizes the lower bound LBy (a) shifts monotonically from minor to intermediate features. This matches the initialization-iis sequential feature amplification discussed in Section 4.2.1 (i.e., the black dots in Figure 3a). Within Regime 2, the our theoretical bound predicts a progression of the most-amplified coordinate from to m.
Lastly, through the cumulative update analysis, characterize the asymptotic behavior of 22-SAM flow for some extreme ranges of a. We prove the following proposition Appendix D.6.4.
Proposition 4.7. Consider og defined in Theorem 4.4. (i) If a < og, then B(t) converges to zero. (ii) If a P P24 ||
7z) [m]l: then B(t) converge in l max-margin direction.
Recall that Theorem 4.2 assumes that the loss vanishes and the limit direction exists. Proposi- tion 4.7(i) shows that for small a in Regime 1, the loss never vanishes. Proposition 4.7(ii) shows that for some a's in Regimes or 3, the limit direction exists and is the & max-margin direction.
5 EXPERIMENTS
Our investigation shows how depth, perturbation geometry, and initialization jointly shape SAM's optimization trajectory. We substantiate these findings with controlled experiments: 2-layer CNNs and linear networks on synthetic banded data, where we systematically vary the dataset construc- tion and metries across architectures (Appendix E.3), as well as multi-point (Appendix D.8.2) and deeper diagonal models (Appendix D.8.3). We also present experiments with practical CNNs trained on MNIST, where we Grad-CAN (Selvaraju al., 2017) visualize which image pixels emphasized (Figure 5 and Appendix E.4). These experiments show that f-MS allocates relatively bigger emphasis weaker/background pixels than GD, qualitatively matching our theory.
(a) GD
(b) f2-SAM
Figure 5: Grad-CAM comparison of GD and SAM on a CNN trained on MNIST, GD focuses on dominant digit pixels, whereas f2-SAM highlights minor background regions.
CONCLUSION
We characterized how network depth changes SAM's implicit bias on linear diagonal networks. For depth 1, SAM preserves GD's implicit bias. For deeper networks (I. 2 2) with -SAM, we derived precise weight trajectories depending on initialization scale and perturbation radius, where each weight coordinate either diverges toward standard basis vector or converges to a finite point. The most interesting regime arises for L 2 with & S-MM while the limit direction converges to the L max mar-mn solution, the finite-time dynamics exhibit sequential feature amplification. where the weight coordinate initially relies on minor coordinates and gradually shifts to larger ones. These observations suggest that implicit bias statements made only in the - oo limit can overlook how the bias emerges, motivating finit-timme perspective.
10

---

Published as a conference paper at ICLR 2026
ETHICS STATEMENT
This work is purely theoretical, analyzing the optimization dynamics and implicit bias of SAM in simplified models. does not involve human subjects, personal data, or sensitive information, and introduces new datasets. Broader impacts are indirect: the results may inform more reliable training and diagnosis of SAM-like methods, but they do not by themselves address safety, fairness, or deployment risks, which must be evaluated in application-speciae settings.
REPRODUCIBILITY STATEMENT
We support the reproducibility of our results by (1) fully specifying the SAM variants and rescaled- flow dynamics, together with complete theoretical statements and proofs; (2) reporting exact exper- imental setups, including the initialization scheme, models, datasets, and the values of step sizes, perturbation radius, and initialization scale; and (3) computing several quantities used in theo- retical simulations in closed form, making the corresponding plots exactly reproducible.
ACKNOWLEDGMENT
We thank Junsoo Oh for the helpful discussions. This work was supported by National Re- search Foundation of Korea (NRF) grant funded by the Korean government (MSIT) (No. RS-2023- 0021 1352) and Institute of Information & communications Technology Planning & Evaluation (ITPP) grant (No. RS-2019-11190077 Artificial Intelligence Graduate School Program (KAIST)) funded by the Korean government (MSIT).
REFERENCES
Atish Agarwala and Yann Dauphin. Sam operates far from home: eigenvalue regularization as a dynamical phenomenon. In International Conference on Machine Learning, pp. 152-168. PMLR, 2023.
Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness- aware aware mini- mization. In International conference machine learning, pp. 639-668 PMLR, 2022.
Maksym Andriushchenko, Dara Bahri, Hossein Mobahi, and Nicolas Flammarion. Sharpness-awann minimization leads to low-rank features. Advances in Neural Information Processing Systems, 36: 47032-47051, 2023.
Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. Advances in neural information processing systems, 32, 2019.
Christina Baek, Zico Kolter, and Aditi Raghunathan. Why is sam robust to label noise? arXiv preprint arXiv:2405.03676 2024.
Dara Bahri, Hossein Mobahi, and Yi Tay. Sharpness-aware minimization improves language model generalization. arXiv preprint arXiv:2110.08529, 2021.
Peter L Bartlett, Philip Long, and Olivier Bousquet. The dynamics of sharpness-awee aware minimization: Bouncing across ravines and drifting towards wide minima. arXiv preprint arXiv:2210.01513, 2022.
Kayhan Behdin and Rahul Mazumder. Sharpness-awar aware minimization: An implicit regularization perspective. arXiv preprint arXIv:2302.11836 2023a.
Kayhan Behdin and Rahul Mazumder. statistical properties of sharpness-aware minimization: Provable guarantees. arXiv preprint 20Xhv:2302.11836, 2023b.
Raphaël Berthier. Incremental learning in diagonal linear networks. Journal of Machine Learning Research, 24(171):1-26 2023.
11

---

Published as a conference paper at ICLR 2026
Xiangning Chen, Cho-Jui Hsich, and Boqing Gong. When vision transformers outperform resnets without pre- trainininng or strong data augmentations. arXiv preprint arXiv:2106.01548 2021.
Zixiang Chen, Junkai Zhang, Yiwen Kou, Xiangning Chen, Cho-Jui Hsieh, and Quanquan Gu. Why does sharpness-aware minimization generalize better than sgd? Advances in neural information processing systems, 36:72325-72376 2023.
Hung- Muing-Hsuua Chou, Carsten Gieshoff, Johannes Maly, and Holger Rauhut. Gradient descent for deep matrix factorization: Dynamics and implicit bias towards low rank. Applied and Computational Harmonic Analyvis, 68: 68:101595, 2024.
Gabriel Clara, Sophie Langer, and Johannes Schmidt-Hieber. Training diagonal linear networks with stochastic sharpness- aware minimization. arXiv preprint arXiv:2503. 1189, 2025.
Yan Dai, Kwangjun Ahn, and Suvrit Sra. The crucial role of normalization in sharpness-awand aware minimization. Advances in Neural Information Processing Systems, 36:67741-67770, 2023.
Li Deng. The mnist database handwritten digit images for machine learning research. JEEE Signal Processing Magazine, 29(6):141-142, 2012.
Mathieu Even, Scott Pesme, Suriya Gunasekar, and Nicolas Flammarion. (s) gd over diagonal linear networks: Implicit bias, large stepsizes and edge of stability, Advances in Neural Information Processing Systems, 36:29406-29448, 2023.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimiza- tion for efficiently improving generalization. arXiv preprint arXiv:2010.01412, 2020.
Jacob Gildenblat and contributors. Pytorch library for cam methods. jacobgil/pytoree warettrrchhhr-d---- cam, 2021.
com/
Daniel Gissin, Shai Shalev-Shwartz, and Amit Danicly. The implicit bias depth: How incremental learning drives generalization. arXiv preprint arXiv: (119999922... 2019.
Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Characterizing implicit bias in terms of optimization geometry. In International Conference Machine Learning, pp. 1832- 1841. PMLR, 2018a.
Suriya Gunasekar, Jason D Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient descent on linear convolutional networks. Advances in neural information processing systems, 31, 2018b.
Sepp Hochreiter and Jürgen Schmidhuber. Simplifying neural nets by discovering flat minima. Advances in neural information processing systems, 1994.
Tom Jacobs and Rebekka Burkholz. Mask in the mirror: Implicit sparsification. arXiv preprint arXiv:2408. 09966, 2024.
Tom Jacobs, Chao Zhou, and Rebekka Burkholz, Mirror, mirror of the flow: How does regularization shape implicit bias? arXiv preprint 1.XXi:22504128833 2025.
Ziwei and Matus Telgarsky. Directional convergence and alignment in deep learning. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neu- ral Information Processing Systems, volume 33, pp. 17176-17186. Curran Associates, Inc., neurips cc/ faaeeerffieeppppppeeeeee wwp.p.pprpp......
2020. URL https://proceeeee
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them. arXiv preprint arXiv: 19l:.02/72200 2019.
Jean Kaddour, Linqing Liu, Ricardo Silva, and Matt Kusner. When do flat minima optimizers work? In Advances Neural Information Processing Systems, 2022a.
Jean Kaddour, Linqing Liu, Ricardo Silva, and Matt Kusner. When do flat minima optimizers work? Advances in Neural Information Processing 3ystems,35:1657777 35:16577-116555, 2022b.
12

---

Published as a conference paper at ICLR 2026
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Pe- ter Tang. On large-batcle training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv:1609.044 04836, 2016.
Hoki Kim, Jinseong Park, Yujin Choi, Woojin Lee, and Jaewook Lee. Exploring the effect of multi- step ascent in sharpness-awars minimization. arXiv preprint arXiv:2302.1018). 2023.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi. Asam: Adaptive sharpness- aware minimization for scale-invariant learning of deep neural networks. In International confer- ence on machine learning, pp. 5905-5914. PMLR, 2021.
Bingcong Li, Liang Zhang, and Niao He. Implicit regularization of sharpness-aware minimization for scale-invariant problems. Advances in neural information processing systems, 37:44444- 44478, 2024a.
Tao Li, Pan Zhou, Zhengbao He, Xinwen Cheng, and Xiaolin Huang. Friendly sharpness-aware minimization. In Proceedings the IEEE/CVF conference on computer vision and pattern recog- nition, pp. 5631-5640, 2024b.
Yong Liu, Siqi Mai, Minhao Cheng, Xiangning Chen, Cho-Jui Hsich, and Yang You. Random sharpness-awan aware minimization. Advances in neural information processing systems, 35:24543- 24556, 2022.
Philip M Long and Peter L Bartlett, Sharpness-aware minimization and the edge of stability. Journal of Machine Learning Research, 25(179):1-20, 2024.
Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks.
arXiv preprint arXiv: /906. 996600599, 2019.
Edward Moroshko, Blake E Woodworth, Suriya Gunasckar, Jason D Lee, Nati Srebro, and Daniel Soudry. Implicit bias in deep linear classification: Initialization scale vs training accuracy. Ad- vances in neural information processing systems, 11:22182-22199 2020.
Mor Shpigel Nacson, Kavya Ravichandran, Nathan Srebro, and Daniel Soudry. Implicit bias of the step size in linear diagonal neural networks. In International Conference on Machine Learning, pp. 16270-16295. PMLR, 2022.
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Baolin Wu, Andrew Y Ng, et al. Reading digits natural images with unsupervised feature learning. In NIPS workshop on deep learning and unsupervised feature learning, volume 2011, pp. 7. Granada, 2011.
Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nati Srebro. Exploring general- ization in deep learning, Advances in neural information processing systems, 30, 2017.
Hristo Papazov, Scott Pesme, and Nicolas Flammarion. Leveraging continuous time to understand momentum when training diagonal linear networks. In International Conference on Artificial Intelligence and Statistics, pp. 3556-3564. PMLR, 2024.
Scott Pesme and Nicolas Flammarion. Saddle-to-saddle dynamics in diagonal linear networks. Ad- vances Neural Information Processing Systems, 36:7475-7505, 2023.
Scott Pesme, Loucas Hilaud-Vivien, and Nicolas Flammarion. Implicit bias of sgd for diagonal linear networks: provable benefit of stochasticity. Advances in Neural Information Processing Systems, 34:29218 000992118-2220,,.. 29230, 2021.
Ramprasaath R Selvaraju, Michael Cogswell, Abhishek Das, Ramakrishna Vedantam, Devi Parikh, and Dhruv Batra. Grad-cam: Visual explanations from deep networks via gradient-based local- ization. In Proceedings of the IEEE international conference on computer vision, pp. 618-626, 2017.
13

---

Published as a conference paper at ICLR 2026
Ramprasaath R. Selvaraju, Michael Cogswell, Abhishek Das, Ramakrishna Vedantam, Devi Parikh, and Dhruv Batra. Grad-cam: Visual explanations from deep networks via gradient-based lo- calization. International Journal of Computer Vision, 128(2)3336335, October 2019. ISSN 1573-1405. doi: 10.
URL http:: /ttp:////x//oo doi. wttp://o..oo..o...oooooooo
s s11263-019 0000-------12---7
Baekrok Shin and Chulhee Yun. Implicit bias and loss plasticity in matrix completion: Depth pro- motes low-rankness. In The Fourteenth International Conference on Learning Representations, 2026.
Dongkuk and Chulhee Yun. Practical sharpness-aware minimization cannot converge all the way to optima. Advances in Neural Information Processing Systems, (6:26190-262222 2023.
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The im- plicit bias of gradient descent on separable data. Journal of Machine Learning Research, 19(70): 1-57, 2018.
Jacob Mitchell Springer, Vaishnavh Nagarajan, and Aditi Raghunathan. Sharpness-awan minimiza- tion enhances feature quality via balanced learning. arXiv preprint arXiv:2405.20453 2024.
Hao Sun, Li Shen, Qihuang Zhong, Liang Ding, Shixiang Chen, Jingwei Sun, Jing Li, Guangzhong Sun, and Dacheng Tao. Adasam: Boosting sharpness-aware minimization with adaptive learning rate and momentum for training deep neural networks. Neural Networks, 169:506-555 -19, 2024.
Gal Vardi. On the implicit bias in deep-learning algorithms. Communications of the ACM, 66(6): 86-93, 2023.
Shuyang Wang and Diego Klabjan. A mirror descent perspective of smoothed sign descent. arXiv preprint arXiv:2410. 141188 2024.
Kaiyue Wen, Tengyu Ma, and Zhiyuan Li. How does sharpness-aware minimization minimize sharpness? arXiv preprint arXiv:2211.05729 2022.
Kaiyue Wen, Zhiyuan Li, and Tengyu Ma. Sharpness minimization algorithms do not only minimize sharpness to achieve better generalization. Advances in Neural Information Processing Systems, 36:1024-1035, 2023.
Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Conference on Learning Theory, pp. 3635-3673. PMLR, 2020.
Chulhee Yun, Shankar Krishnan, and Hossein Mobahi. unifying view on implicit bias in training linear neural networks. arXiv preprint arXiv:2010.02501, 2020.
Jihun Yun and Funho Yang. Riemannian sam: Sharpness-aware minimization riemannian mani- folds. Advances in Neural Information Processing 3 36:65784-658000 2023.
Yihao Zhang, Hangzhou He, Jingyu Zhu, Huanran Chen, Yifei Wang, and Zeming Wei. On the duality between sharpness-aware minimization and adversarial training. arXiv preprint 255 2024.
Zhanpeng Zhou, Mingze Wang, Yuchen Mao, Bingrui Li, and Junchi Yan. Sharpness-aware min- imization efficiently selects flatter minima late in training. arXiv preprint 2400:22410..0333, 2024.
Zhanpeng Zhou, Mingze Wang, Yuchen Mao, Bingrui Li, and Junchi Yan. Sharpness-aware mini- mization efficiently selects flatter minima late in training, 2025. URL httttt//////xi..orr abs/2410.100 10373.
Juntang Zhuang, Boqing Gong, Liangzhe Yuan, Yin Cui, Hartwig Adam, Nicha Dvornek, Sekhar Tatikonda, James Duncan, and Ting Liu. Surrogate gap minimization improves sharpness-aware training. arXiv preprint arXiv:2203.0800 2022.
14

---

Published as a conference paper at ICLR 2026
CONTENTS
Introduction
1.1 Summary of Our Contributions
1.2 Related Work
2
Preliminaries
3
SAM with fo -errturrbati
3.1 Depth-1 Networks..
3.2
Deeper Networks (L 2 2)
SAM with &e-Perturbations: Sequential Feature Amplification
4.1
Asymptotic Behavior on Depth-1 and Depth-2 Networks
4.2
Pre-asymptotic Behavior on Depth-2 Networks
4.2.1
Sequential Feature Amplification
4.2.2
Understanding the Effect of f2-SAM
4.2.3
Analysis of Time-wise Sequential Feature Amplification
8
4.2.4
Analysis of Initialization-wiss Sequential Feature Amplification
5
Experiments
10
6
Conclusion
10
A Figures and Discussions Omitted from Main Text
17
Flow Trajectories GD and SAM
17
A.2 More Discussion on Related Work
17
A.2.1 Recent Work on Implicit Bias Diagonal Linear Networks
17
A.2.2
Comparison with Saddle-to-saddle Dynamics
17
A.2.3 Properties of SAM
18
A.2.4 Depth-Induced Implicit Bias
19
A.3 Derivation Rescaled t, -AM flow
19
A.4 GD and SAMAM do not exhibit sequential feature amplification
19
A.5 Interesting Trajectory Regime 2 of Theorem 4.4
20
Core Lemma for SAM on Depth-1 Networks
21
C SAM with pertuurrbaaso:::: Proof of Section 3
25
C.1 Depth-1 Networks: Proof of Theorem 3.1
25
C.2 Technical Challenges for Multi-Point Datasets ................................
25
C.3 Proof of Theorem 3.2
26
C44 Proof of Corollary 3.5
31
15

---

Published as a conference paper at ICLR 2026
C.5
Finite-time Blow-up
32
C.6 Empirical Verification
33
C.6.1 One-point Case: Discrete Continuous Dynamics
33
C.6.2 Multi-point Case: Persistence One-point Behavior
................................
D SAM with &--erturbattion:: Proof of Section
36
D.1 Depth-1 Networks: Proof of Theorem 4.1
36
D.2 Derivation of f2-SAM flow
36
................................
D.3 Proof of Theorem 4.2
37
D.4 Proofs for Section 4.2.3
................................
40
D.4.1 Recap: Basic Notation
D.4.2
Preliminary Analysis
42
D.4.3 Proof of Lemma 4.3
42
D.4.4 Preliminary Analysis for m. (l) Trajectory Analysis
43
D.4.5 Proof of Theorem 4.4
46
D.5
Extension to deeper diagonal linear networks
53
D.6 Proofs for Section 4.2.4
D.6.1 Derivation of the Dynamics of (tt)
58
D.6.2 Proof of Theorem 4.5
................................
59
D.6.3 Proof of Proposition 4.6
62
D.6.4 Proof of Proposition 4.7
................................
63
D.7
Numerical Evaluation of Theorem 4.5
66
D.8
Empirical Verification
66
D.8.1
One-point Case: Continuous vs. Discrete Dynamics
67
D.8.2
Multi-point Case: Persistence of One-point Behavior
70
D.8.3
Depth-L Case: Persistence of Depth-2 Dynamicss
72
Experiments
74
E.1
Loss Dynamics
74
E.2
Sequential Feature Amplification under Random Initialization
74
E,3
Alternative 2-Layer Models
76
E.3.1
Linear Network
76
E.3.2
Convolutional Neural Network
77
E.4
Grad-CAM
78
E.4.1
MNIST
79
E.4.2
SVHN
81
E.4.3
CIFAR-10
82
16

---

Published as a conference paper at ICLR 2026
DECLARATION OF LLM USAGE
We used Large Language Models (LLMs) solely to aid or polish writing. They did not generate ideas, analyses, conclusions. All LLM-assisted text reviewed and edited by the authors.
A FIGURES AND DISCUSSIONS OMITTED FROM MAIN TEXT
A.1
FLOW TRAIECTORIES OF GD AND SAM
1.4
1.4
1.2
1.2
1.0
1.0
GD
GD
0.8
f.-SAM
B2
0.8
F. SAM
f2M
F2 SAM
0.6
0.6
0.4
0.4
0.2
0.2
0.0
0.0
0.0
0.5
1.0
1.5
0.0
0.5
1.0
1.5
B1
(a) Depth (linear network)
(b) Depth (linear diagonal network)
Figure 6: Trajectories the predictor B(t) E IR² from identical initial conditions under GF, foo- SAM flow and SAM flow on {(,++1)) with μ (1, For SAM, p 1.
A.2 MORE DISCUSSION ON RELATED WORK
A.2.1 RECENT WORK ON IMPLICIT BIAS IN DIAGONAL LINEAR NETWORKS
Jacobs & Burkholz (2024) study continuous sparsification with time-varying weight decay, formu- lating time-dependent Bregman potential that causes the implicit bias to evolve from 62- to -type behavior over the course of training. Wang & Klabjan (2024) study smoothed sign descent on quadratically parameterized regression problem. introducing time varying mirror map. and prove that the resulting limit point is an approximate KKT point of where the stability constant & quantifies the gap to KKT optimality. Papazov et (2024) analyze momentum gradient descent diagonal linear network through a momentum gradient flow, show- ing that newly defined intrinsic parameter determines the optimization trajectory and admits a second order, time varying mirror-flow formulation. Within this framework, they characterize the induced implicit regularization and demonstrate that smaller values of this intrinsic parameter yield more balanced weights and sparser solutions compared to standard gradient flow. Jacobs et al. (2025) extend the mirror flow framework to account for explicit regularization and analyze the evolution of the corresponding Legendre function time, thereby describing how the implicit bias changes in different reparameterizations, including diagonal linear networks. In particular, they track how the implicit bias evolves in terms of its positional bias, bias type, and range shrinking.
objective,
A.2.2 COMPARISON WITH SADDLE- HADDLE-TO-SEE - DYNAMICS
In this section, we provide further details the relation between our work and the saddle-to-saddle dynamics of gradient descent/flow. Pesme & Flammarion (2023) consider diagonal linear networks trained with squared loss in the iimnitssmmll-iiittlizaati hibits incremental, stage-wise learning: the flow undergoes long plateaus near a saddle whose pre- dictor is supported the first & coordinates, then escapes along a low-dimensional "fast escape"
limit. In this regime, gradient flow
17

---

Published as a conference paper at ICLR 2026
manifold to saddle with support on k+1 coordinates, and Sequentiality thus appears dis- crete transitions between saddles with support size & and k+1. In the diagonal setting, complexity is captured by the number of active coordinates, which is constant on each plateau and changes only at these transition times.
In contrast, our work the sequential feature amplification focuses linear diagonal classifier trained with 22-SAM and logistic loss, and on a different notion of complexity: individual coordi- nates (features) ordered by the strength of the teacher signal, from minor major features. In our jumpps, setting, all coordinates are present from the beginning. Instead of coordinate jumps, we track how the coordinate-we wise alignments and margins evolve both over time and as a function of the initialization scale, where by "alignment" we mean the magnitude of the predictor at each coordinate, indicating how strongly the predictor attends to each feature. We show that f2-SAM gives rise to two comple- mentary forms of sequential feature amplification: (i) a time-wise ordering, where alignment with minor features is relatively amplified earlier training and gradually shifts toward major features; and (ii) an initialization sealoeew-sallesseee ordering, where the most-amplified feature over finite training process changes systematically with the initialization scale. In both views, the ordering emerges through continuous evolution of the alignment across coordinates, and sequentiality is captured by which feature is currently most amplified, rather than by discrete activation or deactivation of features.
The mechanisms underlying these two phenomena are conceptually distinct. First, saddle-to-saddle dynamics start from the zero vector and involve successive coordinate activations, where previously inactive coordinates become active over time. Our setting, by contrast, starts from a1 (without taking the limit a + 0), where all coordinates are already active, and the dynamics involve suc- cessive amplification of already-active coordinates. Activation and amplification are fundamentally different: even if saddle-to-sadde dynamics exhibit successive activation, the identity of the most dominant coordinate can remain unchanged, unlike in our setting where dominance itself shifts over time.
Second, the ordering principles differ. In our work, the ordering of amplified coordinates is driven directly by the data geometry, namely the ordering of the signal strengths Py. In saddle-to-saddle dynamics, the progression is governed by dual-thresholding mechanism, tied to when integrated gradients hit constraint boundaries, and does not correspond to a minor-to-major feature progression.
Third, the role of initialization is opposite. Saddle-to-saddle dynamics arise in the vanishing- initialization limit (a - 0). In contrast, we observe sequential feature amplification across a wide range of vaniishiiii initialization scales, and in fact show that increasing induces a clear and systematic amplification ordering. Our phenomenon is therefore not a inallinniiallzzation effect.
Fourth, saddle points play constructive role in our mechanism. Aside from the trivial effect that extremely small initialization can prevent SAM trajectories from escaping the origin, saddle points do not drive the sequential feature amplification characterize. The observed dynamics not mediated by saddle escape.
Finally, the problem setups are fundamentally different. Prior saddle-to-saddle works analyze gression under squared loss, whereas our work studies classification under logistic loss, where the optimization landscape and asymptotic behavior are qualitatively different.
Taken together, these observations indicate that sequential feature amplification is a SAM-specific phenomenon, distinct from known saddle-to-saddle or incremental learning dynamics, and does arise under conventional gradient descent.
A.2.3 PROPERTIES OF SAM
SAM exhibits distinctive valley-bouncing dynamics (Bartletr et al., 2022; Wen et al., 2022) and convergence instability near local minima (Si & Yun, 2023; Kim et al., 2023). SAM prefers low- rank solutions (Andriushchenko et al., 2023), with its normalization term playing a crucial role (Dai et al., 2023).
Previous works (Andriushchenko Flammarion, 2022; Clara et al., 2025) have studied SAM's im- plicit bias in diagonal linear networks. Andriushchenko & Flammarion (2022) analyze 2-layer lin- ear diagonal networks under sparse regression with MSE loss, showing SAM induces better sparsity than gradient descent, but require the small-p assumption. Clara et al. (2025) study SAM dynamics
18

---

Published as a conference paper at ICLR 2026
with noise, proving weight balancing across layers and sharpness minimization, also limited to MSE loss. Our analysis removes the small-p assumption to capture the full perturbation effect and studies logistic loss, revealing distinct implicit bias properties compared the squared loss setting.
A.2.4 DEPTH-INDUCED IMPLICIT BIAS
A line of work shows that increasing depth in overparamettrizee lincarized models qualitatively change the implicit bias of gradient- based training dynamics (Vardi, 2023). matrix problems, deep matrix factorization exhibits creases (Arora et al., 2019; Chou et al., 2024), and recent work further connects depth to low-rau low-rank bias and loss of plasticity in matrix completion (Shin & Yun, 2026). Relatedly, depth also in- behavior in simplified models (Gissin
enhanced tendency toward low-rank solutions depth in-
duce incremental learning dynamics and sparsity al.. 2019). Our results complement these findings by isolating a depth-dependent implicit bias mechanism specific to SAM in diagonal linear networks.
A.3 DERIVATION OF RESCALED 4.-SMM FL.WW
For the dataset {(p, +1)}. the loss function is given as:
(()) = ((3(0),
For each i [Z],, the gradient is
wio ((0) = ( (( (,000)0))))((00))0))000000001110
(1)
(5)
Esi
Then, we have the fp-SAM flow of ( as
i ( (1) - Vo 1 L(((1)) = -=(((((()))
( ih
(1))
Since (u) = 1+exp(u)
0, has the same spatial trajectory (up to reparameterization of time):
0 (t) = م O
w (0))))) = pa @
(O w (t) ( ((())))
This derivation works for any p. not just p and
A.4
GD AND for -AAM DO NOT EXHIBIT SEQUENTIAL FEATURE AMPLIFICATION
0.8
2.00
0.7
1.75
5
initialization Scale
0.6
Initialization Caal
1.50
4
2
0.5
1.25
0.4
1.00
3
xew
0.3
0.75
2
0.2
0.50
0.1
0.25
1
0.0
0.00
0.0
0.5
1.0
1.5
time
time
(a) GF
(b) lx SAAM
Figure 7: Dominant index jt arg max, (t) for GF and lo -AM flow over (t, ar) on D with
(4, 5, 6, 7,8) E 5
19

---

Published as a conference paper at ICLR 2026
A.5
INTERESTING TRAJECTORY IN REGIME 2 OF THEOREM 4.4
0.7
0.7
12
0.6
initialization scale a
10
10
0.5
9
initialization scale
(
0.4
8
7
0.4
Xee
6
6
0.3
0.3
5
0.2
4
3
0.1
0.0
0.0
0.2
0.4
0.6
0.8
1.0
2
time
time
(a) 7 (short horizon)..
(h) = (long horizon).
Figure 8: Dominant index for f2-SAM flow with μ (,, 12). The black line indicates the interesting trajectory.
In Regime 2 of Theorem 4.4, there is also interesting sub-regime that corresponds smaller values of a with the range of Regime 2. Define a critical threshold Aeeitt := <2[| |||] E (a1, Q2).
When 01 < a < Cterit. the trajectory mc(() initially decreases to minimum above HL and then increases. During this decreasing phase, the 22-A flow amplifies coordinates with smaller indices j < (0) than the most-amplified index at initialization jo(0) E arg min, [Ký m,)) enabling an aggressive exploration weaker features before transitioning the standard minor-first-majo last sequential amplification pattern. Along the black path in Figure 8, this manifests as the most- amplified coordinate starting at (4. then stepping down 31 sequentially during the initial decrease, after sufficient time time-stepping back up sequentially toward 3. as m. increases.
20

---

Published as a conference paper at ICLR 2026
B CORE LEMMA FOR SAM ON DEPTH-1 NETWORKS
Although our argument is inspired by the simple proof of Theorem 9 in Soudry et al. (2018), ex- tending that analysis from gradient descent to the SAM flow is far from straightforward. In GD the gradient has a clean exponential form and all coefficients are fixed, which makes the support/non- support decomposition almost immediate.
In contrast, SAM evaluates the gradient at the perturbed point ti((), introducing the dependent factors "y. (t) and the perturbed margins (t), neither of which appear in GD. Controlling these additional terms turns out to be technically delicate: one must show that the SAM-induced coeffi- cients remain uniformly bounded, that the perturbed margins stay within a fixed range, and that the resulting two-variable function (z, 6) admits a uniform upper bound. Only after establishing these new ingredients can the GD-style argument be recovered. The proof below develops these steps and shows that, despite the additional complexity, the SAM flow converges the same &g max mar-mn direction as GD.
Lemma B.1. For almost every dataset which is linearly separable, any perturbation radius and any initialization, consider the linear model f(x) (w, x) trained with logistic loss. For any SAM perturbation of the form
+ E(w)
with perturbation direction E(w) satisfying
||E())l12 S B for some finite constant B co and all w,
the resulting SAM flow converges in lz max-margin direction.
Proof. Let {( bn) C I2d x {+1} be linearly separable dataset, that is, there exists a vector such that
Fn w. > 0 for all n.
As usual in this setting, we absorb the labels into the inputs and assume without loss of generality that all labels are yn In other words, we redefine x,, + v, and work with dataset {2,, J such that
Be, with x, w. > 0 for all n.
For the linear model f(x) = x:' T w. the logistic loss is
em
((w) Σ [ w) f(u) c ((u)) 1 e-
n=1
The SAM flow with perturbation E(w) is the gradient flow mb() un(() e(w))
ab(()) (())((())
(6)
Let (t) = xrj wt) and (t) x,; +6 (t). Then
vC(((tt) = Σ m.,
No (t) mn m0 1m)
(.) Bn
On.
with
Yn (t) e (f)
1 I 2 0.
Because (({)= wv(() = e(ur(t) and [|]())(1))l S B, if the data are bounded, say ||æ,||| $ R. then
(t, (t) m, (t)l | (w(t) w(t)) < B:
for all n, Hence there constant A > 0 such that
0 S (1)))) S for all n,t.
21

---

Published as a conference paper at ICLR 2026
The SAM flow equation 6 can therefore be written
(t)
Nn (tt)emmll0a)
5 5. (t) S A.
(8)
Let w" denote the & max-mnrgin solution
w" arg min || s.t. a, w for all n.
Let S {n F. un }} be the support set. Standard KKT conditions yield coefficients b, for n E S with Enes bn such that
20 = Σ
bn
nEg
Define the residual
r(t) w() w" log t.
Our goal is to show that rt) is bounded. This will imply that
w(t) w log rt))
w
[t)l
| log t + o(log t)
|l
that is, the SAM flow converges in the & max-margin direction. Differentiating and substituting equation obtain
+tt) it)
w
(t)e mm) xnn
w
We split the sum over the support and non-support points:
+())
Σ
(n)e)e (1)()) m, + Σ
mn)e (4) xn
w
1
RES
ngS
For we have n u= 1, 8o
mn,(t) x w(t) x "" logt x., r(t) logt t + xx]
and therefore
be m1) =0.
For é S we have
m. (1)
x. log 20, logtoo, r(t) x w r(
Using w² =E Un rewrite
+(t)
1 1
Σ
bn
In (t)
e- r]((( - xn Σ (t(tt a. ca5r(()
(9)
on
nES
ngS
Consider the squared norm:
1 d
| rr)) Ti(t) Tt(),
2 dt
(() where Ti (t) and (t()) are the contributions of the two terms in equation 9, For the non-support term
(t(t) in equation 9, we have
(1)))
Σ In 20 2. r.)) (()).
ngS
There is margin gap 0 > 0 such that 2. ut > + 0 when n $ S. Then
T. w' (1+0)
22

---

Published as a conference paper at ICLR 2026
and using (t) v A and V2 = $ 1, we have
T2(t) - +
For the support points, write zn = 2; r(1) and define
(1) = (t(()
vn((1) - (1) =a(()) (0),
On
so that
1 Σ bn (t)().. (i(t) = 8 nES
We first justify that the coefficients equation
(t) (t),t/bn /bn remain in fixed compact interval. By
(m, mt) S.
Since
(m(m,(t)-mm(())
(x(t) = 1+e-0,,(
and the denominator satisfies + e (t) 2 1, we obtain the uniform bound
H (t) S e (t), (t)) I <ee
for all n,t.
Thus each "n (t) lies in the compact interval
[0,e
Next, since every b, > for n E S and is finite set, define
bmin min bn 0, nES
bmax max bn. WES
Therefore
(t)
(t)
0 058 S
for all S and all t.
b
bmin
Hence 5, (t) ranges over the compact interval
Emins mmax]=
0. Umin
For each fixed > O, consider the function
(12,8)== == (Be - 1)z.
As - t00 we have (22,6) oC, and therefore 4/(2, 8) attains finite global maximum on IR. Since 5, (L) E [8min; dmax] for all i, there exists constant Cp > O such that
(22 6) < C.
Vz E R, 4o E [smins Smax].
Consequently,
(1) (1),6,)))))))))))))))))) S C2
En S, Vl,
and therefore
Ci (i(1) < +
C. Cpp Σ bn
nES
Combining the two bounds on (T(t),,((t), for sufficiently large t,
1 d
12 (t) +T=(t) < Ci
A
C2
2 dt
for constant C2 > 0.
23

---

Published as a conference paper at ICLR 2026
Integrating from to to gives
||((1]12 s [|r(()))l? 2C2 /.o (" "u''uu- = ||r(())1 202 log (
so
||(()]]==( (lylot)) o(log t).
Since
w(t) un logt r(t),
||((t]] = ologg))
we obtain
(t)
w
o(1),
(t)ll
which proves
w()
w
[t)
|| 1
Thus f2-SAM flow converges in the l max-margin direction any initialization and any fixed >> 0.
24

---

Published as a conference paper at ICLR 2026
C SAM WITH E:o-PERTURBATIONS:: PROOF OF SECTION
C.I DEPTH-1 NETWORKS: PROOF OF THEOREM 3.1
Theorem 3.1. For almost every dataset which is linearty separable, any perturbation radius and any initialization, consider the linear model f(x) = (w, x) trained with logistic loss. Then, SAM flow directionally converges in the E2 max-margin direction.
Proof. Apply Lemma B.1 with E(w) = sign( V (())). Then |E(w)]l2 I pvd for all w, so the conditions of Lemma B.1 hold. Thus, the flow converges to the &g max-margin direction.
Theorem C.I. Consider the linear model f(x) = (w, x)) trained on the dataset Du with loss [(w)- ((w, x)) where (u) < 0 for all Then, GF and Ex -SM flow, starting from any w(O), evolve on the same affine line w()) + span( and have the same spatial trajectory.
Proof. The model ((x) = (w, x)) = w x. The loss is L(wD) E(w 11). The gradient is Vwl(w) ('ww" u) a with (s) 0.
Gradient Descent The GF
ib vww(w) (''www 11) p.
SAM with loo perturbation
The ascent point is
v w + x (w)
w + psign V (ww)
w psign
The equation of -SA flow is
(wv(
= VwL(w - psign(11))
= -'(w"" μ ps) T (1).
-l'(w" μ 1--ppll --ppullii n.
Therefore, they have the same spatial trajectory as:
i n4.
The term -l'(w" 4 plpH(11) is the accelation in terms of since -s'(s) is decreasing in
C.2 TECHNICAL CHALLENGES FOR MULTI-PONS DATASETS
In the multi-point setting, as w(t) diverges the SAM perturbation becomes asymptotically negligi- ble, so SAM and GD share the same long-term behavior. The regime where they differ is precisely when the p-perturbation is non-negligible, but in the multi-point case the resulting gradients (and thus SAM updates) become considerably complex for tractable characterization of the SAM flow in the regime where SAM and GD diverge. This motivates our focus the single-example dataset Dμ = {(1. +1)}. where the SAM dynamics admit a tractable dynamical characterization while still capturing depth-dependent phenomena unique SAM. In Appendix C.6, empirically verify that these behaviors persist under multi point datasets and discrete SAM updates, indicating that our insights extend beyond the single-point setting.
25

---

Published as a conference paper at ICLR 2026
C.3 PROOF OF THEOREM 3.2
Theorem 3.2. For i E [L], suppose w a E IQ. Let (t) follow the rescaled to -AM flow (2) with perturbation radius p the dataset Du. Then, for the j-th coordinate of B(t):
If aa p. then (t) converges to 0 if L is even, to p if is odd.
Ifa; = p. then (t) p² for all z O.
Ifa; > p and
then (1) grows exponentially; 8,(1) ((exp(2)x)))
If aj and L > 2. let = arg paxpa,>0PP Pa p) T, and also let T minkeJ (((t.-2)julo-pp 2 If E J. then 85(() bounded for all < <.
+ T; otherwise, By(1) stays
Proof. Since suppose u (6) (0) E TR. for all i [L], and the dynamics of the linear diagonal network are invariant under any permutation of the layer indices {1,..., L}, we obtain
w (1) w (1) =.... (0)))))))))()))))) w(l)
for all / 2 0.
With for perturbation, the rescaled lo -SA flow (2) becomes
i 100 (t)= = O
(9 (t) + + =2 (0(1)))
Egi
uu (2) (t) + psign(V L((()))))
p の
Exi
Recall the gradient (5)
(n ((((t)) = =((((()))) o
1
E
where (1) Itexp(u)
0. Since also have μ 0 (element-wise), we have
sign( V vi ((O(t))) sign
()
Efi
sign
w(()
where (a) follows from the fact that w (() w()) for all the ODE
[L]. Using this fact again, have
ii(t) = u (t) o
w(t) psign
w(t)
p o
w(l1 -psign
w((
This can be written as coordinate- wise as
و (t) μ4 (w)(() psign (w)) L-1
for j [d].
Divide into three cases:
26

---

Published as a conference paper at ICLR 2026
Case 1: 2
wi)(( My (wp(() - psign (w)).)).
By Lemma C.2, we have
p + (w,(0) - p)etse
if w,(()) > Pt if to,(0) = P:
P
wy(()
p + (w5(0) - p)ett? (L<T), ( T)
if (yy(0) < p.
0
if (0) 0,
where T = log (0)
Then, we have
(
0k > p:
8,(1) ws(() 7 p²
if aj = p;
as t - 00.
if 0x < p:
Case 2: L 2 and even.
(b)(t) = My (w,(t) psign (wy(t)))
By Lemma C.3, have
p
-(L -2) + (wy(0)-p)(-3
1
if (ey(0) > if wy (0) = P.
())(() =
P
p
(-(I. 2)2st + (wc
(t (t<T), ((2T))
if wy (0) <
0
if (wwy(0) = 0,
where :- (p-w;(0)) (1
Then, have
(1--2)ps
af > p, t + ts, - (t -)==)= 2=
8,(t)= wt))²
ph
if 0; p, f - 00, if a, p. f 00,
where 1" = 1/(1.-2)), --)
Case 3: L> 2 and odd.
(b)(t) = my (w,(t) p)'-1
By Lemma C.4, we have
if w p. if wf tp.
wt(()= p+ (-(L- (-(L-2)yugt [
Then, we have
4 (1)
if p,as t", af S p, 1 00,
B,(1) = ut))=
where 1" = 1/1122)j,(mm,,0 2² 2.
These cases of cover possible cases in Theorem 3.2.
The following three lemmas (Lemmas C.2 to C.4) are used in the proof of Theorem 3.2 and corre- spond, respectively, to the three cases.
27

---

Published as a conference paper at ICLR 2026
Lemma C.2. Let μs 0 and p 0. Consider
ri(t) = μ (w(t) - p sign(w(t))).
Then, there exists the solution such that it is absolutely continuous (AC) and satisfies
w(t) wo(0) + S wi(s)ss.
(10)
In particular,
p + (w(0) p)eue
if > P. (f((0) = P. if w() < if
wt)) =
P
p + (w(0) - p)em (t< (t2
where T' = log p p-w(())
Proof. Case 1: wo(0) O. The constant function w(I) 0s AC, and
S. ((0 psign(0)) ds /-¹ dss 0. Thus, Equation (10) holds.
Case 2: w0)) p. The constant function w(t) is AC, and since sign(w(t)) 1, we have
[/ ulp p.1) ds ( Ods 0.
Thus, Equation (10) holds.
Case 3: w() p. At t O, we have (0) u(w(() p) 0. Assume, for contradiction, that there exists t, > 0 with w(t.) = p. Then on 10, t.) we have w(t) p and hence vi(t) = u (ww(t) - p) > 0, so w is strictly increasing on [0, t.). An increasing function cannot reach the smaller value starting from w(0) > f: contradiction. Thus w(t) > p for all t 2 0. On the region {w(t) > p}. sign(w(t)) and the ODE reduces the linear equation
p(w p)-
Then, have
wh(t) 4
w(l)
w(()
ds / μds
w(s) P
w(t) - P
log
μt
w((0) P
wt(( = p+ + (w: - p) eul
This function AC and satisfies Equation (10).
Case 4: 0 < w((0) < Initially sign(w(0)) = 1, so again vi = p(w p) and w(t) = p + (w(0) - p)e
Since (0) - p 0, the function strictly decreasing and reaches exactly once at
1
log w(O) 0.
On [o. T), this solution AC and satisfies Equation (10). Define w(t) for all 2 T. Then, using sign(0) = 0.
(() = w(T) + [ p (0)) ds 0+ [ 0ds 0,
so Equation (10) also holds on 00). The function w is AC on j0, T] and on [T, oo), and it is continuous T, hence it is absolutely continuous.
28

---

Published as a conference paper at ICLR 2026
Lemma C.3. Let μ > 0. p> O, and L is even. Consider
( μ (w((t) psigu((u(t))=
Then, there exists the solution w such that it is absolutely continuous (AC) and satisfies Equa- tion (10). In particular,
p+ + (-(1--2)) (w(0)-p)t<>
ifw(0) P.
w(t)
P
ifw(O) P.
--2 (t<
p (-(1 -2) + (w²(0)-p)= pp)
(t<T), T), (+2
if p.
if 0,
where T = (e-w(0)) (L- 2) (2-2)
(1--2)p
Proof. The proof is similar to the proof of Lemma C.2.
Case 1: (0) The constant function w(t) 0 AC, and
L. m(0 psign 4-1 ds = 1/¹ ['100 1..0%-1 do ds==00
Thus, Equation (10) holds.
Case 2: w00) p. The constant function w(t) p is AC, and since sign(w(t)) 1, we have
L: "lle - p. 4-1 do ' p.0% -- do =0.
Thus, Equation (10) holds.
Case 3: w0) > At = O, we have w0) - (ww(0) - p) p) > O. Assume, for contradiction, that there exists t, > 0 with w(t.) = p. Then on [0, t.) we have w(t) > p and hence ti(t) = ((w(t) - p) > 0, so w strictly increasing on [0, 1.). An increasing function cannot reach the smaller value p starting from w(0) p: contradiction. Thus w(t) > p for all t 2 0. On the region (w(t) > p}. sign(w(t)) and the ODE reduces to
w u(w - D) -)--1
Then, have
(i(
(w(t) -p) 6-1 =M
//' (w(s) -p)
i1())
ds [' uds
1
1
L --2 (wit) - 1.2 (w(() - -p) nt
(wi(t) - p2 = ( (-LL - -2) + (w(0) p)
1
=w(()= + (-17. ((I.-2)pt 2) + (w(0) p- (a)
where (a) follows from w(t) rho 0. This function is AC and satisfies Equation (10).
Case 4: 0 < w((0) < Initially sign(w(()) = 1, again vi = u(w p) 6-1 and (T.
(w(t) - h22 (-(T- 22)tt + (w(0) p)
29

---

Published as a conference paper at ICLR 2026
Since w= p and L is even, we have
w(1) (-(L- 2) + (w(0) b-2
1
6-2
The function w is strictly decreasing and reaches 0 exactly once at
(o - w()) (1--2) ((-))) (1-2)
T
0.
(L 2)24 2)
On [0, 7], this solution AC and satisfies Equation (10). Define w(() := for all 2 T. Then, using (0g =0, 0,
[f u(0 (00pn()))(())) 1-1 ds 0+ [ 0 0,
w(t) wT)
so Equation (10) also holds oo). The function w is AC on [0, and on [T, oo), and it is continuous at t = T, hence it is absolutely continuous.
Lemma C.4. Let p 0. p> 0 and L is odd. Consider
úi(t μ (w)
Then, there exists the solution w such that it is absolutely continuous (AC) and satisfies Equa- tion (10). In particular,
P p+ (-(L- 2) + (wr()) (w] wf((0) +
if = P.
w(t)
Proof. The proof is similar to the proof of Lemma C.2. Case 1: w(0) = p. The constant function wr(t) = p is AC, and
S (ulo p) ds = / d = O.
Thus, Equation (10) holds.
Case 2: w(0) * p. Separate variables:
dw (w p)1-1 dt.
Integrating from 0 to gives
1
L-2 2 (w((l) p) (w((0) - p) L-2 = pt.
Solving for w yields
-1-2
w(t) = p (-(I-2) 2) jut (w(0) 2-2
1
The function is AC and satisfies Equation (10).
30

---

Published as a conference paper at ICLR 2026
C.4 PROOF OF COROLLARY 3.5
Corollary 3.5. Under the assumptions of Theorem 3.2, let S [j: p} and assume S + Ø. If there is a unique maximizing index j == arg maxjes 15{03 p)6-2 then the SAM flow converges the ej direction. In particular, when L we have j arg max Pi.
Proof. Work under the assumptions of Theorem 3.2 and let
5 = [j 0 > p) + # Ø,
= arg max pj(a3 - p)
jES
where the maximizer unique. We prove that the (rescaled) fa- -SAM flow satisfies
(t))
esr
[t)(1)
Case I, 2. By Theorem 3.2, for S,
By(t) 2
whereas for j & S we have either (t) - (if L even) or By(t) = p" when aj = p; in any event these coordinates stay bounded. Since the maximizer is unique and T. 0,
j arg max JES
hence for every l E 5 {5"}
(t)
0 26y -ee)e)
35- (1)
and for k & S we also have (,, (1) + 0 because the denominator grows exponentially while the numerator is bounded. Therefore (tt)/[(((t]ll + e
Case L > 2. By Theorem 3.2, for each E S there is blow-up time
17 (L 2) u4 (a; - p) L-2
and as t ty.
8,00=
j & S. then 3t(t) is bounded (either converging to when L. is even, or equal to pt when aj = p). The uniqueness of j" implies
Eg mininng
min ej and fy. < fg Vk (i*) Jes
Hence, for any fixed f < 17. all coordinates with * j are finite; moreover,
Bx(t)
lim TT..
=0 for every k + 3y-(1)
because (1-(1) + oo while Be(t) remains finite as & < f. Consequently.
lim 3(1) ttc.. [[(())]2 Cj-.
Combining the two cases establishes the claim. In particular, when L we have j* arg maxjes Mij-
31

---

Published as a conference paper at ICLR 2026
C.5 FINITE-TIME BLOW-UP
In the setting of Theorem C.I, the fx-SAM flow evolves independently across coordinates, In the rescaled &x -SAM flow, each coordinate indeed admits finite blow-up time. However, explained in Remark 3.3, the smallest these blow-up times corresponds to forig = oc in the original SAM time scale. Consequently, both the original flow and the rescaled flow terminate this same time and cannot be extended beyond it.
To illustrate this behavior concretely, we provide Figures 9 and 10 using μ (1,2,3,4,5), p and a depth-Z. = 3 network. In the original flow, only one coordinate diverges as forig + oc. As shown Figure 9b, the normalized trajectories (t(t)/l show that the remaining coordi- nates grow much more slowly than the dominant oneinddeed, they remain bounded. Because their growth negligible compared to the blow-up coordinate, their normalized values converge to zero. Thus, in this example, the trajectory converges the direction Es
In contrast, Figure 10a shows that the rescaled foo -AMM flow, each coordinate (t(t) has its own finite blow-up time. However, Theorem 3.2 identifies blow-up time TL-2)a, -p) any J, which is the minimum of these blow-up times only the coordinates in J/ blow up at T, while all remaining coordinates stay bounded. Since this rescaled time T corresponds forig oc, the flow cannot proceed past 7'. In this example, I ~ 0.25.
Because the rescaled system simply a time reparameterization of the original one, the two plots differ only in their z-axis scaling. Before reaching T, the two flows exhibit the same evolution along the y-axis, Indeed, reparameterizing the original trajectory (Figure 9) by r(t) reproduces the same curve as shown in Figure 10 before T.
Original SAM Flow: Bet) trajectories
Original f.-SAM Flow: A(t3/88tt trajectorie
8.10
8:10
7
8,10
8.10
beta j [bbtaal
0.6
8:0001))001)
buta (00
0.5
0.4
3-
2.0
time
(a) (t) trajectory.
(b) Normalized (t()) trajectory.
Figure 9: (t) and normalized (t) trajectory of the original toe -SAM flow.
Rescaled 1.-SAM Flow: 8i(f) traj actories
Rescaled 1.-SAM Flow: 8(tl//8i1] trajectori
10000
A.00
A
8-00
Polt)
8: (0V1
buta jtt
0.6
0.4
2000
0.2
Bime
(a) (t(t) trajectory.
(b) Normalized (,()) trajectory.
Figure 10: (f( and normalized (t) trajectory of the rescaled -SMM flow.
32

---

Published as a conference paper at ICLR 2026
C.6 EMPIRICAL VERIFICATION
Our theoretical analysis (Theorem 3.2 and Corollary 3.5) establishes the behavior of the foc -SAM flow in the one-point setting 4u In this section, we investigate whether these phenomena extend be- yond the idealized one-point regime. We first examine the discrete-time dynamics (GD and discrete -SAM) on the one-point dataset and verify that they exhibit exactly the same trajectory patterns predicted by the continuous-time theory. We then turn to multi-point datasets and demonstrate that the same qualitative behaviors persist in both the continuous- time flows and their discrete counter- parts. Taken together, these experiments empirically confirm that the insights obtained from Dj carry over robustly to multi-point datasets and to practical discrete SAM updates.
For reproducibility, detail the exact initialization used in all experiments. We adopt the layer- wise balanced initialization wx() (0) a for every i E [Z.) consistent with the setup of Theorem 3.2. The black-edged dot in Figures 11 and 13 indicates the initial predictor ( We set ( (0) 3(0) I/L element-wise that ((0) O (+))) (0) holds exactly. For the continuous-time trajectories, approximate the flow using the corresponding discrete updates with small step size 10-3 via an explicit Euler scheme.
C.6.1 ONE-POINT CASE: DISCRETE VS. CONTINUOUS DYNAMICS
To verify that our continuous timee analysis faithfully predicts the behavior of the corresponding discrete algorithms, we repeat the experiments in Figure 2 using exactly the same initializations, SAM radius P. and feature vector μ.. We simulate both the gradient flows (black curves) and their discrete counterparts (blue dots), including GD and discrete &x -SAM updates. As shown below, the discrete trajectories closely trace the qualitative evolution of their continuous-time versions.
4
4
3
3
2
2
1
1
0
0
(a) GD (L
(b) -SM (L 2)
3
3
2
2
1
1
O
0
0
0
2
0
2
B1
B1
(c) GD (T. 3)
(d) l.o SAM (T. =
Figure 11:
Trajectories ((t) under GF. -SM flow (black line), GD, and discrete lxs -SAM
updates (blue dots) for = 2 and μ = 21, For SAM, we set = 1. For GD and discrete -SAM, we step size 0.1.
33

---

Published as a conference paper at ICLR 2026
C.6.2 MULTI-POINT CASE: PERSISTENCE OF ONE-POINT BEHAVIOR
To examine whether the qualitative behaviors identified in the one-point analysis persist on more realistic datasets, we construct random linearly separable binary data by sampling two Gaussian clusters centered at +u and -r as shown in Figure 12. Specifically, we draw
in (+) = + Ex: Yn +1,
x x = - En Yn = -1,
with En o and use N/2 samples per class (with μ = N1, 100, a 0.5).
Figures 11 and 13 show that the same qualitative patterns predicted by our one-point theory-such such as the asymptotic trajectory structure- also emerge clearly in this multi-point setting. Importantly, these behaviors are observed not only in the continuous-time flows but also in their discrete coun- terparts (GD and discrete -SMM). This empirical evidence demonstrates that the phenomena described in Theorem 3.2 and Corollary 3.5 extend robustly beyond the one-point setting to general linearly separable datasets.
+1 class
-1 class
0
-2
-2
2
X1
Figure 12: randomly generated linearly separable dataset used in our multi-poin experiments. We sample two Gaussian clusters centered at +u =11 with isotropic noise (s ~ NN 0.52I2)) and assign labels +1 and accordingly. This dataset used to evaluate whether the one-point phenomena from Theorem 3.2 and Corollary 3.5 persist in the multi-point regime.
34

---

Published as a conference paper at ICLR 2026
4
4
3
3
2
&
2
0
0
B1
B1
(a)
GD (L 2)
(b) fo SAM (L = 2)
4
4
3
3
2
2
1
1
0
0
(c) GD (L
(d) Exc -SAM (L 3)
Trajectories (tt) under GF, -SM flow (black line), GD, and discrete SS- updates (blue dots) for on random multi multi-point dataset in Figure 12. For SAM, we set p 1. For GD and discrete l SAM, we use step size 0.1.
Figure 13:
35

---

Published as a conference paper at ICLR 2026
D SAM WITH 12- 4-PPERTURBATIONS PROOF OF SECTION 4
D.1 DEPTH-1 NETWORKS: PROOF OF THEOREM 4.1
Theorem 4.1. For almost every dataset which linearly separable, any perturbation radius p and any initialization, consider the linear model f(x) = (w, x) trained with logistic loss. Then, S2- flow directionally converges the Ez max-margin direction.
PI Then I|E(w))l2 S for all w, so the conditions of Lemma B.1 hold. Thus, the flow converges the & max-margin direction.
Proof. Apply Lemma B.1 with <(w)
vC(0)
Theorem D.I. Consider the linear model f(x) (w,a) x)) trained the dataset Du with loss [(w) = (( where (u) < 0 for all u. Then, GF and 22-SAM flow, starting from any w0)), evolve on the same affine line w(0) + span { and have the same spatial trajectory.
Proof. The model is f(x) (w, x) = w æ. The loss is L(w) E(w u)) The gradient is Vwl(w) (www μ)) a with (s) <.
Gradient Descent GF is
w = [wl(w) = =w(ww 14) pe.
SAM with & perturbation
The ascent point is
i w + pEz w)
V w
w
+
V m ((
/
w
|1||
The update of &-MA flow is
w V ((w)
1
vwf(ww
"|pa||2
-l'(wn" T P 1i
l''ww" μ p p||pa(2) μ.
Therefore, they have the same spatial trajectory as:
i n.
The term -l'(w" 4 p pl|1a122 is the accelation terms of since -(((s) is decreasing in s.
D.2 DERIVATION OF SAM
Let us get the f2-SAM flow. The gradient is
Von L(0) = (((((()), 1)) V not (3(0),
{((((()), 11) 1 o (1)
for (i, E ((1,2), (2, 1)).
From the gradient, we have
(1) E2 (0) = Vuin L())
1 (
μ ④ w (
|CC()) (0)
-p Vl will
4l H 2lll
-p
ng
36

---

Published as a conference paper at ICLR 2026
for (i, (11, where no (1) - Ite" < 0.
Vll C w + ll o w and (a) follows from
We consider the initialization (0) = w (2) (0) R Then, since the loss function and dynamics are invariant under exchanging w and w(22, we have (t = w = wi() for all 2 O. Therefore, the update w(() by rescaled 22-SAM flow is given
/ (()
tb(t) pa ② (w(()
nig (t)
D.3 PROOF OF THEOREM 4.2
Theorem 4.2. For almost every dataset which is linearly separable, and any perturbation radius
P. consider the linear diagonal network of depth f(x) (wr()) + ur(2), x) trained with logistic loss. Let (wr(1) (1), w(() (n(2)()) (I)) follow the 22-SAM flow with ()))( (0) w (0). Assume (a) the loss vanishes L(w () (t), w (t()(t)) - 0, (b) the predictor (1) == w (t) c wn(2) (t) converges in direction. Then the limit direction of B(t) is the & max-margin direction.
Proof. Let ((on Mn)) n=1 C TRdd x {+13 be a linearly separable dataset, meaning that there exists w, IRd such that
Fn W+ > 0 Vn.
As usual, we absorb the labels into the inputs by redefining on + In so that we may assume 3/- = 1 for all n and
Ew. such that x; w.
Vn.
We consider depth-2 diagonal linear network with parameters wi, W2 Id defining the predictor f(x; w1, w2) (wi w2) Tx TTx,
β == wi w2.
The loss function is logistic:
('(u) e
C(w1, W2)
Ee(BTe
((0 xn),
((u) = log(1 + te
1te
We study the 22MS flow with fixed perturbation radius p >> 0:
wi (1) - -., L( 2 (1), (1)),
rd2(t) = - V m L(w (t), (2)
where
Vw, L(w (t), w2)
(t) (t) + w C( 1 (t), w (2)
i 1,2.
Step 1: Balanced initialization removes layer imbalance. Let =)()== = u (t) w0) (1).
(2)
From the SAM flow and
0c
w Σ
1
xn Ennj
10
ac
w Σ
e' (33 x In.j W3
(1)
Ouy (1)
Ow
(2)
one obtains
N
,y -100=-((011000000000000
G²(t) = Σ (( x..) Kn.j-
Here the factor 1 + (1) arises because the gradients in the SAM update are evaluated at the perturbed parameter
V VL (ww(t)
@(t) wt) +
V ((w((()))
37

---

Published as a conference paper at ICLR 2026
rather than w(t) itself. Since the perturbation has fixed magnitude p while the parameter norm satisfies [t]]]] + oo along any vanishing-los trajectory of a 2-homogeneous model, the relative perturbation decays:
(t) w())ll2
P
0.
(1)l
I2((t)22
Consequently, the gradients (t) and vl(w((t) become asymptotically colinear, and replac- ing the latter by the former introduces only vanishing multiplicative error 1 ++((1) in the imbalance ODE for (t).
Since 2, (0) under balanced initialization and the ODE ()) GG,((),,)))) +o)) is linear with a Lipschitz right-hand side, uniqueness solutions implies 2,(() = 0 for all I. Hence for all
w (t (t) w
/ (t) =: at)
By(t) ay(t)?
Step 2: Predictor ODE.
From the SAM ODE,
(t) = -a ()(t) (
Hence
3,(t) = (()((((()(( = -aa==-2a,ttta0 - (1) o(1)) = -23(()0,((((+000))
Noting that
Vol(B), Σ
(B (n) In.t
n=1
since
N G₁(t) Σ
(B an) In.j
Σ a=1
((B(t) x,) and (1+ o(1)),
we have
G,(t) (L
Hence the coordinate-wise predictor dynamics
() = -1)
become
3,(0 = -00==-2200V,
Writing this in vector form using
we obtain
B(1) - ddaag(((()) (L((())
(11)
Step 3: Geometry induced by the diagonal parameterization. To characterize the optimization geometry associated with the depth-2 diagonal model, we invoke Lemma D.2. The lemma shows that, for the parameterization
B w
w(2
and R(w(") w (2)) (1)
the induced predictor norm is exactly the & norm:
= min w )w
(1) w (2)) = -|]11
Moreover, on the balanced submanifold (1) (2) with f = 02 the lemma establishes that the Riemannian metric induced on predictor space
(u, v).v M()e, M(B) diag(G).
Therefore, the natural- natural-gradient steepest-descent flow with respect to the induced norm || Ilw takes the form
(()) (((((( = - 2 dia)( ()C((())
We next compare this asymptotic steepest-descent flow with the predictor ODE arising from the t 2-MAM dynamics.
38

---

Published as a conference paper at ICLR 2026
Step 4: Asymptotic identification with & steepest descent. Comparing equation 11 with the steepest-descent flow above shows that the SAM predictor dynamics coincide with the & steepest- descent dynamics up multiplicative factor + o(1) and vanishing perturbation. Assumptions (a) and (b) guarantee that these perturbations do not change the limiting direction of ((1)/1]
Step 5: Conclude & max-margin. By the max-margin theorem for steepest descent given norm (Gunasekar et al. (2018a), Thm. 5; extended to logistic loss by Lyu & Li (2019)), any trajectory following l steepest descent and satisfying ((3(t)) + 0 converges in direction the & max- margin solution. Since the SAM predictor dynamics are asymptotically equivalent to & steepest descent, and by (b) the direction limit exists, we obtain
B 3
3 arg min || s.t. B x, 1.
8
Lemma D.2 (Induced Norm and Natural Gradient Metric for Depth-2 Diagonal Models), Consider the depth-2 diagonal parameterizution
B = w we Rd
and the quadratic parameter regularizer
R(w (2)) = 2 (120(1) + 11(0(2111))
Then the induced predictor norm
|| == min R( (2)
(ww =3
satisfies
11@ 113311.
Moreover, the submanifold where w (2) and B 02 the Riemannian metric induced on the predictor space by is
(u, u)N u M()u, M(B) diag(G).
Consequently, the natural steepest- flow w.r.t || Ilv is
= - v = 2 diag( Vg((B).
Proof. Computation of the induced norm. For each coordinate j. the constraint 3, w/) W/u w decouples. If B, the minimum is attained at (w w (0) (0, 0) and equals 0 = |8,1-
For 8g # 0, eliminate uf (2) via w B//ww/" (1) and minimize
<r(w)=
w
w +
2
Differentiation yields (w(w) 833 w -3 whose nonzero roots satisfy - 883. so that |wo] 18,11/2, Substitution gives (ww)) ) = 13,1. Summing over j yields the induced norm
d
1| = Σ 1351 = 11,3111
J=1
(ii) Local parametrization and Jacobian. On the balanced submanifold (1) (2 cdd. the predictor is
8, = o3.
Hence the Jacobian of the map a + / is diagonal:
09;
2a, 8zk-
Oaz
39

---

Published as a conference paper at ICLR 2026
(iii) Riemannian metric induced from R. The regularizer restricted becomes
R(a, a) = ||a|
Thus the parameter-space metric Euclidean on a. For a tangent predictor perturbation d3, the corresponding parameter perturbation is
daj = d3, d3,
20j 2/
Thus the squared parameter differential is
2
d
(d,3)²
||da|
Σ
2
4
j=1
j=1
Therefore the predictor-spee spacee inner product induced by R
d 4
(u, U)N =
48,
==1
Equivalently,
M(8) diag( (3 ,...,,11)
Inverting yields
M = 4 diag (BI..... Ba).
(iv) Removal irrelevant constant factor. Steepest-descent flows are invariant to multiplication of M by any positive scalar constant. Thus M (33) is equivalent, for optimization dynamics, to
((3) = 2 (daag(3;
which is the conventional normalization in the induced-norm literature.
(v) Natural gradient flow. By definition of steepest descent under the induced norm,
e M(3) Vp((B) -2 diag VaL(g).
D.4 PROOFS FOR SECTION 4.2.3
In this section, we provide detailed proofs for the trajectory analysis of SAM flow, with a focus on the roles of the initialization scale a, the perturbation radius p. and the feature vector μt. For notational simplicity, we omit the time dependence (t) when the context is clear.
Assumption D.3. the initial weight parameters are positive and symmetric: (n (0) w (2) ol for some scaling factor
0.
Assumption D.4. the vector l4 has strictly positive, increasing coordinates: /41 <
pdd
(Equivalently, up to fixed permutation we may assume the coordinates are monotone.)
We introduce two auxiliary quantities. Define the normalized weights p,(() = Σ E(-1 and
their moments M²(t)= = E/=n (t).. Using these, we set the thresholds
M2(1) (t) 2Mr(t)
my 2
In the proof, consider E(BB, 14) term, so not only considering the spatial trajectory but full gradient flow without any reparameterization. We define the margins at the current and perturbed parameters as s(t) == (B(t), μt) and s(t) == (ß(t), μ). Set X(t) == je' (s(t)). the slope of the loss with respect to the margin evaluated at the perturbed margin.
40

---

Published as a conference paper at ICLR 2026
D.4.1 RECAP: BASIC NOTATION
Recall the margin (3, p) and the loss [(s) log (1 + exp(-s)). The derivatives of loss with respect to the margin are:
df -(ss) + exp(s)
1
ds
d² (1s)){(-8)00 >8)
ds = =((s)((-
(-))- is the sigmoid function. We define == = (--)) E 1) as the logistic loss slope magnitude. The gradients with respect to the weight parameters, obtained via the chain rule, are:
where ((s) = (1 + exp(-s))
dw (1) = DC ds
dC
Appio
dc
dC ds
Apy u (1)
ds dwp (1)
du (2)
ds dw (2)
The squared norm of the gradient vector
|Vocl1 Σ
X23
(2) 11) (unft)²
(w
X 22n
j=1
((m5))² E==1 13
(w) (w))))
where no
SAM perturbs parameters by taking step of
size p along the normalized gradient direction.
VoL
E2 :
V eeCl2
PP W (2)
(52)
n2
11 (1)
(
The perturbed weight parameters are
(2)
P w (1)
p
(wh1); == u
(1b2); = 10
71
The perturbed By becomes
B ; W
(1) (2)
(1) (2)
PHc
(n (uuf3)) 2)
p² (1) (2) m
+
w
na
p²pu 22
PI1j
10 )
(w (122) 2) )
3g
1
ng
The perturbed margin and loss slope magnitude are
: - p1)
x:= o(-))
Recall that the SAM flow dynamics are given by:
Ô - )=VoC(()
41

---

Published as a conference paper at ICLR 2026
D.4.2 PRELIMINARY ANALYSIS
We first establish a key property of the SAM flow: the balancedness of the weights.
Lemma D.5. Under Assumption D.4, the SAM flow decays the quantity w/ u))(t) (t) ww w/ (t) exponen- tially to zero.
Proof. Define A == w) w/ The SAM dynamics yield
(2)
a = )pgrgs
nt}²) = Eppzioy
The time derivative of A
3 W /
(1) (2)
(2)
(1)
(2)
w
pmy
spz w
Plg w
n2
-ipp
1
pPj Ay.
ne
Since Å is positive and M4 >> 0, it gives exponential decay.
A,(T) A,(0) exp (-r) / (1+
MHg ne
dt
Hence, the quantity (t) - w (t) decays exponentially.
Proposition D.6. Under initialization with wj (0) w/)² (0) and Assumption D.4, the equality
(1) (t) w (2) (t) is preserved for all t 2 0. Furthermore, the sign of wi (t) and w (2) (t) remains unchanged throughout the dynamics.
Proof. With w ()) we²² (0), we have A,( = ( w/ (0) = 0. By Lemma D.5,
(,(t) 0 for all 2 0. Given this balancedness, each coordinate evolves multiplicatively according
to
PMGW (1)
w = )n =
Anz
PHg
no
n2
This differential equation has the unique solution
[- T i(t) (1 Pf4
w (T) (T) w/ (0) exp
My
dt
n1
Since the exponential function is always positive, w (t) and w/ (t) maintain the same sign as their initial values throughout the dynamics.
(t)
D.4.3 PROOF OF LEMMA 4.3
We begin by restating Lemma 4.3.
Lemma 4.3, The rescaled fsM flow (2) (t) (t()));(t) with (t))) = 2/4r (1-
on, no (t)
Proof. Under Assumption D.3 and Assumption D.4, the Proposition D.6 holds, which ensures that w w B3 for all t 2 O. So we have
w ( 2 (w/))) (=f)))² 23,
n
Σ
j=1
42

---

Published as a conference paper at ICLR 2026
The evolution equation for 3, is
83 = 10j (1) w.f + ww 1i)
2X (1 pP.
(12)
no
This yields
T ["x 3 (1-
B,(T) = 3, exp 2pl Jn
PHC
Tig
Let ry 27μu (1- PM4 When rj 0, Bg grows locally exponentially. Otherwise, it decays lo- cally exponentially. The key insight is that each 3, 's growth rate depends on the interaction between the gradient magnitude Ả and the perturbation term PM) This interaction drives SAM's implicit
ne
n
bias.
D.4.4 PRELIMINARY ANALYSIS FOR m²(t) TRAJECTORY ANALYSIS
Before proving Theorem 4.4, we establish some preliminary results that will be used in the proof. Lemma D.7. Under Assumption D.3 and Assumption D.4, the time derivative of me()) is given by
m = ((t)MM(t) (me(t) mm(()).
Proof. Recall that my = 2M4
Ma. where
d
1233z
M, Σ PoKlf
Pj =
(13)
Σ 1= 1r.33
Substituting the definition of Pj, obtain
En ZE, 13
E,
2E, 13333
M2 Σn 12
M1 Σn 17.32
ni
T
Since μ1 <... Pa and pj 2 0 with Σ,P; Pj new expression for me-
1, we have m1 S mH
M2 2M²
< 2
We define
(t) = vs
where == n
(14)
2p"
Taking the time derivative of S, we have
8 = 2
1g.
2
From Lemma 4.3, we have B5 r3 where r3 2 Ui (1- tuting this into the expression for S, we get
2 (py
13 2m
Substi-
no
113
$= 2 11
21
B3
2m
22
= 41
13
11
2m
1=1
13
= 4 Σ
13
2mc
1=1
43

---

Published as a conference paper at ICLR 2026
Recalling that M1 bzm= n9993
2
1398
and M2
we can rewrite the sums as
d
33 = MIS
M2S
13
2
2
1=1
j=1
Therefore, we have
MMS
M25
$ = 4
2. 2m
M2
2AS
Mi
2m
Since me
v3 we have:
2p
$
$
me
2p
2/8
4pv5
Substituting our expression for S:
[M
M2
2m
me
4pvS
\VS
M2
M1
2p
2iTec
M2
Smc
M1
2mo
M2
ÄM
mc
2M
ÄM (m - 11))
Next, derive the time derivative of mi- Lemma D.8. Under Assumption D.3 and Assumption D.4, the time derivative of my is given by
miu 2m
(2m
where D1 == M, M2 and D2 = Mi M4 M2M2.
Proof. Starting from mm
Ma. 2M.
we have
niu MrM1 Ma
2(M2)²
M2
Mr
Mr
2M
Mi
M2
d
Σ 2313
Σ popp
2M1
M1
j=1
j=1
d
Σ
P. (M²
2mHHz))
2M
j=1
Since 33 =r33; where r; = 25 (13
u²
2mm
we can compute
p. = (138,) .r.. (E
(1²33))
(Ein M7
EE 1R
44

---

Published as a conference paper at ICLR 2026
d
Tj Σ
Pj
P
k=1
(11 11g
14
Pj 2λ
Σ ps
2m
2m
Pj 2t ((R) M1)
(43
M2)
2mne
Substituting this into the expression for rig we have
Σ Pj (Ks M1)
(N] M2) (nz 2mz/())
M
2m
We split the sum into two components:
First term:
C1 c1= = Σ Ps (Ky - M1) (A3 22mH/s)
Second term:
C2 Σp; (ng -M2) (n3 2mmH4y))
For the first term,
Ci = Σ Piic 2mm1 Σ Pipl - -M Ersog + 2 Σ P511s
1
M3
2m M2 + M2
M, M3 I'1
M3
M
Mi
Mi
For the second term,
C2 Σ
Pin 2m14 Σ -M2
M2
Epsug + 2ma M2 Σ pslss
M4 2, M3 + 2m
M2
M1 - M2
T2
= M1 Mi
Mi
Mi
Therefore, we have
min
Σ
Pj
(nJ
Mi)
2m (13 Ma) ) (n3 -2mp/4))
Mi
j=1
I'2
Mi
M,
2m
(2m.T1 - -T)).
2 (M1)²
me
Next, we establish key inequalities involving the threshold mi- Proposition D.9. ri 2 and l22 0.
Proof. P1 and I2 are defined in Lemma D.8. M, and pj are defined in Equation 13. Let M, Ej Pı1{ Ep [u]]] By Cauchy-Schward with X = 1/2 and Y = 3/3/2
(Eep [nrT)2 < Epp (ea) Ep [
r Mi M2 M 2 0.
45

---

Published as a conference paper at ICLR 2026
By Cauchy-Schwarr Schwarz with X μ and Y u
(Ep (100 (1=1) S Ep (42)
Ep (k4)
Multiplying the two inequalities gives
Pp [1_]] Ep [n*] s Ep [e)) Ep [1=*]
T2= = M, M2M 2 0.
Proposition D.10. Let mp
2
We have mp for all 0.
Proof. We use same notation as in the proof of Proposition D.9. Let == by Proposition D.9. Then we have
Mr. Mi
rí 2 0 and l'2 2
Ep [12 axt)=] Epppy 2a Eyle + a² Eplu²)
= M1 - 2a.23 + a² a²M2.
Substituting
Ma and multiplying by Mf gives Mi
MP [u²
1 2
MR 2M, M2 + M2.
M,
Since an expectation of a square is nonnegative and M 2 0, it follows that
Mf M4
2M + M 2 0.
Therefore, we have
T2
M2
2
mH
2r,
2M1
D.4.5 PROOF OF THEOREM 4.4
We begin by restating Theorem 4.4 for convenience.
Theorem 4.4. There exists unique aı such that oo := P1
/2
a1
02
Pa and the trajectory of mt(t) falls into one of the following three regimes. V21|aa12
(Regime 1) If < 04,, then m²(t) strictly decreases for all t. 2.2 O and there exists Tj such that for j E [d], (t(t) strictly decreases for all 2 Tí.
(Regime 2) on < a < 02, there exists T2 such that me me(T2) < increases for all T2.
Pa- and me(l) strictly
(Regime If 0z, then me me()) > pa-1tra and Ball) has the largest growth rate for all 2 0.
Proof. Recall from Appendix D.4 that
12
d M, = Σ Popt.
Pj =
Σd 1
1 B
j=1
M2
4 = (((()) 0.
mH 2M
mL
From Lemma D.8, we further have
T2 2r,
P M, M2,
P2:= = M, M2
mp =
Also note that
VS
= n
m.
2p"
46

---

Published as a conference paper at ICLR 2026
From Lemma D.7 and Lemma D.8, we have
me AM1 (me -mi),
TH (M1)² me
(2m-I1 - r2)
We partition Regime into sub-regimes I-a (ox < ag) and 1-b (ao a < an), and similarly parti-
tion Regime 2 into sub-regimes 2-a aa < Al and 2-b ( E < 02
|
The proof is then structured into three distinct cases: Regime 1-a, Regimes 1-b and 2-a treated jointly, and Regimes 2-b and 3 treated jointly.
Regime 1-a. (0) 22llull2 a, so a < on implies mc(0) = ne(0)
2p T = For any
I 12 O. if (4) mR, then m²(() < # < =m((). Hence B(() < O, and therefore mi,(1) < 0. Consequently, for any t 2 0, whenever m²(t) < my, the function m.(-) is strictly decreasing. Since m.(0) < my, we have m.(() (t) < my for all 2 0, and it is strictly decreasing.
Moreover, we have (t) < P1 S My. Therefore.
(1) = 2Î(1)
2m O,
Thus 3,(1) = (t))) < 0, and (t( decays exponentially for all t 2 0.
Regimes 1-h and 2-a. We define A(t) ÅM (t) and B(t) me(t) mmn(t). Then we get the following equalities:
me AB,
B ne - mhH = AB vig.
Let ((t) == exp (-5o A(T)dr) Then:
IB IAB - Iinn
(15)
d
(IB) iB - +IB = -IBB IB = -Imi df
(16)
I(t)B(t)- I(O)B(())
((u))n(u)dd.
(17)
(0) 22 have B(0) < 0 and thus mc( A(0) ((0) < O. so m. initially drifts downward.
ag <a < P /2[ B implies my < mm)(0) ma(0). In this case,
For an initial condition ma E (mL, my)))). let me(: mo) denote the solution with mc(0) ma- We define the floor hitting time as T=())) = inf (t 2 0 m,(t; mo) m.. E [0, 20). and the ceiling hitting time as TC(m)) == inf (t 2 0: B(t; = 0} E (0, oo]. Define the first exit time as (C((m)) Tm(m)))
T(mo)= min (FF(mO),
Now let
F = {mo E (mL, www.mmh(00 (0)) (F())) < Ta(m))) == {mo E (m,, www (0)) :a(m)) < T (m())
We will prove:
and G are open in (m,, (0)).
# 0 and G + 0.
therefore there exists my E (m,, (FUg). and for such mū, we must have Tr(m)) Ta
47

---

Published as a conference paper at ICLR 2026
Since O, we have 3y(0) for all E [d], and from the explicit solution of the 33 dynamics in Lemma 4.3, we have that for each finite time t. (y(t) > 0. Therefore, Py (t) > 0 for all finite and j E [d], which implies that p(t) is not collapsed onto single feature at any finite time. Based on this observation, we make the following Lemma:
Lemma D.II. ())())(m)) < 0 for all mo & F,, and ((Ta(moo) > 0 for all mo E g. In other words, the floor crossing and ceiling crossing are transversal.
TF(ma) cc. Let by tF(mo), so Proof. We first prove for the floor crossing case. mo my. Then we have
( = A(tz)
Here, A(ly) Also p(tz) has positive mass on every Mij hence B(ty) = m - mm((r) 0.
Thus riy mo) < O, in particular (4g) o O. So the intersection with the floor is transversal.
Now prove for the ceiling crossing case. ra(m)) < oc. Let ig == Ta(ma), so B(gg) which means mt(t mtt))). Then we have
B(t,)) = A(t,) mmg(tt)) = (t))(t))
From Proposition D.10, we have my I mp, and the equality holds if and only if p is collapsed onto a single feature. But at any finite time (g. p((g) is not collapsed onto single feature, so mm((g) < mp. Therefore, we get me(lg) - mmp(yy) mm((g) mmp(gy) <
Using the formula
(2mel's - 12) 2 (M))² me
21'1 (mg mp),
2 (M1)²²
and noting T₁ >> O for non-collapsed we conclude
rin(()g) < 0.
Hence,
B(G)) = -in((())) >
So the intersection with the ceiling is transversal.
Next, make lemma for the openness of F and G.
Lemma D.12. The sets F and G are open subsets of (mL, (0)))
Proof. We use continuous dependence of ODE solutions on initial data on compact time intervals.
1. F open. Take any mn F and set 1 == (m(m)). Then ly < do and me (e((y; = m.. By Lemma D.11, the floor crossing is transversal: rh. mo) < 0.
Consider Fit, m) == me(t; mn)-m_ We have F(ty, mm) = 0 and 0 F(t mg) = melt * O. By the Implicit Function Theorem, there exists neighborhood U1 of mo and continuous map m + (m) such that for all m E U1.
me(Tr(m); m)) mi.
(TF(M) - 1 <8
for some
Next we show that the ceiling (gap) cannot be hit before TF (m) for m near mo. without appealing continuity of TG. Since my E F, we have Ta(m)) ty. hence
B(t; mo): < 0 for all t [0, ty].
Moreover, at by we have B(1y; m.. mmlll; m)) 0. By continuity in I, there exists 5u E (0, 6) such that B(t:ma) for all [0, f + 5o)-
48

---

Published as a conference paper at ICLR 2026
Since [0, ty + du] is compact and m-) is continuous and strictly negative it, define the margin
7)
max oo)
B(t: ma) 0,
so that (t; mo) S -1 for all [0, 5o].
By continuous dependence on initial data on [0, ty + do], there exists neighborhood U2 of mn such that for all m U2
sup
B m)) B(t: n/2.
+6)
Hence for all m E U2 and all E [0, 50).
B(t: m) - mo): + + - O. In particular, B(t m) cannot reach 0 before time ty + do, so
(m(m) > 1y + 60- Finally, take m U == U1 n U2 and shrink U1 if needed so that [Tr(m) - 114 < do- Then Tr(m) < y + 60 < <Ta(m), i.e. E F. Therefore F is open.
2. 9 open. Take any my E G and set tg == (mo). Then tg < 0o and B(tg:m)) = 0. By Lemma D.11, the ceiling crossing is transversal: (ty; mo) > O.
Let G(l, m) m() We have G(lg; mo) = and 0,G( ma) = B((;; mo) * 0. By the Implicit Function Theorem, there exists a neighborhood V1 of mg and continuous map m + (mm)) such that for all E V1.
B(Ta(m);:) = 0,
(Ta(m) - 1p <
for some >
Next we show that the floor cannot be hit before Ta(m) for m near mo- Since mo G, we have (y(mo) > fg. hence
melt: mo) > m. for all t E fgl Because (m) is continuous and [0, fgl is compact, define
== min (mo mo) mL) > 0, [g)
E: ::=E0/2.
Then me(4;ma) 2 m. + 2e for all E [0, [g) By continuity in at lg. there exists dg (0,8) such that
melt: mo) 2 mL + for all t E [0, 6o].
By continuous dependence on initial data on [0, fg + 5u]. there exists neighborhood V/2 of mo such that for all m V2.
sup [me m)) me(t: mo) < </2
r rE(0,tt +50
Hence (t: m) 2 my + + on [0, to + 6ol. so in particular TF (m) > ty + do for all m E Vz.
Finally, take in V Vi n V2 and shrink Vi if needed that (Ta(m) tgl 8n- Then Ta(m) < by + 50 < (r(m), i.e. m E G. Therefore g is open.
Next, we make lemma for the no simultaneous finite exit. Lemma D.13 (No simultancous finite exit). There is no mo (mL, such that
TF(mo) = Ta(mo) < co.
Proof. Suppose for contradiction that there exists mg (mi, mm)) and finite time TF(m)) = (mo) < co. Then by definition.
(t) = m. and B(1"; =
Since B())= me(() ml))), the second equality implies mn(t"; m)) me(!"; mi.
49

---

Published as a conference paper at ICLR 2026
However, mz(t) mL 21/2 holds if and only if p(t) collapses onto the first feature (i.e., (t( = 1). Under the uniform positive initialization, we have 3,(0) > for all j. and by the explicit B-dynamics (Lemma 4.3) we have (y(t) > 0 for every finite t. Hence pz(t) > for all at every finite time, so p(t) cannot be collapsed at time '" < co. This contradicts mm(("") = m..
Therefore, such my exists.
Next, we make the non-emptiness of F and G.
Lemma D.14. The set F is non-empty.
Proof. We show that there exists at least one initial condition mo (mL, mp(0)) that hits the floor before the gap closes.
Consider initial condition mo arbitrarily close to my (with mo > mL). At time t = O, we have B(0) = my mn((). Since (0) > my. the initial gap ((0) is strictly negative and bounded away from 0 as m2 + m.. Because A(0) = M (0) > O, we have m = ((0) < O, meaning mt(t) strictly decreases initially. Because the velocity mt((t) continuous and bounded away from zero in a small neighborhood of t = 0, the time needed for the trajectory to travel the distance to the floor scales with the initial distance my my. Thus, we can make the floor hitting time (m)) arbitrarily small by choosing mo sufficiently close to my.
On the other hand, the gap (()) is a continuous function time and starts from a strictly negative value (00) R m.. m (0) < O. By continuity, it cannot reach 0 instantly. Therefore, there exists some uniform constant Eg > such that for all my sufficiently close to mig we have B() < 0 for all e Eo] This implies that the ceiling hitting time satisfies TG 2 Eo for all such mo-
By choosing mo > my close enough to my, we can guarantee that (m)) < so I Ta(mo). This strict incquality implies TF(mo) < Ta (mo), meaning mo F. Hence, F + 0.
Lemma D.15. The set G is non-empty.
Proof. We show there exists at least one initial condition that closes the gap B(t + 0 before reaching the floor. Consider initial condition my mm(()) - 5 with 5 > 0 small. Thus, the initial gap is (0)) = -8, which is strietly negative but arbitrarily close to O.
First, we establish that B(0) is strictly positive and bounded below for small 8. From the corre- sponding ODE, t = 0:
B(0) A(0) rinh (0) A (0).
We evaluate the (0) term Under the uniform positive initialization, 33(0) > 0 for all j. meaning the initial distribution p00) has full support across all features and is not collapsed. Because mh mi if and only if p collapses onto single feature, we have a strict inequality at t = 0:
mp(0) > mn()).
Using the formula for rhnh evaluated at O:
i 2r (0) (mo - (0)).
(0) = (0)² mo
For mo E (0) 60, (0)). we have mo mH mm(0) < mD), which implies mo m0) < (0) 0, mo > my > 0, and (0) > (by Cauchy-Schwart on non-collapsed distribution p). Therefore, there exists a constant co > 0, independent of sufficiently small 8, such that:
m0( (0) 0. Furthermore, i(0)
(0) 2 co > 0.
For 8 chosen small enough, the magnitude of the -A(0)6 term becomes strictly less than /2/ Consequently, we obtain (0))= > co/2> 0.
Next, we bound the time required to close the gap versus the time required to hit the floor. Because B(0)) cn/2 > and the vector field is continuous, there exists short, finite time interval [o, T]
50

---

Published as a conference paper at ICLR 2026
over which Btt) 2 co/4. Starting from B(0) -6, the trajectory must hit the ceiling (i.e., B(tg) = 0) at some time tg bounded by:
46
co/4
Co
We now control how much the center me can drop during this interval [0, tyl- Since r) B(t), and knowing that B(t) [- 0] on this interval while A(t) is bounded by some constant Amax > 0, the total drop is bounded by:
m²(t)) mo ! (o(() mo - / Amax mo
46 Co
Amaxxtt 2 mo
Amax
This reveals that me decreases by at most (6)²) the time it takes the gap to close. Since the initial position me mn(0) 8 is separated from the floor my by (1)) distance (specifically, mn(0) - my > 0), we can choose sufficiently small such that:
me(ty) 2 mo ( mL.
Thus, the trajectory hits the ceiling at time ig before it could possibly reach the floor. This implies Ta(ma) = lg < Tp (ma), meaning mo E G. Hence, 9 & 0.
We have now established the following topological properties:
and G are open subsets of (mi,, mp (0)) (Lemma D.12).
F * 0 (Lemma D.14) and G * @ (Lemma D.15).
Fng by their definitions (a trajectory cannot hit the floor first and the ceiling first simultaneouslyy.
The interval (mL, (0)) is connected topological space.
connected space cannot written as the union of two disjoint, non-empty open sets. Therefore, must have (mL, (0)) (m() (FUg) + D.
Let us pick an initial condition my E (mL, m0) \ (FUg). By the definition of F and G, this means that the trajectory does not hit the floor first in finite time, does it hit the gap (ceiling) first in finite time. It implies that:
(m(m)) = Ta(m)) = cc.
Thus, the trajectory starting at mž never reaches my and never reaches B 0 in finite time. Conse- quently, satisfies:
my < me mb) < mp(t:mp) Vt 2 0.
This boundary trajectory exactly corresponds to the case where the trajectory reaches neither bound- ary, completing the proof of existence at the mi level.
Existence of the Regime 2 Threshold ai. Finally, we map the boundary initial condition mi back into the initialization scale parameter Under the uniform positive initialization, the initial scale determines no /211,u|120, which in turn dictates the initial center m0) via:
m. m.(0) (0) = (0
V21
a.
2p
2p
Because m0(0) is continuous and strictly increasing linear function of a, the existence of mi E (m,, (0)) guarantees the existence corresponding initialization scale an such that (0; m This on is precisely the boundary threshold defining Regime 2.
51

---

Published as a conference paper at ICLR 2026
Uniqueness of the Regime 2 Threshold a1. We now show that within the interval (ag, 02), there can be at most one initialization scale or for which the corresponding trajectory is boundary trajectory, i.c. it never reaches either the floor my or the ceiling mt((t) in finite time.
Suppose for contradiction that there exist two distinct boundary initializations a, < ag in (ag. 42). Let (m² (t), p()) and (ml (t), (t)) denote their trajectories. Since both are boundary tra- jectories, they satisfy
mL <m (1) < mn1 (0) (1) and m.. (1)
mii 11 (1)
1A 2 0.
(18)
In particular, (1) = m (1) mm (((()) (t) 0 for all t. Since the A(t) B(t)) with A(t)
(1))M(() > O, we have rig (0) (1) < and me ( (1) < 0, so both me (1) and me (1) are strictly
decreasing and bounded below by my; hence they converge.
Next, along the boundary corridor we have mt(t) < ma(() & mp(t) (using Proposition D.10), and therefore
A(t) nialtt 2(M (f))² (t) (2mo (((())
((t)r [2())) (t)² me(() (mc(t) mp(()) < 0,
(a) (t) and (b) (t) are also strictly decreasing and bounded below by my, hence both converge
so mi as well.
Since me (() converges while ( = (1) ()) < 0, we must have in! (n) 0, hence
(t) + 0 along each boundary trajectory:
m mH (1) - 0.
(19)
Likewise, since mH (t) converges and vhh
O (t) < 0, we must have (() - 0. Because X(t) > 0,
M1(t) 0, mt) > 0, and (t) 2 0, this forces
(t) mD (-) (t) 0.
(20)
Combining equation 19 and equation 20 yields mH²(1 ( + 0. By the equality case of Proposition D.10, this implies that (t) collapses to a single feature as t - co.
We now rule out collapse any major coordinate k: 1. If mo(t) + px/2 for some k > 1, then the minor feature j = 1 has asymptotic growth rate
(t) = 2 (1
H1 2mm(()) 2p1
M1 PAk
> 0,
so any soollapse-to-majjeee state unstable (the minor coordinate gets re-amplified). Therefore, the only possible boundary limit is the minor-feature collapse equilibrium
E
p(t)
+ e1. m²(t) + my 2
In particular, both boundary trajectories (for an and (s)) must converge to Es.
Now the special structure of the uniform initialization. Under uniform positive initialization, the initial proportions p(0) depend only on (and not on a), so the two trajectories share identical p(0) and differ only their initial m.()):
ma(0; ors) > mo a.).
We examine the stability of E, in the m.ddirection. At E₁ we have p = e1, hence M₁ = Ab and my my. Thus the -dddnamics
m = M (ma-mp)
- mn)
has strictly positive linearization along the m. axis:
0m 0m,
= -((Ei) (E)) -(Ei) > O.
52

---

Published as a conference paper at ICLR 2026
Hence Ej is repelling along the me-direction.
Finally, by uniqueness of solutions to smooth ODEs, trajectories cannot cross in phase space. Since the two trajectories start with the same initial proportions p(0) but different initial height me (0), and since the limiting equilibrium Ej repels along the -ddirrection, it impossible for both trajectories to remain on the boundary corridor equation 18 and still converge to E1.
Therefore, there is at most one boundary initialization scale 01 in (av, a2). Combined with the existence shown above, the threshold d;j is unique.
Since we checked the existence and uniqueness of the boundary threshold aj in Regime 1-b and 2-a, we can now conclude that for < 01, the dynamics goes to floor-first behavior(Regime and for a > 01, the dynamics goes to ceiling-first behavior(Regimme 2).
Regimes 2-b and 3. (0) = al
P / implies mc(0) = no > mn (0). 2p
M
When (0) > mO),), we have (0) > 0 and thus m (0) A(O)B(0) > 0, so me initially
increases. We now show that (t)) > 0 for all t 2 0. Suppose for contradiction that there exists first time r > 0 such that (T)) = 0 (i.e., (T) = mn(T)). Then
B(r) = mi(r)) - m
A(T) min(T)
= min(T)
X(T)
2 (M2(T)) me(T) (2me(T) (T) [2(T)) -T(())))))
Proposition D.10 gives mp(r) 2 mH(T). Therefore, 2mc(r)
l'2(7) I 0, ensuring that B(r) 0. Yet, if 7 represents the first instance where B reaches zero from above, must have
B(r) 0. This contradiction establishes that B(t) for all 2 0, which implies me(t) > m=(()
for Z O. Finally, since A(t) = AMi (1) > and (()) > 0 for 2 0, follows that
(t)(t) = A(t) B(t) > 0
for all t 2 0, proving that mf(t) is strictly increasing. Additionally, रे ag implies m²(0) ? pa-1tHa ensuring that sub-regime 2-b exhibits the dynamics claimed for Regime 2, and that Regime 3 follows its respective theorem statement.
D.5 EXTENSION TO DEEPER DIAGONAL LINEAR NETWORKS
In this section, we extend our analysis to L-layer diagonal linear networks. As the depth increases (L 2), some notational adjustments are necessary.
Recall that the margin is given by
(B,
(1) w 1 O... (L))) 14)
where denotes elementwise (Hadamard) product.
The gradient of the loss L with respect to a particular weight w/ can be computed via the chain rule:
dC dc ds (1
dw (1) Ap
du
ds
kgl
where À is as before, and o indicates multiplication over layers except The squared Euclidean norm of the gradient vector VoL is then
| =
dc
=
w
(k)
du
1=1
1=1
1=1
53

---

Published as a conference paper at ICLR 2026
Accordingly, define
wt (k)
j=1
1m1
The resulting perturbation is:
VoL
62 P
VoLLl2
(e2) (2)
PH4
(1)
no
kyl
Thus, the perturbed weights are given by
(1)
pj
w
The perturbed product then takes the form
B =
Therefore, the ODE for each coordinate is:
aC(0) Ayz (K) Ouly (1)
di
Additionally, we define an assumption on the weight initialization scheme:
Assumption D.16. The weights are initialized symmetrically at = 0, that is, w (0) = w/ (2) (0)
Wjj (L) (0) = (0) for all
Now we show the balancednness-pessrvinng property of the SAM flow.
Lemma D.17. Suppose Assumption 1.16 holds. Then for all t 2 0,
w = w1(() w1) for every 1, j.
Furthermore, the sign of wtyt) is preserved for all 2 0.
Proof. Fix j. Assume that at some time / all weights corresponding to j across the layers are equal,
i.e.,
w ( = m (3 (1) = ... w (1) = w.(().
Then (t) simplifies as follows:
d
(1) = Σ Σ
2
II 1o
j=1 1=1
= ΣΣ
33 (w)(t)(-1)²
5=1
= Σ
Lu3 (w)()))2-2
5=1
54

---

Published as a conference paper at ICLR 2026
Therefore, the perturbed weight for each layer I simplifies to:
1 (t) = w (t)
!1)
PHj
(k)
(t kyd
w,(t) PPj
no (1 wj (t)
which is independent of I. Hence,
10 i ) (t)= = -.. =0 20;
Substituting this into the SAM flow equation yields:
i (t) = (0)))
which is likewise independent of I. Now, for fixed j. consider the L-dimensional vector
uy(t) (1) m (0).....
The SAM dynamics specify the ODE:
ún = F, (t),tt)((t))
where F3 is the vector whose 1-th entry (()) I vi (0) (1). This ODE is locally Lipschitz in Uj, ensuring uniqueness of solutions for given initial conditions.
Consider the one-dimensional diagonal manifold
D, =
x) E R : IR.
if (t) E Dj, then (t) E Dj as well, because all coordinates have the same derivative. So Dj is invariant under the flow.
Since the initial condition u, (0) lies in D, due to symmetric initialization, and the ODE solution is unique, we conclude that uy(() E Dj for all 2 0. Therefore,
w 1 (t) (t) for all I, j, and t 0.
In summary, Assumption D.16 guarantees balancedness at all times for any depth L.
Next, we consider the sign preservation property.
Recall that on the balanced manifold, we may write w (1) (1) for all I, j. and / 2 0, so the per-coordinate dynamics reduce
wy (t)2-1 1-1
(t) = )(t)as ( -
Hj (t(t)
We claim that the sign of wyy( (t) is preserved for all 2 0. To see this, observe that the right-hand side of the ODE is a smooth (in fact, polynomial) function of wy, so it locally Lipschitz in wy for each fixed t. In particular, if at some time T we have w = O, then úi (T) = 0. so wy(t) = 0 for all t is a solution with the same initial value. By uniqueness of solutions ODEs with Lipschitz right-hand side, it follows that once w, reaches zero, it remains identically zero for all future time and cannot cross to the opposite sign. Therefore, if ()) o O, the sign of w,,( (t) preserved for all I 2 0 by continuity; if (0) 0, it remains zero.
In summary, the sign of wy (1) cannot change during the flow.
Utilizing the balancedness prestervvng property, we can now extend the lemma for the depth- L diag- onal network.
55

---

Published as a conference paper at ICLR 2026
Lemma D.18. Under Assumption D.16 and Assumption D.4, the rescaled Eg SAM SAM satisfies, each coordinate j.
di 31(0) (())))-))))()),,(0))
where
L,,5, (t) (1
no(t) 85(0)(4-2)/2) (4-1)
MPj
and
3,(1) w0)(1) no( Σ (1)
Proof. Now define the effective coefficient per coordinate, for general depth L:
(,(1) = II (1
w 3 ())=
Under the balanced ég SAM flow, the coordinate dynamics become:
8,(t) dt (w,(t)²) = (to,(t)(L Lug (1) òp.roō
We first compute the perturbed weight for coordinate j:
PHi E-1
wj (1
(1- MHg m 1/-2 n2
ú w0
w
no
Substituting this into the expression for (1) gives:
MH4 1-22 1-1
3,(t) = L)(t) w (1-
nt w
To express this in terms of 3, w/. note that
wpt-2 137-211 2/1.
wh-2 = g.1-21/1
Therefore, obtain:
8,(1) Li(1), 8,(1)22/12 2,)
(1
(,(1)(12)144 L-1 (1())
PHj
Absorbing X(t) into the time parameter yields the rescaled SAM flow equation:
d dt
Bt(t) = (t)
where
100 (1) = Lp, (1) 2/1 (1-
PHj (t) L-1 (t(
This provides the Depth- generalization of the SAM feature amplification dynamics.
56

---

Published as a conference paper at ICLR 2026
Proposition D.19. Consider the depth-L diagonal network under Assumption D.16 and Assump- tion D.4. Define
B,(() - II w))))),(()( 2x(t) == (1w
(t()) =- L
Σ 2 wx (t)
1=1
1=
and the critical effective scale:
zolt) = no(t)
PL
Then for each time we have
(1) L2,(t) (1 P (0)))(1))
=: (z)(1)).
The function -> de(2) is strictly increasing (0, zc))), strictly decreasing on and possesses a unique interior maximum at = = zo().
In particular, any fixed the coordinate(s) whose effective scale zz(() is closest the peak of dr. ie., near ()), experience the largest instantaneous relative growth rate.
Proof. In rescaled SAM time, we have
PH4 (,(1((22)122 1-1 no(t)
(1) = Lug 3,(1)2-2/12 (1-
where
d (1)() = L Σ
3 (t) 26-2
k=1
Define the effective z-scale by
2z(t) = M1 wy (t)()(-2
Note that
M1 15 3/1- M1 w 2z.
Plugging this yields
L-1
(1)=
where <b(2) Laz (1=(1-10) no( (1=
Define the critical effective scale:
2c(t)= (t)
PL
Consider dr(z) = Lz (1 b-1 where == (!) > O. Its derivative with respect to is:
d()) L(1 ez) 1-2 cz)--2 (1 Lez),
dz
so that:
d((z) 0 for 0 < zc(1),
d((2) when (t),
d()) < for (t) < 2 < (ttp//p.
Therefore, for each fixed t, the function (2(2) is strictly increasing (0, (t)),, strictly decreasing on (ze()),no()))))) and has a unique interior maximum at = =a()).
57

---

Published as a conference paper at ICLR 2026
Unlike the 2epth-2 case, where each Hj is fixed constant and their order remains unchanged throughout training, in the depth-L case the effective quantities 2y(t) are time-dependent and could, in principle, change order the SAM flow evolves. However, the following proposition establishes that the order of 2y(t) is actually preserved throughout the entire SAM trajectory.
Proposition D.20. Under Assumptions D.16 and D.4, the order of the 2,(t) is preserved in the pa then 2t 22(t) ... za(t) for all t 2 0.
depth-L L SAM flow. That is, if f1
Proof. We first compute the ODE satisfied by (t). By definition,
2j = u 2-
Taking the time derivative, we get
2j = U4 (L. 2) u (1/(-33/11
= 14 (I. 2) n²g (3- ( (2p1)
Therefore, the perturbed weight is
w w/ (1 (1-
P14 wi
(L-2)
n1
Also, we get
PJi (L-2) (1-11)
(L (1111)
1 (2L-4)
1
7 wj 106l
= ]
Using w
(4-2) 3i and wj
w
(2L-1)
13 we obtain
22 2 &j (L- 2)
(I-1)
(L-1)
"j
1
PH4 2j
= -2)_)
1
p2j
15
no pÿ
no
Thus, the ODE for (t) can be expressed as
²y(t) = (4-) (11)))) (1- p2, (1)
4-1
=
no nt
Notice that in this expression, the dependence on j appears only through 2, (4); both Ā and no are time-dependent scalars shared across all coordinates. So each 2z()) solves the same scalar non- autonomous ODE,
a(t) f(t (()),
with =(t) (t).
> O, we have 2,(0) 1zaa-2 Since Now at = O. under symmetric positive init wy(
141 < jud and a'-2 > 0, we have 21 (0) < (0) < ... < zz( For this ODE with f is smooth and locally Lipschitz in z, the two different solutions (t) cannot cross each other. If two solutions ever meet (same values at some time), then uniqueness makes them to be identical for all
times. So the order of (t) is preserved for all t 2 0. Thus, we have 21 (t) < 22(t) < .... < za()) for all 2 0.
D.6 PROOFS FOR SECTION 4.2.4
D.6.1 DERIVATION OF THE DYNAMICS OF (())
The dynamics of ((t) = w(t) C (t) is given by
(t) ii()) wtt) + w(t) o it)).
58

---

Published as a conference paper at ICLR 2026
By Equation (3), it is given
1 w(t)
B()) 24 O w(t)
ww())
no (1)
2μ @ (8())
+ ((t)
n²(()
Coordinate-wise, we have the linear equation
B,(() = 2pd (,(()) - - m3
21, (1 (1-p
no (T) (1)))o(7))
Hgj
no(
Therefore, separating variables and integrating, we get
(t
1
2/1
2p no(t)
(t(
B,(s) ds = 1/¹ 224 2p no(s)
ds
By(s)
(2(()
2ppb Srd no(s)
log
3,(0) 2p,t
ds.
Define I(t): = J0 1 ds. Then, the solution is given by no(e)
(1(1) (0) exp (2pz - 2(12
for j E [d].
D.6.2 PROOF OF THEOREM 4.5
Before proving Theorem 4.5, we establish Theorem D.21, which provides lower and upper bounds for I(t) and serves as key ingredient the proof of Theorem 4.5 below.
Theorem D.21. Suppose w11) = (2( = a E Rdd. Let (w))) (1)) and (w (1))e2 follow the rescaled f2-SAM flow (2) reduced to (3) with perturbation radius p and data point μ.. Define Cpa
and C H.a v2d(I]) 14 Then,
1ull
(a) 1(t) tHE log Paua expl- (ocppp--()+-pppp -ppp when I(t)
p(u)+u2]
(b) T(t)= S P pC expl 1ell +)+ -pCr
log
Proof. From the definition of I(t), I(f) fo no(s)
1 ds, we have I'()) ma(1))
Since we suppose ()) = w (0). and the loss function and dynamics are invariant under exchanging (!) and (2), we have ( (t) w (2) w(t) for all t 2 0.
From the definition no (t), have
no (() I (1)) 1
222 wt)lll
2
Σ
3 (t)²
59

---

Published as a conference paper at ICLR 2026
Σ (3,
2
From Equation (4), which is 8y(t) 3,( exp (2 2p12]((t), we have
Σ (0) exp (2 -213 )
nt(
2
and therefore,
(t())
/3 E (0) exp (2pzt
(a) When [(t) 2 p(H1+ p(mi+ for ...... d. holds that
1
(2pt 2p T(t)) (2pzz - 2pp 2(1 mn)(t p(us + (z) 2 0.
Therefore,
('()) =
V 2 1 exp (2pzt - (t)
1 (0) exp (2uu - 2pp
/ exp (11t - pp))((t)
Separating variables and integrating, we get
1
exp(=Pr((())dT S
exp( wip(-
1 13,3(00)
7(0) exp(- (1-pppuudd S
6-'
exp(-H13)) (11ssss
(()
/2Σ2 (s)33,
(cxp( (pp)t(t))) - (pp(-pppI(())) s
(cxp(-Mit) cxp(-H10)
pHd
V (3))3) 11
1 (exp(- (M1t) - 1)
expp(-prG(0) - 1
(0) pMd
/ (3) M1
=cxp(-pp(I(t)) p
11
(oxp(-Mit) 1) +1
2 dd My ())))38,(0))
# - -tp 2 log (pC (exp( (mit) 1) 1)
1
(t)) 2
log PCu,a exp(-xppp +11-Cppp pCp.a)
where (a) holds since 1(0) = 0 from the definition of [tt).
(b) By AM-GM inequality, we have
r'(1
2 d (0);(0) exp (21z1 ((1)
60

---

Published as a conference paper at ICLR 2026
exp (2 2p13 1/d
2d
(II, 1333(00)
1/a
2E / Mit
2p Σク 13 T(t)
2d
cxp
2d
[1 11503)
1/d
exp
(1l 20llall 2 T T(t))
1
V2d
) 1/d [T k5as) (m))))
a (Ilalla
ollxa12 I(t) ((t)
exp
Separating variables and integrating, we get
|leelli dt
exp(-
T()))dT
1// exp
V2d [I] [[]==11100 [1s=))aa)) (1jas))
T(0
1 uddu /-8
1
1/d exp (- ||aell2 (-lleela,
exp(-
ds
((0)
d
22 (Tz= kjas)
d
(exp( p||aa113 ((1) exp(
pllpa
1
1
1/d || (exp( |1el 1) - 1)
d
p
p
2
d
2d [1] M503)
pl I(t)) 2 p
14
[ t) 1) +
exp(
(exp(
d
V2d (IT==1 13jas
1/d ||Aa|/2
d
| d
pT (exa(exp(( |1elll1 t) 1) +
((t) log
d
+1(t) <
log
pC | pCu,a p1,00
Theorem 4.5. Let ai,, a2 be defined in Theorem 4.4 and an be the threshold from there. Suppose tua
01 <a S p
/21 a2. Then, for j E [d], there exists T, such that
9,17,) LB,(a) = exp R2R) ((R, 1) log (T- BaCl
+ log ma/o -())))())))
where Rg == (m)+ > 2, R2 == (Ha-P))/uu and C(R) = R Rlog R - (R 1) log( - 1).
Proof. By the assumption an < a, we have Cpa,
1
We also have
pa
p.
C Cpa=
11
2
#1
#1
#1
1
for all j e [d].
V211u|12a
V21lu|12pa11"
p(p1 + Haa
p(as + pa)
p RR,
1 -Ppuu
1 - Rg-11 11
1
for all j E [d]
pCua
Let T,
1
log
(Ry - 1)) 2 0.
#1
From Theorem D.21, we have
1
I(T))= 2 pnd
log
pC
+ 1 pC
61

---

Published as a conference paper at ICLR 2026
log
PC PCu,a u.a exp (log
pl (1, (by )
L-penn
+ - PC v.a
pC
log
pnd
1-pCua +1 - -1-ppups pCura AT-1
log
ppf
(1 PC suaal (1+ Ri--T
1 pji
1
log
(1 ppp
Rj
Rg-1)
log
2
PCur,
Recall from Equation (4) that
B,(T,) (, exp (2n₁T} 203z for E [d].
Thus, for j [d]. we have
(T(T,)
exp (-2(pa - m5)TT + 1 (13)
Ba(T))
exp
2 Ud Mi
log
ph
(Ry-1) 1) + 2p(12 - 13))(TT))
#1
PC una
PC pCpa p.a
143 - M3 log
1 1 n2
exp
2
Md
MI
Pi log
1-pCpua
(Rl - 11
2
1 1-
exp
Md Hj
M1
Nd + Hi
MI
log
1
pC
Rg
1 log
1- 1-pCpp
PC no
pC
(R,
1 -
1
PCpa
cxp
2R
Rí log
R
log
(Ry -1)
1 PC p
PC pa
R,-1
pao
R2
exp
2R'
Ry log
log
(Ry 1)
poka
1
pin
a
a
exp
2Ry
(R3 1) log 1) - R2 log (R1) - (R2 - 1) log (1
par)
log
po0g
exp
2R
(-C(R)) - -RR (1)-1)log pao
-log
por
a
exp
22R
(R, - 1) (1-1oo/a) log
C(R,)
poo/a
D.6.3 PROOF OF PROPOSITION 4.6
Proposition 4.6. Under the conditions of Theorem 4.5, define ja arg argmax (By and
set of == ag. Then, there exist thresholds of of < < ar S p /2linl]a such that (a) forr (0]-1:07}}
pi+ps for some m S d
62

---

Published as a conference paper at ICLR 2026
Proof. For (ov, P - ) let au/a E (0, 1) and write </lual2
G5(z) = log (B) = 2R, nn, (x),
where
(R(Z) = (R (o) 1- + log 1
1
C(R),
C(R) R Rog R - (R 1) log( R - 1),
and R, = (Hg + Ma)/M1 H1 > 1, R2j = ppg py)/M1 2 O.
(1) Shape of 2R,- We have
R,
R2 (1 x)
OR (x) = zx1 xx)
1 (2((2) (z) =
> 0.
x²
Thus 2R, is strictly convex on (0, 1) and attains its unique minimum at 1/Rj, where Rn (1/R,) = O. Consequently RR (a) 2 for all x and it is strictly increasing on [1/R5, 1). (2) Crossing between adjacent indices. For any j E (1.....d d 1} define
Hz+15(0)) = - G,(x) 2 $ n,+ (x) 2 -)(z)=
(x))
Because R3+1
Rj. have F12,, (1/R++1)
0 and Rn (1/R (1/Rj+1)
0, hence
H5+15 (1/R,+1) <
0. Likewise RR (1/Rj) = 0 and &R. (1/Rg) >
0, giving
H5+1. (1/R,)
> By continuity, Hj+1J has at least one zero x E (1/Rj41, 1/1]l.
To show uniqueness, using the expression for In we obtain
2 x(1 x)) (1 R R2R5)x (R/+1 R2)).
Hz)15(0)=
Since
(Ma px) + pa)
44 1
R2 R2
2
we obtain R/41 R2+1 - R2 = -
ri < O. Its zero occurs at
++ R
11
Ve R¹ Rit1 R2
15+1 + M4
and therefore
(x) for x < Dee
(x) for > Izc-
Hence Hz+ig(x) is strictly increasing up He and strictly decreasing afterward. Since 1/R, uI/(u, + paa) < p1/(p++1 + pz). H,+15 is strictly increasing in the interval (1/R,+1, 1/R,]. 0 and (1
0. this implies that HJ41 crosses zero exactly once in (1/ (1/R441, 1/ R.). Consequently the root æ is unique, with Hz+15( < 0 for x < 2z and Hイ+15(I) > 0 for x > z3.
Because
RR++1) <
(3) Thresholds and staircase structure. As a increases, I = at/ a decreases. Define aī. = a0/ 1 When a crosses a7, the maximizer between indices j and j switches once from to j 1. Because the intervals (1/Rf+1. 1/Ry] are disjoint and ordered, the thresholds satisfy aj aj <
< a S p( + (a)/(v/2lllll for some m d 1.
Thus (a) takes constant values on each interval (aj-1.a3) increasing step by step until the last threshold within the admissible range.
D.6.4 PROOF OF PROPOSITION 4.7
Proposition 4.7. Consider ogg defined in Theorem 4.4. (i) If < on. then B(1) converges to zero. (ii) If a > V2a n:) julli' then B(t) converge in l max-margin direction.
||
63

---

Published as a conference paper at ICLR 2026
Proof. We use Theorem D.21 to prove the theorem.
When w
(() al, have
Cn,a
M1
Mi
11
2 j=1 13
2 " j= 123a
V2llu|12a
|
|
Cp.a= = v2 j=1 (3)) |||||
V2 5= )
eCu.
(i) By the assumption a S oo. have C,,,a
no pa
p. Let T = log
2
From Theorem D.21, we have
1
1 pC pC exp(-xit) + PC pCp.al)
((t) 2 PUi
log
As T, we have
pnua exp(-M11) + +11--uua
PCu.a
+C exp(-1111 + 111ppu
pCu,a
PC CCuia -11
=pCua exp(log
) ++1
pCu.
pCpia
(CCna
-1
=Puaa
+
+ =0
PCuia
Since pCp,o exp(-11() pCu,a strictly decreasing in I, have
PC n.a exp(-Mit) + 1 -pppppp n.a - +0+ as f + T.
Therefore, ((t) too as t - T. Recall from Equation (4) that
(t) (0) exp (2 2pp for E [d].
As + T, we have (t) + 0 for all j [d] since [(t) + too. Therefore, B(t) - 0 as t + T.
(ii) By the assumption a > V2d Ma b we have Tuaa < b.
From Theorem D.21, have
d log
I(t)
pllpel (C exp(- t) + pmna))
For f E oo), have
| t) p 1.
0<1 p n.a S pC p exp(=
d
and as 1 1 0o, we have
||el|l, -lell1, 1)
p expeepppp-p
pC
PC u1a > 0.
d
As f. 00, have
d log
d
I(t) < p1 Therefore, T(() 00 xc.
log
< <000
pCp.a n.ax exp(-
[e|| t) + 1 pC 11.0
p |4 2
pCp,ao
d
64

---

Published as a conference paper at ICLR 2026
Recall from Equation (4) that
5,(t) (0) exp (2μtt - (t) for j E [d.
Thus, for j [d]. have
(t) exp (-2( +1)6+++22222 2p(u3 (1))
Ba(()
As oc, we have 8,10 Ba(0)
-> 0 for all j < d since 1im 1(1) < xc. Therefore, ((() converges to
the direction of ea t +
65

---

Published as a conference paper at ICLR 2026
D.7 NUMERICAL EVALUATION OF THEOREM 4.5
In this section, we provide numerical illustrations of the lower bound LB,(a) derived in Theo- rem 4.5, For several choices of μ., compute the value of
(By(a) == exp (2R' (RR, - log (T0% + log mo/o - ( c(()))
and visualize how much the ratio 3,(1)/3a() must amplified at minimum.
Figure 14 shows that for small in Regime 2 and for la with large spectral gap Pa/ /1, LB,(a) easily exceeds 10. Since this only a lower bound, actual amplification can even larger, indicating that minor-to-inee tontrrtt-ittermediate coordinates can grow by substantially more than the major coordinate.
2,75
2.50
2,25
LBjiA)
1.75
1,50
1.00
(a) /z (1,5,6,7,8)
(b) 1 (1, 4,
10000
20
8000
L8
6000
4000
5
(c) 4 71,3,5,7,9)
(d) μ = (1,2,,,,,,
Figure 14: Numerical evaluation of LB,(a) for various choices of μ.
For reproducibility, describe the numerical procedure used to generate Figure 14. For each choice of μ (with d = dim(H))), we evaluate LBy(a) for all E [d] on a uniform grid of a values. Following the assumptions of Theorem 4.5, first obtain the threshold an specified in Theorem 4.4. We then set E [11, +ua_ using 400 grid points. The quantities ag. R,, Ry, and C(Rg) are computed /2]|
directly from their definitions in Theorems 4.4 and 4.5 using the given μ.. The index E [d] μa. Since the computation is closed-form, no randomness is involved and the plots are exactly reproducible.
corresponds the coordinate ordering μ4
D.8 EMPIRICAL VERIFICATION
Our analysis in Section 4.2 focuses on the one-point setting D,,. begin by verifying that the se- quential feature amplification occurs across multiple choices of in this one-point regime: both the continuous -iteoouus--ieeeeee rescaled flows and the discrete foc -SAM updates exhibit the same coordinate- wise progression, and the loss dynamics follow the theoretical prediction. We then turn multi-point datasets and show that the sequential feature amplification persists this more realistic setting un- der both the rescaled f-SA flow and discrete SAM updates, illustrated in Figure 11. Finally,
66

---

Published as a conference paper at ICLR 2026
confirm that this phenomenon is not limited to depth 2: the same foordinate-wisg progression arises in deeper diagonal networks (general depth L). Taken together, these results demonstrate that the sequential feature amplification is robust and widely recurring behavior: it appears consistently across different l4, across multiple multi-point datasets, across both continuous and discrete SAM dynamics, and across depths I. 2.
To clarify the heatmap visualizations (e.g., Figures 3a and 15 to 23), for each time and initialization scale a, we compute j arg miny (t) and color the grid point (4, a)) according this index. Grid regions where the predictor 33 becomes negligibly small are shown in gray, indicating convergence toward O. We use the threshold (()())l2 I 10 to define gray regions.
Following the visualization style of Figure 3a. we also partition the a-axis into the three regimes defined in Theorem 4.4: Regime 1 (small a), Regime 2 (intermediate a), and Regime (large a). These regime boundaries are indicated by horizontal black dashed lines in heatmap figures.
For reproducibility, we detail the exact initialization used all experiments. As mentioned Sec- tion 4.2, we adopt uniform initialization across coordinates and layers: ())) w (0) al for depth-2 setup and w (0)
w L) (0) a1 for depth-L.. To approximate continuous-time trajectories, we simulate the flow using an explicit Euler scheme with a small step size 10-4
For discrete updates, we use step size of 1] 0.01.
D.8.1 ONE-POINT CASE: CONTINUOUS vs. DISCRETE DYNAMICS
We first verify that sequential feature amplification appears robustly across multiple choices of μ in the one-point setting. To demonstrate that this phenomenon is not limited to the continuous f2-SAM flow, additionally evaluate discrete 22-S updates. Across all tested choices of μ4, the resulting heatmaps closely match the structure in Figure 3a, showing both time-wise and initialization-wise sequential feature amplification. To better visualize the evolution of (3t), we also provide the loss heatmaps over (a, t). In the discrete & f2MSAM case, Regime I often appears unstable and does not become fully gray. This occurs because the relatively large step size causes the trajectory to hover near the origin without collapsing exactly to O. As a result, the predictor norm stays above the gray threshold-ae is not colored gray-yet the loss remains large, revealing that the trajectory still effectively stuck in the vicinity of the origin.
For comparison, we first present the results of GF and discrete GD with μ = (4, 5, 6, 7, 7,8). The behavior is similar across different choices of Both GF and GD consistently recover the major feature, independent of the initialization scale a, and they do not exhibit sequential feature amplifi- cation.
0.8
0.8
3.5
0.6
0.2
0.8
0.9
3.0
(a) GF
0.8
1.2
0.6
Iopatec Teaa
0,0
0.4
0.6
02
0.2
8.8
3.0
0.2
0.6
0.8
0.2
5.4
0.6
1.0
Ummt
Ume
(b) Discrete GD (17 = 0.01)
Figure 15: Dominant index jt a, and logistic loss Dr with μ (4,5,6,7,8).
67

---

Published as a conference paper at ICLR 2026
1. f (4, 6, 7)
0.6
aig mane &. (t
Intialtzation scale
0.5
0.4
0.3
0.2
0.2
0.1
0.0
0.0
0.0
1.0
time
timet
(a) Rescaled 22-S flow
0.6
1 200 Somii & 00
Intialtzation scale
0.5
Intualcaaion
0.4
0.4
0.3
logos
0.2
0.2
0.2
0.1
0.0
0.0
0.0
1.0
0.2
0.6
time
time
(b) Discrete SAM updates 0.01)
Figure 16: Dominant index jt over a, and logistic loss on D with μ (4, 7,8) and p = 1.
2. /a = (1,2,3,4,5
0.8
5
0.6
Intratzstionn scale
0.5
0.6
Iogasiic less
Initialization
0.4
0.4
0.4
0.3
0.2
0.2
0.2
0.1
0.0
0.0
0.5
1.0
1.5
0.0
0.5
1.0
time
(a) Rescaled f2MS flow
0.8
0.8
0.6
inibalization scale
0.5
0.6
logidtic 60
Inbalication
0.4
0.5
1.0
1.0
(b) Discrcte SaMSAM updates (n = 0.01)
Figure 17: Dominant index ji over a, and logistic loss D,, with / 2,,,,,,455) and p
68

---

Published as a conference paper at ICLR 2026
3. μ (1, 7,9)
0.8
0.8
- mag mas (
intialization scale
0.6
0.6
inttalcattion
0.4
0.4
0.0
0.5
1.0
1.5
0.0
0.5
1.0
1.5
(a) Resealed 1A-SAM flow
0.8
0.8
In scale
0.6
0.6
intialccttin
0.4
0.4
0.5
1.0
1.5
0.5
1.0
1.5
time
(b) Discrete SAM updates (n = 0.01)
Figure 18: Dominant index ji over a, and logistic loss D with μa (1,3,5,7,9) 7,9) and p 1.
4. P (1,2,4,8,10)
0.6
0.8
airg mane 20
Inttalizatioo scale
0.8
0.5
0.4
0.2
0.2
0.0
0.5
1.0
1.5
2.0
0.0
0.5
1.0
1.5
2.0
time
(a) Rescaled f2-SAM flow
0.6
0.8
0.8
Infoaiizaiion scale
0.5
Intoalication
0.4
0.2
0.2
0.0
0.5
1.0
1.5
2.0
0.0
0.5
1.0
1.5
2.0
time
(b) Discrete lg SAM updates (n = 0.01)
Figure 19: Dominant index jt over a, and logistic loss on D, with μz (1,2,4,8, and p p= 1.
69

---

Published as a conference paper at ICLR 2026
D.8.2 MULTI-POINT CASE: PERSISTENCE OF ONE-POINT BEHAVIOR
To examine whether the sequential feature amplification identified in the one-point analysis persist in more realistic datasets, we construct random linearly separable binary data by sampling two Gaussian clusters centered + and - ye for various choices of μ. Specifically, draw
(+) =++ En: Vn +1,
x/- - En Yn = 1,
with En ~ (0 and use 2/2 samples per class (with μa = 2), visualization, we plot only the first two dimensions of the dataset in the left panels. The middle panels show the results of the rescaled f2-SAM flow on this dataset, and the right panels show the discrete f--AM updates. Across all choices of multi-point datasets, the same sequential feature amplification behavior observed in the one-point setting persists.
100, a 0.5). For
For comparison, present the results of GF and discrete GD with the multi-point dataset gener- ated with mean μ = (4, 5, 6, 7, 8). The behavior is similar across different choices of μu. As in the one- one-point setting, both GF and GD consistently recover the major feature, independent of the initialization scale a, and they do not exhibit sequential feature amplification.
Updates
class diss
0.8
0.8
intalization scals
10
intaliiation scale
10
0.6
0.6
2
0.4
-2
6-
-2.5
0.0
2.5
0.0
0.0
51
Figure 20: First two dimensions of Dr with = (4, 5, 6, 7, 7,8) and the dominant index jt over a, under GF and discrete GD updates.
70

---

Published as a conference paper at ICLR 2026
class -1 dass
Rescaled 4-SAM Flow
4-SAM Updates
0.8
0.8
intia Piatttttoon scalee
Intilization scales
0.6
0.6
0.4
-2
-0
0.0
0.0
1.0
-5.0
-2.5
0.0
2.5
5.0
(a) p (4,5,6,7,8)
+1 class
Rescaled 5-5AM Flow
Discrete &A-MM Updates
2
indel intiliztionn Scale-
Intalization scales
0.6
0.6
R
0.4
0.4
-1
0.2
0.0
0.5
1.0
1.5
0.0
0.5
1.0
1.5
time
(b) P (1, 4, 5)
class -1 dass
Rescaled 4-5AM Flow
Discrete 4:SAM Updates
0.8
Intelizalion scale-
max A,i
Intia Ization scate
R
0.4
0.4
-2-
0.2
0.2
0.0
0.0
0.5
1.0
1.5
0.0
0.5
1.0
1.5
time t
time
C
X1
(c) = 31,3,5,7,9)
3
Rescaled 1-SAM Flow
Discrete &-SAM Updates
1 dass
intiiiiation scale
0.8
0.8
0
0.6
Intba iratioon
0.6
200
-1
2
-3
0.0
0.0
0.0
1.0
2.0
0.0
1.0
time
time C
(d) μ = (1, 16)
Figure 21: First dimensions of Dr, and the dominant index jt a, under the rescaled fg- SAM flow and discrete fa-SAM updates.
71

---

Published as a conference paper at ICLR 2026
D.8.3 DEPTH-L CASE: PERSISTENCE OF DEPTH-2 DYNAMICSS
We confirm that the sequential feature amplification is limited depth L the same coordinate-wise progression arises in deeper diagonal networks (general depth L). Specifically, we observe GF and rescaled 22-SAM flow on the one-point dataset Du with μ (4, 5, 6. 7, 7) The behavior remains similar across different choices of μ4, multi-point datasets, and under discrete updates. While GF appears to exhibit Regime (being trapped near the origin), does not show the sequential feature amplification, even in the deeper models. However, the rescaled f2-SAM flow clearly demonstrates the sequential feature amplification for general depth L.. Even though Regime appears chaotic, Regime and 3 are distintcly observed. Thus, the sequential feature amplification robustly occurs not only at depth = 2 but also deeper models.
5
1.2
5
1.0
1.0
Intiiltzatin scale
0.8
0.8
0.6
3
0.6
0.4
2
0.4
-1
0.4
time
time
(a) Depth
(b) Depth 3
1.2
5
1.2
1.0
1.0
Inttialization Scale
4
0.8
0.6
0.2
0.2
0.0
0.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
time
time
(c) Depth
(d) Depth 5
Figure 22: Dominant index jt over a, under the GF on Dμ with μ (4,5,6,7,8).
72

---

Published as a conference paper at ICLR 2026
1.2
1.2
1.0
Initialization scale
0.8
Initializattoon
0.6
0.4
0.4
0.2
0.2
0.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
(a) Depth
(b) Depth 3
1.0
1.0
Initaaiiation scale
0.8
0.8
-3
0.6
0.6
0.4
2
0.4
-2
0.2
0.2
time
time
(c) Depth
(d) Depth 5
Figure 23:
Dominant index jt over t under the rescaled SAM flow on D with
(4,5,6,7, 6, 7,8) and p
73

---

Published as a conference paper at ICLR 2026
E EXPERIMENTS
E. LOSS DYNAMICS
For initialization scales in the intermediate regime (Regime 2 in Theorem 4.4), SAM first amplifies minor coordinates and only later focuses on the major ones. This also affects to the training loss curve. As shown in Figure 24, the loss curve of SAM is noticeably flatter than that of GD in the early phase of training.
this experiment, we train the diagonal linear network with full-batch SAM using radius p = 0.5, learning rate 0.05, and 10000 epochs. We fix the initialization scale to a = 0.06 as a representative intermediate value. The data vector is μ = (1, 3, 4, 6), and all other settings follow the default Stagonal-networt. configuration.
To make this precise, we track the dominant index arg max, ry()), where (t) denotes the growth rate of (t) In the early phase, this dominant index corresponds to minor features (coordinates with small μı)), while in the later phase it switches to major features (coordinates with larger μı)). minor features, the loss decreases slowly, leading shifts to major features, the loss drops much faster. In contrast, GD does not exhibit this minor-to- major feature focusing behavior, and its loss decreases more rapidly from the beginning, without such plateau.
When SAM is focusing
plateau; once SAM
Los Curve CC phand 06)
3000
8030
30030
Loss Curve with Deminon (idee 2Wthowd. (40)
07
586 3aunnk
60.0 MGD
J00
Figure 24: Training loss curves of GD (top) and SAM (bottom) on the 2-layer diagonal network in the intermediate initialization regime (ar 0.06). The colored areas correspond regimes where each feature is mostly amplified. Compared to GD, SAM exhibits an early plateau loss curve: in this phase, SAM primarily amplifies minor coordinates, leading slow loss decrease. Once SAM shifts its focus to major coordinates, the loss drops rapidly. GD does not display this minor-to-mapy feature focusing behavior, thereby showing more steadily decreasing loss without such a plateau.
E.2 SEQUENTIAL FEATURE AMPLIFICATION UNDER RANDOM INITIALIZATION
In the main analysis, focused symmetric and layer-wise balanced initialization to obtain clean theoretical characterization. Here, we examine whether the sequential feature amplification phenomenon persists under more general random initialization.
74

---

Published as a conference paper at ICLR 2026
We initialize the two layers independently
w (0), w (0)) ~(0,a"1)))
where the parameter a controls the initialization scale as the standard deviation of the Gaussian distribution.
Figure 25a shows the normalized coordinate trajectories 3, (t(t)/]]((t)] under random initialization 0.65, μ (1, 2, 3, 4, 5, 6), and 0 0.1. In this case, all coordinates except
(Seed 0) for the fourth are sequentially amplified, with activation progressing roughly from the second to the sixth coordinate. Correspondingly, Figure 25b shows that the layer-wise discrepancy (|
w (t) rapidly decays to zero, indicating fast balancing of the two layers.
A qualitatively similar but quantitatively different pattern is observed under a different random seed. In Figure 25c (Seed 1), the sequential amplification begins from the third coordinate and proceeds toward the sixth. Despite this seed-dependent variation in the detailed activation order, the over- all sequential feature amplification phenomenon persists. Moreover, Figure 25d confirms that balancedness property is again achieved rapidly in the early stage of training.
These empirical observations are theoretically supported by Lemma D.5, which shows that even when the layers start from imbalanced initializations, the dynamics drive them toward balanced regime exponentially fast. This explains why the simplified, balanced initialization assumed in the main analysis captures the essential behavior of the training dynamics beyond this restricted setting.
Trajectory of normalized coordinates
Norm of w w
1.0
[w( (0lttllz
0.5
0.20
BeteM
0.0
0.10
0.05
0.00
10
2
4
8
10
time
time
(a) Normalized trajectories (t)) ||2 (Seed 0)
(b) Evolution of |/ (1) (1) (x))(( (t)ll2 (Seed 0)
Trajectory of normalized coordinates
Norm of w(1) w(2)
1.0
0.25
[[w)) w02
0.20
0.5
0.15
(c(C/[PP()
0.10
[x(t\[]pt))
0.05
pmm
-0.5
0.00
2
5
2
5
time
time
(c) Normalized trajectories ()()) (Seed
(d) Evolution of || (1) w (t)l22 (Seed I)
Figure 25: Sequential feature amplification under random initialization in a two-layer diagonal net- work. Rows correspond to different random seeds (Seed 0 and Seed 1), and columns correspond to different plot types (left: normalized coordinate trajectories, right: balancedness).
75

---

Published as a conference paper at ICLR 2026
E.3 ALTERNATIVE 2-LAYER MODELS
To evaluate the generality of our theoretical predictions, we conduct experiments on alternative 2- layer models featuring different parameterizations and metries. all cases, the experimental settings and hyperparameters are chosen to closely match those used in our main theoretical simulations with the diagonal network.
E.3.1 LINEAR NETWORK
We fix small matrix dimension d 5. All inputs are d d matrices. We first draw single random "signal" matrix μ E Raxd with i.d. standard normal entries, and then compute its singular value decomposition (SVD)
# = Up diag(Sp) Vp From this SVD, we construct an orthonormal basis of rank-1 matrices
H4s = u.u," i= i=l....,d d,
where is the i-th column of U2 and i is the i-th row of VTT. These μ4 play the role of "feature directions", analogous the coordinates in the diagonal model.
We use the logistic loss, and the dataset follows the same format as in the diagonal model: we consider the two points (+,,-1) with opposite labels {+1, -13. The 2-layer linear network is
fx(X) (3,X)F ((1) X)F,
with learnable matrices W (((2 E IRddd and effective weight / (())((), Each layer is initially set to the identity matrix, and before training we rescale all layers by a scalar a, so that (0))))) (7(2) (0) al and hence 3(0) I²
For training, we use full-batch SAM with radius = 0.5, learning rate 0.05, and finite train- ing epochs of T 5000. We repeat the experiment over a range of initialization scales, (0.20, ... 0.703.
As our tracking metric, we monitor the normalized squared alignment
(t) = ($(t), pi)
i = 1...... d
[[8(t)
where (() denotes the effective weight at training iteration L.
The results are shown in Figure 26. As plotted in the figure, the dynamics of SAM and GD are qualitatively different. For SAM, when the initialization scale is smaller than 0.225, training does not converge to a solution with sufficiently small loss. Beyond this regime, the initialization scale increases, the dominant singular direction that maximizes the alignment (i.e., arg max; a, (T))) moves from d5 to d1, indicating that SAM sequentially aligns from the minor component to the major component as grows.
33
Figure 26: Alignment the effective weight 3(t) for GD (left) and SAM (right) across initialization scales.
76

---

Published as a conference paper at ICLR 2026
E.3.2 CONVOLUTIONAL NEURAL NETWORK
We consider 2-layer linear convolutional network trained on synthetic dataset built from single image matrix This experiment is designed to probe frequency- wise feature selection under SAM.
We fix an image size d 32 and construct single base image μ E 1xxxx as sum of cosine plane waves with radial frequencies:
K
IcosOs.1 + y sin Ok! + du.).
plx, 1)
wk
WTTk
The experiment uses K different frequency bands, where r& are larget bands, wa are band weights, and ŌR.I. Φx.1 are random orientations and phases for each band. We take re {3,9, 9, 11, 13, 15) and Wi (1.0, 2.0, 3.0, 4.0, 5.0} for all k. We set L 8 for all We then renormalize f to have unit euclidean norm, then shift it slightly to be strictly positive. Next, define the frequency bands by constructing radial masks Mj C {.. dd 1]2 in the fourier domain. Let fi denote the 2D FFT of μ.. The band energy of μ at band k is then given by
Σ (a(m))
MEMM
The bands are sorted by μ4a As we apply low weights low frequency bands when constructing μ,, in this setting, low frequency bands have smaller μa and treated as minor features, and high frequency bands have larger μa and treated as major features.
The utilized model is a depth-2 convolutional network without nonlincarities. For the first convo- lutional layer, we use 3 x 3 convolution with 32 output channels. stride I, and padding For the second convolutional layer, we use same size of kernel, channel size, stride, and padding.
We used realistic gaussian initialization for the weights the convolutional layers. The weights for each layer are independently initialized. Lastly, the final FC layer is linear layer. the input for fc layer is squeezed Id vector, and the output is single logit.
Logistic loss is used, and full-batch training is employed. We use learning rate of 0.03 and 0.1. We train for 6000 epochs.
Band-wise effective weights. To compare with the diagonal model, require band-wise de- composition of the effective weight 3(0) in input space. Since the network is linear, ((0) can be recovered from gradients. At a given parameter vector 6, we consider the empirical margin
(()) = E(ee)) [vfo(z)] = (Jo(s) - fo(-)))-
We compute the gradient of s(8) with respect the input and form "virtual gate" version of B in input space:
V v=s(0))z.n = (V
So.
Amap v) = ((a.) (V (o C x)..
which is proportional to (B(0) O () u. in our linear setting. In practice, this expectation is computed exactly by averaging over æ E (4, -p).
We then take the 2D FFT of Smap, denoted Pmsp. and define the band-wise effective weights by
(x)
Σ
3map (m)
me irc
For each training epoch record the vector
(0.)...., (0:)),
77

---

Published as a conference paper at ICLR 2026
and, in particular, the index of the dominant band
Kdom (1) arg max Bx (0c).
In our initialization-scale experiments, we repeat this procedure over a range of E [0.13, 0.20 and, for each a, track both the dominant band kdom at the end training. This provides a CNN analogue of the feature- behavior observed in the diagonal model, where coordinates replaced by frequency bands.
Figure 27 displays how the final dominant frequency band selected by the CNN varies with the initialization scale a. Consistent with expectations, when trained with SAM, the model emphasizes minor features (i.e., low frequency bands) for small a, and shifts its focus to major features (high frequency bands) as a increases. In contrast, under standard GD, the dominant frequency band remains unchanged regardless of the initialization scale.
CNN Dominant Band Tranditions
CHN Dominant Band Transitions
1240
11:5
03:0
11.00
1334
1300
013
0.4
0.55
0le
In Scalc/c)
Scale fol
Figure 27: Dominant band for GD (top) and SAM (bottom) across gaussian initialization with dif- ferent scales. Each point shows the dominant band (the band that model mostly focuses on) at the end of training; SAM systematically shifts from minor (low-frequency) to major(high--reeqeenyy bands as increases, whereas GD remains insensitive to
E.4 GRAD-CAM
As our theoretical analysis rigorously characterizes the dynamics of SAM in linear diagonal net- works, we extend our empirical investigation to convolutional neural networks (CNNs) to examine whether the same phenomena persist in more realistic architectures. Combining the results for both -SAM and &2-SAM, our theory predicts three practical regimes: for small initialization scale a, SAM collapses toward the origin; for large a, SAM behaves similarly to GD; and for intermediate o, SAM preferentially amplifies minor intermediate features relative to GD.
To examine these predictions in practice, we train depth-2 CNNs with ReLU activations using both SAM and GD. We then apply Grad-CAM (Selvaraju et al., 2019; Gildenblat & contributors, 2021) to visualize which regions of the input image are emphasized by each model. In addition to qualitative visualizations, we compute the average values of pixels whose Grad-CAM activation exceeds threshold (0.5) and plot this quantity as function of the initialization scale a. To characterize the sequential feature amplification as a function of the initialization scale, we rescale the default random initialization by multiplying it by and train the model under this controlled initialization scheme. Unlike the theoretical setting of Theorem 4.5, which assumes a structured initialization, we use randomized initialization with rescaling in practice. In the corresponding figures, indicate collapse-to-origin behavior in green and blow-up behavior in purple.
We conduct experiments on MNIST (Deng, 2012), SVHN (Netzer et al., 2011). and CIFAR-10 (Krizhevsky et al., 2009). Across all datasets, we consistently observe that GD-trained models concentrate dominant, high-intensity pixels, whereas SAM-trained models emphasize lower- intensity, minor pixel regions. These results demonstrate that the distinct feature prioritization mechanism predicted by our theory persists in nonlinear CNN architectures.
78

---

Published as a conference paper at ICLR 2026
E.4.1 MNIST
We first study this phenomenon on MNIST. MNIST has simple structure, where the black back- ground takes the minimum pixel value (0) and the white digit takes the maximum pixel value (I).
We construct subset of 1,000 images whose labels are in 0, 1, and train models using either GD or 2-MMM. After training, we visualize the learned attention patterns using Grad-Can -.M, as shown in Figure 28. We observe that the GD-trained model primarily bases its predictions on the white digit region, whereas the 2--SAM-trainee
model concentrates more strongly on the black background region. Unless otherwise stated, we use learning rate of 0.1, SAM perturbation radius of 0.5, and train for 500 epochs with batch size of 64. We use momentum and no weight decay. For the CNN architecture, we use 3 3 convolutional kernels and do not apply batch normalization or layer normalization.
CruOCAM
CraSCAM
Ind
Qarar
(a) GD
f2-SAM
Figure 28: Grad-CAM comparison between GD and f--SAM on MNIST (labels 0-3).
To study the practical behavior of -SAM, we train models using fee -AM subset of 1,000 MNIST images with labels in {0, 1}. We then visualize the Grad-CAM maps, as shown in Figure 29. We observe bias pattern similar to that of SAM, where the model places greater emphasis background regions corresponding to minor features. We use the same hyperparameters as in the previous experiment: learning rate 0.1, perturbation radius 0.5, training for 500 epochs, and batch size of 64.
Input
Grad-CAM
Overlay
Input
Grad-CAM
Overlay
Input
Grad-CAM
Overlay
Input
Grad-CAM
Overlay
(a) GD
(b) SSA
Figure 29: Grad-CAM comparison between GD and &o -SAM on MNIST (labels 0-1).
We quantify the average values of activated pixels (Grad-CAM 0.5) function of the initialization scale a across different dataset subsets. this experimental setup (Figure 30), we observe that GD consistently concentrates more the white digit region, which be interpreted as the major component the pixel value manner, unless GD fails to minimize the loss because of too large initialization scale. We denote as purple dots where GD blows up. Moreover, observe
79

---

Published as a conference paper at ICLR 2026
three regimes of of SAM. We denote as green dots where too small initialization scale fails to escape near the origin and so the loss is not changed. Here can be seen as Regime 1. After that, SAM concentrates the pixels whose average is almost 0, the background region. This implies SAM concentrating on the minor component of the data more than GD, which can be seen as Regime
2. When GD blows up, SAM also goes out of the trend and almost blows up.
avg/activated pixels)
avg(activated pixels)
0.8
GD
0.6
SAM
GD
SAM
0.4
0.2
0.2
0.0
0.0
0.5
1.0
1.5
2.0
2.5
a
(b) 1K MNIST images with labels and 1.
(a) MNIST with labels 0,1,2,3.
avg(activated pixels)
0.8
avg(activated pixels)
GD
GD
0.6
0.6
SAM
SAM
0.4
0.4
0.2
0.2
0.0
0.0
0.5
1.0
1.5
2.0
2.5
0.5
1.0
1.5
2.0
2.5
(c) Ik MNIST images with labels 0,1,2,3.
(d) Full MNIST 1k subset.
Figure 30: Average number pixels with Grad-CAM activation exceeding 0.5 as function of the initialization scale a, comparing GD and f2-SAM across different MNIST subsets.
-SA exhibits a similar pattern (Figure 31). When a is small, the dynamics collapse toward the origin. For intermediate values of a, loo -AM tends to prioritize minor features more strongly than GD. For sufficiently large a, however, the behavior of loe -SAM deviates from this trend.
avglactivated pixels)
0.6
0.4
SAM
0.2
0.0
0.5
1.0
1.5
2.0
a
Figure 31: Average number pixels with Grad-CAM activation exceeding 0.5 as function of the initialization scale a, comparing GD and -AM on Ik MNIST images with labels 0 and 1.
80

---

Published as a conference paper at ICLR 2026
E.4.2 SVHN
We next study this phenomenon on SVHN. SVHN is more complex than MNIST, as contains both images with dark backgrounds and light digits, well images with light backgrounds and dark digits. Nevertheless, we observe that 22-SAM consistently emphasizes the darker regions of the image.
We construct a subset 1,000 images with labels in {0, 1] and train models using either GD or é2- SAM. We use learning rate of 0.01, SAM perturbation radius of 0.05, and train for 200 epochs.
The images in Figure 32 contain dark digits on light backgrounds. In this case, we observe that SAM concentrates more strongly the digit regions than the background, the digits constitute the minor features in these images. By contrast, the images in Figure 33 contain light digits on dark backgrounds. For these images, SAM concentrates more strongly on the background regions than on the digits, as the background constitutes the minor feature in this setting.
Input
Grad-CAM
Overlay
Input O
Grad-CAM
Overlay
1
0
O
0
Input
Grad-CAM
Overlay
Input 0
Grad-CAM
Overlay
Input
Grad-CAM
Overlay
Input O
Grad-CAM
Overlay
103
Mo
103
103 10
103
Input
Grad-CAM
Overlay
Input 1
Grad-CAM
Overlay
10
11
Input
Grad-CAM
Overlay
Input
Grad-CAM
Overlay
(a) GD
(b) f2-SAM
Figure 32: Grad-CAM comparison between GD and f2 on SVHN (Ik images, labels 0-1) with dark digits and light backgrounds.
81

---

Published as a conference paper at ICLR 2026
Input
Grad-CAM
Overlay
Input 0
Grad-CAM
Overlay
26
26
26
Input O
Grad-CAM
Overlay
Input
Grad-CAM
Overlay
(a) GD
(b) f2-SAM
Figure 33: Grad-CAM comparison between GD and f2-SAM on SVHN (Ik images, labels 0-1) with light digits and dark backgrounds.
Across different values of a, we observe that small a causes f-M to collapse toward the origin, while intermediate a leads f-MS to emphasize minor features with lower pixel intensities as shown in Figure 34, where pixel intensity is computed the average over the three color channels.
avglactivated pixels)
0.6
GD
SAM
0.4
0.0
Figure 34: Average number of activated pixels (Grad-CAM > 0.5) as a function of the initialization scale a, comparing GD and SAMSAM.
E.4.3 CIFAR-10
We also observe the same phenomenon on the CIFAR-10 dataset. We construct a subset of CIFAR- 10 with labels in (0. 1} and train models using learning rate of 0.01, SAM perturbation radius of 0.05, 500 epochs. As shown in Figure 35, small values of lead SAM to emphasize minor features, while larger values of a make the behaviors of GD and SAM increasingly similar.
avg(activated pixels)
0.6
0.4
0.2
GD
SAM
0.0
0
4
Figure 35: Average number of activated pixels (Grad-CAM > 0.5) as function of the initialization scale a, comparing GD and SAM.
82


Reviews:
[{"review_id": "J41lmuMKXp", "reviewer": "Reviewer_SBDC", "summary": "The paper analyzes the implicit bias of sharpness aware minimization (SAM) in the setting of binary classification, linear separable data, and $L$-layer linear diagonal networks. The work shows the implicit bias changes with depth. In addition, unlike GD / GF, the limiting direction of $\\ell_\\infty$ SAM is dependent on initialization.", "questions": "1. For Theorem 4.1 and Theorem 4.2, what are the conditions on $\\rho$ (the radius of SAM)?\n\n2. In Theorem 4.2, could the authors explain what the main obstructions are for proving directional convergence in the special case of one data point?\n\n3. In theorem 3.2, what how do the coordinates of the coefficient $\\beta(t)$ grow in the original flow? In a related vein, what is the rate of convergence of the loss using SAM in the original flow?\n\n4. How would the theorems in section 4 change for $L > 2$?\n\n5. Are the figures generated using rescaled flow or the original flow?", "limitations": "", "rating": 2, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 1, "strengths": "The paper is overall easy to read. In addition, the fact that the implicit bias of SAM changes with depth is interesting.\nA nice result is that $\\ell_\\infty$-SAM's limiting direction is dependent on initialization unlike GD / GF. (That said, for non linear architectures, GD / GF is also dependent on initialization).", "weaknesses": "One weakness is the restrictive setting as the paper deals with L-layer diagonal linear networks and assumes the data is linear separable.\n\nTheorem 4.2 assumes directional convergence of the inner and outer layer as well as their flows which is an extremely strong assumption and something that is highly nontrivial to prove in general.\n\nIn addition, the analysis is restricted to a dataset consisting of one point (that has a monotonic structure with respect to its entries). \n\nFinally, the analysis heavily uses the fact that the depth is 1 or 2 and it seems hard to generalize the analysis to the multi-layer case (L > 2).", "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": "vgJMum6dmo", "reviewer": "Reviewer_gsYm", "summary": "The paper studies the bias (and training trajectory) sharpness-aware minimization in diagonal linear networks for a dataset consisting of one data point and logistic loss. In linear models, no difference between GD and SAM is found. For a larger number of layers, infinity-norm based SAM can exhibit different behavior depending on initialization and the perturbation radius of SAM. $\\ell_2$ based SAM converges to the direction of the $\\ell_1$ max-margin solution, but does so by first having the minor coordinates as the dominant direction and then sequentially moving, from coordinate to coordinate, to the major one.", "questions": "What about more practical experiments for $\\ell_2$ SAM and $L>2$? Do you still observe the sequential feature discovery? Or is this something that only happens in a very narrow regime?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The highlighted rich behaviors that SAM can exhibit is interesting and it is good to have more work studying logistic loss in this setting.  The paper is well written and the experiments on MNIST to some extend backup the theoretical insights.", "weaknesses": "Clearly, the model and data are very specialized. Maybe my main question would the impact of a non-linearity like ReLU would be. This is not covered by the experiments.\n\nNo code was provided, making it more difficult to reproduce results or checking up implementational details of the experiments. For example, I do not think the paper details how exactly the data for Figure 7 is generated (although from the picture one might guess that it could be points drawn from a Gaussian distribution with means (1,2) and (-1,-2)).", "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": "SZrnwgLqJo", "reviewer": "Reviewer_bXyt", "summary": "The paper analyzes the implicit bias of sharpness-aware minimization (SAM) in linearly separable binary classification using (L)-layer linear diagonal networks. For (L=1), both $\\ell_\\infty$-SAM and $\\ell_2$-SAM recover the $\\ell_2$ max-margin solution, matching gradient descent (GD). For depth (L=2), the behavior diverges sharply: (\\ell_\\infty)-SAM’s limit direction is highly initialization-dependent—converging to zero or to any standard basis vector—whereas GD aligns with the data’s dominant coordinate. In contrast, $\\ell_2$-SAM exhibits **sequential feature discovery**, where the predictor first leans on minor coordinates and gradually shifts to larger ones as training progresses or initialization increases. The authors trace this to $\\ell_2$-SAM’s perturbation normalization, which initially amplifies small coordinates before allowing major ones to dominate. Synthetic and real-data experiments corroborate these theoretical findings, revealing depth-sensitive and optimizer-specific implicit biases distinct from GD.", "questions": "See weakness.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Exceptionally clear and easy to follow.\n2. The theoretical analysis is solid and convincingly demonstrates how network depth affects the implicit bias of SAM.", "weaknesses": "It is unclear whether these theoretical findings can inform practical algorithmic improvements—for example, proposing a better SAM-style method.", "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": "BYOsmtI1T6", "reviewer": "Reviewer_uC1X", "summary": "This work investigates the implicit bias of two versions of SAM related to the $L_2$ and $L_{\\infty}$ norm in a classification setting whereas most previous work focusses on the finite mean squared loss regression setting. The analysis is based flow arguments for a data set of size one. They find that the SAM can learn features in different sequences whereas gradient descent would always prefer the major feature first. This provides a potential mechanism to explain the success of SAM. In the appendix their finding is substantiated by an experiment with Gradcam on a vision task.", "questions": "See weaknesses.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "Fully distinguishing the implicit bias induced by SAM and GF is definitely an important open problem. This work reveals a new phenomenon for the SAM optimizer: learning minor features first then major features first. The use of a single sample and the classification setting is original compared to previous works. The work gives good theoretical characterization what happens before reaching the margin. The theory is illustrated for small example in the appendix, indicating that the result might transfer to more complex settings.", "weaknesses": "1) Sequential feature discovery is a known phenomenon for gradient flow at least also known as saddle to saddle dynamics [1]. There are more works showing this for the gradient flow setting. It would be good to reposition the claim in this light as sequential feature discovery for SAM.\n\n2) There is recent work on margin classification for mirror flow and steepest descent flow which cover more ground, how is the analysis for SAM more difficult? Can the authors contrast their work with [2,3]? What would we need to do for an extension or what prevents it?\n\n3) The work is dense making it hard to value each part of the analysis. I would like to propose to put some parts in the appendix completely and provide more details of the current proofs in the main text and make clear where the text starts again and details of the proof are given. For instance highlight how Theorem 4.5 is proven and put the helper lemmas and propositions to the appendix while still highlighting the dynamical description as this important for intuition as well.\n\n4) The experimental validation should be put in the main text as it helps substantiate the claims made. In addition the experimental validation is limited and could be improved to larger scale settings. Additional data analysis of the experiments is also needed. For example rescaling will most likely have changed the final performance of the model. Moreover, hyperparameters play an important role in these type of experiments and the architecture used these would need be mentioned in a table.\n\n6) If SAM learns minor features first can we not see this in the loss curve trajectories in training i.e. GF should learn faster than SAM? Please show a loss plot illustrating this.\n\nMinor points:\nIn appendix D the titles could be improved by using $L_{\\infty}$, also in the text line 2194 it is said L2 sam is used but in the figures it is Linfty, which is it?\n\nTypo line 364 coordinate\n\nFormulation: line 451 \"this is consistent with figure\" should it not be an illustration of or this is substantiated by\n\n[1] Pesme, Scott and Nicolas Flammarion. “Saddle-to-saddle dynamics in diagonal linear networks.” Journal of Statistical Mechanics: Theory and Experiment 2024 (2023): n. pag.\n\n[2] Pesme, Scott et al. “Implicit Bias of Mirror Flow on Separable Data.” ArXiv abs/2406.12763 (2024): n. pag.\n\n[3] Tsilivis, Nikolaos et al. “Flavors of Margin: Implicit Bias of Steepest Descent in Homogeneous Neural Networks.” ArXiv abs/2410.22069 (2024): n. pag.\n\nWhile the work makes an interesting finding the current manuscript may need a substantial rewrite and more experimental illustrations given the limited theory setup.", "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": "rnGallVLGt", "reviewer": "Reviewer_f8R2", "summary": "This paper studies the implicit bias of SAM for linear diagonal networks for a single data point example. While previous works have studied the regression problem, this work focuses on classification and identifies some new phenomena:\n- for $\\ell_\\infty$ SAM with depth at least $2$, the model can converge towards directions that do not maximize the margin\n- for $\\ell_2$ SAM with depth $=2$, the authors identify a novel \"sequential feature discovery\" phenomenon where the network most significant coordinate can switch during training", "questions": "see above", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "The paper is well written and highlights some new unknown phenomena for learning with SAM that are quite interesting. In particular, I found the sequential feature discovery quite cool", "weaknesses": "The main limitation of this work is the single data point assumption. This work indeed solely focuses on a single point dataset example, which is obviously very toyish. I presume extending this work to multiple points is challenging, but I would like to have at least discussions on the challenges when considering multiple points and also simulations illustrating similar phenomena on multiple training examples would be cool.\n\nI also have a question of how surprising/interesting is this sequential feature discovery phenomenon. First, in practice, SAM is often used with very small $\\rho$, so that we can expect the relevant initialisation regime is $\\alpha \\gtrsim \\rho$, so that we would be in the Regime 3. Actually, this sequential feature discovery illustrates how complex and non-monotonic the dynamics can be as soon as we leave the linear parametrization, but does it really say more than that? And in particular, is this property really specific to SAM (and not happening for GD)?\n\nIt is not so clear what Section 4.2.4 (and Proposition 4.7) in particular really says about this sequential discovery feature. In particular, it focuses on **lower bounds of the rates** but it seems we cannot say anything about what coordinate is really the dominant one at some fixed time. In particular, I do not see how these results explain what we see in Figure 3.\n\nI also have a concern/remark about the time reparametrization done in Equation (2). At first sight, I would say the trajectories would only match if $$\\int_{0}^\\infty \\ell'(\\langle \\beta(\\hat{\\theta}(t)),\\mu\\rangle)\\mathrm{d} t = -\\infty.$$\nSo I guess a justification of that would be needed. Also, this time reparametrization leads to finite time blow ups as discussed line 240. However I am not sure how to interpret the result in case of blow-up: the trajectory can be considered until the first (smallest) blow up time and nothing happens afterwards? If that would be the case, this means that only one coordinate converges to $\\infty$ in case of blow up, while all the other ones remain bounded as $t\\to\\infty$. This seems very weird to me, as in general when overfitting in classification, the coordinates would all converge to infinity (although potentially at different rates). Maybe this is made possible by the $\\ell_\\infty$-SAM regularization, but I guess it should at least be discussed as it seems surprising to me.\n\nAnother hidden assumption is that $L(\\theta(0))< L(\\mathbf{0})$ (eg assuming $\\alpha$ has only positive coordinates) which actually seems to play a critical role in the analysis. I think it would be needed to discuss how necessary this assumption is and what would happen otherwise.\n\n----------------------\n# Other remarks:\n- Theorem 4.2 assumes three points (a,b,c), but are these needed? Moreover, the final Proposition 4.7 does not totally explain it: it says that any $j$\n- sequential feature discovery made me think about *incremental learning*. Although this is different, it might be worth to cite the works [1] and [2] that characterize it for diagonal linear networks\n- Figure 3 is very nice. I would have liked a comparison with GD somewhere here.\n- Maybe recall before Theorem 4.5 that $0<\\mu_1<\\ldots<\\mu_d$, as I believe this is where it is really needed (and not before) in the paper\n\n\n----------------\n# References\n\n[1] Berthier, Raphaël. \"Incremental learning in diagonal linear networks.\" Journal of Machine Learning Research 24.171 (2023): 1-26.\n\n[2] Pesme, Scott, and Nicolas Flammarion. \"Saddle-to-saddle dynamics in diagonal linear networks.\" Advances in Neural Information Processing Systems 36 (2023): 7475-7505.", "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)