PaperlayerView exact revision on arXiv

Source-linked AI summary

Beyond the Edge of Stability via Two-step Gradient Updates

Lei Chen, Joan Bruna

arXiv:2206.04172v3cs.LGmath.OCstat.ML

TL;DR

The paper asks when gradient descent can converge beyond local-stability guarantees, where curvature may exceed the usual threshold. It analyzes two-step updates and representative learning problems, establishing conditions for fixed points and reporting teacher alignment and period-2 behavior.

  • Problem

    The paper asks what convergence guarantees GD has when learning rates exceed stability thresholds, including which minimizers it reaches on overparameterized objectives.

    Source: §1 · p. 1, §4.1 · p. 5

  • Method

    The paper analyzes two-step GD updates and higher-order derivative conditions, then applies the analysis to a teacher-student network and matrix factorization.

    Source: §1 · p. 2, §7 · pp. 13–14

  • Results

    Higher-order conditions establish two-step fixed points above the local threshold; the paper also proves teacher-direction alignment and reports matrix-factorization period-2 orbits.

    Source: §1 · p. 2, §7 · p. 14

  • Takeaways & Limitations

    The analyses show that beyond-EoS behavior can be studied in these settings through two-step fixed points and period-2 dynamics.

    Source: §7 · pp. 13–14

    The matrix-factorization observation is limited to a quasi-symmetric version because rigorous proof is difficult and general initialization does not illustrate the sharper-to-flatter behavior well.

    Source: §6.1 · p. 10

Abstract

from arXiv · show

Gradient Descent (GD) is a powerful workhorse of modern machine learning thanks to its scalability and efficiency in high-dimensional spaces. Its ability to find local minimisers is only guaranteed for losses with Lipschitz gradients, where it can be seen as a `bona-fide' discretisation of an underlying gradient flow. Yet, many ML setups involving overparametrised models do not fall into this problem class, which has motivated research beyond the so-called ``Edge of Stability'' (EoS), where the step-size crosses the admissibility threshold inversely proportional to the Lipschitz constant above. Perhaps surprisingly, GD has been empirically observed to still converge regardless of local instability and oscillatory behavior. The incipient theoretical analysis of this phenomena has mainly focused in the overparametrised regime, where the effect of choosing a large learning rate may be associated to a `Sharpness-Minimisation' implicit regularisation within the manifold of minimisers, under appropriate asymptotic limits. In contrast, in this work we directly examine the conditions for such unstable convergence, focusing on simple, yet representative, learning problems, via analysis of two-step gradient updates. Specifically, we characterize a local condition involving third-order derivatives that guarantees existence and convergence to fixed points of the two-step updates, and leverage such property in a teacher-student setting, under population loss. Finally, starting from Matrix Factorization, we provide observations of period-2 orbit of GD in high-dimensional settings with intuition of its dynamics, along with exploration into more general settings.

1 Introduction

The introduction asks when gradient descent can converge with learning rates beyond local-stability guarantees, then presents two-step analysis and representative learning-problem results. Its contributions include higher-order conditions for local fixed points, teacher-neuron alignment, and period-2 behavior in matrix factorization.

1 Introduction

  • The motivating questions are whether GD converges under typical objectives and which minimizers it reaches when overparameterization creates a manifold with potentially different generalization performance.Although a curvature bound gives a convergence guarantee, neural-network training has been observed to continue decreasing loss with curvature at or above 2/η.Source: §1 · pp. 1–2
  • The paper studies non-asymptotic GD beyond EoS and characterizes learning rates above EoS for dynamics that hover around minimizers.This is contrasted with prior asymptotic sharpness-reduction-flow analyses under structural assumptions.Source: §1 · p. 2
  • Higher-order derivative conditions guarantee existence of two-step fixed points near one-dimensional minima for learning rates slightly above 2/λ*.The paper extends the condition to matrix factorization and reports MLP experiments in Appendix B.3.2.Source: §1 · p. 2
  • The passage reports that the example converges to an orbit of period 2 from a more global initialization, rather than through analysis based on a high-order local approximation.Source: §1 · p. 2
  • Under population loss, a two-layer single-neuron ReLU network learns the teacher direction beyond EoS while its two-layer weight norms stably oscillate.The paper also reports empirical evidence from 16-neuron networks in Appendix B.3.1.Source: §1 · p. 2
  • For matrix factorization, the paper reports convergence to period-2 orbits, gives theoretical intuition, and proposes two further models for future analysis.The minimizers form a manifold with different local curvatures.Source: §1 · p. 2

