Distribution Control via Optimal Transport

• Distribution control via optimal transport is the use of cost-minimizing strategies to steer probability measures in multi-agent systems and learning models.
• It leverages mathematical formulations like the Monge and Kantorovich problems, as well as dynamic OT, to achieve provable convergence and efficient distribution shaping.
• Practical applications include multi-agent formation control, federated model aggregation, and generative diffusion models with advanced regularization techniques.

Distribution control via optimal transport (OT) refers to the systemic steering, matching, and regularization of probability measures in multi-agent systems, learned models, and engineered networks using the mathematical and computational machinery of OT theory. This paradigm exploits the cost-minimizing structure of OT to derive principled controllers, scalable algorithms, and theoretical guarantees for mapping or shaping the spatial, feature, or outcome-wise distribution of agents, data points, or decisions in accordance with desired objectives that may encode shape, density, constraints, communication patterns, or social welfare. OT is foundational for tasks such as multi-agent formation control, coverage optimization, federated model aggregation, capacity-constrained network flows, domain adaptation in machine learning, sample generation in diffusion models, density steering in nonlinear systems, and more.

1. Mathematical Foundations and Optimal Transport Formulations

OT establishes the minimum-cost coupling between probability measures, traditionally modeled through the Monge and Kantorovich formulations. The Monge problem seeks a deterministic map $T:X\to Y$ pushing source measure $\mu$ to target measure $\nu$ under cost $c(x,y)$, while the Kantorovich relaxation seeks a transport plan $\pi\in\Pi(\mu,\nu)$ minimizing $\int_{X\times Y} c(x,y) \, d\pi(x,y)$ (Choi et al., 2024). The quadratic cost yields the 2-Wasserstein metric $W_2^2(\mu,\nu)$. Dynamic OT (Benamou–Brenier) provides a fluid-dynamical interpretation, framing OT as the minimization of action among continuous paths $\rho_t$ (Elamvazhuthi et al., 2016). For distribution control, such mathematical machinery enables the explicit shaping of empirical agent measures, coverage densities, or model outputs—for example, by penalizing $W_p^p(\mu_t,\nu)$ within optimal control functionals as in multi-agent shape control (Lin et al., 2022), or constraining agent coverage using the Wasserstein metric (Lee et al., 23 Nov 2025).

2. Multi-Agent Control and Distributed Algorithms

