Modified Loss of Momentum Gradient Descent: Fine-Grained Analysis

We analyze gradient descent with Polyak heavy-ball momentum (HB) to prove that it is exactly plain gradient descent with a modified loss on an exponentially attractive invariant manifold, providing rigorous approximation bounds and deriving continuous modified equations.

How it works

The paper investigates the dynamics of full-batch and mini-batch Polyak [39] heavy-ball momentum (HB) optimization, showing that on an exponentially attractive invariant manifold, the algorithm behaves exactly like plain gradient descent with a modified loss. This is established by defining a function gh(θ) such that the dynamics simplify to:

on an exponentially attractive invariant manifold, HB is plain gradient descent: θ(n+1) = θ(n) − h/ (1 − β)∇L - hβ(1 − β)Ghθ(n).

Approximation and Memoryless Representation

A key contribution is the proof that the algorithm can be approximated by a memoryless iteration with arbitrary precision. For any approximation order R ≥ 2, Theorem 3.1 establishes a global error bound:

sup n∈[0:⌊T /h⌋]∥θ(n) − θ˜(n)∥ = O(h R), where T is any “time” horizon.

This memoryless iteration is defined by the coefficients d j(n)(θ), which can be expressed in a tractable form involving sums over unlabeled rooted trees. Specifically, Theorem 4.1 provides the form of these coefficients:

d m(n) = −β/Xτ∈A˜[m]σ(τ)E(n)τ,1, where E(n)τ,l is defined recursively based on the iteration number n and memory distance variable l.

Modified Equations and Principal Flow

The analysis leads to the derivation of continuous modified equations, which approximate the discrete iterations. Corollary 5.1 finds a set of functions f j(n)(θ) such that their unique continuous solution satisfies:

sup n∈[0:⌊T /h⌋]∥θ(n) − θ(tn)∥ = O(h R), where t˜ is between tn and tn+1.

Furthermore, the analysis of the principal part of these equations yields a principal flow, which approximates the HB dynamics. Corollary 5.8 defines this flow using coefficients z m(n), leading to an approximate continuous equation:

θ˙(t) = X∞ k=0 z k(n+1)h k∇2L(θ(t))k∇L(θ(t)) + NPT for large n.

Combinatorial Insights into Polynomials

A significant finding is the combinatorial structure underlying the memoryless approximation coefficients. Analyzing the form of d m(n)(θ) in the limit as n → ∞ reveals a rich family of polynomials in β:

we prove that it contains Eulerian and Narayana polynomials.

Specifically, Corollary 5.3 shows that coefficients corresponding to chains with m vertices are the rescaled Narayana polynomials, while those corresponding to trees consisting of a root and m-1 leaves are related to the Eulerian polynomials. This forms a rich A˜[m]-parametrized (m ⩾ 1) family of polynomial[s] of β.

Implicit Regularization and Convergence

The modified loss exhibits features interpreted as implicit regularization by memory, which can explain empirical performance. In the stochastic (mini-batch) case, additional implicit regularization by mini-batch noise is identified. The principal iteration for large n is approximately:

θ˜(n+1) = θ˜(n) − h∇Lθ˜(n) − βhgβh∇2Lθ˜(n)∇Lθ˜(n) + NPT.

The convergence of this principal flow is characterized by the generating function g β(x), which is the Stieltjes transform of the standard Marchenko-Pastur law with parameter β, indicating that convergence occurs when∥∇2Lθ˜(n) < (1 − √β) 2/h. This framework provides a general roadmap for analyzing other optimization algorithms with decaying memory.

Conclusion and Generalization

The theoretical results establish that the memoryless iteration is equivalent to plain gradient descent with a modified loss, and the analysis of its coefficients reveals hidden combinatorial structures. These findings are applicable not only to HB but also serve as a framework for studying other numerical methods with decaying memory, such as Adam or Shampoo. The paper concludes by showing that the principal flow coincides with known solutions for specific cases like least-squares regression when initialized on the invariant manifold.

As a fastidious researcher, I have analyzed this paper on Gradient Descent with Polyak Heavy-Ball Momentum (HB). The core contribution is transforming the memory-dependent HB iteration into an exact, high-order memoryless approximation that behaves like plain Gradient Descent (GD) with a modified loss function.

Here are the specific improvements and capabilities you can implement in AI systems:

)

1. The system can be trained using a Modified Loss instead of the raw loss function, which is derived from the underlying invariant manifold dynamics. This modification implicitly regularizes the optimization process based on the momentum parameter β and step size h.

2. The system's training dynamics can be analyzed via a continuous flow (Modified Equation) rather than purely discrete iterations. This allows for better qualitative understanding of convergence, oscillatory behavior, and divergence in high-dimensional loss landscapes.

3. The system can utilize a memoryless approximation that maintains high accuracy even in mini-batch settings by leveraging the derived coefficients (Eulerian and Narayana polynomials). This means the system can be trained using stochastic gradients with guaranteed error bounds related to the order of approximation R, rather than relying solely on second-order approximations.

4. The system's long-term optimization trajectory can be approximated by a Principal Flow, which is a smooth gradient flow (for large n) that captures the essential dynamics while explicitly separating it from higher-order non-principal terms (NPT). This allows for noise and high-frequency fluctuations to be modeled separately.

5. The system's learning rate can be dynamically rescaled by incorporating implicit regularization terms derived from the empirical covariance matrix of per-sample gradients (in the mini-batch case), leading to a more robust and adaptive learning rate schedule that accounts for stochasticity.

)

By implementing these improvements, your AI systems can achieve:

1. The ability to converge reliably in complex, non-convex loss landscapes by using a loss function that is inherently biased towards the solution (implicit regularization).

2. Superior generalization performance on large-scale models by leveraging the high-order memoryless approximation for stochastic gradient descent (SGD) with mini-batches, ensuring that noise does not lead to catastrophic divergence.

3. More efficient training schedules where you can distinguish between the principal learning dynamics and transient fluctuations or higher-order error components, leading to faster convergence toward a stable manifold.

4. The development of novel optimization algorithms by generalizing the framework to other memory-based methods (like Adam or Shampoo), enabling you to propose new optimizers with explicit regularization terms that are theoretically grounded in dynamical systems theory.

We analyze gradient descent with Polyak (1964) heavy-ball momentum (HB) whose fixed momentum hyperparameter β in (0, 1) provides exponential decay of memory. Building on Kovachki and Stuart (2021), we prove that on an exponentially attractive invariant manifold the algorithm is exactly plain gradient descent with a modified loss, provided that the step size h is small enough. Although the modified loss does not admit a closed-form expression, we describe it up to O(h R) -errors for arbitrary finite order R, and prove global (finite "time" horizon) trajectory approximation bounds O(h R). We then conduct a fine-grained analysis of the combinatorics underlying the memoryless approximations of HB, in particular, finding a rich family of polynomials in β hidden inside which include and lie coefficient-wise in between Eulerian and Narayana polynomials. We prove that these polynomials are h-polynomials of certain graph-associahedra. As corollaries of the main results, we derive continuous modified equations of arbitrary finite approximation order (with rigorous bounds) and the principal flow that approximates the HB dynamics, generalizing Rosca et al. (2023). Approximation theorems cover both full-batch and mini-batch HB. The results shed new light on the main features of HB and outline a roadmap for similar analysis of other optimization algorithms.