Towards Understanding Momentum Acceleration in River-Valley Loss Landscape

1 Introduction

• The paper addresses limited theory for momentum in non-convex deep-learning landscapes, asking how momentum and learning rate jointly affect optimization along a river.Its motivation is that classical acceleration guarantees rely on restricted quadratic-like settings and linear convergence.Source: §1 · pp. 1–2

1.1 Our Contributions

• Momentum enlarges the tolerable learning rate by stabilizing oscillations, enabling faster progress along the river; on flat, slow-spinning rivers, the larger rate—not momentum itself—provides most of the speed gain.The analysis treats rotating and interacting flat and sharp eigenspaces in a non-convex river-valley landscape.Source: §1.1 · pp. 2–3
• A new induction argument tracks river- and valley-projected gradients and momentum, while also improving the analyzed tolerable learning rate for vanilla GD over Wen et al. (2025).The coupled river and valley components interfere, complicating magnitude tracking compared with vanilla GD analysis.Source: §1.1 · p. 3
• Synthetic experiments and GPT-Neo pretraining on TinyStories match the predicted learning-rate relationship with momentum in SGD-Momentum and a similar trend with β1 in Adam.The authors report that these results verify their predicted acceleration through larger learning rates.Source: §1.1 · p. 3

1.2 Related Works

• Prior work studies momentum in structured non-convex regimes and large-rate phenomena, while this paper focuses on deterministic stable river tracking and momentum's learning-rate gain.The river-valley perspective is also linked to language-model pretraining and low-loss manifold dynamics.Source: §1.2 · p. 4

2 Preliminaries

• The preliminaries define operator norms, the third-derivative tensor norm, and Oβ(·), which hides factors polynomial in 1/(1 − β).They also specify eigenvalue and eigenvector notation and the floor function.Source: §2 · p. 5

2.1 River-Valley Loss Landscape: Setups and Key Assumptions

• The river-valley model assumes a one-dimensional manifold whose gradient aligns with the flattest Hessian direction, separated from sharper orthogonal directions in a neighborhood with controlled curvature and slow spinning.The assumptions also impose neighborhood regularity and river uniqueness; several scale conditions are chosen for the proof rather than optimized.Source: §2.1 · pp. 5–6
• The reference flow is the Riemannian gradient flow constrained to the river, serving as the trajectory against which iterate proximity and downstream speed are formalized.The paper's x(t) denotes points on the river or this flow, with t its continuous time index.Source: §2.1 · p. 6
• The examples include a non-spinning river and a non-convex river-valley function, with optimization progress characterized by how quickly the trajectory advances while tracking the river.For the non-spinning example, κ = 0.Source: §2.1 · p. 6

2.2 Gradient Descent with Heavy-Ball Momentum

• Vanilla GD uses a learning rate η, while heavy-ball GD updates a momentum state and then the parameters; setting β = 0 recovers vanilla GD.The stated momentum update uses β in [0, 1) and scales the gradient contribution by 1 − β.Source: §2.2 · p. 6

3 Review: What Happens for Vanilla GD?

• GD stays close to the river, and its river-tracking time is nearly proportional to the learning rate when the rate is tolerable.The stated tracking distance is O(κ/γ), and the time-alignment error is O(κ + ηγflat) · ηk.Source: §3 · p. 7
• A larger tolerable learning rate means faster progress along the river, even when oscillations in sharp directions temporarily raise training loss.A subsequent learning-rate decay suppresses those oscillations, so the loss reflects progress made along the river.Source: §3 · p. 7
• The learning-rate threshold is approximately 2/γmax, with sharp-direction curvature and river spinning limiting stable tracking.The section notes that its analysis tightens the bound in Wen et al. (2025) through a sharper analysis of sharp-direction dynamics.Source: §3 · p. 7

4.1 Main Theory: Momentum Acceleration in River-Valley Landscape

• With on-river initialization and zero initial momentum, GD-M stays close to the river and tracks it at a pace proportional to η after a post-transient period.The theorem assumes β ∈ (0, 0.99] and η below ηGD-M max; tracking and speed errors depend on κ, γflat, and β.Source: §4.1 · pp. 7–8

4.2 Theoretical Predictions by Theorem 4.1 and Numerical Verification

