Manifold-Aware Batch Normalization

• Manifold-aware batch normalization extends standard BN operations to Riemannian manifolds by leveraging geometric constructs like the Fréchet mean and tangent space scaling.
• It replaces Euclidean operations with manifold-specific metrics such as the affine-invariant and Bures–Wasserstein methods to robustly normalize structured features like covariance matrices and rotations.
• Practical implementations demonstrate improved convergence, enhanced stability, and superior learning performance in domains such as computer vision, medical imaging, and action recognition.

Manifold-aware Batch Normalization generalizes the principles of Euclidean batch normalization to data and parameters that inhabit non-Euclidean spaces, notably Riemannian manifolds such as the symmetric positive definite (SPD) manifold. This class of normalization layers—embodied by developments such as Riemannian Batch Normalization (RBN), ManifoldNorm, LieBN, GyroBN, and, most recently, Bures–Wasserstein-based methods—systematically leverages the intrinsic geometry to compute batch statistics, normalize representations, and stabilize optimization for manifold-valued features. These frameworks extend centering, scaling, and biasing operations to the manifold context by replacing Euclidean counterparts with the appropriate geometric constructs: Fréchet means (for centering), tangent space scaling (dispersion control), and manifold translations or group actions (biasing). The adoption of manifold-aware normalization is crucial for learning with structured features—e.g., covariance matrices, rotations, and directions—where respecting the data geometry is essential to numerical stability, ability to learn robust representations, and convergence speed.

1. Background and Motivation

Euclidean batch normalization (BN) standardizes activations in vector spaces by subtracting the mean and dividing by the standard deviation across a mini-batch. However, many domains, such as computer vision (covariance descriptors), medical imaging (diffusion tensors), and sequential modeling (correlation trajectories), operate on data naturally lying on Riemannian manifolds. The SPD manifold $\mathcal{S}^{++}_d$ (the set of $d \times d$ symmetric positive definite matrices) is a prototypical example, endowed with a rich family of distances and geometric structures. In these settings, naïvely applying Euclidean normalization can destroy the structural constraints (e.g., positive definiteness), ignore curvature, and undermine numerical stability (Wang et al., 1 Apr 2025). Accordingly, manifold-aware normalization schemes replace Euclidean operations with those respecting the underlying geometry: centering with respect to the Fréchet mean, scaling in the tangent space at an identity or mean point, and biasing via intrinsic translations. These methods aim to stabilize optimization, ensure well-conditioned features, and deliver gains in accuracy and convergence across a range of learning paradigms (Wang et al., 1 Apr 2025, Chen et al., 2024, Chen et al., 8 Sep 2025).

2. Geometric Foundations: Metrics and Statistics on Manifolds

Key to manifold-aware BN is the definition of appropriate means, variances, and transformation maps. For a Riemannian manifold $(\mathcal{M},g)$ with geodesic distance $\mathrm{dist}_g$, the Fréchet mean of a set $\{X_i\}$ is

$$M = \underset{S \in \mathcal{M}}{\arg\min} \sum_{i=1}^N \mathrm{dist}_g^2(X_i, S),$$

with variance $$v^2 = \frac{1}{N}\sum_{i=1}^N \mathrm{dist}_g^2(X_i, M)$$. Classical choices on $\mathcal{S}^{++}_d$ include:

• Affine-Invariant Metric (AIM): $$\delta_{AI}(X,Y) = \| \log(X^{-1/2} Y X^{-1/2}) \|_F$$ (Brooks et al., 2019).
• Log-Euclidean Metric (LEM): Induced by the Euclidean metric after mapping $X \mapsto \log X$ (Chen et al., 2024).
• Bures–Wasserstein Metric (BWM): $$d_{BW}^2(X,Y) = \mathrm{tr}(X) + \mathrm{tr}(Y) - 2\mathrm{tr}\left((X^{1/2} Y X^{1/2})^{1/2}\right)$$ (Wang et al., 1 Apr 2025).

Generalizations such as the Generalized Bures–Wasserstein Metric (GBWM) introduce a learnable SPD parameter $M$, modulating the geometry:

$$d_{GBW}^2(X,Y) = \mathrm{tr}(X M) + \mathrm{tr}(Y M) - 2\mathrm{tr}\left( (M^{1/2} X M^{1/2})^{1/2} \right).$$

Closed-form and iterative algorithms for computing means and variances under these metrics ensure theoretically meaningful and practically computable statistics as the basis for normalization (Wang et al., 1 Apr 2025).

