Template-Based Wasserstein Embeddings

• Template-Based Wasserstein Embeddings are methods that employ reference distributions to derive feature representations using optimal transport geometry.
• They integrate iterative barycenter construction and linearized OT techniques to achieve scalable, computationally efficient dimensionality reduction and metric learning.
• They extend to structured data via Fused Gromov–Wasserstein embeddings, boosting performance in graph classification and archetype extraction tasks.

Template-based Wasserstein embeddings formalize the use of Wasserstein (optimal transport) geometry for statistical inference, dimensionality reduction, and representation learning by defining feature representations with respect to one or more reference distributions, or "templates". This methodology subsumes frameworks based on Wasserstein barycenters for distribution alignment and archetype extraction, linearized optimal transport (OT) embeddings for computationally efficient representation, and template-driven metric learning for graph structures. Template-based embeddings enable both theoretical tractability and scalable implementation by rooting inference in the geometry of the Wasserstein space, and are leveraged in a range of machine learning and data analysis applications.

1. Foundational Principles

Wasserstein metrics, particularly $W_2$, provide a notion of distance between probability measures by quantifying the minimum transportation cost required to morph one distribution into another. Let $\{\mu_j\}_{j=1}^J \subset P_2(\mathbb{R}^d)$ be a finite family of probability measures with weights $\lambda_j \geq 0$, $\sum_{j=1}^J \lambda_j = 1$. The Wasserstein barycenter $\mu^*$ is defined as the Fréchet minimizer:

$$\mu^* = \argmin_{\nu \in P_2(\mathbb{R}^d)} \sum_{j=1}^J \lambda_j\, W_2^2(\nu, \mu_j)$$

where $W_2^2(\nu,\mu_j)$ is given by

$$W_2^2(\nu,\mu_j) = \inf_{\pi \in \Pi(\nu,\mu_j)} \int_{\mathbb{R}^d \times \mathbb{R}^d} \|x - y\|^2\, d\pi(x,y)$$

and $\Pi(\nu, \mu_j)$ denotes the set of all couplings with prescribed marginals. Existence and uniqueness are guaranteed under mild regularity: no regularization is needed and uniqueness holds if at least one measure admits a density (Boissard et al., 2011).

Under a deformation model, observed measures $\mu_j$ can be written as pushforwards $\{\mu_j\}_{j=1}^J \subset P_2(\mathbb{R}^d)$0 of an unknown template $\{\mu_j\}_{j=1}^J \subset P_2(\mathbb{R}^d)$1 by random maps $\{\mu_j\}_{j=1}^J \subset P_2(\mathbb{R}^d)$2 in an admissible family (e.g., gradients of convex functions), centered such that $\{\mu_j\}_{j=1}^J \subset P_2(\mathbb{R}^d)$3. In this setting, the barycenter estimate converges to the template as the sample size grows.

2. Iterative and Linearized Template-Based Embeddings

Iterated Barycenter Construction

Computing the barycenter over many measures is generally intractable directly. The iterated barycenter scheme simplifies computation by induction: starting from $\{\mu_j\}_{j=1}^J \subset P_2(\mathbb{R}^d)$4, each step computes a two-measure barycenter,

$\{\mu_j\}_{j=1}^J \subset P_2(\mathbb{R}^d)$5

Convergence is established both for the true barycenter and the iterated approximation; as the number of input measures increases, concentration inequalities guarantee that the barycenter recovers the underlying template (Boissard et al., 2011).

Linearized Optimal Transport Embeddings

Cloninger–Hamm–Khurana–Moosmüller introduce a linearized OT approach ("LOT Wassmap") to template-based embedding (Cloninger et al., 2023). For a fixed reference (template) measure $\{\mu_j\}_{j=1}^J \subset P_2(\mathbb{R}^d)$6, the optimal Monge map $\{\mu_j\}_{j=1}^J \subset P_2(\mathbb{R}^d)$7 pushing $\{\mu_j\}_{j=1}^J \subset P_2(\mathbb{R}^d)$8 to $\{\mu_j\}_{j=1}^J \subset P_2(\mathbb{R}^d)$9 is computed, and the linearized embedding vector is

