Conditional Wasserstein Gradient

• Conditional Wasserstein Gradient is a concept in optimal transport that computes gradients from conditional distributions while preserving structured data fidelity.
• It employs restricted couplings and dual formulations to stabilize training in models like conditional WGANs, enabling tasks such as domain adaptation and dataset distillation.
• Applications include Bayesian inverse problems and image-to-image translation, where improved metrics like FID, PSNR, and SSIM demonstrate its practical impact.

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 (Bonet et al., 9 Jun 2025, Chemseddine et al., 2023, Chemseddine et al., 2024, Ebenezer et al., 2019).

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 (Chemseddine et al., 2023, Chemseddine et al., 2024).

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 $(Y, X) \sim P_{Y,X}$0 restricts couplings $(Y, X) \sim P_{Y,X}$1 to those supported only on pairs $(Y, X) \sim P_{Y,X}$2 with identical $(Y, X) \sim P_{Y,X}$3 values. The distance is then the infimum over such couplings:

$(Y, X) \sim P_{Y,X}$4

Dual Formulation

For $(Y, X) \sim P_{Y,X}$5, a Kantorovich-type duality gives:

$(Y, X) \sim P_{Y,X}$6

with $(Y, X) \sim P_{Y,X}$7 the set of functions $(Y, X) \sim P_{Y,X}$8 that are bounded and, for each $(Y, X) \sim P_{Y,X}$9, $(Y, Z) \sim Q_{Y,Z}$0 is 1-Lipschitz (Chemseddine et al., 2024, Chemseddine et al., 2023).

Geodesics and Continuity Equation

In the induced metric space, geodesics interpolate between conditionals; under regularity, geodesics $(Y, Z) \sim Q_{Y,Z}$1 are given by pushforward along linear interpolation between optimal pairs, and velocity fields $(Y, Z) \sim Q_{Y,Z}$2 satisfy the continuity equation with no $(Y, Z) \sim Q_{Y,Z}$3-transport:

$(Y, Z) \sim Q_{Y,Z}$4

$(Y, Z) \sim Q_{Y,Z}$5

The velocity $(Y, Z) \sim Q_{Y,Z}$6 is the displacement field given by the conditional optimal transport maps (Chemseddine et al., 2024).

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 $(Y, Z) \sim Q_{Y,Z}$7 is:

$(Y, Z) \sim Q_{Y,Z}$8

where $(Y, Z) \sim Q_{Y,Z}$9 is an approximately optimal critic function (Chemseddine et al., 2023). 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 $A \times B$0 is:

$A \times B$1

and is discretized by forward Euler steps using the WoW (outer) gradient. The explicit update for atoms $A \times B$2 associated with class-conditional measures is:

$A \times B$3

where $A \times B$4 depends on both the outer and inner gradients (Bonet et al., 9 Jun 2025).

4. Algorithmic Schemes and Practical Computation

Conditional WGAN with Gradient Penalty

Training a conditional WGAN requires enforcing 1-Lipschitz continuity of the critic in $A \times B$5 (conditioned on $A \times B$6), typically via the gradient penalty proposed by Gulrajani et al.:

$A \times B$7

where $A \times B$8. The critic and generator are conditioned on $A \times B$9 through input concatenation and parameter-sharing (Ebenezer et al., 2019).

Conditional OT Flow Matching

Numerical schemes for conditional OT flow matching approximate the $$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]$$0-diagonal coupling via a penalized cost:

$$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]$$1

with large $$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]$$2 ensuring almost-perfect label preservation. Empirical minimization of the squared error between model-predicted and target velocities is performed via stochastic gradient descent (Chemseddine et al., 2024).

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 $$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]$$3, updating atomic particles via coupled within-class and between-class interactions (Bonet et al., 9 Jun 2025).

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 (Chemseddine et al., 2024)
Image-to-Image Translation	Stable, conditional mapping with cWGAN + GP	Improved PSNR/SSIM and perceptual quality (Ebenezer et al., 2019)
Dataset Distillation	WoW gradient flows for matching class-conditional mixtures	Achieves synthetic sets matching real data in WoW sense (Bonet et al., 9 Jun 2025)

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 $$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]$$4), label-fidelity and data quality metrics (e.g., FID, PSNR, SSIM) improve, and mode collapse is mitigated (Chemseddine et al., 2024, Ebenezer et al., 2019).

6. Connections to Ordinary and Higher-Order Wasserstein Gradients

A key property is that if the condition variable $$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]$$5 is independent of $$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]$$6 or $$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]$$7, the conditional Wasserstein collapses to the ordinary marginal Wasserstein distance, i.e., $$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]$$8 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 (Bonet et al., 9 Jun 2025).

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 (Ebenezer et al., 2019, Chemseddine et al., 2023).
• Conditional optimal transport flow matching and Bayesian flows for improved posterior alignment and conditional image synthesis (Chemseddine et al., 2024).
• WoW-gradient flows for mixture-matching and dataset distillation across domains or classes (Bonet et al., 9 Jun 2025).

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.