Gradient Descent with Polyak's Momentum Finds Flatter Minima via Large Catapults

2023/11/25 by Phunyaphibarn, Prin, Lee, Junghyun, Wang, Bohan +2 · 1 citation
#FOS: Computer and information sciences #FOS: Mathematics #Machine Learning (cs.LG) #Machine Learning (stat.ML) #Optimization and Control (math.OC)

paper · doi:10.48550/arxiv.2311.15051

Abstract

Although gradient descent with Polyak's momentum is widely used in modern machine and deep learning, a concrete understanding of its effects on the training trajectory remains elusive. In this work, we empirically show that for linear diagonal networks and nonlinear neural networks, momentum gradient descent with a large learning rate displays large catapults, driving the iterates towards much flatter minima than those found by gradient descent. We hypothesize that the large catapult is caused by momentum "prolonging" the self-stabilization effect (Damian et al., 2023). We provide theoretical and empirical support for our hypothesis in a simple toy example and empirical evidence supporting our hypothesis for linear diagonal networks.

Cited by

• Variational Learning Finds Flatter Solutions at the Edge of Stability

Related

• A Modular Analysis of Provable Acceleration via Polyak's Momentum: Training a Wide ReLU Network and a Deep Linear Network
• Leveraging Continuous Time to Understand Momentum When Training Diagonal Linear Networks
• The Marginal Value of Momentum for Small Learning Rate SGD
• Catapults in SGD: spikes in the training loss and their impact on generalization through feature learning
• Gradient descent with momentum --- to accelerate or to super-accelerate?