Optimal Transport Aggregation

• Optimal Transport Aggregation is a framework that leverages mathematical OT to aggregate data by minimizing transport costs under structural constraints.
• It integrates techniques like OT pooling, graph neural networks, and federated model alignment to yield richer representations and improved predictive accuracy.
• The approach offers strong theoretical guarantees including universal approximation, stability, and fairness, making it pivotal for applications in vision, 3D modeling, and distributed learning.

Optimal Transport Aggregation refers to a suite of methodologies that leverage the theory of optimal transport (OT) to combine (aggregate) sets, features, or models while minimizing distortion or cost under a prescribed geometric, probabilistic, or algebraic structure. These approaches generalize classical pooling or averaging by incorporating the OT geometry, providing richer representations, robustness to distributional shifts, and principled solutions to multi-modal aggregation tasks. OT aggregation methods are widely utilized across domains, including graph neural networks, federated learning, 3D representation, semantic segmentation, and economic matching.

1. Mathematical Foundations of Optimal Transport Aggregation

At the core of OT aggregation lies the Kantorovich formulation: given two probability measures $\mu$ on a space $X$ and $\nu$ on $Y$ with cost $c(x, y)$, the OT problem seeks a coupling $\pi$ minimizing the expected cost,

$$\min_{\pi \in \Pi(\mu,\nu)} \int_{X \times Y} c(x, y) \, d\pi(x, y),$$

where $\Pi(\mu,\nu)$ denotes the set of couplings with prescribed marginals. Aggregation within the OT framework generalizes core set operations by enforcing structural or statistical constraints via the transport plan.

Recent extensions accommodate nonlinear aggregation (via barycentric costs), multi-agent objectives, or distributional matching under constraints adapted to the application's geometry or semantics.

2. Prototypical Aggregation Algorithms and Architectures

Various OT aggregation schemes have been instantiated for different data modalities and tasks:

• OT Pooling for Sets and Sequences: OT-based pooling embeds and aggregates input features by solving an OT problem to a trainable reference set, yielding a fixed-size, permutation-invariant representation. The entropic regularization of the transport plan ensures computational tractability, often employing the Sinkhorn algorithm. OT pooling unifies set kernels and constrained attention mechanisms, typically providing better generalization than mean/sum pooling or naïve self-attention for structured, variable-size input collections (Mialon et al., 2020).
• OT Aggregation in Graph Neural Networks: Optimal Transport Graph Neural Networks (OT-GNNs) compute graph-level embeddings by solving a Wasserstein distance between the node feature cloud and a family of learnable prototype point clouds. This procedure is provably universal on point clouds, unlike sum/mean aggregation, and benefits from regularization mechanisms (e.g., noise-contrastive) to prevent prototype collapse. OT-based aggregation captures more nuanced geometric and semantic structure in molecular and relational datasets, yielding both improved smoothness and higher predictive accuracy (Chen et al., 2020).
• Distribution Alignment for Model Aggregation: In federated learning, OT alignment can be used to correct for disparate client participation or domain shifts. For example, FedAVOT formulates the server-side model aggregation as a masked OT problem between the actual client availability distribution and the desired importance distribution, computing aggregation weights with entropic Sinkhorn scaling. Hierarchical aggregation methods such as HFedATM employ filter-wise OT alignment followed by regularized-mean merging, achieving tighter generalization bounds and reducing inter-model divergence (Herlock et al., 17 Sep 2025, Nguyen et al., 7 Aug 2025).
• Global Feature Aggregation in Vision: The SALAD method for visual place recognition replaces soft-assignment of image features to learnable clusters with an OT-based assignment, enforced via entropy-regularized Sinkhorn optimization. A "dustbin" cluster enables the selective removal of low-informative features, further enhancing the quality and robustness of the global descriptor (Izquierdo et al., 2023).
• Distribution Matching in 3DGS Compression: OT aggregation matches dense clouds of 3D Gaussian primitives to a sparse set under the 2-Wasserstein geometry (approximated via the Gelbrich distance), providing state-of-the-art compression without fidelity loss. Efficient block-wise EM algorithms enable scalable merging and densification, enabling large-scale, moment-aware 3D distribution aggregation (Zhao et al., 19 May 2026).

3. Theoretical Guarantees and Universality

OT aggregation offers powerful theoretical properties often absent in traditional means:

• Universal Approximation: Wasserstein-based kernels on point clouds are universal, i.e., any continuous set-function can be approximated by networks using OT aggregation, whereas sum/mean pooling lacks this property due to invariance to permutations and insensitivity to set multiplicities (Chen et al., 2020).
• Stability and Optimality: Entropic regularization confers strict convexity and uniqueness to the transport plan, yielding stable solutions as marginals vary and providing fast convergence rates in projection-based Sinkhorn schemes (Scetbon et al., 2020, Mialon et al., 2020).
• Error Bounds in Federated Model Aggregation: Filter-wise OT alignment and regularized mean aggregation yield exponentially decaying inter-model divergence and strictly tighter generalization bounds compared to naïve averaging—guaranteed across global rounds and valid under mild Lipschitz or Hölder condition on the loss (Nguyen et al., 7 Aug 2025).
• Fairness and Multi-Objective Equitability: Multi-agent or multi-cost OT aggregation (e.g., Equitable OT) achieves simultaneous, balanced satisfaction among competing objectives, reducing to integral probability metrics (IPMs), and links to kernel discrepancies such as the Dudley metric (Scetbon et al., 2020).