3. Manifold-Aware Batch Normalization Architectures

The core pipeline for manifold-aware batch normalization consists of the following principle steps (Wang et al., 1 Apr 2025, Chen et al., 2024, Brooks et al., 2019, Chakraborty, 2020):

1. Metric Normalization (“Whitening”): Precondition each batch element $X_i$ by the inverse square root of $M$ (if $M$ is learnable), resulting in $\hat{X}_i = M^{-1/2}X_i M^{-1/2}$.
2. Centering: Map $\{\hat{X}_i\}$ to zero mean by computing their Fréchet mean $B_b$ and applying parallel transport/logarithmic map compositions to move each element to a canonical base point, typically the identity matrix $I$.
3. Variance Normalization (Scaling): Compute the batch variance $\nu_b^2$ in the tangent space at the canonical point and scale each element in tangent space by $s / \sqrt{\nu_b^2 + \epsilon}$ (where $s$ is a learnable scaling parameter).
4. Biasing: Move the scaled representations from the canonical point to a learnable bias point $G \in \mathcal{S}^{++}_d$ using the exponential map or group action.
5. Unwhitening: Undo the initial metric normalization, often with a matrix power transformation to facilitate additional deformations (e.g., in matrix power $\theta$-GBWBN, apply $(\cdot)^{1/\theta}$) (Wang et al., 1 Apr 2025).

A schematic pseudocode for the full pipeline in the $\theta$-GBWBN layer is:

Hat_Xi = M^{-1/2} X_i^\theta M^{-1/2} Hat_G = M^{-1/2} G^\theta M^{-1/2} B_b = FréchetMean_BW({Hat_Xi}) nu_b^2 = (1/n\theta^2)\sum d_{BW}^2(B_b, Hat_Xi) Xbar_i = PT_{B_b -> I}(Hat_Xi) Xcheck_i = Exp_I[(s/sqrt(nu^2 + epsilon)) * Log_I(Xbar_i)] Xtilde_i = PT_{I -> Hat_G}(Xcheck_i) X_plus_i = (M^{1/2} Xtilde_i M^{1/2})^{1/\theta}

(Wang et al., 1 Apr 2025)

Various design choices in the affine-invariant, Lie-group, and gyrogroup settings admit corresponding instantiations, but the structure—centering, scaling, biasing, “unwhitening”—is universal (Chen et al., 2024, Chen et al., 8 Sep 2025).

4. Robustness, Deformation, and Generalization in SPD Normalization

The practical effectiveness of manifold-aware BN, particularly under BW/GBW/θ-GBW metrics, is rooted in how these metrics respond to conditioning and curvature. The affine-invariant metric can result in gradient instabilities with ill-conditioned inputs, as its dependence is quadratic in eigenvalues near zero. By contrast, the Bures–Wasserstein metric shows linear sensitivity to eigenvalues, handling near-degenerate covariance matrices more stably (Wang et al., 1 Apr 2025):

• BW/GBW metrics: $g_X^{GBW}(S, S)$ involves only the Lyapunov solution, so numerical derivatives remain controlled even for nearly singular $X$.
• Learnable metric $M$ and power deformation $\theta$: $M$ adapts local geometry; $\theta$ interpolates between fully geodesic (GBW, $\theta=1$) and Log-Euclidean ($\theta \rightarrow 0$) regimes, providing further robustness and representational flexibility.

These mechanisms empirically drive features away from the manifold boundary (i.e., prevent extreme condition numbers), facilitating faster convergence and improved generalization. In deep networks for action recognition, EEG, and radar, manifold-aware BN with GBW/θ-GBW results in significant accuracy improvements over both AIM-based and Euclidean baselines, while enhancing the localization and interpretability of gradient maps (Wang et al., 1 Apr 2025).

5. Unified Perspectives: Lie Groups, Homogeneous Spaces, and Gyrogroups

Several recent frameworks extend manifold-aware BN beyond SPD matrices to general manifold classes:

• LieBN: Leverages the Lie group structure, employing group operations for centering (left translation), scaling in Lie algebra (via the exponential map), and biasing by group action. Deformation via power maps and metric parameterization (e.g., LEM, AIM, LCM) recovers and extends previous methods (Chen et al., 2024).
• ManifoldNorm: Abstracts normalization schemes for any homogeneous Riemannian manifold, with centering at the Fréchet mean, tangent space scaling, and bias via group action. Specializes to SPD matrices and spheres as key applications (Chakraborty, 2020).
• GyroBN: Generalizes to any pseudo-reductive gyrogroup, a structure encompassing SPD manifolds (with AIM, LEM, etc.), spheres, hyperbolic spaces, and the Grassmannian. GyroBN replaces addition by gyroaddition, scaling by gyroscaling, and bias by gyrotranslations, recovering LieBN and AIM-based BN as special cases and providing closed-form normalization steps whenever the gyrostructure is explicit (Chen et al., 8 Sep 2025).

