Sharpness-Aware Teleportation on Riemannian Manifolds

TL;DR

This paper introduces a novel sharpness-aware, geometry-aware teleportation method on Riemannian manifolds that improves model robustness and generalization by leveraging intrinsic geometric structures and sharpness-aware optimization techniques.

Contribution

We propose a new teleportation mechanism within Riemannian manifolds that decomposes optimization steps into local orbit transitions and sharpness-aware moves, backed by a theoretical analysis.

Findings

Enhanced generalization on vision benchmarks

Robustness improvements demonstrated across datasets

Theoretical analysis of generalization gap

Abstract

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.