OT is fundamentally leveraged for multi-agent distribution shaping. The MASCOT framework introduces OT-based terms both in the running and terminal cost, with controls $\{u^{(i)}(t)\}$ selected to minimize a composite cost comprising control energy and distributional discrepancy (Earth Mover's Distance) (Lin et al., 2022). Pontryagin-type optimality yields necessary conditions coupling the Hamiltonian, Liouville equation for the evolving empirical measure, and backward adjoint PDEs for co-state variables, with variational derivatives operationalized via OT dual potentials. Direct shooting and entropic OT solvers (Sinkhorn) facilitate numerical solution.

Distributed frameworks push decentralization further. In distributed online optimization, agents iteratively estimate the Kantorovich potential via local primal-dual saddle-point dynamics over the Voronoi graph formed by their positions, and then move by proximal descent along the OT geodesic, converging to the target law (Krishnan et al., 2018). Decentralized multi-agent matching reformulates the global OT objective into sequential local assignment and weight-update rules guaranteeing cyclewise cost descent and scalable feasibility, with only local communication and memory—empirically scaling to thousands of agents (Lee, 2 Jan 2026). In swarm coverage, OT-based coverage control (OTCC) derives gradient flows in position and dual potentials from the Kantorovich dual, generalizing Voronoi-based approaches with improved stability and density reproduction (Inoue et al., 2020).

3. Constraints, Regularization, and Algorithmic Innovations

Distribution control problems routinely feature constraints and regularizations beyond pure OT. Moment-constrained OT fixes one marginal and enforces aggregate properties on the other via generalized moment constraints, with entropic penalties facilitating computational tractability and hard constraint enforcement. Sinkhorn-style dual ascent alternates marginal and moment-matching updates, enabling distributed control applications from EV charging under grid constraints to multi-period steering (Corre et al., 2022).

Partial or adaptive optimal transport (AOT) relaxes the requirement of full mass conservation, instead allowing the transported mass to be determined by the cost structure and noise/outlier presence, vital for robust OT in domain adaptation and biological data analysis. The Kantorovich dual for AOT imposes non-positivity on dual potentials, with entropic regularization and modified Sinkhorn iterations enforcing pointwise marginal inequalities (Yang et al., 7 Mar 2025).

Connectivity and communication constraints are incorporated via smooth penalties or hard constraints in QP reformulations. Fused with Wasserstein objectives, this yields convex optimization programs ensuring spatial coverage while preserving network connectivity (Lee et al., 23 Nov 2025). In large-scale networks, joint OT and minimum-cost flows are formulated as infinite-dimensional LPs and solved via distributed supergradient/ascent methods over network nodes, balancing partitioning and routing under capacity constraints (Laurentin et al., 29 Aug 2025).

4. Model Regularization and Learning via Optimal Transport

OT has emerged as a rigorous backbone for distributional regularization in machine learning, generative modeling, and preference alignment. The AOT framework for LLMs aligns reward distributions by convex relaxations of first-order stochastic dominance, expressed as 1D optimal transport costs with closed-form sorting solutions (Melnyk et al., 2024). The dual formulation admits parametric sample complexity rates and integrates seamlessly into minibatch training pipelines.

Neural OT map learning for generative modeling improves stability and bias by leveraging displacement interpolation (DIOTM) (Choi et al., 2024). Here, the entire OT geodesic (rather than endpoints) is learned, with time-dependent dual potentials tied via Hamilton–Jacobi–Bellman regularization, yielding smoother maps and intermediate distributions useful for morphing and translation.

In federated learning, masked optimal transport aligns client aggregation weights under partial participation, with Sinkhorn scaling computing transport-based weights ensuring unbiased global averaging and $\mathcal O(1/\sqrt{T})$ convergence rates independent of subset sizes (Herlock et al., 17 Sep 2025).

5. Density and Trajectory Control in Nonlinear and Stochastic Systems

OT generalizes naturally to steering densities in nonlinear control systems and stochastic settings. In continuous-time control-affine systems, set-oriented graph methods approximate Liouville evolution and Benamou–Brenier OT via strongly-connected graph generators, enabling the design of low-energy feedback control laws for steering initial densities to targets in complex ambient flows with convex computational guarantees (Elamvazhuthi et al., 2016).

Trajectory-optimized density control with flow matching extends classical endpoint OT by optimizing over full transport paths, embedding collision-avoidance and interaction penalties into running costs. The coupling of flow-matching losses and actor-critic FBSDE training provides collision-free plans in multi-modal domains and constrained passageways (Duan et al., 8 Oct 2025). Entropic regularization is integrated into model predictive control (Sinkhorn MPC), combining OT assignment via Sinkhorn with sequential MPC optimization. This achieves real-time, stable, and cost-effective transport planning with global convergence properties (Ito et al., 2023).

6. Correction of Prior Mismatch and Acceleration in Generative Diffusion Models

Recent advances have clarified that score-based diffusion models intrinsically implement time-dependent OT. When mismatch between forward process termination $p_T$ and reverse process initial distribution $q_T$ occurs (the prior error $W_2(p_T,q_T)$), sampling quality degrades (Wang et al., 2024). Theoretical results prove the probability-flow ODE flow exponentially converges in $L^2$ to the static OT (Monge–Ampère) map as diffusion time increases. Thus, static OT is the most efficient single-step correction for prior mismatch. In practice, learning the OT Brenier map via neural networks provides exact prior correction, yielding zero prior error and faster, higher-quality sample generation—validated on image datasets and class-conditional tasks.

7. Theoretical Guarantees, Stability, and Extensions

Distribution control frameworks built on OT commonly admit convergence, stability, and optimality properties. Saddle-point structures in coupled gradient flows yield Lyapunov stability for multi-agent coverage (Inoue et al., 2020). Strong duality holds in joint OT–flow LPs under generic regularity and convexity (Laurentin et al., 29 Aug 2025). Entropy regularization smooths assignments and enables global convergence via LaSalle invariance principles, local asymptotic stability, and exponential convergence to permutation-matching equilibria in Sinkhorn MPC (Ito et al., 2023).

Extensions include multi-period and multi-marginal settings, integration of risk measures for robust control, scaling of decentralized algorithms, and merging of OT with mean-field or interaction-aware control. Empirical validations demonstrate improved density reproduction, collision-free transport, stability under partial participation, and robustness to outliers/noise in diverse applications.

Summary Table: Specialized Distribution Control via OT

Approach / Domain	Key OT Formulation	Principal Algorithms / Guarantees
Multi-agent shape control (Lin et al., 2022)	Wasserstein cost in running/terminal	Shooting, network simplex, Sinkhorn, Pontryagin-type
Decentralized matching (Lee, 2 Jan 2026)	Kantorovich, sequential LP relaxation	Greedy local assignment, memory-based consensus
Coverage & connectivity (Lee et al., 23 Nov 2025)	Quadratic Wasserstein, barycentric QP	Receding-horizon QP, log-sum-exp penalties
Moment-constrained OT (Corre et al., 2022)	One-sided OT, moment classes	Entropic Sinkhorn-dual iterations
Adaptive OT for domain adaptation (Yang et al., 7 Mar 2025)	Mass-inequality OT (AOT)	Entropic Sinkhorn, marginal clamping
Diffusion model sampling (Wang et al., 2024)	Dynamic/static OT, Monge-Ampère	Brenier map correction, probability-flow ODE
Model aggregation in FL (Herlock et al., 17 Sep 2025)	Masked OT, marginal constraints	Sinkhorn scaling, exact federated averaging

Distribution control via optimal transport thus synthesizes tools from mathematical analysis, convex optimization, distributed systems, neural architectures, and statistical learning to orchestrate ensembles, data, and agent collectives in a manner that is provably cost-efficient, robust to constraints and uncertainty, and scalable to high-dimensional settings.