Navigating loss manifolds via rigid body dynamics: A promising avenue for robustness and generalisation

2025-05-26

AI Summary

Neural network training often suffers from poor generalization and dynamic instability due to ill-conditioned loss landscapes—e.g., sharp minima and fractal structures. To address this, we propose the first optimizer grounded in rigid-body dynamics: parameter updates are modeled as frictionless rolling of a sphere over the loss manifold, with sphere radius serving as a learnable hyperparameter that enables intrinsic multi-scale geometric awareness and naturally suppresses sensitivity to local sharpness. Our method integrates continuous-time dynamical modeling, Riemannian differential geometry, and a gradient correction mechanism. Experiments across diverse architectures and datasets demonstrate substantial improvements in generalization and robustness, marked reduction in overfitting, enhanced training stability, and provision of an interpretable, multi-scale characterization of the loss terrain.

Technology Category

• Search and Optimization: Learning to Search
• Machine Learning: Learning with Manifolds
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Responsible Web: Machine-in-the-loop, human agency and autonomy

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.

Problem

• Navigating high-dimensional loss landscapes with pathological geometry
• Avoiding sharp minima to improve model generalization
• Reducing dependence on loss landscape fine structure

Innovation

• Simulates ball rolling on loss landscape
• Avoids sharp minima via rigid dynamics
• Probes loss geometry at multiple scales