---
title: Riemannian Stochastic Gradient Descent
url: https://www.emergentmind.com/topics/riemannian-stochastic-gradient-descent-rsgd
type: topic
---

# Riemannian Stochastic Gradient Descent

Riemannian Stochastic Gradient Descent (RSGD) encompasses a family of optimization algorithms that generalize the classical stochastic gradient descent paradigm to settings where the model parameters live on a Riemannian manifold rather than flat Euclidean space. This framework is essential when the feasible set or the natural geometry of the problem is non-Euclidean, as in learning on the Stiefel or Grassmann manifold, or when parameter space symmetry/constraints are most naturally enforced with manifold structure. Recent advances have extended sharpness-aware minimization (SAM)—originally formulated for flat, Euclidean spaces—to the Riemannian context, resulting in Riemannian Sharpness-Aware Minimization (RSAM). Notable modern instantiations include Monge SAM (M-SAM) and geometric approaches tailored for learning on manifolds.

## 1. Mathematical Foundations of Riemannian SAM

Let $\mathcal{M} \subset \mathbb{R}^k$ be a $d$-dimensional Riemannian manifold, embedding the parameter set of a model $f_\theta$ with $\theta \in \mathcal{M}$. The loss function, typically sample-averaged,
$$
\mathcal{L}_S(\theta) = \frac{1}{n} \sum_{i=1}^n \ell(f_\theta(x_i), y_i),
$$
is defined for $\theta$ on $\mathcal{M}$. At a given iterate $\theta \in \mathcal{M}$, the *sharpness-aware* objective is
$$
L_{\mathrm{RSAM}}(\theta) = \max_{\Delta \in T_\theta \mathcal{M},\, \|\Delta\|_2 \leq \rho} \; \mathcal{L}_S\left(\mathrm{Exp}_\theta(\Delta)\right),
$$
where $T_\theta \mathcal{M}$ is the tangent space at $\theta$, $\mathrm{Exp}_\theta$ is the Riemannian exponential map (or a computationally tractable retraction $R_\theta$), and $\|\Delta\|_2$ is the standard Euclidean norm in $T_\theta \mathcal{M}$ as a subset of $\mathbb{R}^k$ [2309.17215].

The optimization of this objective requires machinery unique to Riemannian geometry—such as projections onto tangent spaces, manipulation through retraction/exponential maps, and computation of Riemannian gradients and transports.

## 2. RSGD Algorithmic Structure: Teleportation and Descent

Riemannian SAM (RSAM) generalizes SAM via a two-step process at each iteration $t$:

1. **Inner maximization ("Teleportation" step):**
   - Compute the Riemannian gradient $\mathrm{grad}_\theta \mathcal{L}_B(\theta_t)$—the projection of the Euclidean gradient onto $T_{\theta_t} \mathcal{M}$.
   - Solve the adversarial ascent direction in the tangent space, typically via
     $$
     \bar{\Delta} = \rho \frac{\left(\mathrm{grad}_\theta \mathcal{L}_B(\theta_t)\right)^\top D_\theta}{\left\|\left(\mathrm{grad}_\theta \mathcal{L}_B(\theta_t)\right)^\top D_\theta\right\|_2}, \quad
     \Delta_t^* = \mathrm{Proj}_\theta(\bar{\Delta}),
     $$
     where $D_\theta$ is a metric-adjustment matrix (often $I$ or $\operatorname{diag}(|g_i|)$), and $\mathrm{Proj}_\theta$ projects back to the tangent space [2309.17215].
   - Teleport the parameters onto the manifold: $\theta_t^* = R_{\theta_t}(\Delta_t^*)$.

2. **Outer minimization (Riemannian descent):**
   - Compute the Riemannian gradient at the perturbed point.
   - Update by retracting along the negative gradient:
     $$
     \theta_{t+1} = R_{\theta_t}\left(-\eta\, \mathrm{grad}_{\theta_t^*} \mathcal{L}_B(\theta_t^*)\right).
     $$

Complete pseudocode is specified explicitly in [2309.17215]. This framework accommodates both exact and approximate solutions to the inner maximization; the latter (e.g., direct relaxation using $D_\theta$ as identity) is $\sim1.75\times$ faster with negligible loss in empirical accuracy.

## 3. Geometric and Theoretical Properties

RSAM inherits and augments critical geometric properties relative to its Euclidean counterparts:

- **Manifold Consistency:** All updates and perturbations are confined to $\mathcal{M}$, leveraging projection, retraction, and manipulation in tangent spaces.
- **Reparametrization Invariance (M-SAM):** Monge SAM generalizes SAM by introducing a loss-induced Riemannian metric $G(\theta) = I_k + \nabla\ell(\theta)\nabla\ell(\theta)^\top$ (the Monge metric). Adversarial directions and steps are computed with respect to this geometry, yielding invariance under smooth reparametrizations: if $\phi:\eta \mapsto \theta$ is a diffeomorphism, the metric and steps transform covariantly, and the constrained step size is preserved under change of variables [2502.08448].

