Staying on the Manifold: Geometry-Aware Noise Injection

TL;DR

This work argues that perturbations during training should respect the intrinsic geometry of data, not just operate in ambient space. It introduces three geometry-aware noise strategies—tangent-space projection, geodesic noise via the exponential map, and intrinsic Brownian motion—grounded in Riemannian geometry, and extends them to deformed and learned manifolds. Through controlled experiments on Swiss roll, spheroids, and MNIST, the authors demonstrate improved generalisation and robustness on curved manifolds, with Brownian motion often offering a favorable balance of performance and computational efficiency. The framework highlights how manifold-aware augmentation can outperform traditional ambient noise, especially when data lie on highly curved structures, while also outlining limitations related to manifold approximation quality and computational cost.

Abstract

It has been shown that perturbing the input during training implicitly regularises the gradient of the learnt function, leading to smoother models and enhancing generalisation. However, previous research mostly considered the addition of ambient noise in the input space, without considering the underlying structure of the data. In this work, we propose several strategies of adding geometry-aware input noise that accounts for the lower dimensional manifold the input space inhabits. We start by projecting ambient Gaussian noise onto the tangent space of the manifold. In a second step, the noise sample is mapped on the manifold via the associated geodesic curve. We also consider Brownian motion noise, which moves in random steps along the manifold. We show that geometry-aware noise leads to improved generalisation and robustness to hyperparameter selection on highly curved manifolds, while performing at least as well as training without noise on simpler manifolds. Our proposed framework extends to data manifolds approximated by generative models and we observe similar trends on the MNIST digits dataset.

Figures (10)

• Figure 1: Noise injection is a data augmentation technique that can improve generalisation. For a data point () lying on a lower-dimensional manifold, sampling noise in the ambient space () almost surely deviates from the input manifold whereas a sample from a geometry-aware noise process ( ) stays on the manifold and respects the data geometry. Illustration of the biconcave disc that resembles a red blood cell.
• Figure 2: Noise injection strategies with increasing level of conceptual complexity, i.e. ambient space noise (), tangent space noise ( ) and geodesic noise ( ). The Brownian motion strategy is visualised in Figure \ref{['fig:brownian-motion']}.
• Figure 3: Brownian motion from an initial point () in the parameter space (left), and mapped to the manifold (right) via the chart $X$. The endpoint of the Brownian motion on the manifold ( ) acts as the noisy observation.
• Figure 4: The deformation process of the sphere in $\mathbb{R}^3$ for increasing time steps of using a flow field $v_t$.
• Figure 5: Test loss on the SwissRoll as a function of noise intensity $\sigma^2$ for different noise injection strategies. The geometry-aware noise strategies that stay on the manifold, i.e. geodesic noise and Brownian motion noise, show greater robustness to the noise intensity compared to ambient or tangential noise. Our strategies perform at least as well as training without noise (dashed line).

Theorems & Definitions (1)

• Definition 2.1