Multiscaling in Wasserstein Spaces

1. Problem Statement

The paper addresses the lack of multiscale analysis frameworks for sequences of probability measures evolving within Wasserstein spaces ($P_p(\mathbb{R}^d)$). While multiscale methods (like wavelets and subdivision schemes) are ubiquitous in Euclidean signal processing, extending them to the non-linear, geometric setting of optimal transport is challenging.

• The Gap: Existing tools struggle to analyze the evolution of probability distributions (e.g., image sequences, particle clouds, neural network weight distributions) that possess intrinsic geometric structures not captured by linear vector spaces.
• The Goal: To construct a multiscale transform that can decompose sequences of measures into coarse approximations and detail coefficients, enabling tasks such as denoising, anomaly detection, and the quantification of "optimality" (geodesic deviation) in measure flows.

2. Methodology

The authors propose a novel framework that adapts subdivision schemes and multiresolution analysis to the metric geometry of Wasserstein spaces. The methodology relies on three core pillars:

A. Geometric Operators (Addition and Subtraction)

To perform arithmetic on measures, the authors define binary operators analogous to vector addition and subtraction, leveraging the Monge map (optimal transport map) and McCann's interpolants:

• Subtraction ($\ominus$): Defined as $\nu \ominus \mu = T^\nu_\mu - I$, where $T^\nu_\mu$ is the optimal transport map pushing $\mu$ to $\nu$, and $I$ is the identity. This maps the difference to the tangent space of the Wasserstein manifold.
• Addition ($\oplus$): Defined as $\mu \oplus \psi = (I + \psi)_\# \mu$, where $\psi$ is a vector field (tangent vector). This pushes the measure $\mu$ along the field $\psi$.
• Consistency: These operators satisfy $\mu \oplus (\nu \ominus \mu) = \nu$, allowing for perfect reconstruction.

B. The Refinement Operator (Subdivision)

The framework uses McCann's interpolants as the refinement mechanism. For two measures $\mu_i$ and $\mu_{i+1}$, the midpoint is defined as:
$$M(\mu_i, \mu_{i+1}; t) = ((1-t)I + t T^{\mu_{i+1}}_{\mu_i})_\# \mu_i$$
This operator generates a constant-speed geodesic between measures. The elementary subdivision scheme $S$ takes a coarse sequence $\{\mu_i\}$ and produces a refined sequence by keeping original points and inserting midpoints:
$$(S\mu)_{2i} = \mu_i, \quad (S\mu)_{2i+1} = M(\mu_i, \mu_{i+1}; 1/2)$$

C. The Multiscale Transform

The transform decomposes a sequence of measures $\mu^{(J)}$ into a pyramid:

1. Analysis (Decomposition):
   • Downsampling: $D\mu^{(\ell)}$ selects even-indexed measures.
   • Prediction: Apply $S$ to the downsampled sequence to predict the odd-indexed measures.
   • Details: Compute the difference between the actual measures and the prediction: $\psi^{(\ell)} = \mu^{(\ell)} \ominus S\mu^{(\ell-1)}$.
2. Synthesis (Reconstruction):
   • Reconstruct the sequence by adding details back to the prediction: $\mu^{(\ell)} = S\mu^{(\ell-1)} \oplus \psi^{(\ell)}$.

D. The Optimality Number

A novel scalar metric, $\omega$, is introduced to quantify how much a sequence deviates from a geodesic flow:
$$\omega(\mu^{(J)}) = \sum_{\ell=1}^J \|\psi^{(\ell)}\|_1$$
where $\|\cdot\|_1$ is the sum of $L^p$ norms of the detail coefficients. A value of 0 indicates the sequence follows a constant-speed geodesic (optimal transport); higher values indicate irregular dynamics or anomalies.

3. Key Contributions

1. Unified Framework for Continuous and Discrete Measures:

   • The method works for absolutely continuous measures (using Monge maps and density interpolation) and discrete measures (using coupling matrices and linear programming for optimal transport).
   • It handles mixed types and provides specific algorithms for both cases.

2. Theoretical Guarantees:

   • Geometric Decay: The authors prove that if a sequence is sampled from a smooth curve with a finite metric derivative (Lipschitz flow), the detail coefficients decay geometrically across scales ($\|\psi^{(\ell)}\|_\infty \leq C 2^{-\ell}$).
   • Stability: They establish the stability of the inverse transform, proving that small perturbations in detail coefficients result in proportionally small errors in the reconstructed sequence.

3. The Optimality Number:

   • Introduces a scale-invariant scalar to detect anomalies and irregular dynamics. Unlike standard distance metrics, it captures local and global deviations from geodesic behavior across multiple resolutions.

4. Computational Efficiency:

   • The transform scales linearly with the sequence length. For absolutely continuous measures, it reduces to solving Monge-Ampère equations (or entropic regularization), which are computationally tractable.

4. Results and Numerical Experiments

The paper validates the framework through three distinct numerical experiments:

• Gaussian Flows (Denoising & Anomaly Detection):

  • Applied to synthetic sequences of Gaussian measures.
  • Denoising: By thresholding large detail coefficients (which correspond to noise), the authors successfully reconstructed smooth geodesic curves from noisy data.
  • Anomaly Detection: Jump discontinuities in the variance of the Gaussian sequence were successfully localized by spikes in the detail coefficients at specific scales.
  • Optimality: The optimality number $\omega$ was 0 for perfect geodesics and increased linearly as the curve was perturbed away from the geodesic.

• Point Cloud Dynamics (Electromagnetism):

  • Simulated 10 particles evolving under an electric field generated by two charges.
  • The multiscale transform captured the smooth flow of the point cloud.
  • When noise was added to the particle positions, the geometric decay of coefficients vanished, and the optimality number increased significantly, demonstrating the method's sensitivity to non-geodesic perturbations.

• Neural Network Learning Trajectories:

  • Analyzed the evolution of probability distributions over the output classes of a CNN (MNIST) during training.
  • The detail coefficients showed geometric decay, confirming that the learning trajectory in the high-dimensional weight space follows a smooth path.
  • The authors suggest the optimality number could serve as a model-agnostic metric to evaluate the efficiency of different optimizers or learning rates in generative models.

5. Significance

• Bridging Geometry and Data Analysis: This work provides a rigorous mathematical bridge between optimal transport theory and multiscale signal processing, allowing for the analysis of "data as distributions" rather than just vectors.
• New Diagnostic Tools: The Optimality Number offers a new way to diagnose the quality of dynamical systems, generative models, and learning algorithms by measuring their adherence to geodesic principles.
• Versatility: The ability to handle both discrete (point clouds) and continuous (densities) measures makes the framework applicable to a wide range of fields, including computer vision, physics simulations, and deep learning.
• Foundation for Future Work: The paper opens avenues for developing "Wasserstein Wavelets" and using these transforms for compression, interpolation, and regularization in machine learning tasks involving probability distributions.