Sharpness-Aware Teleportation on Riemannian Manifolds
This addresses generalization enhancement in machine learning models, particularly for vision tasks, by combining sharpness-aware minimization with Riemannian geometry insights.
The paper tackles the problem of improving model generalization by introducing a sharpness-aware, geometry-aware teleportation mechanism that decomposes iterations into teleportation within local orbits and sharpness-aware steps between orbits using Riemannian quotient manifolds. The method is evaluated on diverse vision benchmarks with various datasets and Riemannian manifolds.
Recent studies highlight the effectiveness of flat minima in enhancing generalization, with sharpness-aware minimization (SAM) achieving state-of-the-art performance. Additionally, insights into the intrinsic geometry of the loss landscape have shown promise for improving model generalization. Building on these advancements, we introduce a novel sharpness-aware, geometry-aware teleportation mechanism to further enhance robustness and generalization. The core innovation of our approach is to decompose each iteration into a teleportation step within a local orbit and a sharpness-aware step that transitions between different orbits, leveraging the Riemannian quotient manifold. Our approach is grounded in a theoretical framework that analyzes the generalization gap between population loss and worst-case empirical loss within the context of Riemannian manifolds. To demonstrate the effectiveness of our method, we evaluate and compare our algorithm on diverse vision benchmarks with various datasets and Riemannian manifolds.