2 Related Work

Related work studies edge-of-stability dynamics, implicit regularization, balancing, and single-neuron learning. This paper positions its analyses as beyond-EoS results for representative problems, including a two-layer single-neuron ReLU network.

Edge of stability

  • Prior edge-of-stability work reports curvature hovering around 2/η while training loss decreases, with related observations for SGD and a proposed catapult into flatter regions.The passage attributes these observations and arguments to [6], [13], and [21], respectively.Source: §2 · pp. 2–3
  • Concurrent analyses explain unstable convergence through invariant sets [1], multiscale landscapes [24], modified GD or losses [2], and sharpness-constrained trajectories [8].The paper states its learning rate is strictly larger than that of [8], so the manifold proposed there does not exist in its setting.Source: §2 · p. 3

Implicit regularization

  • Existing implicit-bias analyses connect GD to gradient flow or sharpness regularization, but the cited settings do not cover the paper’s interest in learning rates beyond EoS.The passage notes limits for [3] and when η > 1/λ, and that [7] constrains its learning rate below EoS.Source: §2 · p. 3

Balancing effect

  • Prior work studies balancing in homogeneous networks and matrix factorization; this paper extends the matrix-factorization comparison to larger learning rates where perfect balance may occur.The passage contrasts its setting with [30] and connects balancing to flatter solutions in the global-minimum manifold.Source: §2 · p. 3

Learning a single neuron

  • Earlier single-neuron work covers alignment for one-layer students and empirical dynamics for two-layer neurons; this paper analyzes a two-layer single-neuron ReLU network beyond EoS.The passage cites [32], [29], and [28] for the earlier studies.Source: §2 · p. 3

3 Problem Setup

The setup frames GD through curvature-based stability and the observed edge-of-stability regime. It also motivates studying large learning rates in relation to minimizer selection and balancing.

3 Problem Setup

  • A gradient-Lipschitz bound limits η for standard convergence guarantees, but such uniform curvature bounds may fail or be impossible in neural networks.The passage relates this setting to observed curvature crossing 2/η during training.Source: §3 · p. 4
  • With fixed η, neural-network curvature can rise above 2/η, after which training loss may stop decreasing monotonically yet continue decreasing over longer intervals.At the edge of stability, reducing η is followed by curvature quickly rising toward 2/η_s.Source: §3 · p. 4
  • In overparameterized networks, minimizers form a manifold, motivating questions about GD’s implicit bias and whether large learning rates can favor flatter minima.The passage says flatter minima are believed to be preferred and associated with better generalization in cited work.Source: §3 · p. 4
  • Large-learning-rate GD has been observed to produce more balanced matrix-factorization solutions despite imbalanced initialization, though not necessarily perfect balance.The passage contrasts this with gradient flow’s preservation of layer-norm differences.Source: §3 · p. 4
  • The setup motivates studying large-learning-rate GD across low-dimensional and representative models through connections among optimization, loss landscapes, and generalization.Source: §3 · p. 4

4 Stable oscillation on 1-D functions: fixed point of two-step update

The section gives derivative conditions for two-step fixed points beyond the Edge of Stability, then studies convergence in a specific one-dimensional model and its two-dimensional extension.

