RSAM: Learning on manifolds with Riemannian Sharpness-aware Minimization
Tuan Truong Hoang-Phi Nguyen Tung Pham Minh-Tuan Tran Mehrtash Harandi Dinh Phung Trung Le
Nowadays, understanding the geometry of the loss landscape shows promise in
enhancing a model's generalization ability. In this work, we draw upon prior
works that apply geometric principles to optimization and present a novel
approach to improve robustness and generalization ability for constrained
optimization problems. Indeed, this paper aims to generalize the
Sharpness-Aware Minimization (SAM) optimizer to Riemannian manifolds. In doing
so, we first extend the concept of sharpness and introduce a novel notion of
sharpness on manifolds. To support this notion of sharpness, we present a
theoretical analysis characterizing generalization capabilities with respect to
manifold sharpness, which demonstrates a tighter bound on the generalization
gap, a result not known before. Motivated by this analysis, we introduce our
algorithm, Riemannian Sharpness-Aware Minimization (RSAM). To demonstrate
RSAM's ability to enhance generalization ability, we evaluate and contrast our
algorithm on a broad set of problems, such as image classification and
contrastive learning across different datasets, including CIFAR100, CIFAR10,
and FGVCAircraft. Our code is publicly available at
\url{https://t.ly/RiemannianSAM}.