---
title: Conditional Wasserstein Gradient
url: https://www.emergentmind.com/topics/conditional-wasserstein-gradient
type: topic
---

# Conditional Wasserstein Gradient

The conditional Wasserstein gradient is a central concept in optimal transport for learning and generative modeling when conditional distributions or structured data (e.g., labeled datasets, posteriors in inverse problems) are involved. It provides a principled mechanism to compute and optimize gradients derived from conditional Wasserstein distances, particularly via restricted couplings in joint or mixture spaces. This enables the design of tractable dynamics for tasks such as conditional generative modeling, domain adaptation, dataset distillation, and Bayesian inverse problems, and underpins algorithms such as conditional Wasserstein GANs, conditional optimal transport flow matching, and Wasserstein over Wasserstein flows [2506.07534, 2310.13433, 2403.18705, 1903.00395].

## 1. Definition of Conditional Wasserstein Distance

The conditional Wasserstein distance $W_{p,Y}$ quantifies the optimal transport cost between two joint distributions over input-output pairs, but restricts transport plans to only move mass between outputs with the same input condition. For two joint laws $(Y, X) \sim P_{Y,X}$ and $(Y, Z) \sim Q_{Y,Z}$ on $A \times B$:

\[
W_{p,Y}^p(P_{Y,X}, Q_{Y,Z}) = \mathbb{E}_{y \sim P_Y} \left[ W_p^p\big(P_{X | Y = y}, Q_{Z | Y = y}\big) \right]
\]

Here, the expectation is over the conditioning variable $Y$ and $W_p^p$ is the standard $p$-Wasserstein distance between conditionals [2310.13433, 2403.18705].

The set of admissible couplings, denoted $\Gamma_Y$, only allows mass movement on the diagonal in $Y$, ensuring label or condition preservation. This construction is crucial wherever fidelity of conditional structure (e.g., class conditionals, posteriors) must be maintained.

## 2. Primal and Dual Formulations; Geometric Properties

### Primal Formulation

The optimal transport formulation for $W_{p,Y}$ restricts couplings $\alpha$ to those supported only on pairs $\left((y, x_1), (y, x_2)\right)$ with identical $y$ values. The distance is then the infimum over such couplings:

\[
W_{p,Y}(P_{Y,X}, Q_{Y,Z}) = \left( \inf_{\alpha \in \Gamma_Y} \int \| x_1 - x_2 \|^p\, d\alpha \right)^{1/p}
\]

### Dual Formulation

For $p=1$, a Kantorovich-type duality gives:

\[
W_{1,Y}(P_{Y,X}, Q_{Y,Z}) = \sup_{h \in \mathcal F} \left\{ \mathbb{E}_{Y,X}[h(Y,X)] - \mathbb{E}_{Y,Z}[h(Y,Z)] \right\}
\]

with $\mathcal F$ the set of functions $h : A \times B \to \mathbb{R}$ that are bounded and, for each $y$, $x \mapsto h(y,x)$ is 1-Lipschitz [2403.18705, 2310.13433].

### Geodesics and Continuity Equation

In the induced metric space, geodesics interpolate between conditionals; under regularity, geodesics $\mu_t$ are given by pushforward along linear interpolation between optimal pairs, and velocity fields $v_t$ satisfy the continuity equation with no $Y$-transport:

\[
\partial_t \mu_t + \nabla \cdot (v_t\,\mu_t) = 0
\]
\[
\partial_t \mu_t(x | y) + \nabla_x \cdot (\mu_t(x | y)\,v_t(x|y)) = 0
\]

The velocity $v_t(x|y)$ is the displacement field given by the conditional optimal transport maps [2403.18705].

## 3. Conditional Wasserstein Gradients in Optimization

### Gradient Computation

When the target is to minimize a functional defined via conditional Wasserstein distance (e.g., in conditional GANs, flow matching, or gradient flows on random measures), the functional gradient w.r.t. generator parameters or particles is derived from the dual form.

In the conditional WGAN framework, the generator gradient with respect to parameters $\theta$ is:

\[
\nabla_\theta W_{1,Y} \approx -\frac{1}{N} \sum_{i=1}^N \nabla_x h_{\phi^*}(y_i, G_\theta(y_i, z_i)) \cdot \nabla_\theta G_\theta(y_i, z_i)
\]

where $h_{\phi^*}$ is an approximately optimal critic function [2310.13433]. The same formalism applies to more general conditionally-structured flows.

### Gradient Flow PDE

For random measure representations (e.g., mixtures of class-conditionals), the gradient flow induced by a WoW or conditional Wasserstein functional $F$ is:

