Navigating loss manifolds via rigid body dynamics: A promising avenue for robustness and generalisation
Journal:
arXiv
Published Date:
May 26, 2025
Abstract
Training large neural networks through gradient-based optimization requires
navigating high-dimensional loss landscapes, which often exhibit pathological
geometry, leading to undesirable training dynamics. In particular, poor
generalization frequently results from convergence to sharp minima that are
highly sensitive to input perturbations, causing the model to overfit the
training data while failing to generalize to unseen examples. Furthermore,
these optimization procedures typically display strong dependence on the fine
structure of the loss landscape, leading to unstable training dynamics, due to
the fractal-like nature of the loss surface. In this work, we propose an
alternative optimizer that simultaneously reduces this dependence, and avoids
sharp minima, thereby improving generalization. This is achieved by simulating
the motion of the center of a ball rolling on the loss landscape. The degree to
which our optimizer departs from the standard gradient descent is controlled by
a hyperparameter, representing the radius of the ball. Changing this
hyperparameter allows for probing the loss landscape at different scales,
making it a valuable tool for understanding its geometry.
