Paper deep dive
Minor First, Major Last: A Depth-Induced Implicit Bias of Sharpness-Aware Minimization
Chaewon Moon, Dongkuk Si, Chulhee Yun
Intelligence
Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 96%
Last extracted: 3/13/2026, 12:42:46 AM
Summary
The paper investigates the implicit bias of Sharpness-Aware Minimization (SAM) in L-layer linear diagonal networks. It identifies that while SAM matches gradient descent (GD) for linear models (L=1), it exhibits distinct behaviors for L >= 2. Specifically, for L=2, SAM displays 'sequential feature amplification', where the model initially relies on minor coordinates before shifting to major ones, a phenomenon not captured by infinite-time limit analyses. The study provides theoretical characterizations for both l-infinity and l-2 SAM variants.
Entities (5)
Relation Signals (3)
Sharpness-Aware Minimization → differsfrom → Gradient Descent
confidence 98% · in stark contrast to GD, whose limit aligns with the basis vector of the dominant data coordinate.
Sharpness-Aware Minimization → exhibitsimplicitbias → Sequential Feature Amplification
confidence 95% · its finite-time dynamics exhibit a phenomenon we call 'sequential feature amplification'
Linear Diagonal Networks → trainedwith → Logistic Loss
confidence 95% · training L-layer linear diagonal networks on linearly separable binary classification... with logistic loss.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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. For $\ell_\infty$-SAM, the limit direction depends critically on initialization and can converge to $\mathbf{0}$ or to any standard basis vector, in stark contrast to GD, whose limit aligns with the basis vector of the dominant data coordinate. For $\ell_2$-SAM, we show that although its limit direction matches the $\ell_1$ max-margin solution as in the case of GD, its finite-time dynamics exhibit 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.
Tags
Links
- Source: https://arxiv.org/abs/2603.08290v1
- Canonical: https://arxiv.org/abs/2603.08290v1
Trouble viewing inline? Open PDF directly →
Full Text
191,882 characters extracted from source content.
Expand or collapse full text
Published as a conference paper at ICLR 2026 MINOR FIRST, MAJOR LAST: A DEPTH-INDUCED IM- PLICIT BIAS OF SHARPNESS-AWARE MINIMIZATION Chaewon Moon Kim Jaechul Graduate School of AI, KAIST chaewon.moon@kaist.ac.kr Dongkuk Si Mobilint, Inc. dongkuk@mobilint.com Chulhee Yun Kim Jaechul Graduate School of AI, KAIST chulhee.yun@kaist.ac.kr ABSTRACT We study the implicit bias of Sharpness-Aware Minimization (SAM) when train- ing L-layer linear diagonal networks on linearly separable binary classification. For linear models (L = 1), both ℓ ∞ - and ℓ 2 -SAM recover the ℓ 2 max-margin classifier, matching gradient descent (GD). However, for depth L = 2, the be- havior changes drastically—even on a single-example dataset. For ℓ ∞ -SAM, the limit direction depends critically on initialization and can converge to 0 or to any standard basis vector, in stark contrast to GD, whose limit aligns with the basis vector of the dominant data coordinate. For ℓ 2 -SAM, we show that although its limit direction matches the ℓ 1 max-margin solution as in the case of GD, its finite- time dynamics exhibit 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 ℓ 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 analy- ses are insufficient. Synthetic and real-data experiments corroborate our findings. 1INTRODUCTION Modern deep networks often generalize well despite extreme over-parameterization. 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 et 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 et al., 2021; Bahri et 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 (e.g., squared loss), leaving open the case of losses whose infimum lies at 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, ℓ ∞ -SAM 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 ℓ 2 max-margin classi- fier (Soudry et al., 2018). For both ℓ ∞ -SAM and ℓ 2 -SAM, we show that they also align with the same limit direction. Thus, SAM does not change the implicit bias here, as shown in Figure 1a. However, for 2-layer diagonal linear networks, we find that the trajectory of the linear coefficient vector β(t) under both ℓ ∞ - and ℓ 2 -SAM can differ substantially from the maximum ℓ 1 -margin implicit bias of GD (Gunasekar et al., 2018b). In Figure 1b, we consider a toy separable dataset (μ, +1) with μ = (1, 2). In this case, the ℓ 1 max-margin direction is e 2 = (0, 1), the standard basis vector for the major component of μ. As predicted, all GD trajectories and some SAM tra- 1 arXiv:2603.08290v1 [cs.LG] 9 Mar 2026 Published as a conference paper at ICLR 2026 0.00.51.01.5 1 0.0 0.2 0.4 0.6 0.8 1.0 1.2 1.4 2 GD -SAM 2 -SAM (a) Depth 1 (linear network) 0.00.51.01.5 1 0.0 0.2 0.4 0.6 0.8 1.0 1.2 1.4 2 GD -SAM 2 -SAM (b) Depth 2 (linear diagonal network) Figure 1: Trajectories of the predictor β(t) ∈ R 2 from identical initial conditions under discrete GD, ℓ ∞ -SAM and ℓ 2 -SAM on(μ, +1) with μ = (1, 2). We used η = 0.3 and ρ = 1 for SAM. jectories show increasing alignment of β(t) with e 2 . However, for some initializations, we observe that some trajectories of β(t) under ℓ ∞ -SAM and ℓ 2 -SAM instead converge to zero, or even align with e 1 = (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. 1.1SUMMARY OF OUR CONTRIBUTIONS We analyze the optimization trajectory and implicit bias of ℓ ∞ -SAM and ℓ 2 -SAM in L-layer linear diagonal networks trained on linearly separable data with logistic loss. For theoretical analysis, we analyze the evolution of the linear coefficient β(t) of the linear diagonal network under continuous- time versions of SAM, ℓ ∞ -SAM flow and ℓ 2 -SAM 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 a single-example dataset(μ, +1). • Depth 1 (linear). For linear models (L = 1), both ℓ ∞ -SAM flow and ℓ 2 -SAM flow have the same ℓ 2 max-margin implicit bias as GD on linearly separable data; in the single-example setting, we further show that the ℓ ∞ -SAM coincides exactly with the GD trajectory. • Depth L, ℓ ∞ -SAM. For L≥ 2 and ℓ ∞ -SAM flow, we characterize the coordinate-wise trajectory of β(t) determined by the relative scale of each coordinate at initialization and the perturbation radius of ℓ ∞ -SAM (Theorem 3.2). For almost all initializations, β(t) diverges and its limit direc- tion is one of the standard basis vectors e 1 ,...,e d or it converges to a finite point (Corollary 3.5). Compared to GD, the limit direction of ℓ ∞ -SAM becomes more sensitive to initialization. • Depth 2, ℓ 2 -SAM. For L = 2 and ℓ 2 -SAM flow, we first prove that the limit direction (if con- vergent to zero loss) is the ℓ 1 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β(t) initially relies on minor coordinates and gradually shifts to larger ones as t increases or initialization scale grows. We provide a theoretical explanation of both time-wise (Theorem 4.4) and initialization-wise (Theorem 4.5) aspects of the phenomenon. This example shows that fo- cusing only on the t → ∞ 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.2RELATED 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 ℓ 2 max-margin classifier. For linear diagonal networks, gradient descent biases toward sparse predictors (Gunasekar et al., 2018b), with 2-layer models converging to ℓ 1 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 ℓ 2 -type bias, while small initialization favors ℓ 1 -type sparsity (Woodworth et al., 2020; Yun et al., 2020; Moroshko et al., 2020). Stochas- tic gradient descent (SGD)’s noise provides implicit regularization toward sparser solutions (Pesme et al., 2021), amplified at 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 a brief overview in Appendix A.2.1. Prior work on small-initialization GD under squared loss in the same diagonal network setting shows incremental saddle-to-saddle learning dynamics, where coordinates become active in discrete stages as the pre- dictor moves between saddles (Berthier, 2023; Pesme & Flammarion, 2023). We provide a detailed comparison between our setting and these saddle-to-saddle 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 et al., 2024b; Liu et al., 2022; Yun & Yang, 2023; Bahri et al., 2021; Zhuang et al., 2022; Kaddour et al., 2022b). Complementing these empirical findings, theoretical work has analyzed SAM’s opti- mization dynamics, generalization, and implicit bias (Li et al., 2024a; Behdin & Mazumder, 2023b; Zhang et al., 2024; Agarwala & Dauphin, 2023; Wen et al., 2023; Long & Bartlett, 2024; Zhou et al., 2024; Springer et 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 a detailed discussion to Appendix A.2.3. 2PRELIMINARIES Notation.We write the i-th standard basis vector as e i . For n ∈ N, let [n] = 1,· ,n. For a vector v ∈ R d , we denote its coordinates by v = (v 1 ,· ,v d ). For any block vector Z = (z (1) ,...,z (L) )∈ (R d ) L , we denote itsℓ-th block byZ (ℓ) := z (ℓ) ∈ R d . Fora,b∈ R d ,a⊙b denotes the element-wise product; for a collectiona (ℓ) L ℓ=1 , we write J L ℓ=1 a l := a (1) ⊙·⊙a (L) . Model. We consider L-layer linear diagonal networks, a simple family of homogeneous networks widely used for the study of implicit bias (See Section 1.2). Let θ = (w (1) ,...,w (L) )∈ (R d ) L be the parameter vector. For x∈ R d , let the linear coefficient vector β(θ) and output f (x) be β(θ) := K L ℓ=1 w (ℓ) ∈ R d , f (x) :=⟨β(θ),x⟩. Data and Loss. We consider the standard supervised learning setting where a binary classification dataset(x i ,y i ) N i=1 is given. Let the logistic loss be ℓ(u) = log (1 + exp(−u)). Then the training loss function is defined as L(θ) := 1 N P N i=1 ℓ(y i ⟨β(θ),x i ⟩). We write the gradient of L with respect to θ in a block form, as∇L(θ) = (∇ w (1) L(θ),...,∇ w (L) L(θ)). Optimization Algorithms. In this paper, we mainly consider the implicit bias of Sharpness-Aware Minimization (SAM, Foret et al. (2020)) and how depth causes it to deviate from the baseline algorithm, gradient descent (GD). At iteration t, a GD update readsθ(t + 1) := θ(t)−η∇L(θ(t)), where η > 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) +ε p (θ(t)), θ(t + 1) := θ(t)− η∇L( ˆ θ(t)), where the perturbation ε p (θ(t)) is the approximate worst-case direction inside the ℓ p -ball of per- turbation radius ρ > 0: ε p (θ) := arg max ∥ε∥ p ≤ρ ε ⊤ ∇L(θ). We refer to ˆ θ as the ascent point. Since θ = (w (1) ,...,w (L) ) has a block structure, we also write ˆ θ = ˆ w (1) ,..., ˆ w (L) and ε p (θ) = (ε (1) p (θ),...,ε (L) p (θ)) so that we can say ˆ w (i) = w (i) + ε (i) p (θ). For p = 2 and ∞, the perturbation ε p (θ) has clean closed-form solutions: ε 2 (θ) := ρ ∇L(θ) ∥∇L(θ)∥ 2 , ε ∞ (θ) := ρ sign(∇L(θ)), 3 Published as a conference paper at ICLR 2026 and we consider the two variants, referred to as ℓ 2 -SAM when p = 2 and ℓ ∞ -SAM when p = ∞. For p =∞, the maximizer is not unique when a coordinate of the gradient is zero. To make sure that the update is uniquely determined, we adopt the convention sign(0) := 0, applied coordinate-wise. Continuous-time Flows. In the study of optimization algorithms, it is often useful to reduce the original discrete-time updates of an optimizer to a corresponding continuous-time flow. Unless the step size is too large, continuous-time flows offer a good approximation of the discrete-time optimizers, while allowing for clean and simplified analyses. For GD, a common continuous-time counterpart is gradient flow (GF): ̇ θ(τ ) =−∇L(θ(τ )). With gradient flow, the analysis of 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 ̇ θ(τ ) =−∇L( ˆ θ(τ )).(1) Depending on the choice of norm, we will use the terms ℓ ∞ -SAM flow and ℓ 2 -SAM flow to refer to the continuous-time versions of SAM. Figure 6 in Section A.1 plots the trajectory of ℓ ∞ -SAM flow and ℓ 2 -SAM flow under the same setup of Figure 1. We observe that the trajectories stay 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. Rescaled Flows. As shown in Section A.3, for the special case of single-example dataset(μ, +1), the ℓ p -SAM flow (p = 2,∞) of the i-th layer weight follows the same spatial trajectory as the following rescaled ℓ p -SAM flow: ̇ w (i) (t) = μ⊙ K ℓ̸=i w (ℓ) (t) +ε (ℓ) p (θ(t)) ,(2) obtained by taking out the loss derivative −ℓ ′ (⟨β( ˆ θ(t)),μ⟩) > 0 from the original ℓ p -SAM flow. Note that the original ℓ p -SAM flow (1) and the rescaled flow in (2) differ only by a reparameteriza- tion of time. Let θ orig (t orig ) denote the original SAM flow and θ(t) the rescaled flow. Then there exists a strictly increasing map t orig = τ (t) such that θ orig (τ (t)) = θ(t). Applying the chain rule yields the relation dθ dt = dθ orig dτ dτ dt =− ∇L( ˆ θ orig (τ (t))) ℓ ′ (β( ˆ θ(t)) ⊤ μ) , dτ dt =− 1 ℓ ′ (β( ˆ θ(t)) ⊤ μ) . Since ℓ ′ (u)↑ 0 as u→∞, the rescaled flow accelerates time in the large-margin regime. Formally, τ (t) = Z t 0 − 1 ℓ ′ (β( ˆ θ(s)) ⊤ μ) ds. The rescaled flow makes the analysis easier due to the omitted term. Since our goal is to gain a better understanding of the spatial trajectory, we study the rescaled SAM flows in our analysis. Directional Convergence. Let β : [0,T max ) → R d be a trajectory with maximal existence time T max ∈ (0,∞]. We say that β(t) converges in direction or directionally converges if the limit ̄ β ∞ = lim t→T max β(t) ∥β(t)∥ exists. In this case, ̄ β ∞ is called the limit direction of β. 3SAM WITHℓ ∞ -PERTURBATIONS We begin with ℓ ∞ -SAM. For single-example data, its counterpart—rescaled ℓ ∞ -SAM flow—has the nice property that each coordinate evolves independently, enabling an exact characterization of the trajectory for any depth L. 3.1DEPTH-1 NETWORKS We start with the depth-1 case, in which the implicit bias of ℓ ∞ -SAM 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, ℓ ∞ -SAM flow directionally converges in the ℓ 2 max-margin direction. The proof is deferred to Section C.1. Since Theorem 3.1 holds for any ρ, it also recovers the implicit bias of GF. While Theorem 3.1 characterizes the limit direction for almost all linearly separable 4 Published as a conference paper at ICLR 2026 datasets, Theorem C.1 shows that, for the single-example data, the ℓ ∞ -SAM flow follows the same trajectory as GF. The yellow lines in Figure 6a depict the flows. As t → ∞, w(t) converges in direction to the ℓ 2 max-margin direction μ. Hence, when L = 1, GD and ℓ ∞ -SAM share the same bias toward the ℓ 2 max-margin solution, independent of the initialization. 3.2DEEPER NETWORKS (L≥ 2) To isolate the depth-induced implicit bias of SAM from effects of data-point configuration and ob- tain a tractable characterization of the SAM dynamics, we analyze the minimalist separable dataset D μ := (μ, +1) with feature vector μ ∈ R d satisfying 0 < μ 1 < · < μ d ; without loss of generality, we assume this monotone ordering of μ i ’s. The additional technical difficulties that arise in the multi-point setting are deferred to Section C.2. In Section 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 case, for deeper (linear diagonal) networks, the implicit bias of ℓ ∞ -SAM 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≥ 3, we show that the implicit bias of ℓ ∞ -SAM is more sensitive to initialization than GD, in the sense that a 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 i ∈ [L], suppose w (i) (0) = α ∈ R d + . Let w (i) (t) follow the rescaled ℓ ∞ -SAM flow (2) with perturbation radius ρ > 0 on the datasetD μ . Then, for the j-th coordinate of β(t): • If α j < ρ, then β j (t) converges to 0 if L is even, or to ρ L if L is odd. • If α j = ρ, then β j (t) = ρ L for all t≥ 0. • If α j > ρ and L = 2, then β j (t) grows exponentially: β j (t) = Θ(exp(2μ j t)). • If α j > ρ and L > 2, let J := arg max j:α j >ρ μ j (α j − ρ) L−2 , and also let T := min k∈J 1 /(L−2)μ k (α k −ρ) L−2 . If j ∈ J , then β j (t) → ∞ as t → T ; otherwise, β j (t) stays bounded for all t < T . We provide the proof of Theorem 3.2 in Section C.3. The behavior of each coordinate β j (t) is completely determined by whether the initialization α j lies below, at, or above the threshold ρ. In each of these three regimes, β j (t) is monotone in t. Recall that ε ∞ (θ) := ρ sign(∇L(θ)). For D μ , the sign of the gradient (5) is determined coordinate-wise. Thus, the rescaled ℓ ∞ -SAM flow (2) decouples across coordinates, and each β j (t) evolves independently, allowing us to state Theorem 3.2 for each separate trajectory of β j (t). Remark 3.3 (Interpretation of the Finite-time Blow-up). For L > 2, the rescaled ℓ ∞ -SAM flow (2) exhibits finite-time blow-up: some coordinates satisfy β j (t) → ∞ 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 ˆ β(t) ⊤ μ→∞, we have ℓ ′ ( ˆ β(t) ⊤ μ)→ 0 − , and therefore τ (t) = Z t 0 − 1 ℓ ′ ( ˆ β(s) ⊤ μ) ds→∞ as t→ T. Thus, in the original SAM flow, only the coordinates in J diverge as the original time τ (t) → ∞, while all other coordinates remain bounded. Remark 3.4 (Interpretation of Exponential Growth). For L = 2, each coordinate β j (t) with α j > ρ grows exponentially as t→∞. Since τ (t)→∞ as t→∞, divergence occurs on the same infinite- time limit in both the rescaled and original ℓ ∞ -SAM 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 L = 2, all coordinates share the same rescaled time, so their relative growth can be compared. Among the coordinates with α j > ρ, the one with the largest feature weight μ j dominates asymp- totically and the ℓ ∞ -SAM flow therefore converges in that coordinate direction. We formalize these conclusions for general L in the following corollary, characterizing the dominant direction. Corollary 3.5. Under the assumptions of Theorem 3.2, let S :=j : α j > ρ and assume S ̸= ∅. If there is a unique maximizing index j ∗ := arg max j∈S μ j (α j − ρ) L−2 , then the ℓ ∞ -SAM flow converges in the e j ∗ direction. In particular, when L = 2, we have j ∗ := arg max j∈S μ j . 5 Published as a conference paper at ICLR 2026 012 1 0 1 2 2 (a) GD (L = 2) 012 1 0 1 2 2 1 = 2 2 = 2 (b) ℓ ∞ -SAM (L = 2) 012 1 0 1 2 2 (c) GD (L = 3) 012 1 0 1 2 2 1 = 3 2 = 3 (d) ℓ ∞ -SAM (L = 3) Figure 2: Trajectories β(t) from identical initializations under GF and ℓ ∞ -SAM flow with d = 2 and μ = (1, 2). For SAM, ρ = 1. The proof is deferred to Section C.4. When L = 2 and α ∈ R d ++ , setting ρ = 0 in Corollary 3.5 yields S = [d]. Hence, Corollary 3.5 recovers that the GF always aligns in the e d direction—the ℓ 1 max-margin direction—regardless of the initialization. Illustrative Example. Figure 2 shows the trajectories of β(t) under GF and ℓ ∞ -SAM flow with L = 2, 3 and μ = (1, 2). Figure 2a depicts the L = 2, GF case, where GF always aligns in the e 2 direction. For L = 2 and ℓ ∞ -SAM (Figure 2b), the plane (β 1 ,β 2 ) is partitioned by the thresholds β j = α 2 j = ρ 2 . If α 2 > ρ (so 2 ∈ S), the ℓ ∞ -SAM flow shows directional convergence in e 2 (red/blue regions). In the yellow region, 2 /∈ S and 1∈ S, so the limit direction is e 1 —the “minor” feature. If all coordinates satisfy α j < ρ, the flow converges to 0 (purple region), by Theorem 3.2. For L > 2 (Figures 2c and 2d), the blue regions get partitioned once more because large α 1 leads to μ 1 (α 1 − ρ) L−2 > μ 2 (α 2 − ρ) L−2 , leading to directional convergence toward e 1 . Comparing the green regions in Figures 2c and 2d shows that the slope of the boundary between blue and green regions is steeper in ℓ ∞ -SAM flow than that of GF. Considering that initializations in the yellow region also result in the limit direction e 1 , these together indicate that ℓ ∞ -SAM exhibits a greater sensitivity to initialization and stronger implicit bias toward minor features than GD. 4SAM WITHℓ 2 -PERTURBATIONS: SEQUENTIAL FEATURE AMPLIFICATION We now turn to ℓ 2 -SAM, which is the form most commonly used in practice. 4.1ASYMPTOTIC BEHAVIOR ON DEPTH-1 AND DEPTH-2 NETWORKS For depth-1 models, ℓ 2 -SAM converges in the ℓ 2 max-margin direction regardless of initialization, matching the implicit bias of GD and ℓ ∞ -SAM. We prove the following theorem in Section D.1: Theorem 4.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, ℓ 2 -SAM flow directionally converges in the ℓ 2 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 ℓ ∞ -SAM flow follows the same trajectory as GF. For depth-2 models, ℓ 2 -SAM asymptotically converges in the ℓ 1 max-margin direction as the loss converges to 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 Section D.3. Theorem 4.2. For almost every dataset which is linearly separable, and any perturbation radius ρ, consider the linear diagonal network of depth 2, f (x) = ⟨w (1) ⊙ w (2) ,x⟩ trained with logistic loss. Let (w (1) (t),w (2) (t)) follow the ℓ 2 -SAM flow with w (1) (0) = w (2) (0). Assume (a) the loss vanishesL(w (1) (t),w (2) (t)) → 0, (b) the predictor β(t) := w (1) (t)⊙ w (2) (t) converges in direction. Then the limit direction of β(t) is the ℓ 1 max-margin direction. Since Theorems 4.1 and 4.2 holds for any ρ, 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 ℓ 2 -SAM flow for L = 1, and we can check that the trajectories coincide with GD’s, as expected by theory. In the L = 2 case (Figure 6b), the green ℓ 2 -SAM flow curves include ones that (i) drift 6 Published as a conference paper at ICLR 2026 0.00.20.40.60.81.0 time t 0.0 0.2 0.4 0.6 0.8 initialization scale Regime 1 Regime 2 Regime 3 1 2 3 4 5 j = arg max j j ( t ) (a) Dominant index j † := arg max j β j (t) over α and t. 0.00.20.40.60.8 time t 0.2 0.3 0.4 0.5 0.6 j ( t )/ ( t ) 2 j = 1 j = 2 j = 3 j = 4 j = 5 (b) Normalized β j (t)/∥β(t)∥ 2 for α = 0.4. Figure 3: Rescaled ℓ 2 -SAM flow onD μ with μ = (4, 5, 6, 7, 8)∈ R 5 and ρ = 1. toward the origin, and those that (i) initially align with e 1 , a direction orthogonal to the ℓ 1 max- margin direction e 2 . Such behaviors are not explained by Theorem 4.2. Hence, to account for what is observed in Figure 6b, we move on to analyze the dynamics of ℓ 2 -SAM in finite time. 4.2PRE-ASYMPTOTIC BEHAVIOR ON DEPTH-2 NETWORKS We investigate the pre-asymptotic dynamics of ℓ 2 -SAM on depth-2 linear diagonal networks and show that the trajectory exhibits a behavior markedly different from its asymptotic limit. This con- trast highlights the need for a finite-time analysis to understand how the implicit bias of SAM actu- ally emerges. In this section, we study the toy dataset D μ := (μ, +1) with μ ∈ R d satisfying 0 < μ 1 < · < μ d . We further present experiments on multi-point datasets, discrete-time ℓ 2 -SAM, and deeper models (L ≥ 3) in Section D.8, which confirm that the qualitative behaviors identified in the depth-2 single-point ℓ 2 -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 w (1) (0) = w (2) (0) = α1 throughout this subsection. We additionally report similar empirical results under random Gaussian initialization in Section E.2. 4.2.1SEQUENTIAL FEATURE AMPLIFICATION We begin by describing a newly observed and surprising phenomenon of ℓ 2 -SAM—sequential fea- ture amplification. For certain initialization scales α and times t, ℓ 2 -SAM first aligns with minor features; as t increases or as α increases, the dominant coordinate transitions from minor, interme- diate to major features. In contrast, GD selects the major feature regardless of α and t. We visualize this using rescaled ℓ 2 -SAM flow in Figure 3a and show the GF and ℓ ∞ -SAM flow counterparts in Figure 7. To quantify the phenomenon along the two axes—time t and initialization scale α—at each t and α, we track the index j † = arg min j β j (t) and color the grid (t,α) according to j † . Regions where β is negligibly small are shown in gray, indicating convergence to 0. Based on the observations from Figure 3a, we partition the initialization scale α into three regimes. (Regime 1) Starting from any α in this range, the trajectory eventually collapses to 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 a fixed α chosen from this regime and increasing t, there exists the period where the dominant coordinate index j † increases over time, transitioning from minor to major features. As shown in Figure 3b, j † sequentially changes from 1 to 5 over time for α = 0.4. (Regime 3) β 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 α-axis. To see this, consider a fixed slice of time t and navigate through the α- axis: for small α, the predictorβ remains near the origin with no feature discovered. As α grows, the dominant coordinate at t shifts sequentially—β 1 becomes largest first, then β 2 , and so on. However, this is not a fair comparison between trajectories, because Figure 3a is obtained from the rescaled flow; each trajectory (for each α) has a different time scale. 7 Published as a conference paper at ICLR 2026 Nevertheless, we can compare between trajectories if we base our comparison on trajectory-wise maxima. Specifically, in Regime 2 (the sequential feature amplification phase), we define the trajectory-wise most-amplified index, to understand how the initialization scale α affects the “am- plification” of minor components. For each coordinate j, we track the ratio β j (t) /β d (t) over the entire trajectory, and define j ∗ (α) := arg max j max t β j (t) /β d (t) as the coordinate with the greatest max- imum relative amplification. In Figure 3a, for each value of α in Regime 2, we plot the time step that attains the maximum value of β j ∗ (α) (t) /β d (t) in black dots; we can clearly observe that j ∗ (α) increases from the minor index 1 to second-most major index d − 1 in Regime 2. We call this phenomenon initialization-wise sequential feature amplification. 4.2.2UNDERSTANDING THE EFFECT OF ℓ 2 -SAM Before analyzing sequential feature amplification, we describe the rescaled ℓ 2 -SAM flow for depth-2 linear diagonal networks and offer an intuitive explanation of the sequential feature amplification phenomenon. With initialization w (1) (0) = w (2) (0) ∈ R d + , we have w (1) (t) = w (2) (t) =: w(t) for all t≥ 0. Using this, we derive in Section D.2 that the rescaled ℓ 2 -SAM flow for w(t) reads ̇ w(t) = μ⊙ w(t)− ρ μ⊙w(t) n θ (t) , where n θ (t) := q 2∥μ⊙w(t)∥ 2 2 .(3) Compared to the ρ = 0 case, the extra term scales μ⊙ w(t) coordinate-wise by 1− ρ μ n θ (t) < 1. When n θ (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 n θ (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(t). Notice that larger μ j leads to smaller 1−ρ μ j n θ (t) . 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.3ANALYSIS OF TIME-WISE SEQUENTIAL FEATURE AMPLIFICATION We next provide a theoretical account of the time-wise sequential feature amplification. At each time t, we analyze the instantaneous growth rate of each coordinate β j (t), viewed as a function of both t and the initialization scale α. This reveals how the growth behavior of different coordinates evolves across the training trajectory. In particular, we derive a coordinate-wise growth rule of β j (t), in a form analogous to Equation (3). The proof is provided in Appendix D.4.3, and an extension to the L-layer setting—where an analogous growth rate can be derived—is given in Section D.5. Lemma 4.3. The rescaled ℓ 2 -SAM flow (2) is ̇ β j (t) = r j (t)β j (t) with r j (t) := 2μ j 1− ρμ j n θ (t) . By Lemma 4.3, the rate r j (t) controls the instantaneous growth or decay of β j (t). For fixed t, r j (t) is concave quadratic in μ j , maximized at μ j = m c (t) := n θ (t) 2ρ . Hence, indices with μ j closest to m c (t) attain the largest r j (t); coordinates with feature strength μ j nearest to m c (t) are amplified the most, while those farther away may even decay. Consequently, the trajectory of m c (t) dictates the feature-amplification dynamics, and it exhibits three regimes depending on the initialization scale. Recall that 0 < μ 1 <· < μ d . Theorem 4.4. There exists a unique α 1 such that α 0 := ρ μ 1 √ 2∥μ∥ 2 < α 1 < ρ ∥μ∥ 4 4 √ 2∥μ∥ 2 ∥μ∥ 3 3 < α 2 := ρ μ d−1 +μ d √ 2∥μ∥ 2 and the trajectory of m c (t) falls into one of the following three regimes. (Regime 1) If α < α 1 , then m c (t) strictly decreases for all t ≥ 0 and there exists T 1 such that for j ∈ [d], β j (t) strictly decreases for all t≥ T 1 . (Regime 2) If α 1 < α < α 2 , there exists T 2 such that m c (T 2 ) < μ d−1 +μ d 2 and m c (t) strictly increases for all t≥ T 2 . (Regime 3) If α > α 2 , then m c (t) > μ d−1 +μ d 2 , and β d (t) has the largest growth rate for all t≥ 0. The proof of Theorem 4.4 is provided in Appendix D.4.5. Theorem 4.4 identifies three regimes of the m c (t) dynamics, each corresponding to a qualitatively different pattern of feature amplification. Regime 1. m c (t) decreases for all t ≥ 0, and reaches μ 1 2 at time T 1 . Once m c (t) ≤ μ 1 2 , every coordinate satisfies r j (t)≤ 0 by the form of r j (t), and thus β j (t) strictly decreases for all j ∈ [d]. 8 Published as a conference paper at ICLR 2026 Regime 3. When m c (t) > μ d +μ d−1 2 , the closest feature strength to m c (t) is μ d , so β d (t) attains the largest growth rate. This explains why the major feature remains dominant throughout this regime. Regime 2. When m c (T 2 ) < μ d +μ d−1 2 , the closest index j c satisfies j c < d. At this time, the largest growth rate is therefore achieved by the non-major coordinate β j c (T 2 ). Since m c (t) strictly increases for all t ≥ T 2 , 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 m c (t) initially decreases and later increases, leading to a non-monotonic sequential feature amplification phenomenon. We discuss this in Section A.5. Regime 2 also leaves a clear trace in the training loss. ℓ 2 -SAM exhibits an early plateau while it 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.1. 0200040006000 Epoch 0.0 0.2 0.4 0.6 Loss 1 2 3 4 5 6 Dominant Index 0200040006000 Epoch 0.0 0.2 0.4 0.6 Loss 1 2 3 4 5 6 Dominant Index Figure 4: Loss curves of GD (left) and ℓ 2 -SAM (right) on a 2-layer diagonal network in Regime 2 (α = 0.35, μ = (1, 2, 3, 4, 5, 6), ρ = 0.1). Colored regions mark the coordinate with highest growth. 4.2.4ANALYSIS OF INITIALIZATION-WISE 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 initialization-wise sequential feature amplification. In Theorem 4.4, we characterize the range of α (Regime 2) in which sequential feature amplification can occur. Here, we quantify the strength of amplification within Regime 2 as a function of α. Since a coordinate β j (t) can diverge, we assess which feature is amplified—and by how much—via the ratio of the j-th feature to the major feature, β j (t) /β d (t). For a given initialization scale α, we track and bound how large the amplification ratio β j (t) /β d (t) can be along the trajectory. Integrating the rescaled ℓ 2 -SAM flow (3) (derived in Section D.6.1) yields the coordinate ODE β j (t) = β j (0) exp 2μ j t− 2ρμ 2 j I(t) where I(t) :=∫ t 0 1 n θ (s) dsfor j ∈ [d].(4) The behavior of β in (4) is determined by I(t). Recall that n θ (t) controls the behavior of ℓ 2 -SAM in Section 4.2.2 and is used to characterize the instantaneous growth rate in Section 4.2.3. Here, we focus on cumulative updates over time, where the time integral I(t) of 1/n θ becomes decisive. By bounding I(t), we quantify how strongly each feature is amplified relative to the major feature. Theorem 4.5. Let α 0 ,α 2 be defined in Theorem 4.4 and α 1 be the threshold from there. Suppose α 1 < α≤ ρ μ 1 +μ d √ 2∥μ∥ 2 < α 2 . Then, for j ∈ [d], there exists T j such that β j (T j ) β d (T j ) ≥ LB j (α) := exp 2R ′ j (R j − 1) log 1 1− α 0 /α + log 1 α 0 /α − C(R j ) where R j := (μ j +μ d ) /μ 1 > 2, R ′ j := (μ d −μ j ) /μ 1 and C(R) := R logR− (R− 1) log(R− 1). The proof follows from a lower bound on I(t), and is deferred to Section D.6.2. A numerical illustration of LB j (α) for several choices of μ is provided in Section D.7. Theorem 4.5 applies to the small-α portion of Regime 2. For each coordinate j, we select the time T j maximizing β j (t) β d (t) over the entire trajectory, and obtain a nontrivial lower bound LB j (α) 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 coordinate must grow: as shown in Section D.7, LB j (α) often exceeds 10, indicating that the minor to intermediate coordinates can take values more than ten times larger than the major coordinate. Dependence on α. For all α in Regime 2, the ratio α 0 /α lies in (0, 1), so both logarithmic terms in LB j (α) are positive. Since R j > 2, the first logarithmic term dominates the exponent, making LB j (α) grow rapidly as α → α 1 . Thus smaller α in Regime 2 produces stronger amplification as 9 Published as a conference paper at ICLR 2026 shown in Section D.7. This is substantiated by Figure 3a: smaller α in Regime 2 keeps the dynamics aligned with minor-intermediate features for a longer time t, leading to greater amplification. Dependence on Feature Geometry. The coefficients R j and R ′ j increase with the spectral gap μ d /μ 1 , so datasets with larger feature contrast amplify more strongly as shown in Section D.7. Since LB j (α) varies across j, it is natural to ask which coordinate experiences the strongest ampli- fication. Proposition 4.6 identifies the maximizing index j ∗ (α), with the proof in Section D.6.3. Proposition 4.6. Under the conditions of Theorem 4.5, define j ∗ (α) := arg max j∈[d] LB j (α) and set α ∗ 0 := α 0 . Then, there exist thresholds α ∗ 0 < α ∗ 1 < · < α ∗ m ≤ ρ μ 1 +μ d √ 2∥μ∥ 2 for some m ≤ d− 1 such that j ∗ (α) = j for α∈ (α ∗ j−1 ,α ∗ j ]. Proposition 4.6 shows j ∗ (α) monotonically increases sequentially from 1 to m on α ∈ (α 0 ,α ∗ m ]. Namely, as the initialization scale α grows, the index that maximizes the lower bound LB j (α) shifts monotonically from minor to intermediate features. This matches the initialization-wise 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 1 to m. Lastly, through the cumulative update analysis, we characterize the asymptotic behavior of ℓ 2 -SAM flow for some extreme ranges of α. We prove the following proposition in Section D.6.4. Proposition 4.7. Consider α 0 defined in Theorem 4.4. (i) If α < α 0 , then β(t) converges to zero. (i) If α > ρ ∥μ∥ 2 2 √ 2d( Q d i=1 μ i ) 1 /d ∥μ∥ 1 , then β(t) converge in ℓ 1 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 α in Regime 1, the loss never vanishes. Proposition 4.7(i) shows that for some α’s in Regimes 2 or 3, the limit direction exists and is the ℓ 1 max-margin direction. 5EXPERIMENTS 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 construction and metrics across architectures (Section E.3), as well as multi-point (Section D.8.2) and deeper di- agonal models (Section D.8.3). We also present experiments with practical CNNs trained on MNIST, where we use Grad-CAM (Selvaraju et al., 2017) to visualize which image pixels are emphasized (Figure 5 and Section E.4). These experiments show that ℓ 2 -SAM allocates relatively bigger em- phasis to weaker/background pixels than GD, qualitatively matching our theory. (a) GD(b) ℓ 2 -SAM Figure 5: Grad-CAM comparison of GD and ℓ 2 -SAM on a CNN trained on MNIST. GD focuses on dominant digit pixels, whereas ℓ 2 -SAM highlights minor background regions. 6CONCLUSION 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 (L ≥ 2) with ℓ ∞ -SAM, we derived precise weight trajectories depending on initialization scale and perturbation radius, where each weight coordinate either diverges toward a standard basis vector or converges to a finite point. The most interesting regime arises for L = 2 with ℓ 2 -SAM: while the limit direction converges to the ℓ 1 max-margin 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 t → ∞ limit can overlook how the bias emerges, motivating a finite-time 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. It does not involve human subjects, personal data, or sensitive information, and introduces no 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-specific 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 our 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 a National Re- search Foundation of Korea (NRF) grant funded by the Korean government (MSIT) (No. RS-2023- 00211352) and an Institute of Information & communications Technology Planning & Evaluation (IITP) grant (No. RS-2019-I190075, 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, p. 152–168. PMLR, 2023. Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-aware mini- mization. In International conference on machine learning, p. 639–668. PMLR, 2022. Maksym Andriushchenko, Dara Bahri, Hossein Mobahi, and Nicolas Flammarion. Sharpness-aware 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 M Long, and Olivier Bousquet.The dynamics of sharpness-aware minimization: Bouncing across ravines and drifting towards wide minima.arXiv preprint arXiv:2210.01513, 2022. Kayhan Behdin and Rahul Mazumder. Sharpness-aware minimization: An implicit regularization perspective. arXiv preprint arXiv:2302.11836, 2023a. Kayhan Behdin and Rahul Mazumder. On statistical properties of sharpness-aware minimization: Provable guarantees. arXiv preprint arXiv: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 Hsieh, and Boqing Gong. When vision transformers outperform resnets without pre-training 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-Hsu 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 Analysis, 68:101595, 2024. Gabriel Clara, Sophie Langer, and Johannes Schmidt-Hieber. Training diagonal linear networks with stochastic sharpness-aware minimization. arXiv preprint arXiv:2503.11891, 2025. Enea Monzio Compagnoni, Antonio Orvieto, Luca Biggio, Hans Kersting, Frank Norbert Proske, and Aurelien Lucchi.An sde for modeling sam: Theory and insights.arXiv preprint arXiv:2301.08203, 2023. Yan Dai, Kwangjun Ahn, and Suvrit Sra. The crucial role of normalization in sharpness-aware minimization. Advances in Neural Information Processing Systems, 36:67741–67770, 2023. Li Deng. The mnist database of handwritten digit images for machine learning research. IEEE 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. https://github.com/ jacobgil/pytorch-grad-cam, 2021. Daniel Gissin, Shai Shalev-Shwartz, and Amit Daniely. The implicit bias of depth: How incremental learning drives generalization. arXiv preprint arXiv:1909.12051, 2019. Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Characterizing implicit bias in terms of optimization geometry. In International Conference on Machine Learning, p. 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, 7, 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 arXiv:2504.12883, 2025. Ziwei Ji 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, p. 17176–17186. Curran Associates, Inc., 2020.URL https://proceedings.neurips.c/paper_files/paper/2020/ file/c76e4b2fa54f8506719a5c0dc14c2eb9-Paper.pdf. Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them. arXiv preprint arXiv:1912.02178, 2019. Jean Kaddour, Linqing Liu, Ricardo Silva, and Matt Kusner. When do flat minima optimizers work? In Advances in Neural Information Processing Systems, 2022a. 12 Published as a conference paper at ICLR 2026 Jean Kaddour, Linqing Liu, Ricardo Silva, and Matt J Kusner. When do flat minima optimizers work? Advances in Neural Information Processing Systems, 35:16577–16595, 2022b. Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Pe- ter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv:1609.04836, 2016. Hoki Kim, Jinseong Park, Yujin Choi, Woojin Lee, and Jaewook Lee. Exploring the effect of multi- step ascent in sharpness-aware minimization. arXiv preprint arXiv:2302.10181, 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, p. 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 of the IEEE/CVF conference on computer vision and pattern recog- nition, p. 5631–5640, 2024b. Yong Liu, Siqi Mai, Minhao Cheng, Xiangning Chen, Cho-Jui Hsieh, and Yang You. Random sharpness-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:1906.05890, 2019. Edward Moroshko, Blake E Woodworth, Suriya Gunasekar, 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, 33:22182–22193, 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, p. 16270–16295. PMLR, 2022. Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Baolin Wu, Andrew Y Ng, et al. Reading digits in natural images with unsupervised feature learning. In NIPS workshop on deep learning and unsupervised feature learning, volume 2011, p. 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, p. 3556–3564. PMLR, 2024. Scott Pesme and Nicolas Flammarion. Saddle-to-saddle dynamics in diagonal linear networks. Ad- vances in Neural Information Processing Systems, 36:7475–7505, 2023. Scott Pesme, Loucas Pillaud-Vivien, and Nicolas Flammarion. Implicit bias of sgd for diagonal linear networks: a provable benefit of stochasticity. Advances in Neural Information Processing Systems, 34:29218–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, p. 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):336–359, October 2019. ISSN 1573-1405. doi: 10.1007/s11263-019-01228-7. URL http://dx.doi.org/10.1007/ s11263-019-01228-7. Baekrok Shin and Chulhee Yun. Implicit bias and loss of plasticity in matrix completion: Depth pro- motes low-rankness. In The Fourteenth International Conference on Learning Representations, 2026. Dongkuk Si and Chulhee Yun. Practical sharpness-aware minimization cannot converge all the way to optima. Advances in Neural Information Processing Systems, 36:26190–26228, 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-aware minimiza- tion enhances feature quality via balanced learning. arXiv preprint arXiv:2405.20439, 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–519, 2024. Ankit Vani, Frederick Tung, Gabriel L Oliveira, and Hossein Sharifi-Noghabi. Forget sharpness: perturbed forgetting of model biases within sam dynamics. arXiv preprint arXiv:2406.06700, 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.14158, 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, p. 3635–3673. PMLR, 2020. Chulhee Yun, Shankar Krishnan, and Hossein Mobahi. A unifying view on implicit bias in training linear neural networks. arXiv preprint arXiv:2010.02501, 2020. Jihun Yun and Eunho Yang. Riemannian sam: Sharpness-aware minimization on riemannian mani- folds. Advances in Neural Information Processing Systems, 36:65784–65800, 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 arXiv:2402.15152, 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 arXiv:2410.10373, 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 https://arxiv.org/ abs/2410.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.08065, 2022. 14 Published as a conference paper at ICLR 2026 CONTENTS 1 Introduction1 1.1Summary of Our Contributions . . . . . . . . . . . . . . . . . . . . . . . . . . . .2 1.2Related Work . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .2 2 Preliminaries3 3 SAM with ℓ ∞ -Perturbations4 3.1Depth-1 Networks . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .4 3.2Deeper Networks (L≥ 2) . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .5 4 SAM with ℓ 2 -Perturbations: Sequential Feature Amplification6 4.1Asymptotic Behavior on Depth-1 and Depth-2 Networks . . . . . . . . . . . . . .6 4.2Pre-asymptotic Behavior on Depth-2 Networks . . . . . . . . . . . . . . . . . . .7 4.2.1Sequential Feature Amplification . . . . . . . . . . . . . . . . . . . . . .7 4.2.2Understanding the Effect of ℓ 2 -SAM . . . . . . . . . . . . . . . . . . . . .8 4.2.3Analysis of Time-wise Sequential Feature Amplification . . . . . . . . . .8 4.2.4Analysis of Initialization-wise Sequential Feature Amplification . . . . . .9 5 Experiments10 6 Conclusion10 A Figures and Discussions Omitted from Main Text17 A.1 Flow Trajectories of GD and SAM . . . . . . . . . . . . . . . . . . . . . . . . . .17 A.2 More Discussion on Related Work . . . . . . . . . . . . . . . . . . . . . . . . . .17 A.2.1 Recent Work on Implicit Bias in 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 of Rescaled ℓ p -SAM flow . . . . . . . . . . . . . . . . . . . . . . . . .19 A.4 GD and ℓ ∞ -SAM do not exhibit sequential feature amplification . . . . . . . . . .20 A.5 Interesting Trajectory in Regime 2 of Theorem 4.4 . . . . . . . . . . . . . . . . .20 B Core Lemma for SAM on Depth-1 Networks21 C SAM with ℓ ∞ -perturbations: Proof of Section 325 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 C.4 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.1One-point Case: Discrete vs. Continuous Dynamics . . . . . . . . . . . .33 C.6.2Multi-point Case: Persistence of One-point Behavior . . . . . . . . . . . .34 D SAM with ℓ 2 -perturbations: Proof of Section 436 D.1 Depth-1 Networks: Proof of Theorem 4.1 . . . . . . . . . . . . . . . . . . . . . .36 D.2 Derivation of ℓ 2 -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 . . . . . . . . . . . . . . . . . . . . . . . . . . . .41 D.4.2 Preliminary Analysis . . . . . . . . . . . . . . . . . . . . . . . . . . . . .42 D.4.3 Proof of Lemma 4.3 . . . . . . . . . . . . . . . . . . . . . . . . . . . . .42 D.4.4 Preliminary Analysis for m c (t) 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 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .58 D.6.1 Derivation of the Dynamics of β(t) . . . . . . . . . . . . . . . . . . . . .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 E Experiments74 E.1 Loss Dynamics . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .74 E.2 Sequential Feature Amplification under Random Initialization . . . . . . . . . . .74 E.3 Alternative 2-Layer Models . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .76 E.3.1Linear Network . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .76 E.3.2Convolutional Neural Network . . . . . . . . . . . . . . . . . . . . . . . .77 E.4 Grad-CAM . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .78 E.4.1MNIST . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .79 E.4.2SVHN . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .81 E.4.3CIFAR-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, or conclusions. All LLM-assisted text was reviewed and edited by the authors. AFIGURES AND DISCUSSIONS OMITTED FROM MAIN TEXT A.1FLOW TRAJECTORIES OF GD AND SAM 0.00.51.01.5 1 0.0 0.2 0.4 0.6 0.8 1.0 1.2 1.4 2 GD -SAM 2 -SAM (a) Depth 1 (linear network) 0.00.51.01.5 1 0.0 0.2 0.4 0.6 0.8 1.0 1.2 1.4 2 GD -SAM 2 -SAM (b) Depth 2 (linear diagonal network) Figure 6: Trajectories of the predictor β(t) ∈ R 2 from identical initial conditions under GF, ℓ ∞ - SAM flow and ℓ 2 -SAM flow on(μ, +1) with μ = (1, 2). For SAM, ρ = 1. A.2MORE DISCUSSION ON RELATED WORK A.2.1RECENT WORK ON IMPLICIT BIAS IN DIAGONAL LINEAR NETWORKS Jacobs & Burkholz (2024) study continuous sparsification with time-varying weight decay, formu- lating a time-dependent Bregman potential that causes the implicit bias to evolve from ℓ 2 - to ℓ 1 -type behavior over the course of training. Wang & Klabjan (2024) study smoothed sign descent on a quadratically parameterized regression problem, introducing a time varying mirror map. and prove that the resulting limit point is an approximate KKT point of a Bregman-divergence–style objective, where the stability constant ε quantifies the gap to KKT optimality. Papazov et al. (2024) analyze momentum gradient descent on diagonal linear network through a momentum gradient flow, show- ing that a 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 over 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. A.2.2COMPARISON WITH SADDLE-TO-SADDLE DYNAMICS In this section, we provide further details on 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 infinitesimal-initialization limit. In this regime, gradient flow ex- hibits incremental, stage-wise learning: the flow undergoes long plateaus near a saddle whose pre- dictor is supported on the first k coordinates, then escapes along a low-dimensional “fast escape” 17 Published as a conference paper at ICLR 2026 manifold to a saddle with support on k+1 coordinates, and so on. Sequentiality thus appears as dis- crete transitions between saddles with support size k 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 on the sequential feature amplification focuses on a linear diagonal classifier trained with ℓ 2 -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 to major features. In our setting, all coordinates are present from the beginning. Instead of coordinate jumps, we track how the coordinate-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 ℓ 2 -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 in training and gradually shifts toward major features; and (i) an initialization-scale-wise ordering, where the most-amplified feature over a finite training process changes systematically with the initialization scale. In both views, the ordering emerges through a 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 α1 (without taking the limit α → 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-saddle 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 μ j . In saddle-to-saddle dynamics, the progression is governed by a 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 (α → 0). In contrast, we observe sequential feature amplification across a wide range of non-vanishing initialization scales, and in fact show that increasing α induces a clear and systematic amplification ordering. Our phenomenon is therefore not a small-initialization effect. Fourth, saddle points play no 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 we characterize. The observed dynamics are not mediated by saddle escape. Finally, the problem setups are fundamentally different. Prior saddle-to-saddle works analyze re- 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 not arise under conventional gradient descent. A.2.3PROPERTIES OF SAM SAM seeks flatter minima through a two-step update procedure that first perturbs parameters in the direction of steepest ascent before computing the gradient update. This unique optimization strategy has led to the discovery of several distinctive dynamical properties that differ fundamentally from standard gradient descent. SAM exhibits distinctive valley-bouncing dynamics (Bartlett et al., 2022; Wen et al., 2022). Bartlett et al. (2022) demonstrated that SAM oscillates between ravines in quadratic loss landscapes, with update directions naturally aligning with the dominant eigenvector of the Hessian matrix. Building on this geometric intuition, Wen et al. (2022) proved that SAM implicitly follows trajectories that 18 Published as a conference paper at ICLR 2026 minimize the maximum eigenvalue of the Hessian, providing a precise characterization of SAM’s eigenvalue regularization effect. However, SAM’s convergence properties present both opportunities and challenges. Si & Yun (2023) showed that under practical settings, SAM can struggle to converge to local minima, while Kim et al. (2023) identified specific instabilities in SAM dynamics near saddle points. These find- ings highlight the delicate balance between SAM’s beneficial regularization effects and potential optimization difficulties. Andriushchenko et al. (2023) demonstrated that SAM dynamics drive networks toward low-rank feature representations by systematically pruning activations, providing mechanistic insights into how SAM shapes learned representations. Dai et al. (2023) showed that normalization terms in SAM updates play a crucial role in stabilizing training dynamics and preventing gradient vanishing, highlighting the importance of architectural components in SAM’s effectiveness. Compagnoni et al. (2023) derived continuous-time stochastic differential equations for SAM, proving that the dynam- ics are equivalent to SGD on an implicitly regularized loss with Hessian-dependent noise, thereby connecting SAM to principled stochastic optimization theory. Vani et al. (2024) argued that SAM’s ascent perturbation mechanism systematically discards output-exposed biases, offering a perspective on how SAM’s perturbation strategy contributes to improved generalization. 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-ρ assumption. Clara et al. (2025) study SAM dynamics with noise, proving weight balancing across layers and sharpness minimization, also limited to MSE loss. Our analysis removes the small-ρ assumption to capture the full perturbation effect and studies logistic loss, revealing distinct implicit bias properties compared to the squared loss setting. A.2.4DEPTH-INDUCED IMPLICIT BIAS A line of work shows that increasing depth in overparameterized linearized models can qualitatively change the implicit bias of gradient-based training dynamics (Vardi, 2023). In matrix problems, deep matrix factorization exhibits an enhanced tendency toward low-rank solutions as depth in- creases (Arora et al., 2019; Chou et al., 2024), and recent work further connects depth to low-rank bias and loss of plasticity in matrix completion (Shin & Yun, 2026). Relatedly, depth can also in- duce incremental learning dynamics and sparsity-promoting behavior in simplified models (Gissin et al., 2019). Our results complement these findings by isolating a depth-dependent implicit bias mechanism specific to SAM in diagonal linear networks. A.3DERIVATION OF RESCALED ℓ p -SAM FLOW For the dataset(μ, +1), the loss function is given as: L(θ) = ℓ(⟨β(θ),μ⟩). For each i∈ [L], the gradient is ∇ w (i) L(θ) = ℓ ′ ⟨β(θ),μ⟩ ∇ w (i) ⟨β(θ),μ⟩= ℓ ′ ⟨β(θ),μ⟩ μ⊙ K ℓ̸=i w (ℓ) . (5) Then, we have the ℓ p -SAM flow of w (i) as ̇ w (i) (t) =−∇ w (i) L( ˆ θ(t)) =−ℓ ′ ⟨β( ˆ θ(t)),μ⟩ μ⊙ K ℓ̸=i ˆ w (ℓ) (t) . Since ℓ ′ (u) =− 1 1+exp(u) < 0, it has the same spatial trajectory (up to reparameterization of time): ̇ w (i) (t) = μ⊙ K ℓ̸=i ˆ w (ℓ) (t) = μ⊙ K ℓ̸=i w (ℓ) (t) +ε (ℓ) p (θ(t)) . This derivation works for any p, not just p = 2 and p =∞. 19 Published as a conference paper at ICLR 2026 A.4GD AND ℓ ∞ -SAM DO NOT EXHIBIT SEQUENTIAL FEATURE AMPLIFICATION 0.00.51.01.5 time t 0.0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 initialization scale 1 2 3 4 5 arg max j j ( t ) (a) GF 0123 time t 0.00 0.25 0.50 0.75 1.00 1.25 1.50 1.75 2.00 initialization scale 1 2 3 4 5 arg max j j ( t ) (b) ℓ ∞ -SAM Figure 7: Dominant index j † := arg max j β j (t) for GF and ℓ ∞ -SAM flow over (t,α) onD μ with μ = (4, 5, 6, 7, 8)∈ R 5 . A.5INTERESTING TRAJECTORY IN REGIME 2 OF THEOREM 4.4 0.00.20.40.60.81.0 time t 0.0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 initialization scale 1 2 3 4 5 6 7 8 9 10 11 12 arg max j j ( t ) (a) T = 1 (short horizon). 0246 time t 0.0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 initialization scale 1 2 3 4 5 6 7 8 9 10 11 12 arg max j j ( t ) (b) T = 6 (long horizon). Figure 8: Dominant index for ℓ 2 -SAM flow with μ = (1, 2,..., 12). The black line indicates the interesting trajectory. In Regime 2 of Theorem 4.4, there is also an interesting sub-regime that corresponds to smaller values of α with the range of Regime 2. Define a critical threshold α crit := ρ∥μ∥ 4 4 √ 2∥μ∥ 2 ∥μ∥ 3 3 ∈ (α 1 ,α 2 ). When α 1 < α < α crit , the trajectory m c (t) initially decreases to a minimum above μ 1 2 and then increases. During this decreasing phase, the ℓ 2 -SAM flow amplifies coordinates with smaller indices j < j c (0) than the most-amplified index at initialization j c (0) ∈ arg min j |μ j − m c (0)|, enabling an aggressive exploration of weaker features before transitioning to the standard minor-first-major- last sequential amplification pattern. Along the black path in Figure 8, this manifests as the most- amplified coordinate starting at β 4 , then stepping down to β 1 sequentially during the initial decrease, and—after sufficient time—stepping back up sequentially toward β d as m c (t) increases. 20 Published as a conference paper at ICLR 2026 BCORE 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 ˆ w(t), introducing the time–dependent factors γ n (t) and the perturbed margins bm n (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,δ) 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 to the same ℓ 2 max-margin 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 ˆ w = w +ε(w) with a perturbation direction ε(w) satisfying ∥ε(w)∥ 2 ≤ Bfor some finite constant B <∞ and all w, the resulting SAM flow converges in ℓ 2 max-margin direction. Proof. Let(x n ,y n ) N n=1 ⊂ R d ×±1 be a linearly separable dataset, that is, there exists a vector w ∗ such that y n x ⊤ n 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 y n = 1. In other words, we redefinex n ← y n x n and work with a datasetx n N n=1 such that ∃w ∗ with x ⊤ n w ∗ > 0 for all n. For the linear model f (x) = x ⊤ w, the logistic loss is L(w) = N X n=1 ℓ(x ⊤ n w), ℓ(u) = log(1 + e −u ), ℓ ′ (u) =− e −u 1 + e −u . The SAM flow with perturbation ε(w) is the gradient flow ̇ w(t) =−∇L( b w(t)), b w(t) = w(t) +ε(w).(6) Let m n (t) = x ⊤ n w(t) and bm n (t) = x ⊤ n b w(t). Then ∇L( b w(t)) =− N X n=1 e − bm n (t) 1 + e − bm n (t) x n =− N X n=1 γ n (t)e −m n (t) x n , with γ n (t) = e −( bm n (t)−m n (t)) 1 + e − bm n (t) ≥ 0. Because b w(t)− w(t) = ε(w(t)) and∥ε(w(t))∥ 2 ≤ B, if the data are bounded, say∥x n ∥ 2 ≤ R, then |bm n (t)− m n (t)| =|x ⊤ n ( b w(t)−w(t))|≤ BR =: C(7) for all n,t. Hence there is a constant A > 0 such that 0≤ γ n (t)≤ A for all n,t. 21 Published as a conference paper at ICLR 2026 The SAM flow equation 6 can therefore be written as ̇ w(t) = N X n=1 γ n (t)e −m n (t) x n ,0≤ γ n (t)≤ A.(8) Let w ∗ denote the ℓ 2 max-margin solution w ∗ = arg min w ∥w∥ 2 s.t. x ⊤ n w ≥ 1 for all n. Let S = n : x ⊤ n w ∗ = 1 be the support set. Standard KKT conditions yield coefficients b n > 0 for n∈ S with P n∈S b n = 1 such that w ∗ = X n∈S b n x n . Define the residual r(t) = w(t)−w ∗ logt. Our goal is to show that r(t) is bounded. This will imply that w(t) ∥w(t)∥ = w ∗ logt +r(t) ∥w ∗ ∥ logt + o(logt) → w ∗ ∥w ∗ ∥ , that is, the SAM flow converges in the ℓ 2 max-margin direction. Differentiating and substituting equation 8, we obtain ̇ r(t) = ̇ w(t)− w ∗ t = N X n=1 γ n (t)e −m n (t) x n − w ∗ t . We split the sum over the support and non-support points: ̇ r(t) = X n∈S γ n (t)e −m n (t) x n + X n/∈S γ n (t)e −m n (t) x n − w ∗ t . For n∈ S we have x ⊤ n w ∗ = 1, so m n (t) = x ⊤ n w(t) = x ⊤ n w ∗ logt +x ⊤ n r(t) = logt +x ⊤ n r(t), and therefore te −m n (t) = e −x ⊤ n r(t) . For n /∈ S we have e −m n (t) = e −x ⊤ n w ∗ logt−x ⊤ n r(t) = t −x ⊤ n w ∗ e −x ⊤ n r(t) . Using w ∗ = P n∈S b n x n we rewrite ̇ r(t) = 1 t X n∈S b n h γ n (t) b n e −x ⊤ n r(t) − 1 i x n + X n/∈S γ n (t)t −x ⊤ n w ∗ e −x ⊤ n r(t) x n .(9) Consider the squared norm: 1 2 d dt ∥r(t)∥ 2 = r(t) ⊤ ̇ r(t) = T 1 (t) + T 2 (t), where T 1 (t) and T 2 (t) are the contributions of the two terms in equation 9. For the non-support term T 2 (t) in equation 9, we have T 2 (t) = X n/∈S γ n (t)t −x ⊤ n w ⋆ e −x ⊤ n r(t) x ⊤ n r(t). There is a margin gap θ > 0 such that x ⊤ n w ∗ ≥ 1 + θ when n /∈ S. Then t −x ⊤ n w ∗ ≤ t −(1+θ) , 22 Published as a conference paper at ICLR 2026 and using γ n (t)≤ A and∀z e −z z ≤ 1, we have T 2 (t)≤ A t 1+θ . For the support points, write z n (t) = x ⊤ n r(t) and define δ n (t) := γ n (t) b n , ψ n (t) = δ n (t)e −z n (t) − 1 z n (t), so that T 1 (t) = 1 t X n∈S b n ψ n (t). We first justify that the coefficients δ n (t) = γ n (t)/b n remain in a fixed compact interval. By equation 7, |bm n (t)− m n (t)|≤ C. Since γ n (t) = e −( bm n (t)−m n (t)) 1 + e − bm n (t) , and the denominator satisfies 1 + e − bm n (t) ≥ 1, we obtain the uniform bound 0≤ γ n (t)≤ e −( bm n (t)−m n (t)) ≤ e C for all n,t. Thus each γ n (t) lies in the compact interval [0,e C ]. Next, since every b n > 0 for n∈ S and S is a finite set, define b min := min n∈S b n > 0, b max := max n∈S b n . Therefore δ n (t) = γ n (t) b n =⇒0≤ δ n (t)≤ e C b min for all n∈ S and all t. Hence δ n (t) ranges over the compact interval [δ min ,δ max ] = h 0, e C b min i . For each fixed δ > 0, consider the function ψ(z,δ) := (δe −z − 1)z. As z →±∞ we have ψ(z,δ)→−∞, and therefore ψ(z,δ) attains a finite global maximum on R. Since δ n (t)∈ [δ min ,δ max ] for all t, there exists a constant C ψ > 0 such that ψ(z,δ)≤ C ψ ∀z ∈ R, ∀δ ∈ [δ min ,δ max ]. Consequently, ψ n (t) = ψ(z n (t),δ n (t))≤ C ψ ∀n∈ S, ∀t, and therefore T 1 (t)≤ C 1 t , C 1 := C ψ X n∈S b n . Combining the two bounds on T 1 (t),T 2 (t), for sufficiently large t, 1 2 d dt ∥r(t)∥ 2 = T 1 (t) + T 2 (t)≤ C 1 t + A t 1+θ ≤ C 2 t , for some constant C 2 > 0. 23 Published as a conference paper at ICLR 2026 Integrating from t 0 to t gives ∥r(t)∥ 2 ≤∥r(t 0 )∥ 2 + 2C 2 Z t t 0 u −1 du =∥r(t 0 )∥ 2 + 2C 2 log t t 0 , so ∥r(t)∥ = O( p logt) = o(logt). Since w(t) = w ∗ logt +r(t), ∥r(t)∥ = o(logt), we obtain w(t) ∥w(t)∥ = w ∗ ∥w ∗ ∥ + o(1), which proves w(t) ∥w(t)∥ → w ∗ ∥w ∗ ∥ 2 . Thus ℓ 2 -SAM flow converges in the ℓ 2 max-margin direction for any initialization and any fixed ρ > 0. 24 Published as a conference paper at ICLR 2026 CSAM WITHℓ ∞ -PERTURBATIONS: PROOF OF SECTION 3 C.1DEPTH-1 NETWORKS: PROOF OF THEOREM 3.1 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, ℓ ∞ -SAM flow directionally converges in the ℓ 2 max-margin direction. Proof. Apply Lemma B.1 with ε(w) = ρ sign(∇L(θ)). Then∥ε(w)∥ 2 ≤ ρ √ d for all w, so the conditions of Lemma B.1 hold. Thus, the flow converges to the ℓ 2 max-margin direction. Theorem C.1. Consider the linear model f (x) = ⟨w,x⟩ trained on the dataset D μ with loss L(w) = ℓ(⟨w,x⟩) where ℓ ′ (u) < 0 for all u. Then, GF and ℓ ∞ -SAM flow, starting from any w(0), evolve on the same affine line w(0) + spanμ and have the same spatial trajectory. Proof. The model is f (x) = ⟨w,x⟩ = w ⊤ x. The loss is L(w) = ℓ(w ⊤ μ). The gradient is ∇ w L(w) = ℓ ′ (w ⊤ μ)·μ with ℓ ′ (s) < 0. Gradient Descent The GF is ̇ w =−∇ w L(w) =−ℓ ′ (w ⊤ μ)·μ. SAM with ℓ ∞ perturbation The ascent point is ˆ w = w + ρε ∞ (w) = w + ρsign(∇ w L(w)) = w− ρsign(μ). The equation of ℓ ∞ -SAM flow is ̇ w =−∇ w L( ˆ w) =−∇ w L(w− ρsign(μ)) =−ℓ ′ (w ⊤ μ− ρsign(μ) ⊤ μ)·μ =−ℓ ′ (w ⊤ μ− ρ∥μ∥ 1 )·μ. Therefore, they have the same spatial trajectory as: ̇ w = μ. The term−ℓ ′ (w ⊤ μ− ρ∥μ∥ 1 ) is the accelation in terms of t since−ℓ ′ (s) is decreasing in s. C.2TECHNICAL CHALLENGES FOR MULTI-POINT 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 ρ-perturbation is non-negligible, but in the multi-point case the resulting gradients (and thus SAM updates) become considerably complex for a tractable characterization of the SAM flow in the regime where SAM and GD diverge. This motivates our focus on the single-example dataset D μ = (μ, +1), where the SAM dynamics admit a tractable dynamical characterization while still capturing depth-dependent phenomena unique to SAM. In Section C.6, we 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.3PROOF OF THEOREM 3.2 Theorem 3.2. For i ∈ [L], suppose w (i) (0) = α ∈ R d + . Let w (i) (t) follow the rescaled ℓ ∞ -SAM flow (2) with perturbation radius ρ > 0 on the datasetD μ . Then, for the j-th coordinate of β(t): • If α j < ρ, then β j (t) converges to 0 if L is even, or to ρ L if L is odd. • If α j = ρ, then β j (t) = ρ L for all t≥ 0. • If α j > ρ and L = 2, then β j (t) grows exponentially: β j (t) = Θ(exp(2μ j t)). • If α j > ρ and L > 2, let J := arg max j:α j >ρ μ j (α j − ρ) L−2 , and also let T := min k∈J 1 /(L−2)μ k (α k −ρ) L−2 . If j ∈ J , then β j (t) → ∞ as t → T ; otherwise, β j (t) stays bounded for all t < T . Proof. Since we supposew (i) (0) = α∈ R d + for all i∈ [L], and the dynamics of the linear diagonal network are invariant under any permutation of the layer indices1,...,L, we obtain w (1) (t) = w (2) (t) =· = w (L) (t) =: w(t) for all t≥ 0. With ℓ ∞ perturbation, the rescaled ℓ ∞ -SAM flow (2) becomes ̇ w (i) (t) = μ⊙ K ℓ̸=i w (ℓ) (t) +ε (ℓ) ∞ (θ(t)) = μ⊙ K ℓ̸=i w (ℓ) (t) + ρ sign(∇ w (ℓ) L(θ(t))) . Recall the gradient (5) ∇ w (ℓ) L(θ(t)) = ℓ ′ ⟨β(θ(t)),μ⟩ μ⊙ K ℓ̸=i w (ℓ) (t) , where ℓ ′ (u) =− 1 1+exp(u) < 0. Since we also have μ > 0 (element-wise), we have sign(∇ w (ℓ) L(θ(t))) =−sign K ℓ̸=i w (ℓ) (t) = (a) −sign L−1 K ℓ=1 w(t) ! , where (a) follows from the fact that w (i) (t) = w(t) for all i ∈ [L]. Using this fact again, we have the ODE ̇ w(t) = ̇ w (i) (t) = μ⊙ K ℓ̸=i w(t)− ρ sign L−1 K ℓ=1 w(t) !! = μ⊙ L−1 K ℓ=1 w(t)− ρ sign L−1 K ℓ=1 w(t) !!! . This can be written as coordinate-wise as ̇w j (t) = μ j w j (t)− ρ sign w j (t) L−1 L−1 for j ∈ [d]. Divide into three cases: 26 Published as a conference paper at ICLR 2026 Case 1: L = 2. ̇w j (t) = μ j (w j (t)− ρ sign (w j (t))). By Lemma C.2, we have w j (t) = ρ + (w j (0)− ρ)e μ j t if w j (0) > ρ, ρif w j (0) = ρ, ρ + (w j (0)− ρ)e μ j t (t < T ),0 (t≥ T ) if w j (0) < ρ, 0if w j (0) = 0, where T := 1 μ j log ρ ρ−w j (0) . Then, we have β j (t) = w j (t) L → Θ(e 2μ j t ) if α j > ρ, ρ L if α j = ρ, 0if α j < ρ, as t→∞. Case 2: L > 2 and L is even. ̇w j (t) = μ j (w j (t)− ρ sign (w j (t))) L−1 . By Lemma C.3, we have w j (t) = ρ + −(L− 2)μ j t + 1 (w j (0)−ρ) L−2 − 1 L−2 if w j (0) > ρ, ρif w j (0) = ρ, ρ− −(L− 2)μ j t + 1 (w j (0)−ρ) L−2 − 1 L−2 (t < T ),0 (t≥ T ) if w j (0) < ρ, 0if w j (0) = 0, where T := (ρ−w j (0)) −(L−2) −ρ −(L−2) (L−2)μ j . Then, we have β j (t) = w j (t) L → Θ (t ∗ − t) − L L−2 if α j > ρ, as t→ t ∗ , ρ L if α j = ρ, as t→∞, 0if α j < ρ, as t→∞, where t ∗ = 1 /(L−2)μ j (w j (0)−ρ) L−2 Case 3: L > 2 and L is odd. ̇w j (t) = μ j (w j (t)− ρ) L−1 . By Lemma C.4, we have w j (t) = ρif w j (0) = ρ, ρ + −(L− 2)μ j t + 1 (w j (0)−ρ) L−2 − 1 L−2 if w j (0)̸= ρ. Then, we have β j (t) = w j (t) L → ( Θ (t ∗ − t) − L L−2 if α j > ρ, as t→ t ∗ , ρ L if α j ≤ ρ, as t→∞, where t ∗ = 1 /(L−2)μ j (w j (0)−ρ) L−2 . These cases of L cover all 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 μ > 0 and ρ > 0. Consider ̇w(t) = μ (w(t)− ρ sign(w(t))). Then, there exists the solution w such that it is absolutely continuous (AC) and satisfies w(t) = w(0) + Z t 0 ̇w(s)ds.(10) In particular, w(t) = ρ + (w(0)− ρ)e μt if w(0) > ρ, ρif w(0) = ρ, ρ + (w(0)− ρ)e μt (t < T ),0 (t≥ T ) if w(0) < ρ, 0if w(0) = 0, where T := 1 μ log ρ ρ−w(0) . Proof. Case 1: w(0) = 0. The constant function w(t) = 0 is AC, and Z t 0 μ 0− ρ sign(0) ds = Z t 0 0ds = 0. Thus, Equation (10) holds. Case 2: w(0) = ρ. The constant function w(t) = ρ is AC, and since sign(w(t)) = 1, we have Z t 0 μ ρ− ρ· 1 ds = Z t 0 0ds = 0. Thus, Equation (10) holds. Case 3: w(0) > ρ. At t = 0, we have ̇w(0) = μ w(0)− ρ > 0. Assume, for contradiction, that there exists t ⋆ > 0 with w(t ⋆ ) = ρ. Then on [0,t ⋆ ) we have w(t) > ρ and hence ̇w(t) = μ w(t)− ρ > 0, so w is strictly increasing on [0,t ⋆ ). An increasing function cannot reach the smaller value ρ starting from w(0) > ρ: contradiction. Thus w(t) > ρ for all t ≥ 0. On the region w(t) > ρ, sign(w(t)) = 1 and the ODE reduces to the linear equation ̇w = μ(w− ρ). Then, we have ̇w(t) w(t)− ρ = μ ⇒ Z t 0 ̇w(s) w(s)− ρ ds = Z t 0 μds ⇒ log w(t)− ρ w(0)− ρ = μt ⇒w(t) = ρ + w(0)− ρ e μt . This function is AC and satisfies Equation (10). Case 4: 0 < w(0) < ρ. Initially sign(w(0)) = 1, so again ̇w = μ(w− ρ) and w(t) = ρ + w(0)− ρ e μt . Since w(0)− ρ < 0, the function w is strictly decreasing and reaches 0 exactly once at T := 1 μ log ρ ρ− w(0) > 0. On [0,T ], this solution is AC and satisfies Equation (10). Define w(t) := 0 for all t ≥ T . Then, using sign(0) = 0, w(t) = w(T ) + Z t T μ 0− ρ sign(0) ds = 0 + Z t T 0ds = 0, so Equation (10) also holds on [T,∞). The function w is AC on [0,T ] and on [T,∞), and it is continuous at t = T , hence it is absolutely continuous. 28 Published as a conference paper at ICLR 2026 Lemma C.3. Let μ > 0, ρ > 0, and L is even. Consider ̇w(t) = μ (w(t)− ρ sign(w(t))) L−1 . Then, there exists the solution w such that it is absolutely continuous (AC) and satisfies Equa- tion (10). In particular, w(t) = ρ + −(L− 2)μt + 1 (w(0)−ρ) L−2 − 1 L−2 if w(0) > ρ, ρif w(0) = ρ, ρ− −(L− 2)μt + 1 (w(0)−ρ) L−2 − 1 L−2 (t < T ),0 (t≥ T ) if w(0) < ρ, 0if w(0) = 0, where T := (ρ−w(0)) −(L−2) −ρ −(L−2) (L−2)μ . Proof. The proof is similar to the proof of Lemma C.2. Case 1: w(0) = 0. The constant function w(t) = 0 is AC, and Z t 0 μ 0− ρ sign(0) L−1 ds = Z t 0 μ· 0 L−1 ds = 0. Thus, Equation (10) holds. Case 2: w(0) = ρ. The constant function w(t) = ρ is AC, and since sign(w(t)) = 1, we have Z t 0 μ ρ− ρ· 1 L−1 ds = Z t 0 μ· 0 L−1 ds = 0. Thus, Equation (10) holds. Case 3: w(0) > ρ. At t = 0, we have ̇w(0) = μ w(0)− ρ L−1 > 0. Assume, for contradiction, that there exists t ⋆ > 0 with w(t ⋆ ) = ρ. Then on [0,t ⋆ ) we have w(t) > ρ and hence ̇w(t) = μ w(t)− ρ > 0, so w is strictly increasing on [0,t ⋆ ). An increasing function cannot reach the smaller value ρ starting from w(0) > ρ: contradiction. Thus w(t) > ρ for all t ≥ 0. On the region w(t) > ρ, sign(w(t)) = 1 and the ODE reduces to ̇w = μ(w− ρ) L−1 . Then, we have ̇w(t) (w(t)− ρ) L−1 = μ ⇒ Z t 0 ̇w(s) (w(s)− ρ) L−1 ds = Z t 0 μds ⇒− 1 L− 2 1 (w(t)− ρ) L−2 − 1 (w(0)− ρ) L−2 = μt ⇒(w(t)− ρ) L−2 = −(L− 2)μt + 1 (w(0)− ρ) L−2 −1 ⇒ (a) w(t) = ρ + −(L− 2)μt + 1 (w(0)− ρ) L−2 − 1 L−2 , where (a) follows from w(t)− rho > 0. This function is AC and satisfies Equation (10). Case 4: 0 < w(0) < ρ. Initially sign(w(0)) = 1, so again ̇w = μ(w− ρ) L−1 and (w(t)− ρ) L−2 = −(L− 2)μt + 1 (w(0)− ρ) L−2 −1 . 29 Published as a conference paper at ICLR 2026 Since w(0)− ρ < 0 and L is even, we have w(t) = ρ− −(L− 2)μt + 1 (w(0)− ρ) L−2 − 1 L−2 . The function w is strictly decreasing and reaches 0 exactly once at T := (ρ− w(0)) −(L−2) − ρ −(L−2) (L− 2)μ > 0. On [0,T ], this solution is AC and satisfies Equation (10). Define w(t) := 0 for all t ≥ T . Then, using sign(0) = 0, w(t) = w(T ) + Z t T μ 0− ρ sign(0) L−1 ds = 0 + Z t T 0ds = 0, so Equation (10) also holds on [T,∞). The function w is AC on [0,T ] and on [T,∞), and it is continuous at t = T , hence it is absolutely continuous. Lemma C.4. Let μ > 0, ρ > 0 and L is odd. Consider ̇w(t) = μ (w(t)− ρ) L−1 . Then, there exists the solution w such that it is absolutely continuous (AC) and satisfies Equa- tion (10). In particular, w(t) = ρif w(0) = ρ, ρ + −(L− 2)μt + 1 (w(0)−ρ) L−2 − 1 L−2 if w(0)̸= ρ, Proof. The proof is similar to the proof of Lemma C.2. Case 1: w(0) = ρ. The constant function w(t) = ρ is AC, and Z t 0 μ ρ− ρ ds = Z t 0 0ds = 0. Thus, Equation (10) holds. Case 2: w(0)̸= ρ. Separate variables: dw (w− ρ) L−1 = μdt. Integrating from 0 to t gives − 1 L− 2 1 (w(t)− ρ) L−2 − 1 (w(0)− ρ) L−2 = μt. Solving for w yields w(t) = ρ + −(L− 2)μt + 1 (w(0)− ρ) L−2 − 1 L−2 . The function is AC and satisfies Equation (10). 30 Published as a conference paper at ICLR 2026 C.4PROOF OF COROLLARY 3.5 Corollary 3.5. Under the assumptions of Theorem 3.2, let S :=j : α j > ρ and assume S ̸= ∅. If there is a unique maximizing index j ∗ := arg max j∈S μ j (α j − ρ) L−2 , then the ℓ ∞ -SAM flow converges in the e j ∗ direction. In particular, when L = 2, we have j ∗ := arg max j∈S μ j . Proof. Work under the assumptions of Theorem 3.2 and let S :=j : α j > ρ̸= ∅, j ∗ := arg max j∈S μ j (α j − ρ) L−2 , where the maximizer is unique. We prove that the (rescaled) ℓ ∞ –SAM flow satisfies β(t) ∥β(t)∥ 2 −→ e j ∗ . Case L = 2. By Theorem 3.2, for j ∈ S, β j (t) = Θ e 2μ j t , whereas for j /∈ S we have either β j (t) → 0 (if L even) or β j (t) ≡ ρ L when α j = ρ; in any event these coordinates stay bounded. Since the maximizer is unique and L− 2 = 0, j ∗ = arg max j∈S μ j , hence for every k ∈ S\j ∗ , β k (t) β j ∗ (t) = Θ e −2(μ j ∗ −μ k )t −→ 0, and for k /∈ S we also have β k (t)/β j ∗ (t) → 0 because the denominator grows exponentially while the numerator is bounded. Therefore β(t)/∥β(t)∥ 2 → e j ∗ . Case L > 2. By Theorem 3.2, for each j ∈ S there is a blow-up time t ∗ j = 1 (L− 2)μ j (α j − ρ) L−2 , and as t↑ t ∗ j , β j (t) = Θ (t ∗ j − t) −1/(L−2) . If j /∈ S, then β j (t) is bounded (either converging to 0 when L is even, or equal to ρ L when α j = ρ). The uniqueness of j ∗ implies t ∗ j ∗ = min j∈S t ∗ j and t ∗ j ∗ < t ∗ k ∀k ∈ S\j ∗ . Hence, for any fixed t < t ∗ j ∗ , all coordinates with k ̸= j ∗ are finite; moreover, lim t↑t ∗ j ∗ β k (t) β j ∗ (t) = 0for every k ̸= j ∗ , because β j ∗ (t)→∞ while β k (t) remains finite as t < t ∗ k . Consequently, lim t↑t ∗ j ∗ β(t) ∥β(t)∥ 2 = e j ∗ . Combining the two cases establishes the claim. In particular, when L = 2 we have j ∗ = arg max j∈S μ j . 31 Published as a conference paper at ICLR 2026 C.5FINITE-TIME BLOW-UP In the setting of Theorem C.1, the ℓ ∞ -SAM flow evolves independently across coordinates. In the rescaled ℓ ∞ -SAM flow, each coordinate indeed admits a finite blow-up time. However, as explained in Remark 3.3, the smallest of these blow-up times corresponds to t orig = ∞ in the original SAM time scale. Consequently, both the original flow and the rescaled flow terminate at 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), ρ = 1, and a depth-L = 3 network. In the original flow, only one coordinate diverges as t orig → ∞. As shown in Figure 9b, the normalized trajectories β j (t)/∥β(t)∥ show that the remaining coordi- nates grow much more slowly than the dominant one—indeed, they remain bounded. Because their growth is negligible compared to the blow-up coordinate, their normalized values converge to zero. Thus, in this example, the trajectory converges to the direction e 5 . In contrast, Figure 10a shows that in the rescaled ℓ ∞ -SAM flow, each coordinate β j (t) has its own finite blow-up time. However, Theorem 3.2 identifies the blow-up time T = 1 (L−2)μ j (α j −ρ) L−2 for any j ∈ 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 to t orig =∞, the flow cannot proceed past T . In this example, T ≈ 0.25. Because the rescaled system is simply a time reparameterization of the original one, the two plots differ only in their x-axis scaling. Before reaching T , the two flows exhibit the same evolution along the y-axis. Indeed, reparameterizing the original trajectory (Figure 9) by τ (t) reproduces the same curve as shown in Figure 10 before T . 0.00.51.01.52.02.53.0 time 2 3 4 5 6 7 8 9 beta_j(t) Original -SAM Flow: j (t) trajectories 1 (t) 2 (t) 3 (t) 4 (t) 5 (t) (a) β j (t) trajectory. 0.00.51.01.52.02.53.0 time 0.2 0.3 0.4 0.5 0.6 0.7 0.8 beta_j / ||beta|| Original -SAM Flow: j (t)/(t) trajectories 1 (t)/|| (t)|| 2 (t)/|| (t)|| 3 (t)/|| (t)|| 4 (t)/|| (t)|| 5 (t)/|| (t)|| (b) Normalized β j (t) trajectory. Figure 9: β j (t) and normalized β j (t) trajectory of the original ℓ ∞ -SAM flow. 0.000.250.500.751.001.25 time 0 2000 4000 6000 8000 10000 beta_j(t) Rescaled -SAM Flow: j (t) trajectories 1 (t) 2 (t) 3 (t) 4 (t) 5 (t) (a) β j (t) trajectory. 0.000.250.500.751.001.25 time 0.0 0.2 0.4 0.6 0.8 1.0 beta_j / ||beta|| Rescaled -SAM Flow: j (t)/(t) trajectories 1 (t)/|| (t)|| 2 (t)/|| (t)|| 3 (t)/|| (t)|| 4 (t)/|| (t)|| 5 (t)/|| (t)|| (b) Normalized β j (t) trajectory. Figure 10: β j (t) and normalized β j (t) trajectory of the rescaled ℓ ∞ -SAM flow. 32 Published as a conference paper at ICLR 2026 C.6EMPIRICAL VERIFICATION Our theoretical analysis (Theorem 3.2 and Corollary 3.5) establishes the behavior of the ℓ ∞ -SAM flow in the one-point settingD μ . 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 D μ carry over robustly to multi-point datasets and to practical discrete SAM updates. For reproducibility, we detail the exact initialization used in all experiments. We adopt the layer-wise balanced initializationw (i) (0) = α for every i∈ [L], consistent with the setup of Theorem 3.2. The black-edged dot in Figures 11 and 13 indicates the initial predictorβ(0). We setw (i) (0) = β(0) 1/L element-wise so that β(0) = J L i=1 w (i) (0) holds exactly. For the continuous-time trajectories, we approximate the flow using the corresponding discrete updates with a small step size η = 10 −3 via an explicit Euler scheme. C.6.1ONE-POINT CASE: DISCRETE VS. CONTINUOUS DYNAMICS To verify that our continuous-time analysis faithfully predicts the behavior of the corresponding discrete algorithms, we repeat the experiments in Figure 2 using exactly the same initializations, SAM radius ρ, and feature vector μ. We simulate both the gradient flows (black curves) and their discrete counterparts (blue dots), including GD and discrete ℓ ∞ -SAM updates. As shown below, the discrete trajectories closely trace the qualitative evolution of their continuous-time versions. 01234 1 0 1 2 3 4 2 (a) GD (L = 2) 01234 1 0 1 2 3 4 2 (b) ℓ ∞ -SAM (L = 2) 01234 1 0 1 2 3 4 2 (c) GD (L = 3) 01234 1 0 1 2 3 4 2 (d) ℓ ∞ -SAM (L = 3) Figure 11: Trajectories β(t) under GF, ℓ ∞ -SAM flow (black line), GD, and discrete ℓ ∞ -SAM updates (blue dots) for d = 2 and μ = (1, 2). For SAM, we set ρ = 1. For GD and discrete ℓ ∞ -SAM, we use step size η = 0.1. 33 Published as a conference paper at ICLR 2026 C.6.2MULTI-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 +μ and−μ as shown in Figure 12. Specifically, we draw x (+) n = μ +ε n , y n = +1, x (−) n =−μ +ε n , y n =−1, with ε n ∼N (0,σ 2 I d ) and use N/2 samples per class (with μ = (1, 2),N = 100,σ = 0.5). Figures 11 and 13 show that the same qualitative patterns predicted by our one-point theory—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 ℓ ∞ -SAM). 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. 202 x 1 3 2 1 0 1 2 3 x 2 O +1 class -1 class Figure 12: A randomly generated linearly separable dataset used in our multi-point experiments. We sample two Gaussian clusters centered at ±μ = ±(1, 2) with isotropic noise (ε ∼ N (0, 0.5 2 I 2 )) and assign labels +1 and −1 accordingly. This dataset is 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 01234 1 0 1 2 3 4 2 (a) GD (L = 2) 01234 1 0 1 2 3 4 2 (b) ℓ ∞ -SAM (L = 2) 01234 1 0 1 2 3 4 2 (c) GD (L = 3) 01234 1 0 1 2 3 4 2 (d) ℓ ∞ -SAM (L = 3) Figure 13: Trajectories β(t) under GF, ℓ ∞ -SAM flow (black line), GD, and discrete ℓ ∞ -SAM updates (blue dots) for d = 2 on random multi-point dataset in Figure 12. For SAM, we set ρ = 1. For GD and discrete ℓ ∞ -SAM, we use step size η = 0.1. 35 Published as a conference paper at ICLR 2026 DSAM WITHℓ 2 -PERTURBATIONS: PROOF OF SECTION 4 D.1DEPTH-1 NETWORKS: PROOF OF THEOREM 4.1 Theorem 4.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, ℓ 2 -SAM flow directionally converges in the ℓ 2 max-margin direction. Proof. Apply Lemma B.1 with ε(w) = ρ ∇L(θ) ∥∇L(θ)∥ 2 . Then∥ε(w)∥ 2 ≤ ρ for all w, so the conditions of Lemma B.1 hold. Thus, the flow converges to the ℓ 2 max-margin direction. Theorem D.1. Consider the linear model f (x) = ⟨w,x⟩ trained on the dataset D μ with loss L(w) = ℓ(⟨w,x⟩) where ℓ ′ (u) < 0 for all u. Then, GF and ℓ 2 -SAM flow, starting from any w(0), evolve on the same affine line w(0) + spanμ and have the same spatial trajectory. Proof. The model is f (x) = ⟨w,x⟩ = w ⊤ x. The loss is L(w) = ℓ(w ⊤ μ). The gradient is ∇ w L(w) = ℓ ′ (w ⊤ μ)·μ with ℓ ′ (s) < 0. Gradient Descent GF is ̇ w =−∇ w L(w) =−ℓ ′ (w ⊤ μ)·μ. SAM with ℓ 2 perturbation The ascent point is ˆ w = w + ρε 2 (w) = w + ρ ∇ w L(w) ∥∇ w L(w)∥ 2 = w− ρ μ ∥μ∥ 2 . The update of ℓ 2 -SAM flow is ̇ w =−∇ w L( ˆ w) =−∇ w L(w− ρ μ ∥μ∥ 2 ) =−ℓ ′ (w ⊤ μ− ρ μ ⊤ μ ∥μ∥ 2 )·μ =−ℓ ′ (w ⊤ μ− ρ∥μ∥ 2 )·μ. Therefore, they have the same spatial trajectory as: ̇ w = μ. The term−ℓ ′ (w ⊤ μ− ρ∥μ∥ 2 ) is the accelation in terms of t since−ℓ ′ (s) is decreasing in s. D.2DERIVATION OF ℓ 2 -SAM FLOW Let us get the ℓ 2 -SAM flow. The gradient is ∇ w (i) L(θ) = ℓ ′ ⟨β(θ),μ⟩ ∇ w (i) ⟨β(θ),μ⟩ = ℓ ′ ⟨β(θ),μ⟩ μ⊙w (ℓ) for (i,l)∈(1, 2), (2, 1). From the gradient, we have ε (i) 2 (θ) = ρ ∇ w (i) L(θ) ∥∇L(θ)∥ 2 = (a) −ρ μ⊙w (ℓ) p ∥μ⊙w (1) ∥ 2 2 +∥μ⊙w (2) ∥ 2 2 =−ρ μ⊙w (ℓ) n θ 36 Published as a conference paper at ICLR 2026 for (i,l) ∈ (1, 2), (2, 1), where n θ = p ∥μ⊙w (1) ∥ 2 2 +∥μ⊙w (2) ∥ 2 2 and (a) follows from ℓ ′ (u) =− 1 1+e u < 0. We consider the initializationw (1) (0) = w (2) (0)∈ R d + . Then, since the loss function and dynamics are invariant under exchanging w (1) and w (2) , we have w (1) (t) = w (2) (t) =: w(t) for all t ≥ 0. Therefore, the update on w(t) by rescaled ℓ 2 -SAM flow is given as ̇ w(t) = μ⊙ w(t)− ρ μ⊙w(t) n θ (t) . D.3PROOF OF THEOREM 4.2 Theorem 4.2. For almost every dataset which is linearly separable, and any perturbation radius ρ, consider the linear diagonal network of depth 2, f (x) = ⟨w (1) ⊙ w (2) ,x⟩ trained with logistic loss. Let (w (1) (t),w (2) (t)) follow the ℓ 2 -SAM flow with w (1) (0) = w (2) (0). Assume (a) the loss vanishesL(w (1) (t),w (2) (t)) → 0, (b) the predictor β(t) := w (1) (t)⊙ w (2) (t) converges in direction. Then the limit direction of β(t) is the ℓ 1 max-margin direction. Proof. Let(x n ,y n ) N n=1 ⊂ R d ×±1 be a linearly separable dataset, meaning that there exists w ∗ ∈ R d such that y n x ⊤ n w ∗ > 0 ∀n. As usual, we absorb the labels into the inputs by redefining x n ← y n x n , so that we may assume y n = 1 for all n and ∃w ∗ such that x ⊤ n w ∗ > 0 ∀n. We consider a depth-2 diagonal linear network with parametersw 1 ,w 2 ∈ R d , defining the predictor f (x;w 1 ,w 2 ) = (w 1 ⊙w 2 ) ⊤ x = β ⊤ x, β := w 1 ⊙w 2 . The loss function is logistic: L(w 1 ,w 2 ) = N X n=1 ℓ β ⊤ x n , ℓ(u) = log(1 + e −u ), ℓ ′ (u) =− e −u 1 + e −u . We study the ℓ 2 -SAM flow with fixed perturbation radius ρ > 0: ̇ w 1 (t) =−∇ w 1 L( b w 1 (t), b w 2 (t)), ̇ w 2 (t) =−∇ w 2 L( b w 1 (t), b w 2 (t)), where b w i (t) = w i (t) + ρ ∇ w i L(w 1 (t),w 2 (t)) ∥∇ w i L(w 1 (t),w 2 (t))∥ 2 , i = 1, 2. Step 1: Balanced initialization removes layer imbalance. Let z j (t) := w (1) j (t)− w (2) j (t). From the SAM flow and ∂L ∂w (1) j ( b w) = N X n=1 ℓ ′ ( b β ⊤ x n )x n,j bw (2) j , ∂L ∂w (2) j ( b w) = N X n=1 ℓ ′ ( b β ⊤ x n )x n,j bw (1) j , one obtains ̇z j (t) =−G j (t) w (2) j (t)− w (1) j (t) (1 + o(1)), G j (t) = N X n=1 ℓ ′ ( b β ⊤ x n )x n,j . Here the factor 1+o(1) arises because the gradients in the SAM update are evaluated at the perturbed parameter b w(t) = w(t) + ρ ∇L(w(t)) ∥∇L(w(t))∥ 2 , 37 Published as a conference paper at ICLR 2026 rather than at w(t) itself. Since the perturbation has fixed magnitude ρ while the parameter norm satisfies∥w(t)∥ → ∞ along any vanishing-loss trajectory of a 2-homogeneous model, the relative perturbation decays: ∥ b w(t)−w(t)∥ 2 ∥w(t)∥ 2 = ρ ∥w(t)∥ 2 −→ 0. Consequently, the gradients∇L( b w(t)) and∇L(w(t)) become asymptotically colinear, and replac- ing the latter by the former introduces only a vanishing multiplicative error 1 +o(1) in the imbalance ODE for z j (t). Since z j (0) = 0 under balanced initialization and the ODE ̇z j (t) =−G j (t)z j (t)(1 +o(1)) is linear with a Lipschitz right-hand side, uniqueness of solutions implies z j (t)≡ 0 for all t. Hence for all t w (1) j (t) = w (2) j (t) =: a j (t), β j (t) = a j (t) 2 . Step 2: Predictor ODE. From the SAM ODE, ̇a j (t) =−a j (t)G j (t) (1 + o(1)). Hence ̇ β j (t) = 2a j (t) ̇a j (t) =−2a j (t) 2 G j (t)(1 + o(1)) =−2β j (t)G j (t)(1 + o(1)). Noting that ∇ β L(β) j = N X n=1 ℓ ′ (β ⊤ x n )x n,j , since G j (t) = N X n=1 ℓ ′ ( b β ⊤ x n )x n,j = N X n=1 ℓ ′ (β(t) ⊤ x n )x n,j (1 + o(1)), we have G j (t) =∇ β j L(β(t)) (1 + o(1)). Hence the coordinate-wise predictor dynamics ̇ β j (t) =−2β j (t)G j (t) (1 + o(1)) become ̇ β j (t) =−2β j (t)∇ β j L(β(t)) (1 + o(1)). Writing this in vector form using diag(β)∇ β L = (β 1 ∇ β 1 L,...,β d ∇ β d L) ⊤ , we obtain ̇ β(t) =−2 diag(β(t))∇ β L(β(t)) (1 + o(1)).(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 β = w (1) ⊙w (2) and R(w (1) ,w (2) ) = 1 2 ∥w (1) ∥ 2 2 +∥w (2) ∥ 2 2 , the induced predictor norm is exactly the ℓ 1 norm: ∥β∥ N :=min w (1) ⊙w (2) =β R(w (1) ,w (2) ) =∥β∥ 1 . Moreover, on the balanced submanifold w (1) = w (2) = a with β = a ⊙2 , the lemma establishes that the Riemannian metric induced on predictor space is ⟨u,v⟩ N = u ⊤ M (β)v, M (β) = 2 diag(β). Therefore, the natural-gradient steepest-descent flow with respect to the induced norm∥·∥ N takes the form ̇ β(t) =−M (β(t))∇ β L(β(t)) =−2 diag(β(t))∇ β L(β(t)). We next compare this asymptotic steepest-descent flow with the predictor ODE arising from the ℓ 2 -SAM dynamics. 38 Published as a conference paper at ICLR 2026 Step 4: Asymptotic identification with ℓ 1 steepest descent. Comparing equation 11 with the steepest-descent flow above shows that the SAM predictor dynamics coincide with the ℓ 1 steepest- descent dynamics up to a multiplicative factor 1 + o(1) and a vanishing perturbation. Assumptions (a) and (b) guarantee that these perturbations do not change the limiting direction of β(t)/∥β(t)∥ 2 . Step 5: Conclude ℓ 1 max-margin. By the max-margin theorem for steepest descent in a given norm (Gunasekar et al. (2018a), Thm. 5; extended to logistic loss by Lyu & Li (2019)), any trajectory following ℓ 1 steepest descent and satisfying L(β(t)) → 0 converges in direction to the ℓ 1 max- margin solution. Since the SAM predictor dynamics are asymptotically equivalent to ℓ 1 steepest descent, and by (b) the direction limit exists, we obtain ̄ β ∥ β ⋆ , β ⋆ ∈ arg min β ∥β∥ 1 s.t. β ⊤ x n ≥ 1. Lemma D.2 (Induced Norm and Natural Gradient Metric for Depth-2 Diagonal Models). Consider the depth-2 diagonal parameterization β = w (1) ⊙w (2) ∈ R d , and the quadratic parameter regularizer R(w (1) ,w (2) ) := 1 2 ∥w (1) ∥ 2 2 +∥w (2) ∥ 2 2 . Then the induced predictor norm ∥β∥ N :=min w (1) ⊙w (2) =β R(w (1) ,w (2) ) satisfies ∥β∥ N =∥β∥ 1 . Moreover, on the submanifold where w (1) = w (2) = a and β = a ⊙2 , the Riemannian metric induced on the predictor space by R is ⟨u,v⟩ N = u ⊤ M (β)v, M (β) = 2 diag(β). Consequently, the natural-gradient steepest-descent flow w.r.t.∥·∥ N is ̇ β =−M (β)∇ β L(β) =−2 diag(β)∇ β L(β). Proof. (i) Computation of the induced norm. For each coordinate j, the constraint β j = w (1) j w (2) j decouples. If β j = 0, the minimum is attained at (w (1) j ,w (2) j ) = (0, 0) and equals 0 =|β j |. For β j ̸= 0, eliminate w (2) j via w (2) j = β j /w (1) j and minimize φ j (w) := 1 2 w 2 + β 2 j w 2 ! , w ̸= 0. Differentiation yields φ ′ j (w) = w− β 2 j w −3 , whose nonzero roots satisfy w 4 = β 2 j , so that|w| = |β j | 1/2 . Substitution gives φ j (w ⋆ ) =|β j |. Summing over j yields the induced norm ∥β∥ N = d X j=1 |β j | =∥β∥ 1 . (i) Local parametrization and Jacobian. On the balanced submanifold w (1) = w (2) = a ∈ R d , the predictor is β j = a 2 j . Hence the Jacobian of the map a7→ β is diagonal: ∂β j ∂a k = 2a j δ jk . 39 Published as a conference paper at ICLR 2026 (i) Riemannian metric induced from R. The regularizer restricted to a becomes R(a,a) =∥a∥ 2 2 . Thus the parameter-space metric is Euclidean on a. For a tangent predictor perturbation dβ, the corresponding parameter perturbation is da j = dβ j 2a j = dβ j 2 p β j . Thus the squared parameter differential is ∥da∥ 2 2 = d X j=1 dβ j 2 p β j ! 2 = d X j=1 (dβ j ) 2 4β j . Therefore the predictor-space inner product induced by R is ⟨u,v⟩ N = d X j=1 u j v j 4β j . Equivalently, M (β) −1 = 1 4 diag(β −1 1 ,...,β −1 d ). Inverting yields M (β) = 4 diag(β 1 ,...,β d ). (iv) Removal of irrelevant constant factor. Steepest-descent flows are invariant to multiplication of M by any positive scalar constant. Thus M (β) is equivalent, for optimization dynamics, to M (β) = 2 diag(β), which is the conventional normalization in the induced-norm literature. (v) Natural gradient flow. By definition of steepest descent under the induced norm, ̇ β =−M (β)∇ β L(β) =−2 diag(β)∇ β L(β). D.4PROOFS 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 α, the perturbation radius ρ, and the feature vector μ. 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: w (1) (0) = w (2) (0) = α1 for some scaling factor α > 0. Assumption D.4. the vector μ has strictly positive, increasing coordinates: 0 < μ 1 < · < μ d . (Equivalently, up to a fixed permutation we may assume the coordinates are monotone.) We introduce two auxiliary quantities. Define the normalized weights p j (t) := μ 2 j β j (t) P d k=1 μ 2 k β k (t) and their moments M k (t) := P d j=1 μ k j p j (t). Using these, we set the thresholds m L := μ 1 2 , m H (t) := M 2 (t) 2M 1 (t) . In the proof, we consider ℓ(⟨β,μ⟩) 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) := ⟨β(t),μ⟩ and ˆs(t) := ⟨ ˆ β(t),μ⟩. Set ˆ λ(t) := |ℓ ′ (ˆ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.1RECAP: BASIC NOTATION Recall the margin s = ⟨β,μ⟩ and the lossL(s) = log (1 + exp(−s)). The derivatives of the loss with respect to the margin s are: dL ds =−σ(−s) =− 1 1 + exp(s) , d 2 L ds 2 = σ(s)σ(−s) > 0, where σ(s) = (1 + exp(−s)) −1 is the sigmoid function. We define λ := σ(−s) ∈ (0, 1) as the logistic loss slope magnitude. The gradients with respect to the weight parameters, obtained via the chain rule, are: dL dw (1) j := dL ds ds dw (1) j =−λμ j w (2) j , dL dw (2) j := dL ds ds dw (2) j =−λμ j w (1) j . The squared norm of the gradient vector is: ∥∇ θ L∥ 2 = d X j=1 λ 2 μ 2 j w (2) j 2 + w (1) j 2 = λ 2 n 2 θ , where n θ := s P d j=1 μ 2 j w (1) j 2 + w (2) j 2 . SAM perturbs parameters by taking a step of size ρ along the normalized gradient direction. ε 2 := ρ ∇ θ L ∥∇ θ L∥ 2 , (ε 2 ) w (1) j =− ρμ j w (2) j n θ , (ε 2 ) w (2) j =− ρμ j w (1) j n θ . The perturbed weight parameters are ( ˆw 1 ) j := w (1) j − ρμ j w (2) j n θ ,( ˆw 2 ) j := w (2) j − ρμ j w (1) j n θ . The perturbed β j becomes ˆ β j := ˆw (1) j ˆw (2) j = w (1) j w (2) j − ρμ j n θ w (1) j 2 + w (2) j 2 + ρ 2 μ 2 j n 2 θ w (1) j w (2) j = β j 1 + ρ 2 μ 2 j n 2 θ ! − ρμ j n θ w (1) j 2 + w (2) j 2 . The perturbed margin and loss slope magnitude are ˆs :=⟨ ˆ β,μ⟩ = d X j=1 μ j ˆ β j , ˆ λ := σ(−ˆs). Recall that the SAM flow dynamics are given by: ̇ θ =−∇ θ L( ˆ θ). 41 Published as a conference paper at ICLR 2026 D.4.2PRELIMINARY 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 (1) j (t)−w (2) j (t) exponen- tially to zero. Proof. Define ∆ j := w (1) j − w (2) j . The SAM dynamics yield ̇w (1) j = ˆ λμ j ˆw (2) j , ̇w (2) j = ˆ λμ j ˆw (1) j . The time derivative of ∆ j is ̇ ∆ j = ̇w (1) j − ̇w (2) j = ˆ λμ j ˆw (2) j − ˆ λμ j ˆw (1) j = ˆ λμ j w (2) j − ρμ j w (1) j n θ ! − ˆ λμ j w (1) j − ρμ j w (2) j n θ ! =− ˆ λμ j 1 + ρμ j n θ ∆ j . Since ˆ λ is positive and μ j > 0, it gives exponential decay. ∆ j (T ) = ∆ j (0)· exp −μ j Z T 0 ˆ λ 1 + ρμ j n θ dt ! . Hence, the quantity w (1) j (t)− w (2) j (t) decays exponentially. Proposition D.6. Under initialization with w (1) j (0) = w (2) j (0) and Assumption D.4, the equality w (1) j (t) = w (2) j (t) is preserved for all t ≥ 0. Furthermore, the sign of w (1) j (t) and w (2) j (t) remains unchanged throughout the dynamics. Proof. With w (1) j (0) = w (2) j (0), we have ∆ j (0) = w (1) j (0) − w (2) j (0) = 0. By Lemma D.5, ∆ j (t) = 0 for allt≥ 0. Given this balancedness, each coordinate evolves multiplicatively according to ̇w (1) j = ˆ λμ j ˆw (2) j = ˆ λμ j w (1) j − ρμ j w (1) j n θ ! = ˆ λμ j 1− ρμ j n θ w (1) j . This differential equation has the unique solution w (1) j (T ) = w (1) j (0)· exp μ j · Z T 0 ˆ λ(t) 1− ρμ j n θ dt ! . Since the exponential function is always positive, w (1) j (t) and w (2) j (t) maintain the same sign as their initial values throughout the dynamics. D.4.3PROOF OF LEMMA 4.3 We begin by restating Lemma 4.3. Lemma 4.3. The rescaled ℓ 2 -SAM flow (2) is ̇ β j (t) = r j (t)β j (t) with r j (t) := 2μ j 1− ρμ j n θ (t) . Proof. Under Assumption D.3 and Assumption D.4, the Proposition D.6 holds, which ensures that w (1) j = w (2) j = p β j for all t≥ 0. So we have w (1) j 2 + w (2) j 2 = 2β j , n 2 θ = 2 d X j=1 μ 2 j β j . 42 Published as a conference paper at ICLR 2026 The evolution equation for β j is ̇ β j = ̇w (1) j w (2) j + w (1) j ̇w (2) j = 2 ˆ λμ j β j 1− ρμ j n θ .(12) This yields β j (T ) = β j (0)· exp 2μ j Z T 0 ˆ λ 1− ρμ j n θ dt ! . Let r j := 2 ˆ λμ j 1− ρμ j n θ . When r j > 0, β j grows locally exponentially. Otherwise, it decays lo- cally exponentially. The key insight is that each β j ’s growth rate depends on the interaction between the gradient magnitude ˆ λ and the perturbation term ρμ j n θ . This interaction drives SAM’s implicit bias. D.4.4PRELIMINARY ANALYSIS FOR m c (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 m c (t) is given by ̇m c = ˆ λ(t)M 1 (t) (m c (t)− m H (t)). Proof. Recall that m H = M 2 2M 1 , where M r := d X j=1 p j μ r j , p j := μ 2 j β j P d k=1 μ 2 k β k .(13) Substituting the definition of p j , we obtain M 2 = P j μ 4 j β j P k μ 2 k β k = 2 P j μ 4 j β j n 2 θ , M 1 = P j μ 3 j β j P k μ 2 k β k = 2 P j μ 3 j β j n 2 θ . Since μ 1 < · < μ d and p j ≥ 0 with P j p j = 1, we have μ 1 2 ≤ m H = M 2 2M 1 ≤ μ d 2 . We define a new expression for m c . m c (t) = √ S 2ρ ,where S := n 2 θ .(14) Taking the time derivative of S, we have ̇ S = 2 d X j=1 μ 2 j ̇ β j . From Lemma 4.3, we have ̇ β j = r j β j where r j = 2 ˆ λ· μ j 1− ρμ j n θ = 2 ˆ λ· μ j − μ 2 j 2m c . Substi- tuting this into the expression for ̇ S, we get ̇ S = 2 d X j=1 μ 2 j · 2 ˆ λ· μ j − μ 2 j 2m c ! · β j = 4 ˆ λ d X j=1 μ 2 j β j μ j − μ 2 j 2m c ! = 4 ˆ λ d X j=1 μ 3 j β j − μ 4 j β j 2m c ! . 43 Published as a conference paper at ICLR 2026 Recalling that M 1 = 2 P d j=1 μ 3 j β j S and M 2 = 2 P d j=1 μ 4 j β j S , we can rewrite the sums as d X j=1 μ 3 j β j = M 1 S 2 , d X j=1 μ 4 j β j = M 2 S 2 . Therefore, we have ̇ S = 4 ˆ λ M 1 S 2 − M 2 S 2· 2m c = 2 ˆ λS M 1 − M 2 2m c . Since m c = √ S 2ρ , we have: ̇m c = 1 2ρ · ̇ S 2 √ S = ̇ S 4ρ √ S . Substituting our expression for ̇ S: ̇m c = 2 ˆ λS M 1 − M 2 2m c 4ρ √ S = ˆ λ √ S 2ρ M 1 − M 2 2m c = ˆ λm c M 1 − M 2 2m c = ˆ λM 1 m c − M 2 2M 1 = ˆ λM 1 (m c − m H ). Next, we derive the time derivative of m H . Lemma D.8. Under Assumption D.3 and Assumption D.4, the time derivative of m H is given by ̇m H = ˆ λ 2(M 1 ) 2 m c (2m c Γ 1 − Γ 2 ), where Γ 1 := M 1 M 3 − M 2 2 and Γ 2 := M 1 M 4 − M 2 M 3 . Proof. Starting from m H = M 2 2M 1 , we have ̇m H = ̇ M 2 M 1 − M 2 ̇ M 1 2(M 1 ) 2 = 1 2M 1 ̇ M 2 − M 2 M 1 ̇ M 1 = 1 2M 1 d X j=1 ̇p j μ 2 j − M 2 M 1 · d X j=1 ̇p j μ j = 1 2M 1 d X j=1 ̇p j μ 2 j − 2m H μ j . Since ̇ β j = r j β j where r j = 2 ˆ λ μ j − μ 2 j 2m c , we can compute ̇p j = (μ 2 j β j )· r j · P d k=1 μ 2 k β k − (μ 2 j β j )· P d k=1 μ 2 k β k r k P d k=1 μ 2 k β k 2 44 Published as a conference paper at ICLR 2026 = p j r j − d X k=1 p k r k ! = p j · 2 ˆ λ μ j − μ 2 j 2m c ! − d X k=1 p k · μ k − μ 2 k 2m c ! = p j · 2 ˆ λ (μ j − M 1 )− 1 2m c μ 2 j − M 2 . Substituting this into the expression for ̇m H , we have ̇m H = ˆ λ M 1 d X j=1 p j (μ j − M 1 )− 1 2m c (μ 2 j − M 2 ) (μ 2 j − 2m H μ j ). We split the sum into two components: First term: C 1 = X j p j (μ j − M 1 ) μ 2 j − 2m H μ j , Second term: C 2 = X j p j μ 2 j − M 2 μ 2 j − 2m H μ j . For the first term, C 1 = X j p j μ 3 j − 2m H X j p j μ 2 j − M 1 X j p j μ 2 j + 2m H M 1 X j p j μ j = M 3 − 2m H M 2 − M 1 M 2 + 2m H M 2 1 = M 3 − M 2 2 M 1 = M 1 M 3 − M 2 2 M 1 = Γ 1 M 1 . For the second term, C 2 = X j p j μ 4 j − 2m H X j p j μ 3 j − M 2 X j p j μ 2 j + 2m H M 2 X j p j μ j = M 4 − 2m H M 3 − M 2 2 + 2m H M 1 M 2 = M 4 − M 2 M 3 M 1 = M 1 M 4 − M 2 M 3 M 1 = Γ 2 M 1 . Therefore, we have ̇m H = ˆ λ M 1 d X j=1 p j · (μ j − M 1 )− 1 2m c μ 2 j − M 2 μ 2 j − 2m H μ j = ˆ λ M 1 Γ 1 M 1 − Γ 2 2m c M 1 = ˆ λ 2 (M 1 ) 2 m c (2m c Γ 1 − Γ 2 ). Next, we establish a key inequalities involving the threshold m H . Proposition D.9. Γ 1 ≥ 0 and Γ 2 ≥ 0. Proof. Γ 1 and Γ 2 are defined in Lemma D.8. M r and p j are defined in Equation 13. Let M r := P d j=1 p j μ r j = E p μ r j . By Cauchy–Schwarz with X = μ 1/2 and Y = μ 3/2 , E p μ 2 2 ≤ E p [μ] E p μ 3 =⇒Γ 1 = M 1 M 3 − M 2 2 ≥ 0. 45 Published as a conference paper at ICLR 2026 By Cauchy–Schwarz with X = μ and Y = μ 2 , E p μ 3 2 ≤ E p μ 2 E p μ 4 . Multiplying the two inequalities gives E p μ 2 E p μ 3 ≤ E p [μ] E p μ 4 =⇒Γ 2 = M 1 M 4 − M 2 M 3 ≥ 0. Proposition D.10. Let m D := Γ 2 2Γ 1 . We have m D ≥ m H for all t≥ 0. Proof. We use same notation as in the proof of Proposition D.9. Let a := M 2 M 1 . Γ 1 ≥ 0 and Γ 2 ≥ 0 by Proposition D.9. Then we have E p (μ 2 − aμ) 2 = E p [μ 4 ]− 2a E p [μ 3 ] + a 2 E p [μ 2 ] = M 4 − 2aM 3 + a 2 M 2 . Substituting a = M 2 M 1 and multiplying by M 2 1 gives M 2 1 E p h (μ 2 − M 2 M 1 μ) 2 i = M 2 1 M 4 − 2M 1 M 2 M 3 + M 3 2 . Since an expectation of a square is nonnegative and M 2 1 ≥ 0, it follows that M 2 1 M 4 − 2M 1 M 2 M 3 + M 3 2 ≥ 0. Therefore, we have Γ 2 2Γ 1 ≥ M 2 2M 1 = m H . D.4.5PROOF OF THEOREM 4.4 We begin by restating Theorem 4.4 for convenience. Theorem 4.4. There exists a unique α 1 such that α 0 := ρ μ 1 √ 2∥μ∥ 2 < α 1 < ρ ∥μ∥ 4 4 √ 2∥μ∥ 2 ∥μ∥ 3 3 < α 2 := ρ μ d−1 +μ d √ 2∥μ∥ 2 and the trajectory of m c (t) falls into one of the following three regimes. (Regime 1) If α < α 1 , then m c (t) strictly decreases for all t ≥ 0 and there exists T 1 such that for j ∈ [d], β j (t) strictly decreases for all t≥ T 1 . (Regime 2) If α 1 < α < α 2 , there exists T 2 such that m c (T 2 ) < μ d−1 +μ d 2 and m c (t) strictly increases for all t≥ T 2 . (Regime 3) If α > α 2 , then m c (t) > μ d−1 +μ d 2 , and β d (t) has the largest growth rate for all t≥ 0. Proof. Recall from Section D.4 that p j := μ 2 j β j P d k=1 μ 2 k β k , M r := d X j=1 p j μ r j , m H := M 2 2M 1 , m L := μ 1 2 , ˆ λ :=|ℓ ′ (ˆs)| > 0. From Lemma D.8, we further have Γ 1 := M 1 M 3 − M 2 2 ,Γ 2 := M 1 M 4 − M 2 M 3 , m D := Γ 2 2Γ 1 . Also note that m c := √ S 2ρ , S := n 2 θ . 46 Published as a conference paper at ICLR 2026 From Lemma D.7 and Lemma D.8, we have ̇m c = ˆ λM 1 (m c − m H ), ̇m H = ˆ λ 2 (M 1 ) 2 m c (2m c Γ 1 − Γ 2 ). We partition Regime 1 into sub-regimes 1-a (α < α 0 ) and 1-b (α 0 < α < α 1 ), and similarly parti- tion Regime 2 into sub-regimes 2-a α 1 < α < ρ ∥μ∥ 4 4 √ 2∥μ∥ 2 2 ∥μ∥ 3 3 and 2-b ρ ∥μ∥ 4 4 √ 2∥μ∥ 2 2 ∥μ∥ 3 3 < α < α 2 . 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. n θ (0) = √ 2∥μ∥ 2 · α, so α < α 0 implies m c (0) = n θ (0) 2ρ < μ 1 2 = m L . For any t ≥ 0, if m c (t) < m L , then m c (t) < μ 1 2 < m H (t). Hence B(t) < 0, and therefore ̇m c (t) < 0. Consequently, for any t≥ 0, whenever m c (t) < m L , the function m c (·) is strictly decreasing. Since m c (0) < m L , we have m c (t) < m L for all t≥ 0, and it is strictly decreasing. Moreover, we have 2m c (t) < μ 1 ≤ μ j . Therefore, r j (t) = 2 ˆ λ(t)· μ j − μ 2 j 2m c (t) ! < 0, Thus ̇ β j (t) = β j (t)r j (t) < 0, and β j (t) decays exponentially for all t≥ 0. Regimes 1-b and 2-a. We define A(t) := ˆ λM 1 (t) and B(t) := m c (t)− m H (t). Then we get the following equalities: ̇m c = AB, ̇ B = ̇m c − ̇m H = AB− ̇m H . Let I(t) := exp − R t 0 A(τ )dτ . Then: I ̇ B = IAB− I ̇m H ,(15) d dt (IB) = ̇ IB + I ̇ B =−IAB + I ̇ B =−I ̇m H ,(16) I(t)B(t)− I(0)B(0) =− Z t 0 I(u) ̇m H (u)du.(17) n θ (0) = √ 2∥μ∥ 2 · α, so α 0 < α < ρ ∥μ∥ 4 4 √ 2∥μ∥ 2 2 ∥μ∥ 3 3 implies m L < m c (0) < m H (0). In this case, we have B(0) < 0 and thus ̇m c (0) = A(0)B(0) < 0, so m c initially drifts downward. For an initial condition m 0 ∈ (m L ,m H (0)), let m c (·;m 0 ) denote the solution with m c (0) = m 0 . We define the floor hitting time as τ F (m 0 ) := inft≥ 0 : m c (t;m 0 ) = m L ∈ [0,∞], and the ceiling hitting time as τ G (m 0 ) := inft≥ 0 : B(t;m 0 ) = 0∈ [0,∞]. Define the first exit time as τ (m 0 ) := minτ F (m 0 ),τ G (m 0 ). Now let F :=m 0 ∈ (m L ,m H (0)) : τ F (m 0 ) < τ G (m 0 ), G :=m 0 ∈ (m L ,m H (0)) : τ G (m 0 ) < τ F (m 0 ). We will prove: • F andG are open in (m L ,m H (0)). • F ̸=∅ andG ̸=∅. • therefore there exists m ∗ 0 ∈ (m L ,m H (0))\ (F ∪G), and for such m ∗ 0 , we must have τ F (m ∗ 0 ) = τ G (m ∗ 0 ) =∞. 47 Published as a conference paper at ICLR 2026 Since α > 0, we have β j (0) > 0 for all j ∈ [d], and from the explicit solution of the β j dynamics in Lemma 4.3, we have that for each finite time t, β j (t) > 0. Therefore, p j (t) > 0 for all finite t and j ∈ [d], which implies that p(t) is not collapsed onto a single feature at any finite time. Based on this observation, we make the following Lemma: Lemma D.11. ̇m c (τ F (m 0 )) < 0 for all m 0 ∈ F , and ̇ B(τ G (m 0 )) > 0 for all m 0 ∈ G. In other words, the floor crossing and ceiling crossing are transversal. Proof. We first prove for the floor crossing case. τ F (m 0 ) < ∞.Let t f := τ F (m 0 ), so m c (t f ;m 0 ) = m L . Then we have ̇m c (t f ;m 0 ) = A(t f )B(t f ). Here, A(t f ) > 0. Also p(t f ) has positive mass on every μ j , hence B(t f ) = m L − m H (t f ) < 0. Thus ̇m c (t f ;m 0 ) < 0, in particular ̇m c (t f )̸= 0. So the intersection with the floor is transversal. Now prove for the ceiling crossing case. τ G (m 0 ) < ∞. Let t g := τ G (m 0 ), so B(t g ) = 0 which means m c (t g ) = m H (t g ). Then we have ̇ B(t g ) = A(t g )B(t g )− ̇m H (t g ) =− ̇m H (t g ). From Proposition D.10, we have m H ≤ m D , and the equality holds if and only if p is collapsed onto a single feature. But at any finite time t g , p(t g ) is not collapsed onto a single feature, so m H (t g ) < m D . Therefore, we get m c (t g )− m D (t g ) = m H (t g )− m D (t g ) < 0. Using the formula ̇m H = ˆ λ 2 (M 1 ) 2 m c (2m c Γ 1 − Γ 2 ) = ˆ λ 2 (M 1 ) 2 m c 2Γ 1 (m c − m D ), and noting Γ 1 > 0 for non-collapsed p, we conclude ̇m H (t g ) < 0. Hence, ̇ B(t g ) =− ̇m H (t g ) > 0. So the intersection with the ceiling is transversal. Next, we make a lemma for the openness ofF andG. Lemma D.12. The setsF andG are open subsets of (m L ,m H (0)). Proof. We use continuous dependence of ODE solutions on initial data on compact time intervals. 1. F is open. Take any m 0 ∈F and set t f := τ F (m 0 ). Then t f <∞ and m c (t f ;m 0 ) = m L . By Lemma D.11, the floor crossing is transversal: ̇m c (t f ;m 0 ) < 0. Consider F (t,m) := m c (t;m)−m L . We have F (t f ,m 0 ) = 0 and ∂ t F (t f ,m 0 ) = ̇m c (t f ;m 0 )̸= 0. By the Implicit Function Theorem, there exists a neighborhood U 1 of m 0 and a continuous map m7→ τ F (m) such that for all m∈ U 1 , m c (τ F (m);m) = m L , |τ F (m)− t f | < δ for some δ > 0. Next we show that the ceiling (gap) cannot be hit before τ F (m) for m near m 0 , without appealing to continuity of τ G . Since m 0 ∈F , we have τ G (m 0 ) > t f , hence B(t;m 0 ) < 0 for all t∈ [0,t f ]. Moreover, at t f we have B(t f ;m 0 ) = m L − m H (t f ;m 0 ) < 0. By continuity in t, there exists δ 0 ∈ (0,δ] such that B(t;m 0 ) < 0 for all t∈ [0,t f + δ 0 ]. 48 Published as a conference paper at ICLR 2026 Since [0,t f + δ 0 ] is compact and B(·;m 0 ) is continuous and strictly negative on it, we can define the margin η :=−max t∈[0,t f +δ 0 ] B(t;m 0 ) > 0, so that B(t;m 0 )≤−η for all t∈ [0,t f + δ 0 ]. By continuous dependence on initial data on [0,t f +δ 0 ], there exists a neighborhood U 2 of m 0 such that for all m∈ U 2 , sup t∈[0,t f +δ 0 ] |B(t;m)− B(t;m 0 )| < η/2. Hence for all m∈ U 2 and all t∈ [0,t f + δ 0 ], B(t;m)≤ B(t;m 0 ) + η/2≤−η/2 < 0. In particular, B(t;m) cannot reach 0 before time t f + δ 0 , so τ G (m) > t f + δ 0 . Finally, take m in U := U 1 ∩ U 2 and shrink U 1 if needed so that |τ F (m) − t f | < δ 0 . Then τ F (m) < t f + δ 0 < τ G (m), i.e. m∈F . ThereforeF is open. 2. G is open. Take any m 0 ∈ G and set t g := τ G (m 0 ). Then t g < ∞ and B(t g ;m 0 ) = 0. By Lemma D.11, the ceiling crossing is transversal: ̇ B(t g ;m 0 ) > 0. Let G(t,m) := B(t;m). We have G(t g ,m 0 ) = 0 and ∂ t G(t g ,m 0 ) = ̇ B(t g ;m 0 ) ̸= 0. By the Implicit Function Theorem, there exists a neighborhood V 1 of m 0 and a continuous map m 7→ τ G (m) such that for all m∈ V 1 , B(τ G (m);m) = 0, |τ G (m)− t g | < δ for some δ > 0. Next we show that the floor cannot be hit before τ G (m) for m near m 0 . Since m 0 ∈ G, we have τ F (m 0 ) > t g , hence m c (t;m 0 ) > m L for all t∈ [0,t g ]. Because m c (·;m 0 ) is continuous and [0,t g ] is compact, define ε 0 := min t∈[0,t g ] m c (t;m 0 )− m L > 0, ε := ε 0 /2. Then m c (t;m 0 ) ≥ m L + 2ε for all t ∈ [0,t g ]. By continuity in t at t = t g , there exists δ 0 ∈ (0,δ] such that m c (t;m 0 )≥ m L + ε for all t∈ [0,t g + δ 0 ]. By continuous dependence on initial data on [0,t g +δ 0 ], there exists a neighborhood V 2 of m 0 such that for all m∈ V 2 , sup t∈[0,t g +δ 0 ] |m c (t;m)− m c (t;m 0 )| < ε/2. Hence m c (t;m)≥ m L + ε/2 on [0,t g + δ 0 ], so in particular τ F (m) > t g + δ 0 for all m∈ V 2 . Finally, take m in V := V 1 ∩ V 2 and shrink V 1 if needed so that |τ G (m) − t g | < δ 0 . Then τ G (m) < t g + δ 0 < τ F (m), i.e. m∈G. ThereforeG is open. Next, we make a lemma for the no simultaneous finite exit. Lemma D.13 (No simultaneous finite exit). There is no m 0 ∈ (m L ,m H (0)) such that τ F (m 0 ) = τ G (m 0 ) <∞. Proof. Suppose for contradiction that there exists m 0 ∈ (m L ,m H (0)) and a finite time t ⋆ := τ F (m 0 ) = τ G (m 0 ) <∞. Then by definition, m c (t ⋆ ;m 0 ) = m L and B(t ⋆ ;m 0 ) = 0. Since B(t) = m c (t)− m H (t), the second equality implies m H (t ⋆ ;m 0 ) = m c (t ⋆ ;m 0 ) = m L . 49 Published as a conference paper at ICLR 2026 However, m H (t) = m L = μ 1 /2 holds if and only if p(t) collapses onto the first feature (i.e., p 1 (t) = 1). Under the uniform positive initialization, we have β j (0) > 0 for all j, and by the explicit β-dynamics (Lemma 4.3) we have β j (t) > 0 for every finite t. Hence p j (t) > 0 for all j at every finite time, so p(t) cannot be collapsed at time t ⋆ <∞. This contradicts m H (t ⋆ ) = m L . Therefore, no such m 0 exists. Next, we make the non-emptiness ofF andG. Lemma D.14. The setF is non-empty. Proof. We show that there exists at least one initial condition m 0 ∈ (m L ,m H (0)) that hits the floor before the gap closes. Consider an initial condition m 0 arbitrarily close to m L (with m 0 > m L ). At time t = 0, we have B(0) = m 0 − m H (0). Since m H (0) > m L , the initial gap B(0) is strictly negative and bounded away from 0 as m 0 → m L . Because A(0) = ˆ λ(0)M 1 (0) > 0, we have ̇m c (0) = A(0)B(0) < 0, meaning m c (t) strictly decreases initially. Because the velocity ̇m c (t) is 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 m 0 − m L . Thus, we can make the floor hitting time τ F (m 0 ) arbitrarily small by choosing m 0 sufficiently close to m L . On the other hand, the gap B(t) is a continuous function of time and starts from a strictly negative value B(0) ≈ m L − m H (0) < 0. By continuity, it cannot reach 0 instantly. Therefore, there exists some uniform constant ε 0 > 0 such that for all m 0 sufficiently close to m L , we have B(t) < 0 for all t∈ [0,ε 0 ]. This implies that the ceiling hitting time satisfies τ G (m 0 )≥ ε 0 for all such m 0 . By choosing m 0 > m L close enough to m L , we can guarantee that τ F (m 0 ) < ε 0 ≤ τ G (m 0 ). This strict inequality implies τ F (m 0 ) < τ G (m 0 ), meaning m 0 ∈F . Hence,F ̸=∅. Lemma D.15. The setG 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 an initial condition m 0 = m H (0)−δ with δ > 0 small. Thus, the initial gap is B(0) =−δ, which is strictly negative but arbitrarily close to 0. First, we establish that ̇ B(0) is strictly positive and bounded below for small δ. From the corre- sponding ODE, at t = 0: ̇ B(0) = A(0)B(0)− ̇m H (0) =−A(0)δ− ̇m H (0). We evaluate the − ̇m H (0) term. Under the uniform positive initialization, β j (0) > 0 for all j, meaning the initial distribution p(0) has full support across all features and is not collapsed. Because m H = m D if and only if p collapses onto a single feature, we have a strict inequality at t = 0: m D (0) > m H (0). Using the formula for ̇m H evaluated at t = 0: ̇m H (0) = ˆ λ(0) 2 (M 1 (0)) 2 m 0 2Γ 1 (0) (m 0 − m D (0)). For m 0 ∈ (m H (0)− δ 0 ,m H (0)), we have m 0 < m H (0) < m D (0), which implies m 0 − m D (0) < m H (0) − m D (0) < 0. Furthermore, ˆ λ(0) > 0, M 1 (0) > 0, m 0 > m L > 0, and Γ 1 (0) > 0 (by Cauchy-Schwarz on a non-collapsed distribution p). Therefore, there exists a constant c 0 > 0, independent of sufficiently small δ, such that: − ̇m H (0)≥ c 0 > 0. For δ chosen small enough, the magnitude of the −A(0)δ term becomes strictly less than c 0 /2. Consequently, we obtain ̇ B(0)≥ c 0 /2 > 0. Next, we bound the time required to close the gap versus the time required to hit the floor. Because ̇ B(0) ≥ c 0 /2 > 0 and the vector field is continuous, there exists a short, finite time interval [0,T ] 50 Published as a conference paper at ICLR 2026 over which ̇ B(t) ≥ c 0 /4. Starting from B(0) = −δ, the trajectory must hit the ceiling (i.e., B(t g ) = 0) at some time t g bounded by: t g ≤ δ c 0 /4 = 4δ c 0 . We now control how much the center m c can drop during this interval [0,t g ]. Since ̇m c (t) = A(t)B(t), and knowing that B(t)∈ [−δ, 0] on this interval while A(t) is bounded by some constant A max > 0, the total drop is bounded by: m c (t g ) = m 0 + Z t g 0 ̇m c (t)dt≥ m 0 − Z t g 0 A max δdt = m 0 − A max δt g ≥ m 0 − A max δ 4δ c 0 . This reveals that m c decreases by at most O(δ 2 ) in the time it takes the gap to close. Since the initial position m 0 = m H (0)− δ is separated from the floor m L by anO(1) distance (specifically, m H (0)− m L > 0), we can choose δ sufficiently small such that: m c (t g )≥ m 0 −O(δ 2 ) > m L . Thus, the trajectory hits the ceiling at time t g before it could possibly reach the floor. This implies τ G (m 0 ) = t g < τ F (m 0 ), meaning m 0 ∈G. Hence,G ̸=∅. We have now established the following topological properties: • F andG are open subsets of (m L ,m H (0)) (Lemma D.12). • F ̸=∅ (Lemma D.14) andG ̸=∅ (Lemma D.15). • F ∩G = ∅ by their definitions (a trajectory cannot hit the floor first and the ceiling first simultaneously). • The interval (m L ,m H (0)) is a connected topological space. A connected space cannot be written as the union of two disjoint, non-empty open sets. Therefore, we must have (m L ,m H (0))\ (F ∪G)̸=∅. Let us pick an initial condition m ∗ 0 ∈ (m L ,m H (0))\ (F ∪G). By the definition ofF andG, this means that the trajectory does not hit the floor first in finite time, nor does it hit the gap (ceiling) first in finite time. It implies that: τ F (m ∗ 0 ) = τ G (m ∗ 0 ) =∞. Thus, the trajectory starting at m ∗ 0 never reaches m L and never reaches B = 0 in finite time. Conse- quently, it satisfies: m L < m c (t;m ∗ 0 ) < m H (t;m ∗ 0 ) ∀t≥ 0. This boundary trajectory exactly corresponds to the case where the trajectory reaches neither bound- ary, completing the proof of existence at the m 0 level. Existence of the Regime 2 Threshold α 1 . Finally, we map the boundary initial condition m ∗ 0 back into the initialization scale parameter α. Under the uniform positive initialization, the initial scale determines n θ (0) = √ 2∥μ∥ 2 α, which in turn dictates the initial center m c (0) via: m c (0) = n θ (0) 2ρ = √ 2∥μ∥ 2 2ρ α. Because m c (0) is a continuous and strictly increasing linear function of α, the existence of m ∗ 0 ∈ (m L ,m H (0)) guarantees the existence of a corresponding initialization scale α 1 such that m c (0;α 1 ) = m ∗ 0 . This α 1 is precisely the boundary threshold defining Regime 2. 51 Published as a conference paper at ICLR 2026 Uniqueness of the Regime 2 Threshold α 1 . We now show that within the interval α∈ (α 0 ,α 2 ), there can be at most one initialization scale α 1 for which the corresponding trajectory is a boundary trajectory, i.e. it never reaches either the floor m L or the ceiling m H (t) in finite time. Suppose for contradiction that there exist two distinct boundary initializations α a < α b in (α 0 ,α 2 ). Let (m (a) c (t),p (a) (t)) and (m (b) c (t),p (b) (t)) denote their trajectories. Since both are boundary tra- jectories, they satisfy m L < m (a) c (t) < m (a) H (t) and m L < m (b) c (t) < m (b) H (t) ∀t≥ 0.(18) In particular, B (·) (t) := m (·) c (t) − m (·) H (t) < 0 for all t. Since ̇m c = A(t)B(t) with A(t) = ˆ λ(t)M 1 (t) > 0, we have ̇m (a) c (t) < 0 and ̇m (b) c (t) < 0, so both m (a) c (t) and m (b) c (t) are strictly decreasing and bounded below by m L ; hence they converge. Next, along the boundary corridor we have m c (t) < m H (t)≤ m D (t) (using Proposition D.10), and therefore ̇m H (t) = ˆ λ(t) 2(M 1 (t)) 2 m c (t) 2m c (t)Γ 1 (t)− Γ 2 (t) = ˆ λ(t)Γ 1 (t) (M 1 (t)) 2 m c (t) m c (t)− m D (t) < 0, so m (a) H (t) and m (b) H (t) are also strictly decreasing and bounded below by m L , hence both converge as well. Since m (·) c (t) converges while ̇m (·) c (t) = A (·) (t)B (·) (t) < 0, we must have ̇m (·) c (t) → 0, hence B (·) (t)→ 0 along each boundary trajectory: m (·) c (t)− m (·) H (t)→ 0.(19) Likewise, since m (·) H (t) converges and ̇m (·) H (t) < 0, we must have ̇m (·) H (t) → 0. Because ˆ λ(t) > 0, M 1 (t) > 0, m c (t) > 0, and Γ 1 (t)≥ 0, this forces m (·) c (t)− m (·) D (t)→ 0.(20) Combining equation 19 and equation 20 yields m (·) H (t)− m (·) D (t) → 0. By the equality case of Proposition D.10, this implies that p (·) (t) collapses to a single feature as t→∞. We now rule out collapse to any major coordinate k > 1. If m c (t) → μ k /2 for some k > 1, then the minor feature j = 1 has asymptotic growth rate r 1 (t) = 2μ 1 1− μ 1 2m c (t) → 2μ 1 1− μ 1 μ k > 0, so any “collapse-to-major” state is unstable (the minor coordinate gets re-amplified). Therefore, the only possible boundary limit is the minor-feature collapse equilibrium E 1 : p(t)→ e 1 , m c (t)→ m L = μ 1 2 . In particular, both boundary trajectories (for α a and α b ) must converge to E 1 . Now use the special structure of the uniform initialization. Under uniform positive initialization, the initial proportions p(0) depend only on μ (and not on α), so the two trajectories share identical p(0) and differ only in their initial m c (0): m c (0;α b ) > m c (0;α a ). We examine the stability of E 1 in the m c -direction. At E 1 we have p = e 1 , hence M 1 = μ 1 and m H = m L . Thus the m c -dynamics ̇m c = ˆ λM 1 (m c − m H ) has strictly positive linearization along the m c axis: ∂ ̇m c ∂m c E 1 = ˆ λ(E 1 )M 1 (E 1 ) = ˆ λ(E 1 )μ 1 > 0. 52 Published as a conference paper at ICLR 2026 Hence E 1 is repelling along the m c -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 m c (0), and since the limiting equilibrium E 1 repels along the m c -direction, it is impossible for both trajectories to remain on the boundary corridor equation 18 and still converge to E 1 . Therefore, there is at most one boundary initialization scale α 1 in (α 0 ,α 2 ). Combined with the existence shown above, the threshold α 1 is unique. Since we checked the existence and uniqueness of the boundary threshold α 1 in Regime 1-b and 2-a, we can now conclude that for α < α 1 , the dynamics goes to floor-first behavior(Regime 1) and for α > α 1 , the dynamics goes to ceiling-first behavior(Regime 2). Regimes 2-b and 3. n θ (0) = √ 2∥μ∥ 2 ·α, so α > ρ ∥μ∥ 4 4 √ 2∥μ∥ 2 2 ∥μ∥ 3 3 implies m c (0) = n θ (0) 2ρ > m H (0). When m c (0) > m H (0), we have B(0) > 0 and thus ̇m c (0) = A(0)B(0) > 0, so m c initially increases. We now show that B(t) > 0 for all t ≥ 0. Suppose for contradiction that there exists a first time τ > 0 such that B(τ ) = 0 (i.e., m c (τ ) = m H (τ )). Then ̇ B(τ ) = ̇m c (τ )− ̇m H (τ ) = A(τ )B(τ )− ̇m H (τ ) = 0− ̇m H (τ ) =− ˆ λ(τ ) 2 (M 1 (τ )) 2 m c (τ ) (2m c (τ )Γ 1 (τ )− Γ 2 (τ )). Proposition D.10 gives m D (τ ) ≥ m H (τ ). Therefore, 2m c (τ )Γ 1 (τ )− Γ 2 (τ ) ≤ 0, ensuring that ̇ B(τ ) > 0. Yet, if τ represents the first instance where B reaches zero from above, we must have ̇ B(τ )≤ 0. This contradiction establishes that B(t) > 0 for all t≥ 0, which implies m c (t) > m H (t) for all t≥ 0. Finally, since A(t) = ˆ λM 1 (t) > 0 and B(t) > 0 for all t≥ 0, it follows that ̇m c (t) = A(t)B(t) > 0 for all t ≥ 0, proving that m c (t) is strictly increasing. Additionally, α≷ α 2 implies m c (0)≷ μ d−1 +μ d 2 , ensuring that sub-regime 2-b exhibits the dynamics claimed for Regime 2, and that Regime 3 follows its respective theorem statement. D.5EXTENSION 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 s =⟨β,μ⟩ = D w (1) ⊙w (2) ⊙·⊙w (L) , μ E , where⊙ denotes elementwise (Hadamard) product. The gradient of the loss L with respect to a particular weight w (l) j can be computed via the chain rule: dL dw (l) j = dL ds · ds dw (l) j =−λμ j Y k̸=l w (k) j , where λ is as before, and k ̸= l indicates multiplication over all layers except l. The squared Euclidean norm of the gradient vector∇ θ L is then ∥∇ θ L∥ 2 = d X j=1 L X l=1 dL dw (l) j ! 2 = λ 2 d X j=1 L X l=1 μ 2 j Y k̸=l w (k) j 2 . 53 Published as a conference paper at ICLR 2026 Accordingly, we define n θ := v u u u t d X j=1 L X l=1 μ 2 j Y k̸=l w (k) j 2 . The resulting perturbation is: ε 2 := ρ ∇ θ L ∥∇ θ L∥ 2 , (ε 2 ) w (l) j =− ρμ j n θ Y k̸=l w (k) j . Thus, the perturbed weights are given by ˆw (l) j := w (l) j − ρμ j n θ Y k̸=l w (k) j . The perturbed product then takes the form ˆ β j := L Y l=1 ˆw (l) j . Therefore, the ODE for each coordinate is: ̇w (l) j =− ∂L( ˆ θ) ∂w (l) j = ˆ λμ j Y k̸=l ˆw (k) j . Additionally, we define an assumption on the weight initialization scheme: Assumption D.16. The weights are initialized symmetrically at t = 0, that is, w (1) j (0) = w (2) j (0) = · = w (L) j (0) = w j (0) for all j. Now we show the balancedness-preserving property of the SAM flow. Lemma D.17. Suppose Assumption D.16 holds. Then for all t≥ 0, w (l) j (t) = w j (t) for every l, j. Furthermore, the sign of w j (t) is preserved for all t≥ 0. Proof. Fix j. Assume that at some time t all weights corresponding to j across the layers are equal, i.e., w (1) j (t) = w (2) j (t) =· = w (L) j (t) = w j (t). Then n 2 θ (t) simplifies as follows: n 2 θ (t) = d X j=1 L X l=1 μ 2 j Y k̸=l w (k) j (t) 2 = d X j=1 L X l=1 μ 2 j w j (t) L−1 2 = d X j=1 Lμ 2 j (w j (t)) 2L−2 . 54 Published as a conference paper at ICLR 2026 Therefore, the perturbed weight for each layer l simplifies to: ˆw (l) j (t) = w (l) j (t)− ρμ j n θ (t) Y k̸=l w (k) j (t) = w j (t)− ρμ j n θ (t) w j (t) L−1 , which is independent of l. Hence, ˆw (1) j (t) = ˆw (2) j (t) =· = ˆw (L) j (t) =: ˆw j (t). Substituting this into the SAM flow equation yields: ̇w (l) j (t) = ˆ λ(t)μ j ˆw j (t) L−1 , which is likewise independent of l. Now, for a fixed j, consider the L-dimensional vector u j (t) := w (1) j (t),w (2) j (t),...,w (L) j (t) . The SAM dynamics specify the ODE: ̇u j (t) = F j (u j (t),θ(t)), where F j is the vector whose l-th entry is ˆ λ(t)μ j Q k̸=l ˆw (k) j (t). This ODE is locally Lipschitz in u j , ensuring uniqueness of solutions for given initial conditions. Consider the one-dimensional diagonal manifold D j := (x,...,x)∈ R L : x∈ R . if u j (t) ∈ D j , then ̇u j (t) ∈ D j as well, because all coordinates have the same derivative. SoD j is invariant under the flow. Since the initial condition u j (0) lies inD j due to symmetric initialization, and the ODE solution is unique, we conclude that u j (t)∈D j for all t≥ 0. Therefore, w (l) j (t) = w j (t)for all l, 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 (l) j (t) = w j (t) for all l, j, and t ≥ 0, so the per-coordinate dynamics reduce to ̇w j (t) = ˆ λ(t)μ j w j (t)− ρ μ j n θ (t) w j (t) L−1 L−1 . We claim that the sign of w j (t) is preserved for all t ≥ 0. To see this, observe that the right-hand side of the ODE is a smooth (in fact, polynomial) function of w j , so it is locally Lipschitz in w j for each fixed t. In particular, if at some time τ we have w j (τ ) = 0, then ̇w j (τ ) = 0, so w j (t)≡ 0 for all t≥ τ is a solution with the same initial value. By uniqueness of solutions to ODEs with Lipschitz right-hand side, it follows that once w j reaches zero, it remains identically zero for all future time and cannot cross to the opposite sign. Therefore, if w j (0)̸= 0, the sign of w j (t) is preserved for all t≥ 0 by continuity; if w j (0) = 0, it remains zero. In summary, the sign of w j (t) cannot change during the flow. Utilizing the balancedness-preserving 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 ℓ 2 SAM flow satisfies, for each coordinate j, d dt β j (t) = r (L) j (t)β j (t), where r (L) j (t) = Lμ j β j (t) (1−2/L) 1− ρμ j n θ (t) β j (t) (L−2)/L (L−1) , and β j (t) = w j (t) L , n θ (t) = L d X k=1 μ 2 k w k (t) (2L−2) . Proof. Now define the effective coefficient per coordinate, for general depth L: β j (t) := L Y l=1 w (l) j (t) = w j (t) (L) . Under the balanced ℓ 2 SAM flow, the coordinate dynamics become: ̇ β j (t) = d dt w j (t) L = Lw j (t) (L−1) ̇w j (t) = Lw (L−1) j ˆ λμ j ˆw (L−1) j . We first compute the perturbed weight for coordinate j: ˆw j = w j − ρμ j n θ w L−1 j = w j 1− ρμ j n θ w L−2 j . Substituting this into the expression for ̇ β j (t) gives: ̇ β j (t) = L ˆ λ(t)μ j w 2L−2 j 1− ρμ j n θ (t) w L−2 j L−1 . To express this in terms of β j = w L j , note that w 2L−2 j = β 2−2/L j , w L−2 j = β (L−2)/L j . Therefore, we obtain: ̇ β j (t) = L ˆ λ(t)μ j β j (t) 2−2/L 1− ρμ j n θ (t) β j (t) (L−2)/L L−1 . Absorbing ˆ λ(t) into the time parameter yields the rescaled SAM flow equation: d dt β j (t) = r (L) j (t)β j (t), where r (L) j (t) := Lμ j β j (t) 1−2/L 1− ρμ j n θ (t) β j (t) (L−2)/L L−1 . This provides the Depth-L 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 β j (t) := L Y l=1 w (l) j (t) = w j (t) L , z j (t) := μ j w j (t) L−2 , n 2 θ (t) := L d X k=1 μ 2 k w k (t) (2L−2) , and the critical effective scale: z c (t) := n θ (t) ρL . Then for each time t, we have r (L) j (t) = Lz j (t) 1− ρ n θ (t) z j (t) L−1 =: φ t (z j (t)). The function z 7→ φ t (z) is strictly increasing on (0,z c (t)), strictly decreasing on (z c (t),n θ (t)/ρ), and possesses a unique interior maximum at z = z c (t). In particular, at any fixed t, the coordinate(s) whose effective scale z j (t) is closest to the peak of φ t , i.e., near z c (t), experience the largest instantaneous relative growth rate. Proof. In rescaled SAM time, we have r (L) j (t) = Lμ j β j (t) 1−2/L 1− ρμ j n θ (t) β j (t) (L−2)/L L−1 , where n 2 θ (t) = L d X k=1 μ 2 k w k (t) 2L−2 . Define the effective z-scale by z j (t) := μ j w j (t) L−2 . Note that μ j β (L−2)/L j = μ j w L−2 j = z j . Plugging this yields r (L) j (t) = φ t (z j (t)),where φ t (z) := Lz 1− ρ n θ (t) z L−1 . Define the critical effective scale: z c (t) := n θ (t) ρL . Consider φ t (z) = Lz (1− cz) L−1 , where c = ρ n θ (t) > 0. Its derivative with respect to z is: d dz φ t (z) = L (1− cz) L−2 (1− Lcz), so that: • φ ′ t (z) > 0 for 0 < z < z c (t), • φ ′ t (z) = 0 when z = z c (t), • φ ′ t (z) < 0 for z c (t) < z < n θ (t)/ρ. Therefore, for each fixed t, the function z 7→ φ t (z) is strictly increasing on (0,z c (t)), strictly decreasing on (z c (t),n θ (t)/ρ), and has a unique interior maximum at z = z c (t). 57 Published as a conference paper at ICLR 2026 Unlike the depth-2 case, where each μ j is a fixed constant and their order remains unchanged throughout training, in the depth-L case the effective quantities z j (t) are time-dependent and could, in principle, change order as the SAM flow evolves. However, the following proposition establishes that the order of z j (t) is actually preserved throughout the entire SAM trajectory. Proposition D.20. Under Assumptions D.16 and D.4, the order of the z j (t) is preserved in the depth-L SAM flow. That is, if μ 1 <· < μ d , then z 1 (t) < z 2 (t) <· < z d (t) for all t≥ 0. Proof. We first compute the ODE satisfied by z j (t). By definition, z j = μ j w L−2 j , Taking the time derivative, we get ̇z j = μ j (L− 2)w (L−3) j ̇w j = μ j (L− 2)w (L−3) j ˆ λμ j ˆw (L−1) j Therefore, the perturbed weight is ˆw j = w j 1− ρμ j n θ w (L−2) j . Also, we get w (L−3) j ˆw (L−1) j = w (2L−4) j 1− ρμ j n θ w (L−2) j (L−1) . Using w (L−2) j = z j μ j and w (2L−4) j = z 2 j μ 2 j , we obtain ̇z j = (L− 2) ˆ λμ 2 j z 2 j μ 2 j 1− ρμ j n θ z j μ j (L−1) = (L− 2) ˆ λz 2 j 1− ρz j n θ (L−1) . Thus, the ODE for z j (t) can be expressed as ̇z j (t) = f (t,z j (t)) := (L− 2) ˆ λz j (t) 2 1− ρz j (t) n θ (t) L−1 . Notice that in this expression, the dependence on j appears only through z j (t); both ˆ λ and n θ (t) are time-dependent scalars shared across all coordinates. So each z j (t) solves the same scalar non- autonomous ODE, ̇z(t) = f (t,z(t)), with z(t) = z j (t). Now at t = 0, under symmetric positive init w j (0) = α > 0, we have z j (0) = μ j α L−2 . Since μ 1 < · < μ d and α L−2 > 0, we have z 1 (0) < z 2 (0) < · < z d (0). For this ODE with f is smooth and locally Lipschitz in z, the two different solutions z j (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 z j (t) is preserved for all t ≥ 0. Thus, we have z 1 (t) < z 2 (t) < · < z d (t) for all t≥ 0. D.6PROOFS FOR SECTION 4.2.4 D.6.1DERIVATION OF THE DYNAMICS OF β(t) The dynamics of β(t) = w(t)⊙w(t) is given by ̇ β(t) = ̇ w(t)⊙w(t) +w(t)⊙ ̇ w(t). 58 Published as a conference paper at ICLR 2026 By Equation (3), it is given as ̇ β(t) = 2μ⊙w(t)⊙ w(t)− ρ μ⊙w(t) n θ (t) = 2μ⊙ β(t)− ρ μ⊙β(t) n θ (t) . Coordinate-wise, we have the linear equation ̇ β j (t) = 2μ j β j (t)− ρ μ j β j (t) n θ (t) = 2μ j β j (t) 1− ρ μ j n θ (t) . Therefore, separating variables and integrating, we get ̇ β j (t) β j (t) = 2μ j − 2ρ μ 2 j n θ (t) ⇒ Z t 0 ̇ β j (s) β j (s) ds = Z t 0 2μ j − 2ρ μ 2 j n θ (s) ! ds ⇒ log β j (t) β j (0) = 2μ j t− 2ρμ 2 j Z t 0 1 n θ (s) ds. Define I(t) := R t 0 1 n θ (s) ds. Then, the solution is given by β j (t) = β j (0) exp 2μ j t− 2ρμ 2 j I(t) for j ∈ [d]. D.6.2PROOF 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 a key ingredient in the proof of Theorem 4.5 below. Theorem D.21. Suppose w (1) = w (2) = α ∈ R d . Let (w (1) (t)) t≥0 and (w (2) (t)) t≥0 follow the rescaled ℓ 2 -SAM flow (2) reduced to (3) with perturbation radius ρ and data point μ. Define C μ,α = μ 1 q 2 P d j=1 μ 2 j α 2 j and C μ,α = ∥μ∥ 2 2 √ 2d( Q d j=1 μ j α j ) 1 /d ∥μ∥ 1 . Then, (a) I(t)≥ 1 ρμ 2 1 log 1 ρC μ,α exp(−μ 1 t)+1−ρC μ,α when I(t) t ≥ 1 ρ(μ 1 +μ 2 ) , (b) I(t)≤ d ρ∥μ∥ 2 2 log 1 ρC μ,α exp(− ∥μ∥ 1 d t)+1−ρC μ,α . Proof. From the definition of I(t), I(t) := R t 0 1 n θ (s) ds, we have I ′ (t) = 1 n θ (t) . Since we suppose w (1) (0) = w (2) (0), and the loss function and dynamics are invariant under exchanging w (1) and w (2) , we have w (1) (t) = w (2) (t) =: w(t) for all t≥ 0. From the definition of n θ (t), we have n θ (t) = q ∥μ⊙w (1) (t)∥ 2 2 +∥μ⊙w (2) (t)∥ 2 2 = q 2∥μ⊙w(t)∥ 2 2 = v u u u t 2 d X j=1 μ 2 j w j (t) 2 59 Published as a conference paper at ICLR 2026 = v u u u t 2 d X j=1 μ 2 j β j (t) . From Equation (4), which is β j (t) = β j (0) exp 2μ j t− 2ρμ 2 j I(t) , we have n θ (t) = v u u u t 2 d X j=1 μ 2 j β j (0) exp 2μ j t− 2ρμ 2 j I(t) , and therefore, I ′ (t) = 1 r 2 P d j=1 μ 2 j β j (0) exp 2μ j t− 2ρμ 2 j I(t) . (a) When I(t) t ≥ 1 ρ(μ 1 +μ 2 ) ≥ 1 ρ(μ 1 +μ j ) for j = 2,...,d, it holds that (2μ j t− 2ρμ 2 j I(t))− (2μ 1 t− 2ρμ 2 1 I(t)) = 2(μ j − μ 1 )(t− ρ(μ j + μ 1 )I(t))≥ 0. Therefore, I ′ (t) = 1 q 2 P d j=1 μ 2 j β j (0) exp 2μ j t− 2ρμ 2 j I(t) ≤ 1 q 2 P d j=1 μ 2 j β j (0) exp (2μ 1 t− 2ρμ 2 1 I(t)) = 1 q 2 P d j=1 μ 2 j β j (0) exp (μ 1 t− ρμ 2 1 I(t)) Separating variables and integrating, we get exp(−ρμ 2 1 I(t))dI ≤ 1 q 2 P d j=1 μ 2 j β j (0) exp(−μ 1 t)dt ⇒ Z I(t) I(0) exp(−ρμ 2 1 u)du≤ Z t 0 1 q 2 P d j=1 μ j (s) 2 β j (0) exp(−μ 1 s)ds ⇒− 1 ρμ 2 1 (exp(−ρμ 2 1 I(t))− exp(−ρμ 2 1 I(0)))≤− 1 q 2 P d j=1 μ j (s) 2 β j (0) 1 μ 1 (exp(−μ 1 t)− exp(−μ 1 0)) ⇒ (a) 1 ρμ 2 1 (exp(−ρμ 2 1 I(t))− 1)≥ 1 q 2 P d j=1 μ j (s) 2 β j (0) 1 μ 1 (exp(−μ 1 t)− 1) ⇒ exp(−ρμ 2 1 I(t))≥ ρ μ 1 q 2 P d j=1 μ j (s) 2 β j (0) (exp(−μ 1 t)− 1) + 1 ⇒− ρμ 2 1 I(t)≥ log ρC μ,α (exp(−μ 1 t)− 1) + 1 ⇒I(t)≥ 1 ρμ 2 1 log 1 ρC μ,α exp(−μ 1 t) + 1− ρC μ,α , where (a) holds since I(0) = 0 from the definition of I(t). (b) By AM-GM inequality, we have I ′ (t) = 1 q 2 P d j=1 μ 2 j β j (0) exp 2μ j t− 2ρμ 2 j I(t) 60 Published as a conference paper at ICLR 2026 ≤ 1 r 2d Q d j=1 μ 2 j β j (0) exp 2μ j t− 2ρμ 2 j I(t) 1/d = 1 s 2d Q d j=1 μ 2 j β j (0) 1/d exp 2 P d j=1 μ j d t− 2ρ P d j=1 μ 2 j d I(t) = 1 r 2d Q d j=1 μ 2 j α 2 j 1/d exp 2∥μ∥ 1 d t− 2ρ∥μ∥ 2 2 d I(t) = 1 √ 2d Q d j=1 μ j α j 1/d exp ∥μ∥ 1 d t− ρ∥μ∥ 2 2 d I(t) Separating variables and integrating, we get exp(− ρ∥μ∥ 2 2 d I(t))dI ≤ 1 √ 2d Q d j=1 μ j α j 1/d exp − ∥μ∥ 1 d t dt ⇒ Z I(t) I(0) exp(− ρ∥μ∥ 2 2 d u)du≤ Z t 0 1 √ 2d Q d j=1 μ j α j 1/d exp − ∥μ∥ 1 d s ds ⇒− d ρ∥μ∥ 2 2 (exp(− ρ∥μ∥ 2 2 d I(t))− exp(− ρ∥μ∥ 2 2 d I(0)))≤− 1 √ 2d Q d j=1 μ j α j 1/d d ∥μ∥ 1 (exp(− ∥μ∥ 1 d t)− 1) ⇒ exp(− ρ∥μ∥ 2 2 d I(t))≥ ρ ∥μ∥ 2 2 √ 2d Q d j=1 μ j α j 1/d ∥μ∥ 1 (exp(− ∥μ∥ 1 d t)− 1) + 1 ⇒− ρ ∥μ∥ 2 2 d I(t)≥ log ρC μ,α (exp(− ∥μ∥ 1 d t)− 1) + 1 ⇒I(t)≤ d ρ∥μ∥ 2 2 log 1 ρC μ,α exp(− ∥μ∥ 1 d t) + 1− ρC μ,α ! . Theorem 4.5. Let α 0 ,α 2 be defined in Theorem 4.4 and α 1 be the threshold from there. Suppose α 1 < α≤ ρ μ 1 +μ d √ 2∥μ∥ 2 < α 2 . Then, for j ∈ [d], there exists T j such that β j (T j ) β d (T j ) ≥ LB j (α) := exp 2R ′ j (R j − 1) log 1 1− α 0 /α + log 1 α 0 /α − C(R j ) where R j := (μ j +μ d ) /μ 1 > 2, R ′ j := (μ d −μ j ) /μ 1 and C(R) := R logR− (R− 1) log(R− 1). Proof. By the assumption α 0 < α 1 < α, we have C μ,α = α 0 ρα < 1 ρ . We also have C μ,α = μ 1 √ 2∥μ∥ 2 α ≥ μ 1 √ 2∥μ∥ 2 ρα (2) μ = μ 1 ρ(μ 1 + μ d ) ≥ μ 1 ρ(μ j + μ d ) = 1 ρ R j for all j ∈ [d]. ⇒ 1− ρC μ,α ρC μ,α = 1 ρC μ,α − 1 < R j − 1for all j ∈ [d]. Let T j := 1 μ 1 log ρC μ,α 1−ρC μ,α (R j − 1) ≥ 0. From Theorem D.21, we have I(T j )≥ 1 ρμ 2 1 log 1 ρC μ,α exp(−μ 1 T j ) + 1− ρC μ,α 61 Published as a conference paper at ICLR 2026 = 1 ρμ 2 1 log 1 ρC μ,α exp log 1−ρC μ,α ρC μ,α (R j −1) + 1− ρC μ,α = 1 ρμ 2 1 log 1 1−ρC μ,α R j −1 + 1− ρC μ,α = 1 ρμ 2 1 log 1 (1− ρC μ,α ) 1 + 1 R j −1 = 1 ρμ 2 1 log 1 (1− ρC μ,α ) R j R j −1 = 1 ρμ 2 1 log 1− 1 R j 1− ρC μ,α ! . Recall from Equation (4) that β j (T j ) = β j (0) exp 2μ j T j − 2ρμ 2 j I(T j ) for j ∈ [d]. Thus, for j ∈ [d], we have β j (T j ) β d (T j ) = exp −2(μ d − μ j )T j + 2ρ(μ 2 d − μ 2 j )I(T j ) = exp −2 μ d − μ j μ 1 log ρC μ,α 1− ρC μ,α (R j − 1) + 2ρ(μ 2 d − μ 2 j )I(T j ) ≥ exp −2 μ d − μ j μ 1 log ρC μ,α 1− ρC μ,α (R j − 1) + 2 μ 2 d − μ 2 j μ 2 1 log 1− 1 R j 1− ρC μ,α !! = exp 2 μ d − μ j μ 1 μ d + μ j μ 1 log 1− 1 R j 1− ρC μ,α ! − log ρC μ,α 1− ρC μ,α (R j − 1) !! = exp 2R ′ j R j log 1− 1 R j 1− ρC μ,α ! − log ρC μ,α 1− ρC μ,α (R j − 1) !! = exp 2R ′ j R j log R j −1 R j 1− ρα 0 α − log ρα 0 α 1− ρα 0 α (R j − 1) = exp 2R ′ j (R j − 1) log(R j − 1)− R j log(R j )− (R j − 1) log 1− ρα 0 α − log ρα 0 α = exp 2R ′ j −C(R j )− (R j − 1) log 1− ρα 0 α − log ρα 0 α = exp 2R ′ j (R j − 1) log 1 1− ρα 0 /α + log 1 ρα 0 /α − C(R j ) D.6.3PROOF OF PROPOSITION 4.6 Proposition 4.6. Under the conditions of Theorem 4.5, define j ∗ (α) := arg max j∈[d] LB j (α) and set α ∗ 0 := α 0 . Then, there exist thresholds α ∗ 0 < α ∗ 1 < · < α ∗ m ≤ ρ μ 1 +μ d √ 2∥μ∥ 2 for some m ≤ d− 1 such that j ∗ (α) = j for α∈ (α ∗ j−1 ,α ∗ j ]. 62 Published as a conference paper at ICLR 2026 Proof. For α∈ (α 0 ,ρ μ 1 +μ d √ 2∥μ∥ 2 ), let x = α 0 /α∈ (0, 1) and write G j (x) = log LB j (α) = 2R ′ j Φ R j (x), where Φ R (x) = (R− 1) log 1 1− x + log 1 x − C(R), C(R) = R logR− (R− 1) log(R− 1), and R j = (μ j + μ d )/μ 1 > 1, R ′ j = (μ d − μ j )/μ 1 ≥ 0. (1) Shape of Φ R j . We have Φ ′ R j (x) = R j x− 1 x(1− x) ,Φ ′ R j (x) = R j − 1 (1− x) 2 + 1 x 2 > 0. Thus Φ R j is strictly convex on (0, 1) and attains its unique minimum at x = 1/R j , where Φ R j (1/R j ) = 0. Consequently Φ R j (x) ≥ 0 for all x and it is strictly increasing on [1/R j , 1). (2) Crossing between adjacent indices. For any j ∈1,...,d− 1 define H j+1,j (x) = G j+1 (x)− G j (x) = 2 R ′ j+1 Φ R j+1 (x)− R ′ j Φ R j (x) . Because R j+1 > R j , we have Φ R j+1 (1/R j+1 ) = 0 and Φ R j (1/R j+1 ) > 0, hence H j+1,j (1/R j+1 ) < 0.Likewise Φ R j (1/R j ) = 0 and Φ R j+1 (1/R j ) > 0, giving H j+1,j (1/R j ) > 0. By continuity, H j+1,j has at least one zero x ∗ j ∈ (1/R j+1 , 1/R j ]. To show uniqueness, using the expression for Φ ′ R j , we obtain H ′ j+1,j (x) = 2 x(1− x) (R ′ j+1 R j+1 − R ′ j R j )x− (R ′ j+1 − R ′ j ) . Since R ′ k R k = (μ d − μ k )(μ k + μ d ) μ 2 1 = μ 2 d − μ 2 k μ 2 1 , we obtain R ′ j+1 R j+1 − R ′ j R j = μ 2 j −μ 2 j+1 μ 2 1 < 0. Its zero occurs at x c = R ′ j+1 − R ′ j R ′ j+1 R j+1 − R ′ j R j = μ 1 μ j+1 + μ j , and therefore H ′ j+1,j (x) > 0 for x < x c , H ′ j+1,j (x) < 0 for x > x c . Hence H j+1,j (x) is strictly increasing up to x c and strictly decreasing afterward. Since 1/R j = μ 1 /(μ j + μ d ) ≤ μ 1 /(μ j+1 + μ j ), H j+1,j is strictly increasing in the interval (1/R j+1 , 1/R j ]. Because H j+1,j (1/R j+1 ) < 0 and H j+1,j (1/R j ) > 0, this implies that H j+1,j crosses zero exactly once in (1/R j+1 , 1/R j ). Consequently the root x ∗ j is unique, with H j+1,j (x) < 0 for x < x ∗ j and H j+1,j (x) > 0 for x > x ∗ j . (3) Thresholds and staircase structure. As α increases, x = α 0 /α decreases. Define α ∗ j = α 0 /x ∗ j . When α crosses α ∗ j , the maximizer between indices j and j + 1 switches once from j to j + 1. Because the intervals (1/R j+1 , 1/R j ] are disjoint and ordered, the thresholds satisfy α ∗ 0 < α ∗ 1 < · < α ∗ m ≤ ρ(μ 1 + μ d )/( √ 2∥μ∥ 2 ) for some m≤ d− 1. Thus j ∗ (α) takes constant values on each interval (α ∗ j−1 ,α ∗ j ], increasing step by step until the last threshold within the admissible range. D.6.4PROOF OF PROPOSITION 4.7 Proposition 4.7. Consider α 0 defined in Theorem 4.4. (i) If α < α 0 , then β(t) converges to zero. (i) If α > ρ ∥μ∥ 2 2 √ 2d( Q d i=1 μ i ) 1 /d ∥μ∥ 1 , then β(t) converge in ℓ 1 max-margin direction. 63 Published as a conference paper at ICLR 2026 Proof. We use Theorem D.21 to prove the theorem. When w (1) (0) = w (2) (0) = α1, we have C μ,α = μ 1 q 2 P d j=1 μ 2 j α 2 = μ 1 q 2 P d j=1 μ 2 j α = μ 1 √ 2∥μ∥ 2 α = α 0 α C μ,α = ∥μ∥ 2 2 √ 2d( Q d j=1 μ j α) 1 /d ∥μ∥ 1 = ∥μ∥ 2 2 √ 2d( Q d j=1 μ j ) 1 /d α∥μ∥ 1 (i) By the assumption α≤ α 0 , we have C μ,α = α 0 ρα ≥ 1 ρ . Let T := 1 μ 1 log ρC μ,α ρC μ,α −1 ≥ 0. From Theorem D.21, we have I(t)≥ 1 ρμ 2 1 log 1 ρC μ,α exp(−μ 1 t) + 1− ρC μ,α . As t→ T , we have ρC μ,α exp(−μ 1 t) + 1− ρC μ,α →ρC μ,α exp(−μ 1 T ) + 1− ρC μ,α =ρC μ,α exp(log ρC μ,α − 1 ρC μ,α ) + 1− ρC μ,α =ρC μ,α ρC μ,α − 1 ρC μ,α + 1− ρC μ,α = 0. Since ρC μ,α exp(−μ 1 t) + 1− ρC μ,α is strictly decreasing in t, we have ρC μ,α exp(−μ 1 t) + 1− ρC μ,α → 0 + as t→ T. Therefore, I(t)→ +∞ as t→ T . Recall from Equation (4) that β j (t) = β j (0) exp 2μ j t− 2ρμ 2 j I(t) for j ∈ [d]. As t→ T , we have β j (t)→ 0 for all j ∈ [d] since I(t)→ +∞. Therefore, β(t)→ 0 as t→ T . (i) By the assumption α > ρ ∥μ∥ 2 2 √ 2d( Q d i=1 μ i ) 1 /d ∥μ∥ 1 , we haveC μ,α < 1 ρ . From Theorem D.21, we have I(t)≤ d ρ∥μ∥ 2 2 log 1 ρC μ,α exp(− ∥μ∥ 1 d t) + 1− ρC μ,α ! . For t∈ [0,∞), we have 0 < 1− ρ C μ,α ≤ ρC μ,α exp(− ∥μ∥ 1 d t) + 1− ρC μ,α < 1. and as t→∞, we have ρC μ,α exp(− ∥μ∥ 1 d t) + 1− ρC μ,α → 1− ρC μ,α > 0. As t→∞, we have I(t)≤ d ρ∥μ∥ 2 2 log 1 ρC μ,α exp(− ∥μ∥ 1 d t) + 1− ρC μ,α ! → d ρ∥μ∥ 2 2 log 1 1− ρC μ,α <∞. Therefore, I(t) <∞ as t→∞. 64 Published as a conference paper at ICLR 2026 Recall from Equation (4) that β j (t) = β j (0) exp 2μ j t− 2ρμ 2 j I(t) for j ∈ [d]. Thus, for j ∈ [d], we have β j (t) β d (t) = exp −2(μ d − μ j )t + 2ρ(μ 2 d − μ 2 j )I(t) . As t→∞, we have β j (t) β d (t) → 0 for all j < d since lim t→∞ I(t) <∞. Therefore, β(t) converges to the direction of e d as t→∞. 65 Published as a conference paper at ICLR 2026 D.7NUMERICAL EVALUATION OF THEOREM 4.5 In this section, we provide numerical illustrations of the lower bound LB j (α) derived in Theo- rem 4.5. For several choices of μ, we compute the value of LB j (α) := exp 2R ′ j (R j − 1) log 1 1− α 0 /α + log 1 α 0 /α − C(R j ) and visualize how much the ratio β j (t)/β d (t) must be amplified at minimum. Figure 14 shows that for small α in Regime 2 and for μ with a large spectral gap μ d /μ 1 , LB j (α) easily exceeds 10. Since this is only a lower bound, the actual amplification can be even larger, indicating that minor-to-intermediate coordinates can grow by substantially more than the major coordinate. 0.350.400.450.500.550.60 1.00 1.25 1.50 1.75 2.00 2.25 2.50 2.75 LB j ( ) j = 1 j = 2 j = 3 j = 4 j = 5 (a) μ = (4, 5, 6, 7, 8) 0.350.400.450.500.55 2 4 6 8 LB j ( ) j = 1 j = 2 j = 3 j = 4 j = 5 (b) μ = (1, 2, 3, 4, 5) 0.400.450.500.55 0 5 10 15 20 25 LB j ( ) j = 1 j = 2 j = 3 j = 4 j = 5 (c) μ = (1, 3, 5, 7, 9) 0.400.450.500.550.600.65 0 2000 4000 6000 8000 10000 LB j ( ) j = 1 j = 2 j = 3 j = 4 j = 5 (d) μ = (1, 2, 4, 8, 16) Figure 14: Numerical evaluation of LB j (α) for various choices of μ. For reproducibility, we describe the numerical procedure used to generate Figure 14. For each choice ofμ (withd = dim(μ)), we evaluate LB j (α) for allj ∈ [d] on a uniform grid ofα values. Following the assumptions of Theorem 4.5, we first obtain the threshold α 1 specified in Theorem 4.4. We then set α∈ h α 1 ,ρ μ 1 +μ d √ 2∥μ∥ 2 i using 400 grid points. The quantities α 0 , R j , R ′ j , and C(R j ) are computed directly from their definitions in Theorems 4.4 and 4.5 using the given μ. The index j ∈ [d] corresponds to the coordinate ordering μ 1 < · < μ d . Since the computation is closed-form, no randomness is involved and the plots are exactly reproducible. D.8EMPIRICAL VERIFICATION Our analysis in Section 4.2 focuses on the one-point settingD μ . We begin by verifying that the se- quential feature amplification occurs across multiple choices of μ in this one-point regime: both the continuous-time rescaled flows and the discrete ℓ ∞ -SAM updates exhibit the same coordinate-wise progression, and the loss dynamics follow the theoretical prediction. We then turn to multi-point datasets and show that the sequential feature amplification persists in this more realistic setting un- der both the rescaled ℓ 2 -SAM flow and discrete ℓ 2 -SAM updates, as illustrated in Figure 11. Finally, 66 Published as a conference paper at ICLR 2026 we confirm that this phenomenon is not limited to depth 2; the same coordinate-wise progression arises in deeper diagonal networks (general depth L). Taken together, these results demonstrate that the sequential feature amplification is a robust and widely recurring behavior: it appears consistently across different μ, across multiple multi-point datasets, across both continuous and discrete SAM dynamics, and across depths L≥ 2. To clarify the heatmap visualizations (e.g., Figures 3a and 15 to 23), for each time t and initialization scaleα, we computej † = arg min j β j (t) and color the grid point (t,α) according to this index. Grid regions where the predictor β becomes negligibly small are shown in gray, indicating convergence toward 0. We use the threshold∥β(t)∥ 2 ≤ 10 −2 to define gray regions. Following the visualization style of Figure 3a, we also partition the α–axis into the three regimes defined in Theorem 4.4: Regime 1 (small α), Regime 2 (intermediate α), and Regime 3 (large α). These regime boundaries are indicated by horizontal black dashed lines in heatmap figures. For reproducibility, we detail the exact initialization used in all experiments. As mentioned in Sec- tion 4.2, we adopt a uniform initialization across coordinates and layers: w (1) (0) = w (2) (0) = α1 for depth-2 setup andw (1) (0) =· = w (L) (0) = α1 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 a step size of η = 0.01. D.8.1ONE-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 ℓ 2 -SAM flow, we additionally evaluate discrete ℓ 2 -SAM updates. Across all tested choices ofμ, 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 β(t), we also provide the loss heatmaps over (α,t). In the discrete ℓ 2 -SAM case, Regime 1 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 0. As a result, the predictor norm stays above the gray threshold—so it is not colored gray—yet the loss remains large, revealing that the trajectory is 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, 8). The behavior is similar across different choices of μ. Both GF and GD consistently recover the major feature, independent of the initialization scale α, and they do not exhibit sequential feature amplifi- cation. 0.00.20.40.60.81.0 time t 0.0 0.2 0.4 0.6 0.8 initialization scale 1 2 3 4 5 j = arg max j j ( t ) 0.00.20.40.60.81.0 time t 0.0 0.2 0.4 0.6 0.8 initialization scale 0.0 0.1 0.2 0.3 0.4 0.5 0.6 logistic loss (a) GF 0.00.20.40.60.81.0 time t 0.0 0.2 0.4 0.6 0.8 initialization scale 1 2 3 4 5 j = arg max j j ( t ) 0.00.20.40.60.81.0 time t 0.0 0.2 0.4 0.6 0.8 1.0 1.2 initialization scale 0.0 0.1 0.2 0.3 0.4 0.5 0.6 logistic loss (b) Discrete GD (η = 0.01) Figure 15: Dominant index j † over α,t and logistic loss onD μ with μ = (4, 5, 6, 7, 8). 67 Published as a conference paper at ICLR 2026 1. μ = (4, 5, 6, 7, 8) 0.00.20.40.60.81.0 time t 0.0 0.2 0.4 0.6 0.8 initialization scale 1 2 3 4 5 j = arg max j j ( t ) 0.00.20.40.60.81.0 time t 0.0 0.2 0.4 0.6 0.8 initialization scale 0.0 0.1 0.2 0.3 0.4 0.5 0.6 logistic loss (a) Rescaled ℓ 2 -SAM flow 0.00.20.40.60.81.0 time t 0.0 0.2 0.4 0.6 0.8 initialization scale 1 2 3 4 5 j = arg max j j ( t ) 0.00.20.40.60.81.0 time t 0.0 0.2 0.4 0.6 0.8 initialization scale 0.0 0.1 0.2 0.3 0.4 0.5 0.6 logistic loss (b) Discrete ℓ 2 -SAM updates (η = 0.01) Figure 16: Dominant index j † over α,t and logistic loss onD μ with μ = (4, 5, 6, 7, 8) and ρ = 1. 2. μ = (1, 2, 3, 4, 5) 0.00.51.01.5 time t 0.0 0.2 0.4 0.6 0.8 initialization scale 1 2 3 4 5 j = arg max j j ( t ) 0.00.51.01.5 time t 0.0 0.2 0.4 0.6 0.8 initialization scale 0.0 0.1 0.2 0.3 0.4 0.5 0.6 logistic loss (a) Rescaled ℓ 2 -SAM flow 0.00.51.01.5 time t 0.0 0.2 0.4 0.6 0.8 initialization scale 1 2 3 4 5 j = arg max j j ( t ) 0.00.51.01.5 time t 0.0 0.2 0.4 0.6 0.8 initialization scale 0.0 0.1 0.2 0.3 0.4 0.5 0.6 logistic loss (b) Discrete ℓ 2 -SAM updates (η = 0.01) Figure 17: Dominant index j † over α,t and logistic loss onD μ with μ = (1, 2, 3, 4, 5) and ρ = 1. 68 Published as a conference paper at ICLR 2026 3. μ = (1, 3, 5, 7, 9) 0.00.51.01.5 time t 0.0 0.2 0.4 0.6 0.8 initialization scale 1 2 3 4 5 j = arg max j j ( t ) 0.00.51.01.5 time t 0.0 0.2 0.4 0.6 0.8 initialization scale 0.0 0.1 0.2 0.3 0.4 0.5 0.6 logistic loss (a) Rescaled ℓ 2 -SAM flow 0.00.51.01.5 time t 0.0 0.2 0.4 0.6 0.8 initialization scale 1 2 3 4 5 j = arg max j j ( t ) 0.00.51.01.5 time t 0.0 0.2 0.4 0.6 0.8 initialization scale 0.0 0.1 0.2 0.3 0.4 0.5 0.6 logistic loss (b) Discrete ℓ 2 -SAM updates (η = 0.01) Figure 18: Dominant index j † over α,t and logistic loss onD μ with μ = (1, 3, 5, 7, 9) and ρ = 1. 4. μ = (1, 2, 4, 8, 16) 0.00.51.01.52.0 time t 0.0 0.2 0.4 0.6 0.8 1.0 initialization scale 1 2 3 4 5 j = arg max j j ( t ) 0.00.51.01.52.0 time t 0.0 0.2 0.4 0.6 0.8 1.0 initialization scale 0.0 0.1 0.2 0.3 0.4 0.5 0.6 logistic loss (a) Rescaled ℓ 2 -SAM flow 0.00.51.01.52.0 time t 0.0 0.2 0.4 0.6 0.8 1.0 initialization scale 1 2 3 4 5 j = arg max j j ( t ) 0.00.51.01.52.0 time t 0.0 0.2 0.4 0.6 0.8 1.0 initialization scale 0.0 0.1 0.2 0.3 0.4 0.5 0.6 logistic loss (b) Discrete ℓ 2 -SAM updates (η = 0.01) Figure 19: Dominant index j † over α,t and logistic loss onD μ with μ = (1, 2, 4, 8, 16) and ρ = 1. 69 Published as a conference paper at ICLR 2026 D.8.2MULTI-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 at +μ and−μ for various choices of μ. Specifically, we draw x (+) n = μ +ε n , y n = +1, x (−) n =−μ +ε n , y n =−1, with ε n ∼ N (0,σ 2 I d ) and use N/2 samples per class (with μ = (1, 2),N = 100,σ = 0.5). For visualization, we plot only the first two dimensions of the dataset in the left panels. The middle panels show the results of the rescaled ℓ 2 -SAM flow on this dataset, and the right panels show the discrete ℓ 2 -SAM updates. Across all choices of multi-point datasets, the same sequential feature amplification behavior observed in the one-point setting persists. For comparison, we 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 μ. As in the one-point setting, both GF and GD consistently recover the major feature, independent of the initialization scale α, and they do not exhibit sequential feature amplification. 5.02.50.02.55.0 x 1 6 4 2 0 2 4 6 x 2 O +1 class -1 class 0.00.20.40.60.81.0 time t 0.0 0.2 0.4 0.6 0.8 initialization scale GF 1 2 3 4 5 j = arg max j j ( t ) 0.00.20.40.60.81.0 time t 0.0 0.2 0.4 0.6 0.8 initialization scale Discrete GD Updates 1 2 3 4 5 j = arg max j j ( t ) Figure 20: First two dimensions ofD μ with μ = (4, 5, 6, 7, 8) and the dominant index j † over α,t under GF and discrete GD updates. 70 Published as a conference paper at ICLR 2026 5.02.50.02.55.0 x 1 6 4 2 0 2 4 6 x 2 O +1 class -1 class 0.00.20.40.60.81.0 time t 0.0 0.2 0.4 0.6 0.8 initialization scale Rescaled 2 -SAM Flow 1 2 3 4 5 j = arg max j j ( t ) 0.00.20.40.60.81.0 time t 0.0 0.2 0.4 0.6 0.8 initialization scale Discrete 2 -SAM Updates 1 2 3 4 5 j = arg max j j ( t ) (a) μ = (4, 5, 6, 7, 8) 21012 x 1 3 2 1 0 1 2 3 x 2 O +1 class -1 class 0.00.51.01.5 time t 0.0 0.2 0.4 0.6 0.8 initialization scale Rescaled 2 -SAM Flow 1 2 3 4 5 j = arg max j j ( t ) 0.00.51.01.5 time t 0.0 0.2 0.4 0.6 0.8 initialization scale Discrete 2 -SAM Updates 1 2 3 4 5 j = arg max j j ( t ) (b) μ = (1, 2, 3, 4, 5) 21012 x 1 4 2 0 2 4 x 2 O +1 class -1 class 0.00.51.01.5 time t 0.0 0.2 0.4 0.6 0.8 initialization scale Rescaled 2 -SAM Flow 1 2 3 4 5 j = arg max j j ( t ) 0.00.51.01.5 time t 0.0 0.2 0.4 0.6 0.8 initialization scale Discrete 2 -SAM Updates 1 2 3 4 5 j = arg max j j ( t ) (c) μ = (1, 3, 5, 7, 9) 21012 x 1 3 2 1 0 1 2 3 x 2 O +1 class -1 class 0.00.51.01.52.0 time t 0.0 0.2 0.4 0.6 0.8 1.0 initialization scale Rescaled 2 -SAM Flow 1 2 3 4 5 j = arg max j j ( t ) 0.00.51.01.52.0 time t 0.0 0.2 0.4 0.6 0.8 1.0 initialization scale Discrete 2 -SAM Updates 1 2 3 4 5 j = arg max j j ( t ) (d) μ = (1, 2, 4, 8, 16) Figure 21: First two dimensions ofD μ and the dominant index j † over α,t under the rescaled ℓ 2 - SAM flow and discrete ℓ 2 -SAM updates. 71 Published as a conference paper at ICLR 2026 D.8.3DEPTH-L CASE: PERSISTENCE OF DEPTH-2 DYNAMICSS We confirm that the sequential feature amplification is not limited to depth L = 2; the same coordinate-wise progression arises in deeper diagonal networks (general depth L). Specifically, we observe GF and rescaled ℓ 2 -SAM flow on the one-point dataset D μ with μ = (4, 5, 6, 7, 8). The behavior remains similar across different choices of μ, multi-point datasets, and under discrete updates. While GF appears to exhibit Regime 1 (being trapped near the origin), it does not show the sequential feature amplification, even in the deeper models. However, the rescaled ℓ 2 -SAM flow clearly demonstrates the sequential feature amplification for general depth L. Even though Regime 1 appears chaotic, Regime 2 and 3 are distintcly observed. Thus, the sequential feature amplification robustly occurs not only at depth L = 2 but also in deeper models. 0.00.20.40.60.81.0 time t 0.0 0.2 0.4 0.6 0.8 1.0 1.2 initialization scale 1 2 3 4 5 j = arg max j j ( t ) (a) Depth 2 0.00.20.40.60.81.0 time t 0.0 0.2 0.4 0.6 0.8 1.0 1.2 initialization scale 1 2 3 4 5 j = arg max j j ( t ) (b) Depth 3 0.00.20.40.60.81.0 time t 0.0 0.2 0.4 0.6 0.8 1.0 1.2 initialization scale 1 2 3 4 5 j = arg max j j ( t ) (c) Depth 4 0.00.20.40.60.81.0 time t 0.0 0.2 0.4 0.6 0.8 1.0 1.2 initialization scale 1 2 3 4 5 j = arg max j j ( t ) (d) Depth 5 Figure 22: Dominant index j † over α,t under the GF onD μ with μ = (4, 5, 6, 7, 8). 72 Published as a conference paper at ICLR 2026 0.00.20.40.60.81.0 time t 0.0 0.2 0.4 0.6 0.8 1.0 1.2 initialization scale 1 2 3 4 5 j = arg max j j ( t ) (a) Depth 2 0.00.20.40.60.81.0 time t 0.0 0.2 0.4 0.6 0.8 1.0 1.2 initialization scale 1 2 3 4 5 j = arg max j j ( t ) (b) Depth 3 0.00.20.40.60.81.0 time t 0.0 0.2 0.4 0.6 0.8 1.0 1.2 initialization scale 1 2 3 4 5 j = arg max j j ( t ) (c) Depth 4 0.00.20.40.60.81.0 time t 0.0 0.2 0.4 0.6 0.8 1.0 1.2 initialization scale 1 2 3 4 5 j = arg max j j ( t ) (d) Depth 5 Figure 23:Dominant index j † over α,t under the rescaled ℓ 2 -SAM flow on D μ with μ = (4, 5, 6, 7, 8) and ρ = 1. 73 Published as a conference paper at ICLR 2026 EEXPERIMENTS E.1LOSS 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. In this experiment, we train the diagonal linear network with full-batch SAM using radius ρ = 0.5, learning rate 0.05, and 10000 epochs. We fix the initialization scale to α = 0.06 as a representative intermediate value. The data vector is μ = (1, 2, 3, 4, 5, 6), and all other settings follow the default diagonal-network configuration. To make this precise, we track the dominant index arg max j r j (t), where r j (t) denotes the growth rate of β j (t). In the early phase, this dominant index corresponds to minor features (coordinates with small μ j ), while in the later phase it switches to major features (coordinates with larger μ j ). When SAM is focusing on minor features, the loss decreases slowly, leading to a plateau; once SAM 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. Figure 24: Training loss curves of GD (top) and SAM (bottom) on the 2-layer diagonal network in the intermediate initialization regime (α = 0.06). The colored areas correspond to 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 to slow loss decrease. Once SAM shifts its focus to major coordinates, the loss drops rapidly. GD does not display this minor-to-major feature focusing behavior, thereby showing a more steadily decreasing loss without such a plateau. E.2SEQUENTIAL FEATURE AMPLIFICATION UNDER RANDOM INITIALIZATION In the main analysis, we focused on a symmetric and layer-wise balanced initialization to obtain a 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 as w (1) (0),w (2) (0)∼N (0,α 2 I), where the parameter α controls the initialization scale as the standard deviation of the Gaussian distribution. Figure 25a shows the normalized coordinate trajectories β j (t)/∥β(t)∥ 2 under random initialization (Seed 0) for α = 0.65, μ = (1, 2, 3, 4, 5, 6), and ρ = 0.1. In this case, all coordinates except 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 (1) (t)− w (2) (t)∥ 2 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 the 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 a 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. 0246810 time t 0.5 0.0 0.5 1.0 j ( t )/ ( t ) Trajectory of normalized coordinates 1 (t)/(t) 2 (t)/(t) 3 (t)/(t) 4 (t)/(t) 5 (t)/(t) 6 (t)/(t) (a) Normalized trajectories β j (t)/∥β(t)∥ 2 (Seed 0) 0246810 time t 0.00 0.05 0.10 0.15 0.20 0.25 w (1) ( t ) w (2) ( t ) 2 Norm of w (1) w (2) (b) Evolution of∥w (1) (t)−w (2) (t)∥ 2 (Seed 0) 012345 time t 0.5 0.0 0.5 1.0 j ( t )/ ( t ) Trajectory of normalized coordinates 1 (t)/(t) 2 (t)/(t) 3 (t)/(t) 4 (t)/(t) 5 (t)/(t) 6 (t)/(t) (c) Normalized trajectories β j (t)/∥β(t)∥ 2 (Seed 1) 012345 time t 0.00 0.05 0.10 0.15 0.20 0.25 w (1) ( t ) w (2) ( t ) 2 Norm of w (1) w (2) (d) Evolution of∥w (1) (t)−w (2) (t)∥ 2 (Seed 1) 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.3ALTERNATIVE 2-LAYER MODELS To evaluate the generality of our theoretical predictions, we conduct experiments on alternative 2- layer models featuring different parameterizations and metrics. In 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.1LINEAR NETWORK We fix a small matrix dimension d = 5. All inputs are d×d matrices. We first draw a single random “signal” matrix μ ∈ R d×d with i.i.d. standard normal entries, and then compute its singular value decomposition (SVD) μ = U μ diag(S μ )V ⊤ μ . From this SVD, we construct an orthonormal basis of rank-1 matrices μ i = u i v ⊤ i , i = 1,...,d, where u i is the i-th column of U μ and v ⊤ i is the i-th row of V ⊤ μ . These μ i play the role of “feature directions”, analogous to 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+μ,−μ with opposite labels+1,−1. The 2-layer linear network is f θ (X) = ⟨β,X⟩ F = ⟨W (1) W (2) ,X⟩ F , with learnable matrices W (1) ,W (2) ∈ R d×d and effective weight β = W (1) W (2) . Each layer is initially set to the identity matrix, and before training we rescale all layers by a scalar α, so that W (1) (0) = W (2) (0) = αI and hence β(0) = α 2 I . For training, we use full-batch SAM with radius ρ = 0.5, learning rate 0.05, and a finite train- ing epochs of T = 5000. We repeat the experiment over a range of initialization scales, α ∈ 0.20, 0.21,..., 0.70. As our tracking metric, we monitor the normalized squared alignment a i (t) = ⟨β(t),μ i ⟩ 2 F ∥β(t)∥ 2 F , i = 1,...,d, where β(t) denotes the effective weight at training iteration t. 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, as the initialization scale increases, the dominant singular direction that maximizes the alignment (i.e., arg max i a i (T )) moves from σ 5 to σ 1 , indicating that SAM sequentially aligns from the minor component to the major component as α grows. Figure 26: Alignment of the effective weight β(t) for GD (left) and SAM (right) across initialization scales. 76 Published as a conference paper at ICLR 2026 E.3.2CONVOLUTIONAL NEURAL NETWORK We consider a 2-layer linear convolutional network trained on a synthetic dataset built from a single image matrix μ. This experiment is designed to probe frequency-wise feature selection under SAM. We fix an image size d = 32 and construct a single base image μ ∈ R 1×d×d as a sum of cosine plane waves with radial frequencies: μ(x,y) = K X k=1 w k L k X l=1 cos wπr k xcosθ k,l + y sinθ k,l d + φ k,l , The experiment uses K = 5 different frequency bands, where r k are target bands, w k > 0 are band weights, and θ k,l , φ k,l are random orientations and phases for each band. We take r k ∈ 3, 9, 11, 13, 15 and w k = 1.0, 2.0, 3.0, 4.0, 5.0 for all k. We set L k = 8 for all k. We then renormalize μ to have unit euclidean norm, then shift it slightly to be strictly positive. Next, we define the frequency bands by constructing radial masks M k ⊂ 0,· ,d − 1 2 in the fourier domain. Let ˆμ denote the 2D FFT of μ. The band energy of μ at band k is then given by μ k = X m∈M k |ˆμ(m)| 2 . The bands are sorted by μ k . As we apply low weights to low frequency bands when constructing μ, in this setting, low frequency bands have smaller μ k and treated as minor features, and high frequency bands have larger μ k and treated as major features. The utilized model is a depth-2 convolutional network without nonlinearities. For the first convo- lutional layer, we use 3× 3 convolution with 32 output channels, stride 1, and padding 1. For the second convolutional layer, we use same size of kernel, channel size, stride, and padding. We used realistic gaussian initialization for the weights of the convolutional layers. The weights for each layer are independently initialized. Lastly, the final FC layer is a linear layer. the input for fc layer is squeezed 1d vector, and the output is a 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, we require a band-wise de- composition of the effective weight β(θ) in input space. Since the network is linear, β(θ) can be recovered from gradients. At a given parameter vector θ, we consider the empirical margin s(θ) = E (x,y) [yf θ (x)] = 1 2 (f θ (μ)− f θ (−μ)). We compute the gradient of s(θ) with respect to the input and form a “virtual gate” version of β in input space: ∇ x s(θ)| x,y = y (∇ x f θ (x)). So, β map (u,v) = E (x,y) h (∇ x f θ (x)⊙ x) u,v i , which is proportional to (β(θ)⊙μ) u,v in our linear setting. In practice, this expectation is computed exactly by averaging over x∈μ,−μ. We then take the 2D FFT of β map , denoted b β map , and define the band-wise effective weights by β k (θ) = X m∈M k b β map (m) 2 . For each training epoch t we record the vector (β 1 (θ t ),...,β K (θ t )), 77 Published as a conference paper at ICLR 2026 and, in particular, the index of the dominant band k dom (t) = arg max k β k (θ t ). In our initialization-scale experiments, we repeat this procedure over a range of α ∈ [0.13, 0.20] and, for each α, track both the dominant band k dom at the end of training. This provides a CNN analogue of the feature-selection behavior observed in the diagonal model, where coordinates are replaced by frequency bands. Figure 27 displays how the final dominant frequency band selected by the CNN varies with the initialization scale α. Consistent with expectations, when trained with SAM, the model emphasizes minor features (i.e., low frequency bands) for small α, and shifts its focus to major features (high frequency bands) as α increases. In contrast, under standard GD, the dominant frequency band remains unchanged regardless of the initialization scale. 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-frequency) bands as α increases, whereas GD remains insensitive to α. E.4GRAD-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 α, SAM collapses toward the origin; for large α, SAM behaves similarly to GD; and for intermediate α, SAM preferentially amplifies minor to 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 a threshold (0.5) and plot this quantity as a function of the initialization scale α. 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, we 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 on 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.1MNIST We first study this phenomenon on MNIST. MNIST has a simple structure, where the black back- ground takes the minimum pixel value (0) and the white digit takes the maximum pixel value (1). We construct a subset of 1,000 images whose labels are in 0, 1, 2, 3 and train models using either GD or ℓ 2 -SAM. After training, we visualize the learned attention patterns using Grad-CAM, 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–trained model concentrates more strongly on the black background region. Unless otherwise stated, we use a learning rate of 0.1, a SAM perturbation radius of 0.5, and train for 500 epochs with a batch size of 64. We use no momentum and no weight decay. For the CNN architecture, we use 3× 3 convolutional kernels and do not apply batch normalization or layer normalization. (a) GD(b) ℓ 2 -SAM Figure 28: Grad-CAM comparison between GD and ℓ 2 -SAM on MNIST (labels 0–3). To study the practical behavior of ℓ ∞ -SAM, we train models using ℓ ∞ -SAM on a subset of 1,000 MNIST images with labels in0, 1. We then visualize the Grad-CAM maps, as shown in Figure 29. We observe a bias pattern similar to that of ℓ 2 -SAM, where the model places greater emphasis on 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 a batch size of 64. Input 1Grad-CAMOverlayInput 1Grad-CAMOverlay Input 0Grad-CAMOverlay (a) GD Input 0Grad-CAMOverlay (b) ℓ ∞ -SAM Figure 29: Grad-CAM comparison between GD and ℓ ∞ -SAM on MNIST (labels 0–1). We now quantify the average values of activated pixels (Grad-CAM > 0.5) as a function of the initialization scale α across different dataset subsets. In this experimental setup (Figure 30), we observe that GD consistently concentrates more on the white digit region, which can be interpreted as the major component in 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, we 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 on the pixels whose average is almost 0, so 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. 024 0.0 0.2 0.4 0.6 0.8 avg(activated pixels) GD SAM (a) MNIST with labels 0,1,2,3. 0.51.01.52.02.5 0.0 0.2 0.4 0.6 avg(activated pixels) GD SAM (b) 1k MNIST images with labels 0 and 1. s 0.51.01.52.02.5 0.0 0.2 0.4 0.6 0.8 avg(activated pixels) GD SAM (c) 1k MNIST images with labels 0,1,2,3. 0.51.01.52.02.5 0.0 0.2 0.4 0.6 avg(activated pixels) GD SAM (d) Full MNIST 1k subset. Figure 30: Average number of pixels with Grad-CAM activation exceeding 0.5 as a function of the initialization scale α, comparing GD and ℓ 2 -SAM across different MNIST subsets. ℓ ∞ -SAM exhibits a similar pattern (Figure 31). When α is small, the dynamics collapse toward the origin. For intermediate values of α, ℓ ∞ -SAM tends to prioritize minor features more strongly than GD. For sufficiently large α, however, the behavior of ℓ ∞ -SAM deviates from this trend. 0.51.01.52.0 0.0 0.2 0.4 0.6 avg(activated pixels) GD SAM Figure 31: Average number of pixels with Grad-CAM activation exceeding 0.5 as a function of the initialization scale α, comparing GD and ℓ ∞ -SAM on 1k MNIST images with labels 0 and 1. 80 Published as a conference paper at ICLR 2026 E.4.2SVHN We next study this phenomenon on SVHN. SVHN is more complex than MNIST, as it contains both images with dark backgrounds and light digits, as well as images with light backgrounds and dark digits. Nevertheless, we observe that ℓ 2 -SAM consistently emphasizes the darker regions of the image. We construct a subset of 1,000 images with labels in0, 1 and train models using either GD or ℓ 2 - SAM. We use a learning rate of 0.01, a 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 on the digit regions than the background, as 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. (a) GD(b) ℓ 2 -SAM Figure 32: Grad-CAM comparison between GD and ℓ 2 -SAM on SVHN (1k images, labels 0–1) with dark digits and light backgrounds. 81 Published as a conference paper at ICLR 2026 (a) GD(b) ℓ 2 -SAM Figure 33: Grad-CAM comparison between GD and ℓ 2 -SAM on SVHN (1k images, labels 0–1) with light digits and dark backgrounds. Across different values of α, we observe that small α causes ℓ 2 -SAM to collapse toward the origin, while intermediate α leads ℓ 2 -SAM to emphasize minor features with lower pixel intensities as shown in Figure 34, where pixel intensity is computed as the average over the three color channels. 0123 0.0 0.2 0.4 0.6 avg(activated pixels) GD SAM Figure 34: Average number of activated pixels (Grad-CAM > 0.5) as a function of the initialization scale α, comparing GD and ℓ 2 -SAM. E.4.3CIFAR-10 We also observe the same phenomenon on the CIFAR-10 dataset. We construct a subset of CIFAR- 10 with labels in0, 1 and train models using a learning rate of 0.01, a SAM perturbation radius of 0.05, for 500 epochs. As shown in Figure 35, small values of α lead SAM to emphasize minor features, while larger values of α make the behaviors of GD and SAM increasingly similar. 024 0.0 0.2 0.4 0.6 avg(activated pixels) GD SAM Figure 35: Average number of activated pixels (Grad-CAM > 0.5) as a function of the initialization scale α, comparing GD and ℓ 2 -SAM. 82