$\lambda_j \geq 0$0

with $\lambda_j \geq 0$1 the Kantorovich dual potential. The embedding map is

$\lambda_j \geq 0$2

The linearized Wasserstein (LOT) distance is

$\lambda_j \geq 0$3

which is isometric to $\lambda_j \geq 0$4 for compatible deformations and provides controlled approximation error otherwise. This embedding avoids $\lambda_j \geq 0$5 complexity by computing only $\lambda_j \geq 0$6 template-to-data maps.

3. Template-Based Embeddings for Structured Data

Template-driven Wasserstein frameworks have been extended to attributed graphs via Fused Gromov–Wasserstein (FGW) distances (Vincent-Cuaz et al., 2022). Let $\lambda_j \geq 0$7 denote an attributed graph with adjacency matrix $\lambda_j \geq 0$8, node features $\lambda_j \geq 0$9, and weights $\sum_{j=1}^J \lambda_j = 1$0, and $\sum_{j=1}^J \lambda_j = 1$1 be a template graph. The FGW distance is defined as

$\sum_{j=1}^J \lambda_j = 1$2

Here, $\sum_{j=1}^J \lambda_j = 1$3 encodes feature distances, $\sum_{j=1}^J \lambda_j = 1$4 encodes adjacency differences, and $\sum_{j=1}^J \lambda_j = 1$5 trades off feature vs. structure matching. Embedding a graph $\sum_{j=1}^J \lambda_j = 1$6 as a vector of FGW distances to $\sum_{j=1}^J \lambda_j = 1$7 templates,

$\sum_{j=1}^J \lambda_j = 1$8

yields a discriminative representation for classification tasks. All template parameters and $\sum_{j=1}^J \lambda_j = 1$9 can be learned end-to-end via backpropagation and Frank–Wolfe optimization.

4. Theoretical Guarantees and Error Bounds

Barycentric Convergence

Under a centered random deformation model, the Wasserstein barycenter $\mu^*$0 of $\mu^*$1 observed measures has

$\mu^*$2

As $\mu^*$3, $\mu^*$4 almost surely.

A nonasymptotic concentration bound guarantees for any $\mu^*$5,

$\mu^*$6

if the deformations $\mu^*$7 are bounded (Boissard et al., 2011).

LOT Embedding Approximation

For linearized OT embeddings, the discrepancy between true ($\mu^*$8) and linearized ($\mu^*$9) distances satisfies

$$\mu^* = \argmin_{\nu \in P_2(\mathbb{R}^d)} \sum_{j=1}^J \lambda_j\, W_2^2(\nu, \mu_j)$$0

where $$\mu^* = \argmin_{\nu \in P_2(\mathbb{R}^d)} \sum_{j=1}^J \lambda_j\, W_2^2(\nu, \mu_j)$$1 quantifies departure from template compatibility.

Sampling and regularization contribute additional $$\mu^* = \argmin_{\nu \in P_2(\mathbb{R}^d)} \sum_{j=1}^J \lambda_j\, W_2^2(\nu, \mu_j)$$2 error terms, where $$\mu^* = \argmin_{\nu \in P_2(\mathbb{R}^d)} \sum_{j=1}^J \lambda_j\, W_2^2(\nu, \mu_j)$$3 is the number of samples per measure and $$\mu^* = \argmin_{\nu \in P_2(\mathbb{R}^d)} \sum_{j=1}^J \lambda_j\, W_2^2(\nu, \mu_j)$$4 from $$\mu^* = \argmin_{\nu \in P_2(\mathbb{R}^d)} \sum_{j=1}^J \lambda_j\, W_2^2(\nu, \mu_j)$$5. In manifold settings, the overall embedding error in $$\mu^* = \argmin_{\nu \in P_2(\mathbb{R}^d)} \sum_{j=1}^J \lambda_j\, W_2^2(\nu, \mu_j)$$6 dimensions is $$\mu^* = \argmin_{\nu \in P_2(\mathbb{R}^d)} \sum_{j=1}^J \lambda_j\, W_2^2(\nu, \mu_j)$$7, with $$\mu^* = \argmin_{\nu \in P_2(\mathbb{R}^d)} \sum_{j=1}^J \lambda_j\, W_2^2(\nu, \mu_j)$$8 controlling manifold deviation and $$\mu^* = \argmin_{\nu \in P_2(\mathbb{R}^d)} \sum_{j=1}^J \lambda_j\, W_2^2(\nu, \mu_j)$$9 the embedding/sampling error (Cloninger et al., 2023).

