Research map / Optimization / Implicit bias and generalization
Effects of Momentum in Implicit Bias of Gradient Flow for Diagonal Linear Networks
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Characterizes the implicit bias of continuous-time heavy-ball and Nesterov momentum flows on 2-layer diagonal linear networks, showing momentum adds an initialization-mitigation effect plus an explicit gradient-initialization dependence, so generalization vs. GD depends on hyperparameters.
Problem
Whether adding momentum (heavy-ball and Nesterov) to gradient descent changes the implicit bias of training overparameterized regression models such as diagonal linear networks, and when this helps generalization for sparse regression.
Approach
Derives O(eta) continuous-time approximations of HB and NAG with constant momentum as a damped second-order ODE, then shows for diagonal linear networks their limit point minimizes a hyperbolic-entropy-plus-linear objective with an effectively reduced initialization scale and an extra linear term depending on the initial gradient, unlike plain linear regression where momentum does not change impli
Key results
- HB and NAG flows induce an extra initialization-mitigation effect (reduced effective initialization scale), which improves sparse-regression generalization over GF when initialization is unbiased and the gradient-initialization term is small.
- With biased initialization, HB and NAG generalize better than GD when the initialization is highly biased, but the explicit gradient-dependent term can erase this benefit when the initialization is only slightly biased.
- Decreasing the learning rate or momentum factor makes HB/NAG implicit bias converge to that of GD.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use HB flow instead of gradient flow | Diagonal linear networks · Sparse regression with unbiased initialization and small gradient initialization | squared parameter distance to ground truth | Helped |
| Use NAG flow instead of gradient flow | Diagonal linear networks · Sparse regression with unbiased initialization and small gradient initialization | squared parameter distance to ground truth | Helped |
| Use HB or NAG flow instead of gradient flow when the initialization-gradient dependence is not small | Diagonal linear networks · Sparse regression | generalization performance | Hurt |
| Decrease the dataset scaling ε (ablation) | HB and NAG flow · Sparse regression with unbiased initialization | squared parameter distance to ground truth | Helped |
| Increase the initialization scale (ablation) | Diagonal linear networks trained with GD, SGD, HB, or NAG · Sparse regression with unbiased initialization | squared parameter distance to ground truth | Mixed |
| Use HB, NAG, or SGD instead of GD | Diagonal linear networks · Sparse regression with biased initialization and fixed initialization scale | squared parameter distance to ground truth | Helped |
| Increase the extent of the unbiased part of initialization (ablation) | Diagonal linear networks trained with HB or NAG · Sparse regression with biased initialization | squared parameter distance to ground truth | Hurt |
| Decrease the learning rate (ablation) | HB and NAG · Diagonal linear networks | squared parameter distance to ground truth | No clear effect |
| Decrease the momentum factor (ablation) | HB and NAG · Diagonal linear networks | squared parameter distance to ground truth | No clear effect |
Surprising findings
- Momentum changes the implicit bias of gradient flow for diagonal linear networks, unlike for single-layer linear regression where it does not.
- Momentum's explicit dependence on the initial gradient can make HB and NAG generalize worse than GD when initialization is only slightly biased.
Experimental setup
Datasets: synthetic Gaussian sparse regression (n=40, d=100, 5 nonzero ground-truth components)
Benchmarks: synthetic over-parameterized sparse regression
Metrics: squared distance D(theta, theta*) between learned and ground
Baselines: gradient descent gradient flow stochastic gradient descent heavy-ball Nesterov accelerated gradient
Builds on: Woodworth et al. 2020 diagonal linear network implicit bias Pesme et al. 2021 stochastic gradient flow implicit bias Kovachki and Stuart 2021 modified equations for momentum Papazov et al. 2024 heavy-ball diagonal linear networks
Closest papers by meaning
- Implicit Bias of the Step Size in Linear Diagonal Neural Networks (ICML 2022 · similarity 0.91)
- (S)GD over Diagonal Linear Networks: Implicit bias, Large Stepsizes and Edge of Stability (NeurIPS 2023 · similarity 0.89)
- Implicit regularization in Heavy-ball momentum accelerated stochastic gradient descent (ICLR 2023 · similarity 0.89)
- Mirror, Mirror of the Flow: How Does Regularization Shape Implicit Bias? (ICML 2025 · similarity 0.87)
- Implicit Bias of Gradient Descent for Non-Homogeneous Deep Networks (ICML 2025 · similarity 0.87)
- Implicit Bias of Mirror Flow for Shallow Neural Networks in Univariate Regression (ICLR 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.