Fast and Scalable Multi-Agent Distribution Matching via Partitioned Optimal Transport
Abstract
This paper presents a scalable optimal-transport-based framework for terminal distribution matching in multi-agent systems. While optimal transport provides a natural way to measure distributional mismatch and assign agents to a desired spatial distribution, global discrete transport can become computationally expensive for large-scale systems. We address this bottleneck by partitioning agents and target samples into spatially corresponding blocks and solving smaller local transport problems. Under a mass-balance condition, the resulting restricted coupling remains feasible for the global problem and provides an upper bound on the Wasserstein cost. The local assignments generate target locations for finite-horizon agent control, applicable to both linear and nonlinear dynamics. By alternating local assignment and control, we establish a cycle-to-cycle descent guarantee for the resulting transport surrogate. The proposed framework therefore enables scalable terminal distribution matching while retaining a rigorous connection to the Wasserstein objective. The technical soundness of the proposed results is validated through simulations.
Related
- Multiscale Supervised Unbalanced Optimal Transport Flow Matching
- Regularized Coupling Maps on Matrix Polytopes: From Assignment to Graph Matching and Optimal Transport
- Expected Batch Optimal Transport Plans and Consequences for Flow Matching
- SWAP: A Scalable, Warm-Startable, Anytime Permutation Solver for Optimal Transport
- Coarse-Graining Optimal Transport for Rollout-Free Post-Training of Flow Matching
BibTeX
@misc{lee2026fast,
title = {{Fast and Scalable Multi-Agent Distribution Matching via Partitioned Optimal Transport}},
author = {Kooktae Lee and Ruchika Singh},
howpublished = {arXiv},
year = {2026},
month = sep,
doi = {10.48550/arxiv.2609.38960},
eprint = {2609.38960},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2609.38960},
}
RIS
TY - UNPB
TI - Fast and Scalable Multi-Agent Distribution Matching via Partitioned Optimal Transport
AU - Kooktae Lee
AU - Ruchika Singh
PY - 2026
DA - 2026/09/30
PB - arXiv
DO - 10.48550/arxiv.2609.38960
UR - https://arxiv.org/abs/2609.38960
AB - This paper presents a scalable optimal-transport-based framework for terminal distribution matching in multi-agent systems. While optimal transport provides a natural way to measure distributional mismatch and assign agents to a desired spatial distribution, global discrete transport can become computationally expensive for large-scale systems. We address this bottleneck by partitioning agents and target samples into spatially corresponding blocks and solving smaller local transport problems. Under a mass-balance condition, the resulting restricted coupling remains feasible for the global problem and provides an upper bound on the Wasserstein cost. The local assignments generate target locations for finite-horizon agent control, applicable to both linear and nonlinear dynamics. By alternating local assignment and control, we establish a cycle-to-cycle descent guarantee for the resulting transport surrogate. The proposed framework therefore enables scalable terminal distribution matching while retaining a rigorous connection to the Wasserstein objective. The technical soundness of the proposed results is validated through simulations.
ER -