Conditional Lagrangian Optimal Transport

• Conditional Lagrangian Optimal Transport is a framework that extends classical optimal transport by integrating Lagrangian mechanics and conditional probability to form geodesics in conditional measure spaces.
• Its methodology minimizes a Lagrangian action under mass conservation and conditional boundary constraints, producing unique dynamic transport paths.
• Applications include conditional generative modeling, time series imputation, surrogate modeling, and hyperparameter inference, enabling simulation-free, data-adaptive implementations.

Conditional Lagrangian Optimal Transport (CLOT) synthesizes concepts from optimal transport (OT), Lagrangian mechanics, and conditional probability to solve transport problems under conditioning constraints. It is characterized by the minimization of Lagrangian action subject to both mass conservation and explicit conditional structure, enabling dynamic, data-adaptive transport plans as geodesics in measure space. This framework generalizes classical OT with path-dependent (and often data-dependent) costs and encompasses both deterministic and stochastic variants, providing both a theoretical foundation and practical methodologies for conditional generative modeling, time series imputation, surrogate modeling under varying parameters, and more.

1. Theoretical Foundations and Formulation

Conditional Lagrangian OT extends the classical Monge-Kantorovich framework by introducing a dynamic, action-minimizing path formulation with explicit conditioning on certain coordinates or variables. Formally, given a source distribution $\rho_0(x)$ and a target distribution $\rho_T(x)$, both conditional on observed entries or parameters ($x^{\text{cond}}$), the goal is to find a time-dependent density $\rho(t, x \mid x^{\text{cond}})$ and a velocity field $v(t, x \mid x^{\text{cond}})$ minimizing the action

$$A[v] = \int_0^T \int_{\mathbb{R}^d} \frac{1}{2} \|v(t,x)\|^2 \rho(t,x) \, dx \, dt$$

subject to the continuity equation $$\partial_t \rho(t,x) + \nabla \cdot [\rho(t,x) v(t,x)] = 0$$ and hard conditioning in the boundary constraints $\rho(0, \cdot) = \rho_0(\cdot | x^{\text{cond}})$, $\rho(T, \cdot) = \rho_T(\cdot | x^{\text{cond}})$.

The action-minimizing curve—subject to these constraints—is uniquely characterized as a geodesic in the conditional Wasserstein space. The Euler–Lagrange equations (augmented with potential terms or density bias when appropriate) provide necessary conditions for optimality, coupling the Hamilton–Jacobi (HJ) PDE for the scalar potential $\phi$ and the mass conservation PDE for the evolving density. This yields a system analogous to those in classical mechanics but tailored to probability flows under conditional structure (Qian et al., 2024, Amad et al., 2 Mar 2026, Kerrigan et al., 2024, Barton et al., 2018).

2. Euler–Lagrange System, Geodesics, and Generalization of OT Distances

The core of CLOT lies in the associated Euler–Lagrange system. Introducing a Lagrange multiplier $\rho_T(x)$0 for the continuity equation, stationarity yields:

• $\rho_T(x)$1: $\rho_T(x)$2
• $\rho_T(x)$3: $\rho_T(x)$4.

Setting $\rho_T(x)$5 gives the Hamilton–Jacobi equation: $\rho_T(x)$6. Together, $\rho_T(x)$7 solve the coupled geodesic system in the $\rho_T(x)$8-Wasserstein conditional metric (Qian et al., 2024). More generally, with a Lagrangian $\rho_T(x)$9 (potentially density- or data-dependent), the path cost becomes: $x^{\text{cond}}$0 defining a conditional Wasserstein-type distance $x^{\text{cond}}$1, with stationary solutions $x^{\text{cond}}$2 given by the appropriate Euler–Lagrange equations. This introduces a flexible family of OT distances, where geodesics (i.e., optimal plans) follow minimal-action paths adapted to background density, conditioning, or other structural bias (Wang et al., 4 Nov 2025, Amad et al., 2 Mar 2026).

3. Dynamic, Conditional, and Triangular Flow Structures

A defining aspect of Conditional Lagrangian OT is its explicit attention to conditionality and triangularity. The particle or measure evolution is performed so that conditioning variables—such as observed time series entries, hyperparameters, or simulator parameters—are held fixed or treated in a non-transported fashion. The velocity field $x^{\text{cond}}$3 is typically block-diagonal or strictly triangular, with nonzero evolution only in the unconditioned, "missing," or latent subspace.

The dynamic aspect is often formalized via the generalized (conditional) Benamou–Brenier theorem: $x^{\text{cond}}$4 subject to triangular structure and conditional continuity equations. Each conditional variable $x^{\text{cond}}$5 defines an independent optimal transport geodesic for the transport variable $x^{\text{cond}}$6 (e.g., missing data, simulator outputs), and the global evolution is assembled as a product over $x^{\text{cond}}$7 (Kerrigan et al., 2024, Zeghal et al., 28 Oct 2025, Alfonso et al., 2023). This key structure enables simulation-free, amortized algorithms where only the unobserved or to-be-inferred components are transported.