4. Computational Aspects: Scalability and Algorithms

Efficient computation is realized via advances in numerical OT and large-scale, structure-aware aggregation algorithms:

Method	Core Algorithm	Complexity / Scalable Trick
OT Pooling (sets/seqs)	Sinkhorn scaling	$O(TNM)$ with T iterations; reference set reduces scaling (Mialon et al., 2020)
OT-GNN aggregation	Earth Mover/Sinkhorn	$O(n^3)$ for small $X$0 (exact); $X$1 for Sinkhorn, GPU batching (Chen et al., 2020)
Distribution matching in 3DGS	Block-wise EM, partitioning	$X$2; distributes computation across KD-tree blocks (Zhao et al., 19 May 2026)
Model aggregation in FL	Sinkhorn on participation matrix	$X$3; mask restricts support, exact weights (Herlock et al., 17 Sep 2025)
Visual descriptor aggregation	Sinkhorn assignment	$X$4 per Sinkhorn pass; dual assignment constraints (Izquierdo et al., 2023)

Entropic OT (Sinkhorn) has become the standard engine for these algorithms, enabling approximate but highly parallelizable and GPU-compatible optimization.

5. Application Domains and Empirical Impact

OT aggregation is foundational in several advanced application areas:

• Graph-level and set-level prediction: OT-GNN achieves strictly better regression accuracy and smoother latent spaces than classical pooling methods in molecular property prediction benchmarks (e.g., RMSE on ESOL drops from 0.635 to 0.605; AUC on BACE from 0.865 to 0.873) (Chen et al., 2020).
• Visual place recognition: The SALAD framework achieves recall@1 improvements of up to 7.6–12.4 percentage points compared to prior NetVLAD-based single-stage systems, operating at <3 ms per image (Izquierdo et al., 2023).
• 3D representation compression: MMGS attains 10× compression of 3DGS with a marginal increase (+0.44 dB PSNR) relative to baseline, outperforming all ablation and prior compression techniques (Zhao et al., 19 May 2026).
• Open-vocabulary semantic segmentation: OV-COAST's OT-guided cost aggregation improves mIoU by +1.72 percentage points over CAT-Seg and +4.94 over SAN-B on the MESS benchmark (Gandhamal et al., 4 Jun 2025).
• Federated learning under participation bias: FedAVOT yields $X$5 convergence regardless of client fraction and eliminates aggregation bias from client selection skew, with performance (loss and accuracy) nearly matching full-participation baselines even when only two clients are active per round (Herlock et al., 17 Sep 2025).
• Multi-agent and fair division: Equitable OT achieves balanced allocation among heterogeneous agents and generalizes classical transport to fairness-aware settings (Scetbon et al., 2020).

6. Extensions: Multi-agent, Weak, and Barycentric Aggregation

Several lines of research generalize OT aggregation beyond classical settings:

• Multi-agent/multi-cost aggregation: "Equitable OT" involves minimizing the worst-case cost among multiple transports, or equivalently maximizing the minimum agent utility, unifying fair division and robust OT. Duality theory extends to semi-infinite programs and entropic regularization, enabling efficient computation and connection to well-known IPMs (Scetbon et al., 2020).
• Barycentric and weak OT: When the cost depends on the barycenter or nonlinear functional of the marginal (rather than pairwise matches), as in economic matching, aggregation is formulated via barycentric weak OT. Primal and dual mirror-descent algorithms based on the KL divergence enable efficient computation, and application-specific interpretations elucidate firm/worker matching in labor markets (Paty et al., 2022).
• Grand-canonical OT aggregation: Aggregates all N-marginal optimal transport costs into a single problem allowing for variable-size ensembles (e.g., in statistical mechanics and DFT), extending aggregation to fluctuating collections and mixed population sizes (Marino et al., 2022).

7. Challenges, Open Problems, and Theoretical Limits

Key challenges and limitations are recognized and addressed in the literature:

• Scalability: Full-rank OT scales quadratically or cubically, but low-rank, partitioned, or block-wise schemes (e.g., hierarchical refinement, block-wise EM, masked OT) achieve linear or log-linear scaling at minor cost to resolution, as in HiRef for scalable full-rank assignment (Halmos et al., 4 Mar 2025).
• Universality versus collapse: While OT-based kernels are universal, practical training necessitates regularization to avoid structural collapse (e.g., prototype collapse in OT-GNN), typically addressed via noise-contrastive or spectral penalties (Chen et al., 2020).
• Robustness to distributional shift: OT-guided aggregation aligns user/model/feature distributions under participation or domain shift, but integrating these methods with highly heterogeneous architectures, asynchronous regimes, or non-Euclidean costs remains an open research area (Nguyen et al., 7 Aug 2025, Herlock et al., 17 Sep 2025).
• Nonlinear aggregation and domain-specific modeling: Weak OT and barycentric costs generalize aggregation for settings where simple pairwise matching does not capture task-specific value (e.g., aggregating nonlinear skills in economics, or features under nonlinear composition in biology or vision) (Paty et al., 2022).

OT aggregation thereby constitutes both a highly general mathematical framework and a suite of computational techniques central to contemporary machine learning, signal processing, distributed optimization, and economic modeling. Its ongoing development is characterized by advances in scalable solvers, theoretical understanding of universality and fairness, and deeper integration into application-specific pipelines spanning vision, language, federated model training, and scientific modeling.