Revisiting Global Pooling through the Lens of Optimal Transport

TL;DR

This paper introduces a unified global pooling framework based on optimal transport, specifically unbalanced optimal transport, and proposes a new layer called UOT-Pooling that improves neural network performance across various tasks.

Contribution

The paper develops a novel global pooling method using optimal transport theory, unifies existing methods, and introduces a new stable algorithmic approach for neural network pooling layers.

Findings

UOT-Pooling can replicate existing pooling methods.

UOT-Pooling achieves better performance in multiple tasks.

The Bregman ADMM algorithm enhances numerical stability.

Abstract

Global pooling is one of the most significant operations in many machine learning models and tasks, whose implementation, however, is often empirical in practice. In this study, we develop a novel and solid global pooling framework through the lens of optimal transport. We demonstrate that most existing global pooling methods are equivalent to solving some specializations of an unbalanced optimal transport (UOT) problem. Making the parameters of the UOT problem learnable, we unify various global pooling methods in the same framework, and accordingly, propose a generalized global pooling layer called UOT-Pooling (UOTP) for neural networks. Besides implementing the UOTP layer based on the classic Sinkhorn-scaling algorithm, we design a new model architecture based on the Bregman ADMM algorithm, which has better numerical stability and can reproduce existing pooling layers more effectively. We test our UOTP layers in several application scenarios, including multi-instance learning, graph classification, and image classification. Our UOTP layers can either imitate conventional global pooling layers or learn some new pooling mechanisms leading to better performance.