The inverse variance-flatness relation in stochastic gradient descent is critical for finding flat minima

Area of Science:

• Deep Learning
• Machine Learning Theory
• Statistical Physics

Background:

• Stochastic Gradient Descent (SGD) is a cornerstone of deep learning, yet its mechanism for finding generalizable solutions in high-dimensional spaces remains unclear.
• Understanding the interplay between SGD learning dynamics and the geometry of the loss function landscape is crucial for improving generalization.
• Existing theories often lack a direct link between the dynamics of SGD and the properties of the loss landscape it navigates.

Purpose of the Study:

• To investigate the relationship between SGD learning dynamics and the loss function landscape in deep neural networks.
• To characterize the nature of SGD's movement in high-dimensional weight spaces and its connection to landscape features.
• To develop a theoretical framework explaining SGD's preference for flat minima and its implications for generalization and continual learning.

Main Methods:

• Principal Component Analysis (PCA) was employed to analyze the low-dimensional dynamics of SGD in the weight space.
• A phenomenological theory of SGD was developed, utilizing statistical properties of minibatch loss functions.
• An algorithm incorporating landscape-dependent constraints was designed to address catastrophic forgetting.

Main Results:

• SGD dynamics were shown to follow a low-dimensional drift-diffusion process.
• A significant inverse relationship was discovered between weight variance and loss landscape flatness across PCA directions.
• Both SGD's noise strength (temperature) and correlation time were found to be inversely proportional to landscape flatness.

Conclusions:

• SGD acts as a landscape-dependent annealing algorithm, effectively preferring flatter minima due to its dynamics.
• The findings provide a theoretical basis for why SGD generalizes well and suggest a mechanism for mitigating catastrophic forgetting.
• This work offers a new framework for understanding learning dynamics and developing more effective deep learning algorithms.