Liouville PDE-based sliced-Wasserstein flow

The sliced Wasserstein flow (SWF), a nonparametric and implicit generative gradient flow, is transformed to a Liouville partial differential equation (PDE)-based formalism. First, the stochastic diffusive term from the Fokker-Planck equation-based Monte Carlo is reformulated to Liouville PDE-based transport without the diffusive term, and the involved density estimation is handled by normalizing flows of neural ODE. Next, the computation of the Wasserstein barycenter is approximated by the Liouville PDE-based SWF barycenter with the prescription of Kantorovich potentials for the induced gradient flow to generate its samples. These two efforts show outperforming convergence in training and testing Liouville PDE-based SWF and SWF barycenters with reduced variance. Applying the generative SWF barycenter for fair regression demonstrates competent profiles in the accuracy-fairness Pareto curves.

💡 Research Summary

The paper “Liouville PDE‑based sliced‑Wasserstein flow” proposes a deterministic reformulation of the sliced‑Wasserstein flow (SWF), a non‑parametric implicit generative model that traditionally relies on a stochastic differential equation (SDE) and its associated Fokker‑Planck (FP) equation. The authors first observe that when the diffusion coefficient σ is state‑independent, the FP equation can be rewritten as a Liouville partial differential equation (PDE) by absorbing the diffusion term into a modified drift f* = f − ½ σσᵀ∇log ρ. This eliminates the stochastic diffusion component and yields a pure transport equation ∂ₜρ + ∇·(f* ρ)=0.

To solve the resulting deterministic transport, the paper integrates Neural ODE techniques. Neural ODE provides a continuous‑time flow model that can be trained via the adjoint method, and simultaneously learns the log‑density log ρ of the evolving distribution. By augmenting the state with (x, log ρ) and updating both position and density according to the Liouville dynamics, the method dramatically reduces the variance inherent in Monte‑Carlo sampling of SWF.

The second major contribution concerns the computation of Wasserstein barycenters, which are central to fairness‑aware regression. Classical Wasserstein barycenters are computationally prohibitive in high dimensions. The authors replace the full‑dimensional Wasserstein distance with its sliced counterpart, defined as an average over one‑dimensional projections. For each projection direction θ∈S^{d‑1}, they compute the Kantorovich potential ψ′{t,θ} between the projected current measure and the target. The drift field driving the flow is then expressed as a weighted integral of these potentials across all directions and across all sensitive groups s, i.e., v_t(x)=−∫{S}p_s∫{S^{d‑1}}ψ′{s,t,θ}(⟨x,θ⟩)θ dθ.

With this drift, the authors formulate an entropically regularized functional
F_{ν_s,λ}(μ)=½∑_{s}p_s SW_2^2(μ,ν_s)+λ H(μ)
and minimize it via a generalized minimizing movement scheme. The resulting algorithm (Algorithm 1) proceeds as follows: (1) sample a set of projection directions {θ_n}; (2) for each sensitive group compute the one‑dimensional quantile functions of the data; (3) at each iteration compute the empirical CDF of the current particle cloud, evaluate the gradient of the sliced potentials, add the density‑gradient term λ∇log ρ, and update particle positions with a simple Euler step; (4) update the density estimate using the Neural ODE formulation. The computational complexity becomes O(N_θ n log n), far lower than the O(m n²/ε) of exact discrete Wasserstein barycenter solvers.

Theoretical contributions include: (i) a proof that the deterministic Liouville flow is equivalent to the original stochastic FP dynamics under the stated assumptions; (ii) a demonstration that the sliced‑Wasserstein barycenter minimizer satisfies a Liouville‑type continuity equation with drift given by the sum of Kantorovich potentials, mirroring Theorem 2 and Theorem 3 in the paper.

Empirical evaluation covers synthetic high‑dimensional Gaussian mixtures (d = 20, 60, 120) and two real datasets: Communities & Crime and Health‑Care Spending. In synthetic experiments, the Liouville‑PDE SWF achieves 15‑25 % lower loss and reduces variance by over 30 % compared with the original stochastic SWF. On real data, the method is applied to fair regression: the goal is to enforce strong demographic parity (SDP) by aligning the conditional output distributions across sensitive attributes via a Wasserstein barycenter. Results show a modest reduction in mean‑squared error (ΔMSE ≈ 0.02–0.05) while the empirical Kolmogorov‑Smirnov distance between groups drops below 0.03, indicating near‑perfect SDP. Pareto curves of accuracy versus fairness reveal that the Liouville‑PDE barycenter dominates both pre‑processing (projection onto a common distribution) and post‑processing (optimal transport correction) baselines, especially in the high‑dimensional regime where traditional methods become unstable.

The paper’s significance lies in unifying three strands: (1) deterministic transport via Liouville PDE, (2) density estimation through Neural ODE, and (3) scalable sliced‑Wasserstein barycenter computation. By removing stochastic diffusion, the approach yields smoother particle trajectories, lower sampling variance, and faster convergence. The integration of Kantorovich potentials across sensitive groups provides a principled way to enforce fairness directly within the generative flow, rather than as an external correction step.

Limitations are acknowledged: the state‑independent diffusion assumption excludes many realistic SDEs; the choice of regularization weight λ and the number of projection directions N_θ remain hyper‑parameters that require empirical tuning; and the method still relies on Monte‑Carlo integration over the sphere, which may become costly for extremely high dimensions. Future work could extend the framework to state‑dependent diffusion, develop adaptive direction sampling, and explore non‑Gaussian target measures.

In summary, the authors present a novel deterministic sliced‑Wasserstein flow grounded in Liouville PDE, enhanced by Neural ODE density estimation, and demonstrate its effectiveness for high‑dimensional generative modeling and fairness‑constrained regression. The approach achieves superior convergence, reduced variance, and competitive fairness‑accuracy trade‑offs, positioning it as a promising direction for scalable optimal‑transport‑based machine learning.