These unified approaches establish rigorous control of intrinsic batch statistics (mean and variance) and guarantee that normalization commutes with group-based data transformations, supporting efficient and robust learning throughout the spectrum of non-Euclidean architectures.

6. Theoretical and Empirical Properties

Manifold-aware BN layers maintain key theoretical properties of their Euclidean analogues:

• Statistical Invariance: The normalization pipeline shifts the batch mean to the bias point and rescales variance by $s^2$ (the learnable scaling parameter) (Chen et al., 8 Sep 2025, Chen et al., 2024).
• Gradient Compatibility: All steps (exponentials, logarithms, Lyapunov maps, matrix powers) admit closed-form or differentiable-backprop implementations, leveraging the Daleckiĭ–Kreĭn formula and Lyapunov equation derivatives (Wang et al., 1 Apr 2025).
• Well-posedness and Convergence: In the context of network optimization, the inclusion of manifold-aware normalization can be recast as altering the Riemannian metric of the parameter space, resulting in gradient flows or even Wasserstein-gradient flows in the mean-field limit (Ma et al., 2021). Theoretical guarantees have been established for global minimization and well-posedness under standard regularity and convexity conditions (Ma et al., 2021).

Empirically, manifold-aware BN consistently improves prediction accuracy, stabilizes training, and enhances sample efficiency in diverse domains: EEG classification, skeleton action recognition, radar target identification, and high-dimensional medical imaging (Wang et al., 1 Apr 2025, Chen et al., 2024, Brooks et al., 2019, Chakraborty, 2020, Chen et al., 8 Sep 2025). Speedups arise from better-conditioned representations and, in gyro-structured settings, from closed-form normalization operations.

7. Applications, Implementation, and Future Directions

Manifold-aware BN is now a core component in SPDNet-type models, geometric deep learning for human action recognition, EEG-based paradigms, radar analysis, and manifold convolutional architectures for connectomics. Implementation best practices include:

• Efficient eigendecompositions and use of structure to minimize $O(d^3)$ costs for moderate $d$ (Brooks et al., 2019).
• Parallelization of batch-centric computations (Fréchet means, exp/log maps) (Wang et al., 1 Apr 2025).
• Hyperparameter tuning for deformation (e.g., θ in θ-GBW, $(\alpha, \beta)$ in LEM/LCM) (Chen et al., 2024).
• Robust handling of numerical stability near the boundary of the SPD manifold (clamping eigenvalues, ridge regularization) (Brooks et al., 2019, Chakraborty, 2020).

Emerging directions include generalized normalization on more exotic homogeneous spaces (e.g., flag manifolds, correlation manifolds), domain-specific normalization strategies (domain-specific momentum BN), and theoretical investigation of interactions between manifold geometry, normalization, and information propagation in ultra-deep architectures (Chen et al., 8 Sep 2025, Chen et al., 2024). Theoretical frameworks based on gyrovector spaces and Wasserstein geometry suggest broader unification and potential for further advances in normalization under geometric constraints.

---

Summary Table: Key Manifold-Aware BN Variants

Approach	Geometry/Metric	Center/Scale/Bias Mechanics
RBN/AIM (Brooks et al., 2019)	SPD, Affine-Invariant	Fréchet mean, tangent scaling, PT
LieBN (Chen et al., 2024)	Lie group (e.g. SPD)	Group mean, Lie algebra scaling, action
ManifoldNorm (Chakraborty, 2020)	Homogeneous manifolds	Fréchet mean, tangent/E-channel scale, group bias
GBWBN (Wang et al., 1 Apr 2025)	SPD, GBW/θ-GBW	Learnable M, BW mean, matrix power
GyroBN (Chen et al., 8 Sep 2025)	Pseudo-reductive gyrogroups	Gyro-barycenter, gyroscaling, gyrotranslation

References: (Wang et al., 1 Apr 2025, Chen et al., 2024, Brooks et al., 2019, Chakraborty, 2020, Chen et al., 8 Sep 2025, Ma et al., 2021).