• At η = 0.4, GD-M with β = 0.9 damps transverse oscillations and advances along the river, while GD continues oscillating without sustained progress.The numerical experiment uses the synthetic non-convex river-valley loss and starts near the river.Source: Figure 2 · p. 8, §4.2 · p. 9
• Empirical learning-rate thresholds nearly follow ηGD max · (1 + β)/(1 − β + 0.05), with the 0.05 adjustment attributed to river spinning and curvature.The empirical threshold markers and theoretical curve are shown in Figure 2(c).Source: Figure 2 · p. 8, §4.2 · p. 9
• In the flat, slowly spinning regime, momentum mainly accelerates river progress by stabilizing larger learning rates, not by directly increasing tracking speed.For sharper or faster-spinning rivers, large β can hinder close tracking; β very close to 1 is specifically cautioned against.Source: §4.2 · pp. 8–9
• For a fixed β, final loss decreases approximately linearly as η rises; at a fixed tolerable η, changing β has little effect on final loss, while β = 0.999 does not converge.Larger β permits a larger tolerable learning rate.Source: §4.2 · p. 9

5 Proof Outline, Key Challenges, and Technical Novelty

• The proof projects gradients and momentum onto flat and sharp directions, then defines river projection through the limiting flow of an ODE.A linear interpolation of discrete iterates yields a continuous path whose river projections define an induced trajectory and absolute time map.Source: §5 · pp. 9–10

5.1 Proof Outline and Key Challenges

• The proof first bounds sharp-direction dynamics to prevent explosion, then shows iterates stay close to the river and river-time grows nearly proportionally to η.The challenge is that river spinning makes gradient and momentum projections interfere across the two directions, requiring careful induction.Source: §5.1 · p. 10

5.2 Key Lemmas and Technical Novelties

• The analysis shows that gradient and momentum dynamics are dominated by the flat direction, while sharp-direction dynamics set the maximum tolerable learning rate.A delicate induction tracks projected gradient and momentum magnitudes; compared with vanilla GD, momentum enlarges the maximum learning rate by (1 + β)/(1 − β).Source: §5.2 · pp. 10–11
• Under the theorem's conditions, the interpolated momentum trajectory has a well-defined river projection and remains close to its projected counterpart.This is the close-tracking step used to establish the river-following guarantee.Source: §5.2 · p. 11
• After a transient, the absolute time map has derivative within ϵ(β) · η of η, making accumulated river-time progress nearly linear in the learning rate.The error factor is ϵ(β) = Oβ(κ + ηγflat).Source: §5.2 · p. 11

• The study trains GPT-Neo on TinyStories with SGD-M and Adam, varying learning rate and momentum, then decays the rate during the final 15% of training.The decay is used to reduce valley oscillations and reveal progress along the river.Source: §6 · p. 12
• For SGD-M, the largest tolerable learning rate increases with β, supporting the predicted scaling trend rather than an exact realization of the idealized formula.The fit parameters are descriptive corrections for non-asymptotic effects, not constants predicted by Theorem 4.1.Source: §6 · p. 12, Figure 3 · p. 13
• At η = 1, SGD-M without momentum diverges while β = 0.9 or 0.95 converges to low losses; at η = 0.1, changing β has only marginal effect.This fixed-grid comparison supports acceleration mainly through stabilizing larger learning rates, not substantially changing river-direction speed at a fixed rate.Source: §6 · p. 12, Table 1 · p. 13
• For Adam, the largest tolerable learning rate generally increases with β1, but decreases as β1 approaches 1.The adjusted scaling fit summarizes this observed trend and is not a theoretical prediction for Adam.Source: §6 · p. 12, Figure 4 · p. 13
• Figure 5 compares three learning rates across training and validation curves; at the two larger rates, Adam with β1 ∈{0.9, 0.95} is more stable.Its columns vary η; the top panels show batch losses and moving averages, while the bottom panels show validation losses at three checkpoints.Source: §6 · p. 12, Figure 5 · p. 13
• With momentum, larger learning rates give lower training and validation losses than a small learning rate, supporting acceleration through a more aggressive stable rate.The reported comparison also finds lower training and validation losses than Adam with β1 = 0.Source: §6 · p. 13

7 Conclusions

• In river-valley landscapes, momentum stabilizes oscillations across sharp directions, enabling larger learning rates and faster progress along the low-loss river.The theory is validated by synthetic experiments and language-model training; extending the analysis to adaptive optimizers is left for future work.Source: §7 · p. 13

A Further Related Works

• Prior heavy-ball analyses emphasize near-quadratic or PL settings, while other work studies network architectures, small-learning-rate momentum, implicit bias, and gradient-history variants.Examples include Wang et al. (2021), Plattner (2022), Cattaneo et al. (2024), and Chen et al. (2023).Source: A Further Related Works · p. 22
• Related river-valley studies describe valley-floor motion, vanishing gradients in sharp directions, and river-valley structure in language-model pretraining.The passage cites Xing et al. (2018), Davis et al. (2024), and Wen et al. (2025).Source: A Further Related Works · p. 22