\[
\partial_t P_t + \mathrm{Div}_{\mathcal{P}_2}\big(P_t \, \nabla_W F(P_t)\big) = 0
\]

and is discretized by forward Euler steps using the WoW (outer) gradient. The explicit update for atoms $x_{i, k}^c$ associated with class-conditional measures is:

\[
x_{i, k+1}^c = x_{i, k}^c - \tau\, \nabla_W F(P_k)(\mu_k^c)(x_{i, k}^c)
\]

where $\nabla_W F(P_k)(\mu)$ depends on both the outer and inner gradients [2506.07534].

## 4. Algorithmic Schemes and Practical Computation

### Conditional WGAN with Gradient Penalty

Training a conditional WGAN requires enforcing 1-Lipschitz continuity of the critic in $x$ (conditioned on $y$), typically via the gradient penalty proposed by Gulrajani et al.:

\[
\mathbb{E}_{(\hat x, y)} \left( \|\nabla_{\hat x} D(\hat x, y)\|_2 - 1 \right)^2
\]

where $\hat x = \alpha x + (1 - \alpha) G(y)$. The critic and generator are conditioned on $y$ through input concatenation and parameter-sharing [1903.00395].

### Conditional OT Flow Matching

Numerical schemes for conditional OT flow matching approximate the $Y$-diagonal coupling via a penalized cost:

\[
d_\beta^p\big((y_1, x_1), (y_2, x_2)\big) = \|x_1 - x_2\|^p + \beta\,\|y_1 - y_2\|^p
\]

with large $\beta$ ensuring almost-perfect label preservation. Empirical minimization of the squared error between model-predicted and target velocities is performed via stochastic gradient descent [2403.18705].

### Wasserstein-over-Wasserstein (WoW) Gradient Flows

For datasets modeled as mixtures over class-conditionals, the WoW gradient flows optimize functionals (e.g., MMDs with Sliced-Wasserstein kernels) over the space $\mathcal{P}(\mathcal{P}_2)$, updating atomic particles via coupled within-class and between-class interactions [2506.07534].

## 5. Applications and Empirical Significance

| Application Area         | Conditional Gradient’s Role                                | Empirical Finding (as reported)              |
|-------------------------|-------------------------------------------------------------|-----------------------------------------------|
| Bayesian Inverse Problems | Posterior sampling by conditional generative modeling      | Lower Sinkhorn/FID distances, class fidelity [2403.18705] |
| Image-to-Image Translation | Stable, conditional mapping with cWGAN + GP               | Improved PSNR/SSIM and perceptual quality [1903.00395]  |
| Dataset Distillation     | WoW gradient flows for matching class-conditional mixtures | Achieves synthetic sets matching real data in WoW sense [2506.07534] |

The conditional Wasserstein gradient ensures that generative or transport flows preserve conditional structure (labels, classes, input-output mappings) even in high-dimensional or inverse settings. Experiments demonstrate that as the conditional structure is more strictly enforced (e.g., via increasing penalty $\beta$), label-fidelity and data quality metrics (e.g., FID, PSNR, SSIM) improve, and mode collapse is mitigated [2403.18705, 1903.00395].

## 6. Connections to Ordinary and Higher-Order Wasserstein Gradients

A key property is that if the condition variable $Y$ is independent of $X$ or $Z$, the conditional Wasserstein collapses to the ordinary marginal Wasserstein distance, i.e., $W_{p,Y} = W_p$ on the marginals.

The framework unifies several perspectives: it generalizes ordinary Wasserstein gradients by coupling with conditional structure, extends to higher-order mixture models (random measures), and admits Riemannian structure enabling explicit geodesics, gradient flows, and discrete optimization in mixture spaces [2506.07534].

A plausible implication is that conditional Wasserstein gradients are the natural tool for generative modeling and inference wherever structured or hierarchical data is essential, and their mathematical structure enables tractable algorithms and theoretical analysis.

## 7. Related Work and Ongoing Directions

Conditional Wasserstein gradients underpin several modern algorithmic paradigms:

- Conditional WGANs and their gradient-penalty variants for conditional generation and image restoration [1903.00395, 2310.13433].
- Conditional optimal transport flow matching and Bayesian flows for improved posterior alignment and conditional image synthesis [2403.18705].
- WoW-gradient flows for mixture-matching and dataset distillation across domains or classes [2506.07534].

Ongoing extensions focus on scalable relaxations of conditional couplings, tractable kernel choices (e.g., Sliced-Wasserstein), and adaptation to high-dimensional and infinite-dimensional spaces.

These frameworks are converging toward providing principled, computationally feasible tools for conditional, structured, and distributional learning in modern machine learning.

Source: https://www.emergentmind.com/topics/conditional-wasserstein-gradient