4. Algorithmic Schemes: Flow Matching and Surrogate Construction

Algorithmic realizations of CLOT frameworks frequently use flow-matching techniques. Instead of solving coupled PDEs directly, velocity networks $x^{\text{cond}}$8 are trained to match ground-truth velocities along geodesics obtained via mini-batch OT, conditional Monge maps, or explicit endpoint couplings. The standard pipeline is:

• Draw conditioning data $x^{\text{cond}}$9 (or $\rho(t, x \mid x^{\text{cond}})$0), and sample source and target variables with conditioning held fixed.
• Solve a conditional OT problem (often using entropic/mini-batch regularization) to obtain endpoint pairs.
• For randomly interpolated points along the geodesic, regress $\rho(t, x \mid x^{\text{cond}})$1 to the analytic velocity $\rho(t, x \mid x^{\text{cond}})$2.
• At inference, integrate the learned velocity field, possibly augmented with potential-derived corrections, to transport samples from source to target conditional distributions.

Surrogate models parameterize these dynamics through neural networks (potentially with FiLM layers or spline decoders), learning both the cost geometry (Riemannian metric $\rho(t, x \mid x^{\text{cond}})$3) and potential terms (density or VAE-based) (Amad et al., 2 Mar 2026, Qian et al., 2024). This produces fast, simulation-free, and amortized inference engines that can adapt to arbitrary conditioning and generalize across observed anchor points or parameter values.

5. Potentials, Density Bias, and Variance Reduction

Modern CLOT approaches often enhance pure-geodesic transport with potentials or density-adaptive bias to regularize or guide flows toward data-supported regions or task-specific structures. Notable schemes include:

• Potential augmentation: addition of a gradient-of-potential term $\rho(t, x \mid x^{\text{cond}})$4 (e.g., VAE reconstruction error) to base velocity, resulting in dynamics $\rho(t, x \mid x^{\text{cond}})$5.
• Density bias potentials: inclusion of terms such as $\rho(t, x \mid x^{\text{cond}})$6 in the Lagrangian to penalize paths through low-density regions, favoring manifold-adherent geodesics (Amad et al., 2 Mar 2026, Wang et al., 4 Nov 2025).
• Variance reduction via Rao–Blackwellization: potential-derived corrections reduce variance in endpoint updates without bias, enabling tighter, more data-consistent imputation or sampling (Qian et al., 2024).

Such mechanisms are empirically observed to yield lower-variance paths, improved empirical coverage, and robust performance in sparse, high-dimensional, or ill-posed conditional generation tasks.

6. Applications and Empirical Performance

The CLOT framework underpins a diverse array of conditional inference and mapping tasks:

• Time series imputation: CLWF achieves sharp, simulation-free imputation of missing values by transporting noise toward data-consistent completions under clamped observation constraints, outperforming classical diffusion approaches in convergence (Qian et al., 2024).
• Hyperparameter trajectory inference: CLOT surrogates accurately capture output distributions of neural networks across a hyperparameter spectrum, combining learned Riemannian geometry and density bias to interpolate across sparse anchors for RL reward weights, quantile regression, and dropout parameters (Amad et al., 2 Mar 2026).
• Bridging simulation domains: COT-FM enables conditioning-preserving corrections from Lagrangian to particle-mesh cosmological simulators without paired data, matching true posteriors at both summary and pixel levels and restoring frequentist coverage (Zeghal et al., 28 Oct 2025).
• Conditional generative modeling: Simulation-free flows leveraging triangular structure allow for infinite-dimensional conditional generation and Bayesian inverse problems, with fully amortized, likelihood-free sampling (Kerrigan et al., 2024, Alfonso et al., 2023).
• Clustering and matching: Density-dependent Lagrangians promote clustering aligned with data geometry, correctly recovering cluster structure and matchings in synthetic experiments through action-minimizing, density-adaptive paths (Wang et al., 4 Nov 2025).

Empirical benchmarks consistently validate the improved interpolation, robustness to sparse conditioning, and computational efficiency of the CLOT family compared to classical diffusion and Euclidean-cost OT approaches.

7. Connections to Classical and Stochastic Control, Uniqueness, and Extensions

Conditional Lagrangian OT generalizes static OT via its dynamic, path-wise formulation, incorporating results from control theory, convex analysis, and mean field games. Classical results—including geodesic uniqueness, existence under Tonelli/Mikami–Thieullen hypotheses, and connection to Monge–Ampère equations—carry over, often adapted to the conditional, block-triangular setting (Barton et al., 2018, Hindawi et al., 2011). The dynamic action can be instantiated with linear-quadratic costs, density-adaptive metrics, or stochastic action to recover standard, density-aware, or mean field game-interpreted formulations.

Infinite-dimensional generalization is established by leveraging Hilbertian geometric and functional-analytic structures, supporting high-dimensional inverse problems and function-space generative modeling (Kerrigan et al., 2024). Algorithmic strategies are notably simulation-free, relying only on sampling, OT matching, and integration of learned conditional velocities, conferring practical scalability to real-world, high-dimensional conditional inference problems.