---
title: Template-Based Wasserstein Embeddings
url: https://www.emergentmind.com/topics/template-based-wasserstein-embeddings
type: topic
---

# Template-Based Wasserstein Embeddings

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 [1111.5927].

Under a deformation model, observed measures $\mu_j$ can be written as pushforwards $\mu_j = T_{j\#}\mu$ of an unknown template $\mu$ by random maps $T_j$ in an admissible family (e.g., gradients of convex functions), centered such that $\mathbb{E} T_j(x) = x$. 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 $\nu_1 = \mu_1$, each step computes a two-measure barycenter,

\[
\nu_j = \argmin_{\nu \in P_2(\mathbb{R}^d)} \left\{ \left(\sum_{i=1}^{j-1} \lambda_i\right) W_2^2(\nu, \nu_{j-1}) + \lambda_j W_2^2(\nu, \mu_j) \right\}
\]

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 [1111.5927].

### Linearized Optimal Transport Embeddings

Cloninger–Hamm–Khurana–Moosmüller introduce a linearized OT approach ("LOT Wassmap") to template-based embedding [2302.07373]. For a fixed reference (template) measure $\mu_0$, the optimal Monge map $T_{\mu_0}^{\mu_k}$ pushing $\mu_0$ to $\mu_k$ is computed, and the linearized embedding vector is

\[
u_k(x) := T_{\mu_0}^{\mu_k}(x) - x = \nabla\phi_k(x) - x
\]

with $\phi_k$ the Kantorovich dual potential. The embedding map is

\[
\mathrm{Emb}: \mu_k \mapsto u_k \in L^2(\mu_0; \mathbb{R}^n)
\]

The linearized Wasserstein (LOT) distance is

\[
d_{\mathrm{LOT}}(\mu_i, \mu_j) := \|u_i - u_j\|_{L^2(\mu_0)}
\]

which is isometric to $W_2$ for compatible deformations and provides controlled approximation error otherwise. This embedding avoids $O(N^2)$ complexity by computing only $N$ 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 [2205.15733]. Let $G=(V, E, X, h)$ denote an attributed graph with adjacency matrix $C$, node features $X$, and weights $h$, and $T$ be a template graph. The FGW distance is defined as

\[
\mathrm{FGW}_p^\alpha(G, T) = \min_{\Pi \in \mathcal{U}(h, h')} (1 - \alpha) \langle C_X, \Pi \rangle + \alpha \langle C_A, \Pi \otimes \Pi \rangle
\]

Here, $C_X$ encodes feature distances, $C_A$ encodes adjacency differences, and $\alpha$ trades off feature vs. structure matching. Embedding a graph $G$ as a vector of FGW distances to $K$ templates,

\[
\phi(G) = [d_1, d_2, \dots, d_K]^\top,\quad d_k = \mathrm{FGW}_p^\alpha(G, T_k)
\]

yields a discriminative representation for classification tasks. All template parameters and $\alpha$ 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_B$ of $J$ observed measures has

\[
\mu_B = \left(\frac{1}{J} \sum_{j=1}^J T_j\right)_\# \mu,\qquad W_2(\mu_B, \mu) = \left\|\frac{1}{J} \sum_{j=1}^J T_j - \mathrm{id} \right\|_{L^2(\mu)}
\]

As $J \rightarrow \infty$, $\mu_B \to \mu$ almost surely.

A nonasymptotic concentration bound guarantees for any $\varepsilon > 0$,

\[
\Pr\{W_2(\mu_B, \mu) \geq \varepsilon\} \leq 2 \exp\left(-J \frac{\varepsilon^2}{M^2 (1 + c\,\varepsilon/M)}\right)
\]

if the deformations $T_j$ are bounded [1111.5927].

### LOT Embedding Approximation

For linearized OT embeddings, the discrepancy between true ($W_2$) and linearized ($d_{\mathrm{LOT}}$) distances satisfies

\[
|W_2(\mu_i, \mu_j) - d_{\mathrm{LOT}}(\mu_i, \mu_j)| \leq C \varepsilon^\alpha
\]

where $\varepsilon$ quantifies departure from template compatibility.

Sampling and regularization contribute additional $O_p(k^{-r}) + O_p(m^{-1/2})$ error terms, where $k$ is the number of samples per measure and $m$ from $\mu_0$. In manifold settings, the overall embedding error in $d$ dimensions is $O(\tau_1+\tau_2)$, with $\tau_1$ controlling manifold deviation and $\tau_2$ the embedding/sampling error [2302.07373].

## 5. Computational Aspects and Implementation

Computational tractability is achieved via several algorithmic innovations:

- **Iterative barycenters** reduce multilateral OT barycenter computation to $O(J)$ two-measure problems, each solvable with algorithms such as Benamou–Brenier ($O(n^3)$) or Sinkhorn regularized OT ($O(n^2)$).
- **Template-based LOT** embeddings avoid $O(N^2)$ pairwise OT solves by mapping all measures to a fixed template and employing principal component analysis or SVD for dimensionality reduction [2302.07373].
- **FGW embedding layers** exploit conditional gradient Frank–Wolfe optimization, with cost per graph–template pair $O(n^2 m + m^2 n)$, and computational parallelism across templates [2205.15733].

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 [1111.5927].
- **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 [2302.07373].
- **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 [2205.15733].

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_p$ for $p \neq 2$), 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 $d > 1$ 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 [1111.5927, 2205.15733, 2302.07373].

Source: https://www.emergentmind.com/topics/template-based-wasserstein-embeddings