B Further Discussions

• The analysis treats deterministic gradients; systematic study of stochastic GD-momentum in river-valley landscapes is left for future work.The discussion motivates the deterministic setting as more relevant when batch size is large.Source: B Further Discussions · p. 23
• The normalized first-moment parametrization matches Adam's first-moment update and yields an ODE whose gradient coefficient differs from the common unnormalized heavy-ball form.The continuous-time model views iterates as samples of a curve with spacing h.Source: B Further Discussions · p. 23
• After momentum burn-in, the river-direction update has the same first-order speed as vanilla GD at learning rate η; momentum instead stabilizes sharp directions, permitting a larger η and faster river progress.The transient momentum term is removed using a lower bound on k in Theorem 4.1.Source: B Further Discussions · pp. 23–24
• The second-order ODE intuition is local to a very flat, slowly spinning river segment, while the theorem tracks a changing reference flow under its stated assumptions.The reference flow captures projected-gradient variation along the river, and the proof controls the difference from its tangent direction.Source: B Further Discussions · pp. 23–24

Define the post-transient iteration

• After k⋆, GD-momentum stays close to the river and its reference-flow time differs from ηk by at most ϵ(β)η(k−k⋆).The theorem assumes initialization on the river with zero initial momentum.Source: Define the post-transient iteration · pp. 25–26
• When flat-direction constraints are inactive and scale separation holds, I(β) is nonempty and its upper endpoint is approximately 2(1 + β)/(1 − β) · 1/γmax.The discussion also cautions against choosing β too close to 1.Source: Define the post-transient iteration · p. 26
• Figure 6's proof graph orders endpoint induction bounds before extending them to interpolated trajectories and deriving river proximity and near-linear reference-flow time.At each iteration, the induction conditions update in the order (I1), (I3), (I2), then (I4).Source: Define the post-transient iteration · p. 26, Figure 6 · p. 27

C.2.1 Main Result and Proof Outline

• Under Assumptions 2.1 and 2.2 and Theorem C.1’s parameter conditions, Lemma C.4 establishes flat-direction dominance for momentum iteration initialized on the river with zero momentum.The lemma states its conclusions hold for every iteration k and interpolation parameter τ ∈ [0, 1].Source: §C.2.1 · p. 28
• The proof first establishes endpoint bounds by induction for sharp- and flat-direction gradients and momentum, then extends them across each interpolation segment.The induction maintains that the sharp direction does not explode, the flat-direction gradient keeps its magnitude, flat-direction momentum resembles a moving average of flat-direction gradients, and the flat gradient dominates sharp-direction gradient and momentum.Source: §C.2.1 · pp. 28–29

C.2.2 Lemmas for Induction Argument

• Supporting lemmas provide a uniform momentum bound when iterates remain in U and zero initial momentum, and analyze sharp directions using parameter-conditioned matrix bounds.Lemma C.7 introduces orthogonal matrices and positive definite matrices to bound the sharp-direction dynamics for each iteration.Source: §C.2.2 · p. 35
• Lemma C.9 provides the first induction step under Assumptions 2.1 and 2.2, using the preceding induction conditions and a dummy step at k = −1.The iteration starts on the river with m0 = 0, and η ∈ I1 ∩ I2 ∩ I3.Source: §C.2.2 · p. 45
• For k ≥ 1, if the stated induction conditions hold, the conclusion holds for every τ ∈ [0, 1]; it also holds at k = 0 when (C.9) holds at k = −1 and (C.8) at k = 0.Source: §C.2.2 · p. 47
• Lemma C.11 closes the next induction step when η lies in I1 ∩ I2 ∩ I3 ∩ I4 and the stated induction conditions hold through step k.The proof sets the induction residual using cend := 3B(β)κ3/4 and bounds it by the sharp-direction contraction factor.Source: §C.2.2 · pp. 49–50
• The induction hypothesis makes the corresponding residual nonpositive, completing the proof of Lemma C.11.The text also says the induction hypothesis includes the dummy initial step when k=0.Source: §C.2.2 · p. 52
• Lemma C.12 states that induction condition (C.7) holds at step k when q_k≤1/2, using 1/(1−x)≤1+2x for 0≤x≤1/2.Source: §C.2.2 · p. 52

C.3.1 Proof of Lemma 5.2

