--- title: Riemannian Sharpness-Aware Minimization (RSAM) url: https://www.emergentmind.com/topics/riemannian-sharpness-aware-minimization-rsam type: topic --- # Riemannian Sharpness-Aware Minimization (RSAM) Riemannian Sharpness-Aware Minimization (RSAM) encompasses a family of methodologies generalizing sharpness-aware optimization by incorporating notions of local or intrinsic geometry, particularly Riemannian metrics and manifolds, into the minimization of loss sharpness. These advances address the limitations of Euclidean, parameterization-dependent sharpness-aware methods, and provide frameworks suited for constrained or geometrically structured parameter spaces. Two main strands have emerged in recent literature: geometric reparameterization-invariant sharpness-aware minimization, and intrinsic manifold-constrained RSAM. ## 1. Objectives and Conceptual Framework RSAM extends the conventional Sharpness-Aware Minimization (SAM)—which minimizes the worst-case loss within a neighborhood of parameters—by reformulating neighborhoods, distances, and gradients using Riemannian geometry instead of Euclidean space. This yields methods that optimize parameters constrained to Riemannian manifolds or adapt the adversarial perturbation direction using a Riemannian metric intrinsically tied to the loss geometry. Given a manifold $\mathcal{M}\subset\mathbb{R}^k$ with an intrinsic (often problem-induced) dimension $d$, and empirical risk $L_\mathcal{S}(\theta) = \frac{1}{n} \sum_{i=1}^n \ell(f_\theta(x_i),y_i)$, the generic RSAM objective is \[ \min_{\theta\in\mathcal{M}} \; \max_{\theta'\in\mathcal{B}_\theta(\rho)} L_\mathcal{S}(\theta') \] with $\mathcal{B}_\theta(\rho)$ the Riemannian ball of radius $\rho$ centered at $\theta$ [2309.17215]. ## 2. Geometric Foundations RSAM requires tools from differential geometry: - **Riemannian Metric:** Inner product $\langle\cdot,\cdot\rangle_\theta$ defined on tangent space $T_\theta\mathcal{M}$ at each $\theta$. - **Retraction:** Map $R_\theta: T_\theta\mathcal{M} \to \mathcal{M}$ approximating exponential map, with $R_\theta(0)=\theta$ and $DR_\theta(0)=\mathrm{Id}$. - **Projection:** For embedded submanifolds, the ambient Euclidean gradient $\nabla L$ is projected to $T_\theta\mathcal{M}$ to yield the Riemannian gradient $\operatorname{grad}L(\theta)$ via an orthogonal projection [2309.17215]. This geometric machinery defines neighborhoods and gradient-based optimization steps. For statistical learning on quotient manifolds or the Stiefel manifold (e.g., orthogonality constraints on weight matrices), these operations are realized in closed form (e.g., QR retraction, symmetric projection). ## 3. Algorithmic Procedures RSAM algorithms generally decompose each iteration into - **Inner maximization:** ("teleportation" step) Find the worst-case perturbation within a Riemannian ball by first-order Taylor expansion, leading to choosing the maximal-ascent direction of the Riemannian gradient and mapping it back to the manifold via retraction: \[ \epsilon^* = \operatorname{Proj}_\theta(\bar{\epsilon}), \quad \bar{\epsilon} = \rho \cdot (g^T D_\theta) / \| g^T D_\theta \|_2, \] followed by $\theta^* = R_\theta(\epsilon^*)$, where $g = \operatorname{grad} L_\mathcal{S}(\theta)$, $D_\theta$ encodes local metric scaling. - **Outer update:** ("sharpness-aware descent") Apply gradient descent to the perturbed point, mapped back via retraction: \[ \theta_{t+1} = R_\theta\big(-\eta\,\operatorname{grad} L_\mathcal{S}(\theta^*)\big) \] [2309.17215]. In Riemannian parameterizations induced by the loss geometry (as in Monge SAM [2502.08448]), the metric is $G(\theta) = I_K + \nabla \ell(\theta)\nabla \ell(\theta)^T$, making both the perturbation norm and direction adapt to the local slope. ## 4. Reparameterization-Invariant and Loss-Induced (Monge) Metrics A key RSAM approach is to define the Riemannian metric from the embedding of the parameter manifold into the loss surface, yielding reparameterization invariance. The pullback metric is \[ G(\theta) = I_K + \nabla \ell(\theta)\nabla \ell(\theta)^T \] with the adversarial step \[ \delta^* = \rho \cdot G(\theta)^{-1} \nabla \ell(\theta) \Big/ \sqrt{ \nabla \ell(\theta)^T G(\theta)^{-1}\nabla \ell(\theta) } \] This approach, termed Monge SAM (M-SAM), interpolates smoothly between SAM and vanilla gradient descent (GD): as $\|\nabla \ell\| \rightarrow 0$, $G(\theta)\rightarrow I$, recovering SAM; as $\|\nabla \ell\|\rightarrow\infty$, the step vanishes, recovering GD [2502.08448]. This metric yields a closed-form, invariant adversarial step and markedly improves robustness to hyperparameters as well as saddle-point escape properties. ## 5. Theoretical Guarantees Theoretical analyses of RSAM provide generalization bounds leveraging Riemannian neighborhoods and PAC-Bayes concentration, typically tightening the dependence on the parameter space's intrinsic dimension $d$ rather than the ambient parameter count $k$. Specifically, for RSAM with retraction $R_\theta$ and ball radius $\rho$: \[ L_\mathcal{D}(\theta) \leq \max_{\theta'\in\mathcal{B}_\theta(\rho)} L_\mathcal{S}(\theta') + C_\rho\,\rho^2 + \sqrt{\frac{d + \log(n/\delta)}{n-1}} \] with $C_\rho$ reflecting geometric aspects of the retraction (e.g., $C_\rho = 1 + \sqrt{2}/2$ for Stiefel/QR) [2309.17215]. For Monge SAM, invariance under reparameterization holds formally: for every smooth chart $\phi$, the pullback metric guarantees that the adversarial step and corresponding descent trajectory commute with coordinate changes [2502.08448]. ## 6. Practical Implementations and Empirical Evaluation RSAM algorithms: - Maintain the computational cost of SAM (two forward/backward passes), with modest additional overhead (typically $<10\%$), even when enforcing geometric constraints such as orthogonality [2309.17215, 2502.08448]. - Require essentially the same hyperparameters as Euclidean SAM: perturbation radius $\rho$, learning rate $\alpha$. - Induce significantly improved generalization and train–test stability, including flatter minima (as indicated by smaller Hessian eigenvalues relative to SAM), larger test performance gains ($+1$–$3\%$ in supervised classification, $+1$–$5\%$ in contrastive learning on vision benchmarks), and greater robustness to hyperparameter misspecification. Empirical protocols utilize benchmarks such as CIFAR-10, CIFAR-100, FGVCAircraft, and architectures including ResNet-34/50, constrained via Stiefel manifold projections. | Model/Task | CE+SGD | CE+SAM | CE+RSAM | SupCon+SGD | SupCon+SAM | SupCon+RSAM | |----------------------|--------|--------|---------|------------|------------|-------------| | CIFAR-10 (ResNet-34) | – | – | – | – | – | – | | CIFAR-100 | +1–3% | +1–3% | +1–3% | +1–5% | +1–5% | +1–5% | Observed sharpness of RSAM-trained models, as measured by the maximal Hessian eigenvalue, is consistently reduced relative to SAM (e.g., $10.9$ vs $13.3$) [2309.17215]. ## 7. Connections, Limitations, and Ongoing Directions RSAM provides a unifying geometric perspective on sharpness-aware training, generalizing previous approaches in loss-induced (Monge) Riemannian metrics [2502.08448], manifold-constrained learning [2309.17215], and the pursuit of reparameterization invariance. Notably, the term "RSAM" has been used in Euclidean, randomized smoothing (random-SAM) contexts as well [2401.08060]; however, those variants do not employ Riemannian geometry but rather exploit stochastic perturbation in Euclidean space. Open research avenues include extending RSAM to more complex quotient geometries, exploring spectral regularization (e.g., Rényi sharpness), and applying RSAM in domains with inherent geometric constraints (e.g., equivariant networks, structured parameterizations). In summary, Riemannian Sharpness-Aware Minimization advances sharpness-aware optimization by adapting adversarial perturbations and robust descent steps to the intrinsic geometry of the parameter space or loss landscape, yielding improved generalization, stability, and invariance properties compared to classical approaches [2309.17215, 2502.08448]. Source: https://www.emergentmind.com/topics/riemannian-sharpness-aware-minimization-rsam