4.1 Existence of fixed points

  • For a 1-D differentiable function near a local minimum, if f'''(x̄) ≠ 0 and 3(f'''(x̄))² − f''(x̄)f''''(x̄) > 0, suitable nearby initialization admits a period-2 orbit at a learning rate beyond EoS.Source: §4.1 · p. 5
  • When f''' vanishes, the sine function still permits stable oscillation because its first nonzero higher derivative at the minimum is f''''<0.Source: §4.1 · pp. 5–6
  • For squared loss on g, mild derivative conditions yield a two-step fixed point, and composing qualifying base functions preserves the condition.The result covers a broad family of nonlinear least-squares models.Source: §4.1 · p. 6

4.2 Convergence to fixed points

  • For f(x)=1/4(x^2−μ)^2, GD from 0<x0<√μ converges to a period-2 orbit, except for a measure-zero initialization.The section distinguishes this global convergence result from the local Taylor-based condition; the latter is recovered as K approaches 1.Source: §4.2 · pp. 6–7
  • In the 2-D factorization loss, |x−y| decays while xy>μ for 1<K<1.5; below 1.121, both cases converge to the same orbit.Source: §4.2 · pp. 6–7

5 On a two-layer single-neuron homogeneous network

For a two-layer single-neuron ReLU network trained on population loss, the analysis establishes convergence beyond EoS under specified conditions. The dynamics approach a period-2 orbit, with the layer imbalance decreasing in the final phase.

5 On a two-layer single-neuron homogeneous network

  • The network uses a ReLU student neuron trained by GD on population loss, with inputs uniform on the unit sphere and labels generated by a teacher neuron.Source: §5 · p. 7
  • On the global-minimum manifold, the largest Hessian eigenvalue measures layer imbalance, linking sharper minima to more unequal layer parameters.Source: §5 · p. 8
  • Under Theorem 4's initialization and η=Kd with K∈(1,1.1], the model converges to a period-2 orbit with wy=0 and (v,wx)∈γK.Source: §5 · p. 8
  • If the learning rate decays after sufficient oscillations, the trajectory approaches the flattest minimum; a 16-neuron experiment also shows convergence near it.Source: §5 · p. 8
  • Training has three phases: the angle first decreases, then wy enters a basin and decays exponentially, after which the dynamics approach the 2-D case and imbalance decreases.The model ultimately converges to the period-2 orbit described in the 1-D analysis.Source: §5 · p. 9

6 Matrix Factorization and beyond

The section examines period-2 and higher-period behavior in matrix factorization, then considers whether existing models explain those observations. It also proposes broader models for stable oscillations in more complicated settings.

6 Matrix Factorization and beyond

  • The section extends earlier beyond-EoS results to matrix factorization, asking whether period-2 orbits occur in higher dimensions and what models can explain the oscillations.It also asks what happens as the learning rate increases.Source: §6 · p. 9

6.1 Observations from Matrix Factorization

  • Matrix factorization has a period-2 orbit reached from a non-measure-zero set of initializations under the stated learning-rate conditions.The minima form a manifold, and the orbit is described through singular-value decompositions of the factors.Source: §6.1 · p. 9
  • Near an unbalanced minimum, quasi-symmetric matrix factorization converges to a period-2 orbit with error O(ϵ) under the stated learning-rate and initialization conditions.The analysis establishes the necessary one-dimensional condition around minima, while the authors note that a rigorous proof of the broader observation is difficult.Source: §6.1 · pp. 9–10
  • η(1/α^2 + σ2^2α^2) < 2 is a sufficient condition for the quasi-symmetric bound, and the grid-search boundary approximately follows it.The authors caution against extending this empirical trend beyond the theorem.Source: Figure 1 · p. 11, §6.1 · p. 11
  • For ησ1^2 in (1.23, 1.28), GD converges to a period-4 orbit; above 1.28 it is chaotic, while below 1 it does not oscillate.Across most reported cases, the balancing effect holds: σmax(Y) = σmax(Z).Source: §6.1 · p. 11, Figure 2 · p. 12

6.2 Implications for more complicated settings

  • The separable eigendirection model cannot explain matrix-factorization dynamics, and the manifold proposed in [8] may not exist beyond EoS.For η < 1, the manifold can exist but be limited, leaving the trajectory’s projection unchanged in early steps.Source: §6.2 · p. 12
  • The proposed extension adds pairwise interactions to separable functions; larger |E| permits fewer Hessian eigenvalues beyond 2/η, while E = ∅ recovers the earlier model.The example uses squared norm and inner-product terms to imitate matrix factorization.Source: §6.2 · pp. 11, 13

7 Conclusions

The conclusion frames the work as an analysis of gradient descent beyond local stability through two-step updates. It summarizes results for low-dimensional functions, a teacher-student network, and matrix factorization, and proposes models for later study.

7 Conclusions

  • High-order derivative conditions guarantee the existence of fixed points of two-step updates in the low-dimensional setting.The paper investigates large step sizes crossing the local-stability threshold through two-step rather than one-step updates.Source: §7 · pp. 13–14
  • Under population loss, a two-layer single-neuron ReLU network converges to align with the teacher neuron.For matrix factorization, the necessary one-dimensional condition holds around any minimum.Source: §7 · p. 14
  • The paper reports period-2 orbits in matrix factorization, gives theoretical intuition for their dynamics, and proposes two models compatible with the observations.Source: §7 · p. 14

A.1 On a 2-D function

The appendix extends the analysis of unstable GD to an imbalanced two-parameter factorization, matrix factorization, and higher-dimensional experiments. Across the supported settings, GD can balance parameters and settle into stable oscillations beyond the edge of stability.

A.1 On a 2-D function

  • On the global-minima manifold, curvature is larger when parameters are more imbalanced, with smallest curvature λ1 = 2µ at x = y = √µ; if x,y stay positive, |x−y| decays to 0 where xy > µ for 1<K<1.5.Source: §A.1 · p. 18

A.2.1 Asymmetric Case: 1D function at the minima

  • In asymmetric matrix factorization, the leading Hessian eigenvector defines a 1-D cross-section whose minima satisfy the stable-oscillation condition.The authors state that this 1-D condition holds around any minimum of two-layer matrix factorization, while arbitrary-direction convergence near minima requires separate analysis.Source: §A.2.1 · pp. 19–20

B.1 Proven Settings

  • The quartic and sine test functions both exhibit stable oscillations around local minima when the learning rate is slightly above the threshold.The numerical examples use η=1.01, above the threshold of 1.Source: §B.1 · p. 20
  • In the teacher-student experiment, the small second-layer weight decays before imbalance falls, after which the dynamics reduce to a balanced 1-D period-2 orbit.The experiment uses 1,000 sphere samples; residual imbalance and weight are attributed to the difference between empirical and population loss.Source: §B.1 · p. 20
  • Symmetric and quasi-symmetric matrix-factorization runs converge to the same period-2 orbit of leading singular values despite different initial scaling.The quasi-symmetric run uses α=0.8 and a learning rate 1.02 times the EoS threshold.Source: §B.1 · p. 20, Figure 6 · p. 21

B.2 2-D function

  • For the 2-D factorization, η=1.05 yields a period-2 orbit and η=1.25 a period-4 orbit, with |x−y| shrinking in both cases.The period-4 case lies outside the stated orbit theorem’s range but within the range for balance.Source: §B.2 · p. 21

B.3 High dimension and MNIST

  • The high-dimensional experiments report that GD beyond EoS drives toward flatter minima and behaves similarly to the low-dimensional dynamics.Source: §B.3 · p. 21

B.3.1 2-layer high-dim homogeneous ReLU NNs with planted teacher neurons

The high-dimensional teacher-student experiment tests whether large-rate GD moves from sharper minima toward flatter ones, while the MNIST experiments examine the proposed oscillation conditions in deeper MLPs. The reported results support flatter-minimum selection and show alignment and derivative-condition patterns across the tested depths.

B.3.1 2-layer high-dim homogeneous ReLU NNs with planted teacher neurons

  • The high-dimensional teacher-student experiment compares learning rates 0.5 and 2.6 from the same initialization, with the latter beyond an EoS threshold around 2.5.Source: §B.3.1 · p. 23
  • In the high-dimensional teacher-student setting, GD beyond EoS moves from sharper minima toward flatter minima and then oscillates stably there.Reducing the learning rate from 2.6 to 0.5 after 100 epochs reaches a nearby minimum flatter than the run using 0.5 throughout.Source: §B.3.1 · pp. 22–23

B.3.2 3, 4, 5-layer non-homogeneous MLPs on MNIST

  • For 3-layer MNIST MLPs, during EoS epochs the gradient closely aligns with the top Hessian eigenvector, and the required derivative expression is positive at most line-search minima.The few negative values occur outside the EoS regime and may reflect fourth-derivative approximation error; the authors report the same arguments for 4- and 5-layer networks.Source: §B.3.2 · p. 24

C Proof of Theorem 1

The proof uses a local third- and fourth-derivative condition to establish a two-step fixed point for GD at a learning rate above the local stability threshold.

C Proof of Theorem 1

  • If f^(3)(x̄)≠0 and 3[f^(3)]^2−f''f^(4)>0 at a local minimum, a nearby starting point admits a learning rate for which two GD steps return to it.Source: C Proof of Theorem 1 · p. 28
  • The learning rate lies above 2/f''(x̄), with the stated upper bound 2/(f''(x̄)−ϵf^(3)(x̄)).Source: C Proof of Theorem 1 · p. 28
  • The proof locates this rate by continuity of the two-step displacement as η varies.Source: C Proof of Theorem 1 · p. 28

D Proof of Lemma 1

The lemma gives derivative-based conditions for a nearby two-step return, with separate cases for odd and even orders of the first nonzero derivative beyond the second. It also applies to the local minima of sine.

D Proof of Lemma 1

  • For odd k, if ϵf^(k)(x̄)>0 and f^(k+1)(x̄)<0, a learning rate in the stated interval yields a two-step return.Source: D Proof of Lemma 1 · p. 29
  • For even k, if f^(k)(x̄)<0, a learning rate in the stated interval yields a two-step return.Source: D Proof of Lemma 1 · p. 29
  • The lemma implies stable oscillation around local minima of sin(x).At such a minimum, the first nonzero derivative beyond f'' is the fourth derivative and is negative.Source: D Proof of Lemma 1 · p. 30

E Proof of Prop 1

The section applies the derivative condition for stable oscillation to squared-residual losses and elementary examples, then extends the result to composed functions.

E Proof of Prop 1

  • For f(x) = (g(x) − y)^2, nonzero g′ at the minimum and g′g(3) < 6[g′′]^2 ensure the condition for stable period-2 oscillation.The proposition states that either Theorem 1 or Lemma 1 supplies the condition.Source: E Proof of Prop 1 · p. 30
  • For nontrivial minima where g′ ≠ 0 and g is nonlinear nearby, the derivative test requires 6[g′′]^2 > g′g(3); when g′′ = 0, the stated case instead uses Lemma 1 with k = 4.At such minima, f′′ = 2[g′]^2 and f(3) = 6g′′g′.Source: E Proof of Prop 1 · pp. 30–31
  • For f(x) = (sin(x) − y)^2, stable oscillation is stated around x̄ = arcsin(y) when y ∈ (−1, 1).Source: E Proof of Prop 1 · p. 31
  • The section gives tanh, exponential, logarithmic, and sigmoid squared-residual losses as examples allowing stable oscillation around their local minima.Source: E Proof of Prop 1 · p. 31
  • Power losses f(x) = (x^α − y)^2 and the quartic loss (x^2 − 1)^2 are also stated to allow stable oscillation around local minima.Source: E Proof of Prop 1 · p. 31
  • If f and g each satisfy the stable-oscillation conditions at the relevant minima, then g(f(x)) also allows stable oscillation around x = x̄.Source: E Proof of Prop 1 · p. 31

F Proof of Theorem 2

The section proves existence and global convergence of a period-2 orbit for one-dimensional gradient descent, and establishes decay of factor imbalance under stated positivity and initialization conditions. It concludes by computing max_η(g(z) − μ), identified as the final result.

F Proof of Theorem 2

  • For K in (1, 1.5], the proof establishes a unique period-2 orbit and shows that initializations in (0, √µ) converge to it, except for a measure-zero initialization converging to √µ.Existence and uniqueness are obtained from a third-order equation; convergence is shown by trapping iterates in an invariant interval and proving convergence there.Source: F Proof of Theorem 2 · pp. 32–35

G Proof of Theorem 3

  • For 1 < K < 1.5, if x and y remain positive, |x − y| strictly decreases while xy > µ and converges to 0 on those points.The proof derives monotone decrease and rules out a positive limiting difference.Source: G Proof of Theorem 3 · pp. 35, 38

H Proof of Lemma 2

  • Under the stated Theorem 3 initialization, with x0y0 > µ and m = |y0 − x0|/√µ, the lemma asserts that both factors remain positive throughout gradient descent.The proof seeks positivity-preserving bounds using the initial factor difference and the preceding imbalance-decay result.Source: H Proof of Lemma 2 · pp. 38–39
  • The proof concludes by computing max_η(g(z) − μ), identified as the final result.Source: H Proof of Lemma 2 · p. 40

I Proof of Theorem 4

This section establishes convergence properties of the stated two-stage dynamics, including a bounded basin and decay of $w_y$, and develops matrix-factorization results on stable oscillation at minima, convergence to a period-2 orbit, and transitions to higher-period orbits.

I Proof of Theorem 4

  • With teacher w̃ = [1, 0], η = Kd for K ∈ (1, 1.1], and the stated initialization, the proof places the variables in a bounded basin after four iterations, where w_y decays at a linear rate of at most 0.97.The initialization is w(0) = v(0) = ε ∈ (0, 0.10] with ⟨w(0), w̃⟩ ≥ 0.Source: I Proof of Theorem 4 · p. 40

Stage 1.

  • In Stage 1, w_y grows more slowly than w_x and v w_x, so its growth and the stage duration can be bounded while controlling v w_x and v − w_x.The proof also shows w_y/w_x ≤ 2.7 and decreasing, and derives bounds on w_y and v − w_x at the stage transition.Source: I Proof of Theorem 4 · p. 41, Stage 1 · pp. 44, 46

Stage 2.

  • In Stage 2, the proof iteratively computes bounds on v w_x, v − w_x, and w_y; after eight iterations they enter and remain in a bounded basin where w_y decays linearly.The bounds are computed using the relatively narrow learning-rate range 1 < K ≤ 1.1.Source: I Proof of Theorem 4 · p. 41, Stage 2 · pp. 47–48
  • The range-based mapping method makes bounds feasible by mapping source intervals to covering target intervals, though this sacrifices tightness.The proof later notes that interval splitting and merging can combine tightness with efficiency.Source: Stage 2 · p. 48
  • The generated interval Is has wy = 0, vwx ∈ [0.181, 1.5], and piecewise bounds on v − wx depending on whether vwx is at most 1.For vwx ≤ 1, v − wx ∈ [−0.735, 0.23]; for vwx > 1, v − wx ∈ [−0.474, 0.148].Source: Stage 2 · pp. 51–52
  • The expanded interval generated by If preserves vwx ∈ [0.181, 1.5] and the piecewise bounds on v − wx, with the lower bound −0.472 when vwx > 1.When vwx ≤ 1, the bound is [−0.735, 0.23]; when vwx > 1, it is [−0.472, 0.148].Source: Stage 2 · p. 55
  • For Cases I–II in Stage 2, wy decays with a linear ratio of at most 0.97 (or 1 − 0.030K).The bound applies after the iterates enter the expanded interval with the stated bounds on vwx and v − wx.Source: Stage 2 · p. 55
  • For Case III, wy ≤ 0.078 in the first step of Stage 2, and the case therefore converges to Is.Source: Stage 2 · p. 55
  • In the Stage 2 proof, the step count needed to reach vwx > 0.181 is bounded within the bound for Cases (I, II).Source: Stage 2 · p. 56

J Proof of Matrix Factorization

  • The matrix-factorization proof focuses on the minimum manifold XY = C, aiming to establish the 1D stable-oscillation condition there and explain movement toward flatter minima.Source: J Proof of Matrix Factorization · p. 56

J.1 Asymmetric Case: 1D function at the minima

  • At a matrix-factorization minimum XY = C, the leading Hessian eigenvector has the stated form and its 1D cross-section satisfies the stable-oscillation condition under the theorem’s singular-value assumptions.The assumptions include unique top singular values for X and Y and a nonzero inner product between the corresponding singular vectors.Source: §J.1 · p. 56

J.2 Quasi-symmetric case: walk towards flattest minima

  • For quasi-symmetric matrix factorization, small perturbations near a minimum lead GD to an approximate period-2 orbit, while transverse components decay exponentially.The orbit approximation has error O(ϵ); the proof notes that control of error terms in non-asymptotic analysis remains missing.Source: §J.2 · pp. 58, 66, 67

L Illustration of period-2 and period-4 orbits

M Discussions

  • The discussion provides a roadmap of the theoretical results and considers implications of the low-dimensional results for more complex models, alongside high-dimensional experiments.Source: M Discussions · p. 68

M.1 Connections between theoretical results

The section connects stable one-dimensional oscillations to two-layer and high-dimensional settings, where gradient descent can converge toward flatter solutions. It also relates the local geometric condition to broader regression models and reports supporting empirical observations.

M.1 Connections between theoretical results

  • The one-dimensional analysis supplies a local condition for stable oscillation in regression models, and a composition rule extends it to more complicated models.The condition and composition rule are stated in Propositions 1 and 2.Source: §M.1 · p. 68
  • In the two-layer single-neuron ReLU model, the student direction converges to the teacher direction as wy → 0, reducing the dynamics to the two-parameter case.The two-parameter argument then implies convergence toward a balanced situation.Source: §M.1 · p. 68, §M.2 · p. 69

M.2 Implications from low-dimension to high-dimension

  • The one-dimensional geometric condition is proved at minima for non-trivial matrix factorization and is reported around minima of 3-, 4-, and 5-layer ReLU MLPs on MNIST.The matrix-factorization result is Theorem 6; the MLP observations appear in Figures 9(d), 10(d), and 11(d).Source: §M.2 · p. 69
  • The high-dimensional discussion connects matrix-factorization trajectories to the leading Hessian eigenvector and reports near-perfect gradient alignment with it for MNIST MLPs.The theoretical intuition concerns quasi-symmetric matrix factorization; the alignment is reported empirically in Figures 9(c), 10(c), and 11(c).Source: §M.2 · p. 69
  • In the two-parameter model, x and y converge to the same value after oscillations of xy, and smaller |x − y| corresponds to a flatter minimizer.The model shares the convergence behavior of the one-dimensional analysis.Source: §M.2 · p. 69
  • From sharp initialization, quasi-symmetric matrix factorization moves toward the flattest minima, where the top singular values of Y and Z are equal.The authors describe the high-dimensional dynamics as an observation with theoretical intuition.Source: §M.2 · p. 69
  • For 16-neuron student-teacher ReLU networks, learning-rate decay after beyond-EoS oscillations brings the model close to the flattest minima.The observation is reported in Figure 8 and Appendix B.3.1.Source: §M.2 · p. 69
  • For 3-, 4-, and 5-layer MNIST MLPs, larger learning rates drive convergence toward flatter minima.The result is shown in Figure 9(b).Source: §M.2 · p. 69

Related papers

Loading 2206.04172v3…