5. Computational Aspects and Implementation

Computational tractability is achieved via several algorithmic innovations:

• Iterative barycenters reduce multilateral OT barycenter computation to $W_2^2(\nu,\mu_j)$0 two-measure problems, each solvable with algorithms such as Benamou–Brenier ($W_2^2(\nu,\mu_j)$1) or Sinkhorn regularized OT ($W_2^2(\nu,\mu_j)$2).
• Template-based LOT embeddings avoid $W_2^2(\nu,\mu_j)$3 pairwise OT solves by mapping all measures to a fixed template and employing principal component analysis or SVD for dimensionality reduction (Cloninger et al., 2023).
• FGW embedding layers exploit conditional gradient Frank–Wolfe optimization, with cost per graph–template pair $W_2^2(\nu,\mu_j)$4, and computational parallelism across templates (Vincent-Cuaz et al., 2022).

Common OT libraries (e.g., POT, GeomLoss, Sinkhorn) are suitable for prototypical implementation. Regularization (e.g., Gaussian smoothing) can be used to ensure uniqueness and numerical stability.

6. Applications and Empirical Findings

Template-based Wasserstein embeddings find application in template estimation, unsupervised and supervised learning on distributional or structured inputs, and dimensionality reduction in Wasserstein space:

• Template estimation: Wasserstein barycenters provably recover unknown template measures from deformed samples under structural assumptions (Boissard et al., 2011).
• Dimensionality reduction: LOT Wassmap reveals low-dimensional manifold structure in Wasserstein space, with substantial computational savings and provable guarantees even in the presence of sampling noise (Cloninger et al., 2023).
• Graph classification: FGW-based template embeddings achieve state-of-the-art accuracy on benchmarks (e.g., MUTAG, PTC, ENZYMES) as single layers or composed with GNNs. They outperform classic GNNs and kernel methods in tasks requiring sensitivity to higher-order structure. The number of templates required is typically small (4–6), and joint learning of template weights and trade-off parameters further improves performance (Vincent-Cuaz et al., 2022).

Experimental findings in the referenced works emphasize theoretical validation and controlled studies (toy datasets, structured synthetic data). Large-scale real-world or image-based scenarios are identified as open for further empirical exploration.

7. Extensions and Limitations

Extensions include alternative OT metrics ($W_2^2(\nu,\mu_j)$5 for $W_2^2(\nu,\mu_j)$6), parametric deformation models, and geodesic PCA in Wasserstein space. For graph data, additional regularizations (e.g., entropy) can be incorporated. A common limitation is the curse of dimensionality in high-dimensional ambient spaces; while linearization and entropic regularization ameliorate some costs, practical performance at scale is yet to be established. Barycentric associativity fails in $W_2^2(\nu,\mu_j)$7 unless further structure is imposed; thus, care is required in application beyond compatible deformation regimes.

Template-based Wasserstein embeddings thus synthesize archetypal modeling, geometry, and scalable computation, enabling principled statistical learning in settings where distributional structure is fundamental (Boissard et al., 2011, Vincent-Cuaz et al., 2022, Cloninger et al., 2023).