• Under Theorem C.1’s conditions, the trajectory has a well-defined river projection throughout interpolation, and its distance from that projection is controlled.The lemma also relates this distance to the gradient in the sharp direction, which is dominated by the gradient in the flat direction.Source: §C.3.1 · pp. 52–53
• An induction keeps each interpolated iterate inside the region where the river projection exists, using a uniform momentum bound and neighborhood containment.The proof extends the property from one interpolation interval to the next.Source: §C.3.1 · p. 53
• The proof bounds the distance to the projection using the sharp-direction gradient, then shows that this gradient is dominated by its flat-direction counterpart.This connects closeness to the river with the gradient’s concentration in the flat direction.Source: §C.3.1 · p. 53

C.3.2 Proof of Lemma 5.3

• After the stated iteration threshold, the reference-flow time derivative satisfies |S′(t) − η| ≤ ϵ(β) · η almost everywhere on interpolation intervals.The bound applies when β ≤ 0.99 and η lies in I(β) under Theorem C.1’s conditions.Source: §C.3.2 · p. 54
• The proof aims to bound the difference between the river tangent and the projected momentum update by O(κ · ∥u(t)∥₂), making the vectors similar.The tangent is compared at reference-flow time S(t).Source: §C.3.2 · p. 54
• The proof decomposes the flat-direction tracking error into river-tangent alignment, change in the flat gradient over time, and mismatch between projected momentum and flat gradient.These are handled as separate terms in the proof.Source: §C.3.2 · pp. 54–55

D.1 Analysis for Eigen-Decomposition

• The eigen-decomposition analysis gives a stability test for T, including the condition ηλ < 2(1 + β)/(1 − β) for its spectral radius to remain below one.It connects this condition to the discrete-time Lyapunov stability criterion.Source: §D.1 · p. 58

D.2 Basics of the River

• The projection ODE flow converges to a point on the river, with movement along the flow decaying exponentially.The existence result is stated under Assumptions 2.1 and 2.2.Source: §D.2 · p. 59

D.3 Analysis of the River and the Projections

• The river-and-projection lemmas characterize river spinning, tangent directions, and how projections change and respond to perturbations.They also include results on the Jacobian of projection to the river and gradients in the flat direction.Source: §D.3 · pp. 59–61

D.4 Matrix Inequalities

• The matrix-inequality subsection supplies eigenvalue perturbation and Davis–Kahan results for symmetric matrices.The stated results concern eigenvalues and eigenspaces.Source: §D.4 · pp. 61–62

E.1 Formal Statement of Theorem 3.1

• Under the stated assumptions and learning-rate bound, GD stays close to the river and its progress time is nearly proportional to η.The theorem bounds river distance by O(κ · g_max/γ) and the time deviation by O(κ + ηγ_flat)ηk.Source: §E.1 · p. 63

E.2 Proofs for Improved Analysis

• A sharper estimate for the sharp directions supplies the new learning-rate condition; the remaining projection, trajectory-tracking, and time-reparametrization arguments follow the earlier proof without momentum-specific terms.The proof separately bounds gradient norms in the sharp and flat directions before combining those bounds.Source: §E.2 · pp. 63–65
• The flat-direction gradient norm dominates under the stated assumptions and learning-rate bound.The lemma applies for any iteration k and τ ∈ [0, 1].Source: §E.2 · p. 63

F.1 Further Experiment Details

• Training losses are recorded every 100 steps, full-validation losses every 5,000 steps, and experiments use one NVIDIA H100 GPU.Plots show faint recorded losses, 900-step moving averages, and insets for steps 30,000–40,000.Source: §F.1 · p. 67

F.2 More Experiment Results

• The experiments plot training and validation loss curves for five learning rates and four momentum settings.The left and right panels show training and validation loss, respectively.Source: §F.2 · p. 67

F.2.1 Fix the learning rate and change the momentum

• With learning rates fixed, momentum enables more stable and faster training, with a larger effect at higher learning rates.The comparison plots training and validation losses across momentum values 0, 0.85, 0.9, and 0.95 for learning rates from 0.0001 to 0.0008.Source: §F.2.1 · p. 67, Figure 7 · p. 67, Figure 8 · p. 68

F.2.2 Fix the momentum and change the learning rate

• At fixed momentum, larger learning rates yield lower training and validation losses during early training, while larger momentum makes the curves more stable, especially for validation loss.The comparisons use momentum values 0, 0.85, 0.9, and 0.95 and learning rates from 1 × 10^-4 to 8 × 10^-4.Source: §F.2.2 · p. 67, Figure 12 · p. 69, Figure 13 · p. 69