# Leveraging Continuous Time to Understand Momentum When Training Diagonal Linear Networks

arXiv: https://arxiv.org/abs/2403.05293 · HTML: https://mlsanity.com/paper/2403.05293/

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Builds on 34 of its 53 references that are in the map.
- First posted on arXiv 8 Mar 2024.

## Builds on (most cited first, 30 of 34)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [Understanding deep learning requires rethinking generalization](https://mlsanity.com/paper/1611.03530/) (2016)
- [Implicit Regularization in Deep Matrix Factorization](https://mlsanity.com/paper/1905.13655/) (2019)
- [Kernel and Rich Regimes in Overparametrized Models](https://mlsanity.com/paper/2002.09277/) (2020)
- [Characterizing Implicit Bias in Terms of Optimization Geometry](https://mlsanity.com/paper/1802.08246/) (2018)
- [A Differential Equation for Modeling Nesterov's Accelerated Gradient Method: Theory and Insights](https://mlsanity.com/paper/1503.01243/) (2015)
- [Algorithmic Regularization in Learning Deep Homogeneous Models: Layers are Automatically Balanced](https://mlsanity.com/paper/1806.00900/) (2018)
- [Momentum Improves Normalized SGD](https://mlsanity.com/paper/2002.03305/) (2020)
- [An Improved Analysis of Stochastic Gradient Descent with Momentum](https://mlsanity.com/paper/2007.07989/) (2020)
- [A Variational Perspective on Accelerated Methods in Optimization](https://mlsanity.com/paper/1603.04245/) (2016)
- [Implicit Bias of SGD for Diagonal Linear Networks: a Provable Benefit of Stochasticity](https://mlsanity.com/paper/2106.09524/) (2021)
- [Shape Matters: Understanding the Implicit Bias of the Noise Covariance](https://mlsanity.com/paper/2006.08680/) (2020)
- [Implicit Regularization for Optimal Sparse Recovery](https://mlsanity.com/paper/1909.05122/) (2019)
- [Understanding the Acceleration Phenomenon via High-Resolution Differential Equations](https://mlsanity.com/paper/1810.08907/) (2018)
- [Accelerated Gradient Descent Escapes Saddle Points Faster than Gradient Descent](https://mlsanity.com/paper/1711.10456/) (2017)
- [Global convergence of the Heavy-ball method for convex optimization](https://mlsanity.com/paper/1412.7457/) (2014)
- [On the insufficiency of existing momentum schemes for Stochastic Optimization](https://mlsanity.com/paper/1803.05591/) (2018)
- [Towards understanding how momentum improves generalization in deep learning](https://mlsanity.com/paper/2207.05931/) (2022)
- [Implicit Bias of Gradient Descent on Reparametrized Models: On Equivalence to Mirror Descent](https://mlsanity.com/paper/2207.04036/) (2022)
- [Exponentiated Gradient Meets Gradient Descent](https://mlsanity.com/paper/1902.01903/) (2019)
- [Continuous Time Analysis of Momentum Methods](https://mlsanity.com/paper/1906.04285/) (2019)
- [Convergence of a Stochastic Gradient Method with Momentum for Non-Smooth Non-Convex Optimization](https://mlsanity.com/paper/2002.05466/) (2020)
- [The Two Regimes of Deep Network Training](https://mlsanity.com/paper/2002.10376/) (2020)
- [Label noise (stochastic) gradient descent implicitly solves the Lasso for quadratic parametrisation](https://mlsanity.com/paper/2206.09841/) (2022)
- [Implicit regularization in Heavy-ball momentum accelerated stochastic gradient descent](https://mlsanity.com/paper/2302.00849/) (2023)
- [From Averaging to Acceleration, There is Only a Step-size](https://mlsanity.com/paper/1504.01577/) (2015)
- [Accelerated Linear Convergence of Stochastic Momentum Methods in Wasserstein Distances](https://mlsanity.com/paper/1901.07445/) (2019)

## Built on by (most cited first, 9 of 9)

- [The AdEMAMix Optimizer: Better, Faster, Older](https://mlsanity.com/paper/2409.03137/) (2024)
- [Alternating Gradient Flows: A Theory of Feature Learning in Two-layer Neural Networks](https://mlsanity.com/paper/2506.06489/) (2025)
- [Implicit Bias of Mirror Flow on Separable Data](https://mlsanity.com/paper/2406.12763/) (2024)
- [Never Saddle for Reparameterized Steepest Descent as Mirror Flow](https://mlsanity.com/paper/2603.02064/) (2026)
- [Optimization Insights into Deep Diagonal Linear Networks](https://mlsanity.com/paper/2412.16765/) (2024)
- [Precise Dynamics of Diagonal Linear Networks: A Unifying Analysis by Dynamical Mean-Field Theory](https://mlsanity.com/paper/2510.01930/) (2025)
- [Gradient Flow Dynamics and Implicit Bias of Diagonal Linear Networks under Infinitesimal Initialization](https://mlsanity.com/paper/2607.12332/) (2026)
- [Heavy-Ball Momentum Method in Continuous Time and Discretization Error Analysis](https://mlsanity.com/paper/2506.14806/) (2025)
- [HORST: Composing Optimizer Geometries for Sparse Transformer Training](https://mlsanity.com/paper/2605.21104/) (2026)

## Abstract (arXiv)

In this work, we investigate the effect of momentum on the optimisation trajectory of gradient descent. We leverage a continuous-time approach in the analysis of momentum gradient descent with step size γ and momentum parameter β that allows us to identify an intrinsic quantity λ = frac γ  (1 - β)2  which uniquely defines the optimisation path and provides a simple acceleration rule. When training a 2-layer diagonal linear network in an overparametrised regression setting, we characterise the recovered solution through an implicit regularisation problem. We then prove that small values of λ help to recover sparse solutions. Finally, we give similar but weaker results for stochastic momentum gradient descent. We provide numerical experiments which support our claims.