# Regularized Optimal Transport Layers for Generalized Global Pooling Operations

arXiv: https://arxiv.org/abs/2212.06339 · HTML: https://mlsanity.com/paper/2212.06339/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 58 of its 92 references that are in the map.
- First posted on arXiv 13 Dec 2022.

## Builds on (most cited first, 30 of 58)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [How Powerful are Graph Neural Networks?](https://mlsanity.com/paper/1810.00826/) (2018)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Deep Sets](https://mlsanity.com/paper/1703.06114/) (2017)
- [Attention-based Deep Multiple Instance Learning](https://mlsanity.com/paper/1802.04712/) (2018)
- [Automatic differentiation in machine learning: a survey](https://mlsanity.com/paper/1502.05767/) (2015)
- [Hierarchical Graph Representation Learning with Differentiable Pooling](https://mlsanity.com/paper/1806.08804/) (2018)
- [TUDataset: A collection of benchmark datasets for learning with graphs](https://mlsanity.com/paper/2007.08663/) (2020)
- [Deep Equilibrium Models](https://mlsanity.com/paper/1909.01377/) (2019)
- [Network In Network](https://mlsanity.com/paper/1312.4400/) (2013)
- [OptNet: Differentiable Optimization as a Layer in Neural Networks](https://mlsanity.com/paper/1703.00443/) (2017)
- [Differentiable Convex Optimization Layers](https://mlsanity.com/paper/1910.12430/) (2019)
- [Optimal Transport for Domain Adaptation](https://mlsanity.com/paper/1507.00504/) (2015)
- [Self-Attention Graph Pooling](https://mlsanity.com/paper/1904.08082/) (2019)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [Wasserstein Auto-Encoders](https://mlsanity.com/paper/1711.01558/) (2017)
- [Sparse Sinkhorn Attention](https://mlsanity.com/paper/2002.11296/) (2020)
- [DeeperGCN: All You Need to Train Deeper GCNs](https://mlsanity.com/paper/2006.07739/) (2020)
- [Order Matters: Sequence to sequence for sets](https://mlsanity.com/paper/1511.06391/) (2015)
- [Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration](https://mlsanity.com/paper/1705.09634/) (2017)
- [Fast Computation of Wasserstein Barycenters](https://mlsanity.com/paper/1310.4375/) (2013)
- [Optimal transport mapping via input convex neural networks](https://mlsanity.com/paper/1908.10962/) (2019)
- [Adversarial Graph Augmentation to Improve Graph Contrastive Learning](https://mlsanity.com/paper/2106.05819/) (2021)
- [Learning with a Wasserstein Loss](https://mlsanity.com/paper/1506.05439/) (2015)
- [Learning Latent Permutations with Gumbel-Sinkhorn Networks](https://mlsanity.com/paper/1802.08665/) (2018)
- [Optimal Transport for structured data with application on graphs](https://mlsanity.com/paper/1805.09114/) (2018)
- [ASAP: Adaptive Structure Aware Pooling for Learning Hierarchical Graph Representations](https://mlsanity.com/paper/1911.07979/) (2019)
- [Sinkformers: Transformers with Doubly Stochastic Attention](https://mlsanity.com/paper/2110.11773/) (2021)

## Built on by (most cited first, 2 of 2)

- [Importance Sparsification for Sinkhorn Algorithm](https://mlsanity.com/paper/2306.06581/) (2023)
- [Preserving Geometric Integrity in Graph Prompting via Measure-Constrained Optimal Transport](https://mlsanity.com/paper/2609.23547/) (2026)

## Abstract (arXiv)

Global pooling is one of the most significant operations in many machine learning models and tasks, which works for information fusion and structured data (like sets and graphs) representation. However, without solid mathematical fundamentals, its practical implementations often depend on empirical mechanisms and thus lead to sub-optimal, even unsatisfactory performance. In this work, we develop a novel and generalized global pooling framework through the lens of optimal transport. The proposed framework is interpretable from the perspective of expectation-maximization. Essentially, it aims at learning an optimal transport across sample indices and feature dimensions, making the corresponding pooling operation maximize the conditional expectation of input data. We demonstrate that most existing pooling methods are equivalent to solving a regularized optimal transport (ROT) problem with different specializations, and more sophisticated pooling operations can be implemented by hierarchically solving multiple ROT problems. Making the parameters of the ROT problem learnable, we develop a family of regularized optimal transport pooling (ROTP) layers. We implement the ROTP layers as a new kind of deep implicit layer. Their model architectures correspond to different optimization algorithms. We test our ROTP layers in several representative set-level machine learning scenarios, including multi-instance learning (MIL), graph classification, graph set representation, and image classification. Experimental results show that applying our ROTP layers can reduce the difficulty of the design and selection of global pooling -- our ROTP layers may either imitate some existing global pooling methods or lead to some new pooling layers fitting data better. The code is available at urlhttps://github.com/SDS-Lab/ROT-Pooling.