- **Generalization Bound:** Under compactness, $K$-Lipschitz loss, and controlled retraction error, RSAM yields a generalization bound
  $$
  \mathcal{L}_{\mathcal{D}}(\theta) \leq \max_{\theta' \in \mathcal{B}_\theta(\rho)} \mathcal{L}_S(\theta') + K C_\theta (\rho + \varepsilon)^2 + \sqrt{ \frac{d + \ln(n/\delta)}{n-1} },
  $$
  where $d = \dim(\mathcal{M})$. The dependence on intrinsic dimension $d$ (as opposed to ambient $k$ in Euclidean SAM) underpins improved statistical guarantees in manifold-constrained problems [2309.17215].

- **Critical Point Behavior:** In RSAM and especially in M-SAM, the method is less prone than Euclidean SAM to become trapped at suboptimal saddle points, due to its step size automodulation. The effective radius
  $$
  \tilde{\rho}_{\text{M-SAM}} = \frac{\rho}{\sqrt{1 + \|\nabla \ell(\theta)\|^2}}
  $$
  endows M-SAM with self-damping properties that make it more robust to hyperparameter choices and gradient magnitude [2502.08448].

## 4. Implementation Details and Computational Overhead

All Riemannian SAM-type methods require no explicit computation of Hessians or high-rank matrix inverses. Key operations per iteration include:

- One forward and one backward pass for the base point.
- Computation of Riemannian (projected) gradients and norm/scaling.
- One additional forward-backward for the adversarial direction.
- Retraction/exponential-map computations, and (optionally) tangent-space projections.

The runtime increase relative to standard SAM is marginal: RSAM incurs $\sim6\%$ additional overhead per epoch over Euclidean SAM (both $\sim2 \times$ slower than vanilla SGD, due to the double-backprop requirement). All empirical evidence in the literature uses tractable choices of retraction and projection, such as the identity map in unconstrained $\mathbb{R}^k$, or projection-based retraction for embedded manifolds [2309.17215], [2502.08448].

## 5. Empirical Performance and Benchmarks

Empirical investigations show RSAM and its Monge metric variant outperform standard SGD and Euclidean SAM in multiple settings, particularly:

- **Supervised Classification (ResNet50, CIFAR-100):** RSAM achieves 77.78% top-1 accuracy compared to 75.04% for SAM and 74.62% for SGD at identical hyperparameters [2309.17215].

- **Contrastive Pretraining (SupCon + RSAM):** Linear evaluation after SupCon pretraining and RSAM achieves 81.62% accuracy (vs. 76.73% for SAM and 75.29% for SGD) [2309.17215].

- **Robustness to Hyperparameter Choices:** On CIFAR-10, both SAM and M-SAM can escape local minima inaccessible to SGD, but only M-SAM avoids catastrophic divergence with large $\rho$, demonstrating its conservative step size adaptation [2502.08448].

- **Multi-modal Representation Alignment (CLIP Fine-tuning):** On WIT/MS-COCO, M-SAM achieved higher mutual-kNN similarity ($S_{\mathrm{kNN}} \approx 0.446$) than SAM ($\approx 0.405$) or SGD/Adam ($\approx 0.388$), exhibiting less sensitivity to $\rho$ [2502.08448].

- **Ablations:** Empirical evaluations demonstrate that approximate inner maximization for adversarial perturbations is nearly as accurate as exact projection, and that choice of metric-adjustment matrix has limited effect on final performance. RSAM is robust to auxiliary constraints, such as orthogonality in autoencoders, where Euclidean regularization strategies struggle [2309.17215].

## 6. Limitations and Open Directions

Current limitations of Riemannian SAM-type methods include the following:

- **Approximate Inner Maximization:** The practical implementations use heuristics for the inner maximization in the tangent space; more accurate manifold-specific solvers (e.g., geodesic searches) remain an open area [2309.17215].
- **Extension to Quotient Manifolds:** While existing RSAM treats embedded manifolds, extension to quotient structures (e.g., the Grassmann manifold) with explicit parallel transport between tangent spaces poses technical challenges.
- **Second-order Corrections:** Incorporating Riemannian analogues of curvature- or sharpness-aware corrections beyond first-order remains unexplored [2309.17215].
- **Empirical Overheads:** Although the per-iteration cost is only marginally higher than Euclidean SAM, the overall training cost remains a barrier compared to vanilla SGD, particularly for large-scale applications.

## 7. Comparison of Methodological Variants

The following table summarizes core distinctions and similarities between Monge SAM, RSAM, and classical SAM, as discussed in the referenced works:

| Method             | Geometry        | Reparametrization Invariant | Manifold Support            |
|--------------------|----------------|-----------------------------|-----------------------------|
| SAM                | Euclidean      | No                          | Flat $\mathbb{R}^k$ only    |
| Monge SAM (M-SAM)  | Loss-induced   | Yes                         | Any model, $\mathbb{R}^k$   |
| RSAM               | Manifold       | Yes (by construction)       | General $\mathcal{M}\subset\mathbb{R}^k$ |

*Monge SAM is a special case of manifold-aware sharpness-aware minimization where the metric is induced by the loss graph in ambient space [2502.08448], while RSAM is a general framework for optimizing over arbitrary Riemannian manifolds [2309.17215].*

---

Source: https://www.emergentmind.com/topics/riemannian-stochastic-gradient-descent-rsgd