Regularized Optimal Transport Layers for Generalized Global Pooling Operations

Most existing models empirically leverage simple pooling operations like add-pooling, mean-pooling, and max-pooling for convenience, whose practical performance is often sub-optimal. Many efforts have been made to achieve better pooling performance, which can be coarsely categorized into two strategies. The first strategy is applying multiple simple pooling operations jointly, e.g., concatenating the output of mean-pooling with that of max-pooling [11, 12]. In [10], the mixed mean-max pooling and its structured variants leverage mixture models of mean-pooling and max-pooling to improve pooling results. Recently, a generalized norm-based pooling (GNP) is proposed in [28, 29]. It can imitate max-pooling and mean-pooling under different settings. By learning its parameters, the GNP achieves a mixture of max-pooling and mean-pooling in a nonlinear way.

The second strategy is designing pooling operations with cutting-edge neural network architectures and empowering them with additional feature transformation and extraction abilities. The early work following this strategy includes the Network-in-Network in [13] and the Set2Set in [14], which integrate neural networks like classic multi-layer perceptrons (MLPs) and recurrent neural networks (RNNs) into pooling operations. Recently, attention-pooling and its gated version merge multiple instances with the help of different self-attentive mechanisms [2]. Based on a similar idea, the dynamic-pooling in [1] applies an iterative adjustment step to improve the self-attentive mechanism. More recently, the DeepSet in [15], the SetTransformer in [16], and the prototype-oriented set representer in [30] leverage and modify advanced transformer modules [31] to achieve sophisticated pooling operations. Besides the above global pooling methods, some attempts have been made to leverage graph structures to achieve pooling operations, e.g., the DiffPooling in [3], the ASAP in [32], and the self-attentive graph pooling (SAGP) in [33].

Different from the above methods, we study the design of pooling operation through the lens of computational optimal transport [34] and propose a novel and solid algorithmic framework to unify many representative global pooling methods. The neural network-based implementations of our framework can be interpreted as solving a regularized optimal transport problem with learnable parameters. As a result, instead of designing and selecting global pooling operations empirically, we can apply our ROTP layers to approximate suitable global pooling layers automatically according to observed data or achieve new pooling mechanisms with better performance.

Optimal transport (OT) theory [35] has proven to be useful in machine learning tasks, e.g., distribution matching [36, 37, 38], data clustering [39, 40], and generative modeling [41, 42]. Given the samples of two distributions, the discrete OT problem aims at learning a joint distribution of the samples (a.k.a., the optimal transport plan) and indicating the correspondence between them accordingly. This discrete OT problem is a linear programming problem [43]. By adding an entropic regularizer [23], the problem becomes strictly convex and can be solved efficiently by the Sinkhorn scaling in [44, 45]. Along this direction, the logarithmic stabilized Sinkhorn scaling algorithm [46, 47] and the proximal point method [48] make efforts to suppress the numerical instability of the classic Sinkhorn scaling algorithm and solve the entropic OT problem robustly. The Greenkhorn algorithm [49] provides a stochastic Sinkhorn scaling algorithm for batch-based optimization. When the marginal distributions of the optimal transport plan are unreliable or unavailable, the variants of the original OT problem, e.g., the partial OT [45, 50] and the unbalanced OT [46, 24] are considered, and the algorithms focusing on these variants are developed accordingly. Besides the Sinkhorn scaling algorithm, some other algorithms are developed, e.g., the Bregman ADMM [26, 51, 27], the smoothed semi-dual algorithm [52], and the conditional gradient algorithm [53, 54].

Recently, some attempts have been made to design neural networks to imitate the Sinkhorn-based algorithms of OT problems, such as the Gumbel-Sinkhorn network [55], the sparse Sinkhorn attention model [56], the Sinkhorn autoencoder [57], and the Sinkhorn-based transformer [58]. Focusing on pooling layers, some OT-based solutions have been proposed as well. In [59], an OT-based feature aggregation method called OTK is proposed. Following a similar idea, a differentiable expectation-maximization pooling method is proposed in [60], whose implementation is based on solving an entropic OT problem. More recently, the pooling methods based on sliced Wasserstein distance [61] are proposed, e.g., the sliced-Wasserstein embedding (SWE) in [62] and the WEGL in [63]. However, these methods only consider solving the pooling problems in specific tasks rather than developing a generalized pooling framework. Moreover, they focus on the optimal transport in the sample space and are highly dependent on Sinkhorn-based algorithms and random projections, which ignore the potentials of other algorithms.

Essentially, our ROTP layers can be viewed as new members of deep implicit layers [19, 64, 21] (or equivalently, called declarative neural networks [22]), which achieve global pooling by solving optimization problems. Many implicit layers have been proposed, e.g., the input convex neural networks (ICNNs) [65] and the deep equilibrium (DEQ) network [21], and so on. From the viewpoint of ordinary differential equations (ODEs), the DEQ network reformulates the ResNet [6] as a single implicit layer whose feed-forward computation corresponds to fixed point iterations. The OptNet in [66] provides a toolbox to implement convex optimization problems as neural network layers. At the same time, some researchers contributed to this field from the viewpoint of signal processing, implementing the iteration optimization steps of compressive sensing by neural networks [67, 68]. In general, the learning of such optimization-driven neural networks is based on auto-differentiation (AD) [69], which often owns high time and space complexity. Fortunately, for the layers corresponding to convex optimization problems, their gradients often have closed-form solutions, and their backpropagation steps are efficient [22].

Note that the study of implicit layers has been interactive with computational optimal transport methods. For example, an optimal transport model based on the input convex neural network has been proposed in [70]. More recently, the work in [71, 72] develops the Sinkhorn-scaling algorithm as implicit layers, whose backward computation is achieved in a closed form. Focusing on the design and the learning of global pooling operations, our generalized pooling framework provides a new optimization-driven solution.

Denote $\mathcal{X}_{D}=\{\bm{X}\in\mathbb{R}^{D\times N}|N\in\mathbb{N}\}$ as the space of sample sets, where each set $\bm{X}=[\bm{x}_{1},...,\bm{x}_{N}]\in\mathbb{R}^{D\times N}$ contains $N$ $D$-dimensional samples. In practice, $\bm{X}$ can be used to represent the instances in a bag, the node embeddings of a graph, or the local visual features in an image. Following the work in [29, 28, 73], we assume the input data $\bm{X}$ to be nonnegative in this study. This assumption is generally reasonable because the input data are often processed by non-negative activations, like ReLU, Sigmoid, Softplus, and so on. For some pooling methods, e.g., the max-pooling and the generalized norm pooling [28], the non-negativeness of input data is even necessary.

A global pooling operation, denoted as $f:\mathcal{X}_{D}\mapsto\mathbb{R}^{D}$, maps each set to a single vector and ensures the output is permutation-invariant, i.e., $f(\bm{X})=f(\bm{X}_{\pi})$ for $\bm{X},\bm{X}_{\pi}\in\mathcal{X}_{D}$, where $\bm{X}_{\pi}=[\bm{x}_{\pi(1)},...,\bm{x}_{\pi(N)}]$ and $\pi$ is an arbitrary permutation. As aforementioned, many pooling methods have been proposed to achieve this aim. Typically, the mean-pooling takes the average of the input vectors as its output, i.e., $f(\bm{X})=\frac{1}{N}\sum_{n=1}^{N}\bm{x}_{n}$. Another popular pooling operation, max-pooling, concatenates the maximum of each dimension as its output, i.e., $f(\bm{X})=\|_{d=1}^{D}\max_{n}\{x_{dn}\}_{n=1}^{N}$, where $x_{dn}$ is the $d$-th element of $\bm{x}_{n}$ and “$\|$” represents the concatenation operator. The attention-pooling in [2] outputs the weighted summation of the input vectors, i.e., $f(\bm{X})=\bm{X}\bm{a}_{\bm{X}}$, where $\bm{a}_{\bm{X}}\in\Delta^{N-1}$ is a vector on the $(N-1)$-Simplex. The attention-pooling leverages a self-attention mechanism to derive $\bm{a}_{\bm{X}}$ from the input $\bm{X}$, $i.e.$, $\bm{a}_{\bm{X}}=\text{softmax}(\bm{w}^{T}\text{tanh}(\bm{VX}))^{T}$.

The element $x_{dn}$ of the $\bm{X}$ represents the $d$-th feature of the $n$-th sample. From the perspective of statistical signal processing, it can be treated as the signal corresponding to a pair of the sample index and the feature dimension. Given all sample indices and feature dimensions, we can define a joint distribution for them, denoted as $\bm{P}=[p_{dn}]\in[0,1]^{D\times N}$, and the element $p_{dn}$ indicates the significance of the signal $x_{dn}$. Accordingly, the above global pooling methods yield the following generalized formulation that calculates and concatenates the conditional expectations of $x_{dn}$’s for $d=1,...,D$:

where $\odot$ is the Hadamard product, $\text{diag}(\cdot)$ converts a vector to a diagonal matrix, and $\bm{1}_{N}$ represents the $N$-dimensional all-one vector. $\bm{P1}_{N}=\bm{p}$ is the distribution of feature dimensions. $\text{diag}^{-1}(\bm{p})\bm{P}=\tilde{\bm{P}}=[p_{n|d}]$ normalizes the rows of $\bm{P}$, and the $d$-th row leads to the distribution of sample indices conditioned on the $d$-th feature dimension.

Based on the generalized formulation in (1), we can find that the above global pooling methods apply different mechanisms to derive the $\bm{P}$. Given $\bm{X}$, the mean-pooling treats each element evenly and $\bm{P}=[\frac{1}{DN}]$. The max-pooling sets $\bm{P}\in\{0,\frac{1}{D}\}^{D\times N}$ and $p_{dn}=\frac{1}{D}$ if and only if $n=\arg\max_{m}\{x_{dm}\}_{m=1}^{N}$. The attention-pooling derives $\bm{P}$ as a learnable rank-one matrix parameterized by the input $\bm{X}$, i.e., $\bm{P}=\frac{1}{D}\bm{1}_{D}\bm{a}^{T}_{\bm{X}}$. All these operations set the marginal distribution of feature dimensions to be uniform, i.e., $\bm{p}=\bm{P1}_{N}=[\frac{1}{D}]$. For the other marginal distribution $\bm{q}=\bm{P}^{T}\bm{1}_{D}$, some pooling methods impose specific constraints, e.g., $\bm{q}=\frac{1}{N}\bm{1}_{N}$ for the mean-pooling and $\bm{q}=\bm{a}_{\bm{X}}$ for the attention-pooling, while the max-pooling makes $\bm{q}$ unconstrained. In the following content, we will show that in general scenarios, we can relax the constraints to some regularization terms when the marginal distributions have some uncertainties.

The above analysis indicates that we can unify typical pooling operations in an interpretable algorithmic framework based on the expectation-maximization principle. In particular, the pooling operation in (1) is determined by the joint distribution $\bm{P}$. To keep the fused data as informative as possible, we would like to obtain a joint distribution that maximizes the expectations in (1), which leads to the following optimization problem:

where $\langle\cdot,\cdot\rangle$ represents the inner product of matrices. $\bm{p}=[p_{d}]\in\Delta^{D-1}$ and $\bm{q}\in\Delta^{N-1}$ are predefined distributions for feature dimensions and sample indices, respectively. Accordingly, the marginal distributions of $\bm{P}$ are restricted to be $\bm{p}$ and $\bm{q}$, i.e., $\bm{P}\in\Pi(\bm{p},\bm{q})=\{\bm{P}\geq\bm{0}|\bm{P}\bm{1}_{N}=\bm{p},\bm{P}^{T}\bm{1}_{D}=\bm{q}\}$.

As shown in (2), the objective function is the weighted summation of the expectations conditioned on different feature dimensions, which leads to the expectation of all $x_{dn}$’s. Here, we have connected the global pooling operation to the theory of optimal transport — (2) is a classic optimal transport problem [35], which learns the optimal joint distribution $\bm{P}^{*}$ to maximize the expectation of $x_{dn}$. Plugging $\bm{P}^{*}$ into (1) leads to a global pooling result of $\bm{X}$.

However, solving (2) directly often leads to undesired pooling results because of the following three reasons:

$i)$ The nature of linear programming. Solving (2) is time-consuming and always leads to a sparse $\bm{P}^{*}$ because it is a constrained linear programming problem. A sparse $\bm{P}^{*}$ tends to filter out some weak but possibly-informative signals in $\bm{X}$, which may have negative influences on downstream tasks.

$ii)$ The uncertainty of marginal distributions. Solving (2) requires us to set the marginal distributions $\bm{p}$ and $\bm{q}$ in advance. However, such exact marginal distributions are often unavailable or unreliable.¹¹1That is why most existing pooling methods either assume the marginal distributions to be uniform or make them unconstrained.

$iii)$ The lack of structural information. The optimal transport problem in (2) did not consider the structural relations among samples and those among features. However, real-world samples, like the node embeddings of a graph and the instances in a bag, are non-i.i.d. in general, and their features can be correlated. The state-of-the-art pooling methods, especially those for graph representation [74, 75, 3, 32], often take such structural information into account.

According to the above analysis, to make the optimal transport-based pooling applicable, we need to $i)$ improve the smoothness of the optimization problem, $ii)$ take the uncertainties of the marginal distributions into account, and $iii)$ leverage the sample-level and feature-level structural information hidden in the input data. Therefore, we extend (2) to a regularized optimal transport (ROT) problem:

where $\Omega=\{\bm{P}>\bm{0}|\bm{1}_{D}^{T}\bm{P}\bm{1}_{N}=1\}$. In (3), we introduce the following three regularizers, each of which solves one of the above three challenges.

$i)$ Smoothness regularization. $\text{R}(\bm{P})$ is a regularizer of $\bm{P}$, which is used to improve the smoothness of the objective function. Typically, we can set $\text{R}(\bm{P})$ as the negative entropy of $\bm{P}$ ($\text{R}(\bm{P})=\langle\bm{P},\log\bm{P}-\bm{1}\rangle=\sum_{d,n}p_{dn}(\log p_{dn}-1)$) [23] or the quadratic regularizer of $\bm{P}$ ($\text{R}(\bm{P})=\|\bm{P}\|_{F}^{2}$) [52]. The parameter $\alpha_{1}$ controls the significance of $\text{R}(\bm{P})$.

$ii)$ Marginal prior regularization. Instead of imposing strict constraints, we leverage two KL divergence terms in (3) to penalize the difference between the marginals of $\bm{P}$ and the predefined prior distributions (denoted as $\bm{p}_{0}$ and $\bm{q}_{0}$, respectively). Here, $\text{KL}(\bm{a}|\bm{b})=\langle\bm{a},\log\bm{a}-\log\bm{b}\rangle-\langle\bm{a}-\bm{b},\bm{1}\rangle$ represents the KL-divergence between $\bm{a}$ and $\bm{b}$. The strength of these two terms is controlled by the weights $\alpha_{2}$ and $\alpha_{3}$, respectively. This regularization helps us achieve a trade-off between the utilization of prior information and the robustness to its uncertainty.

$iii)$ Gromov-Wasserstein discrepancy-based structural regularization. Given $\bm{X}$, we construct its feature-level and sample-level covariance matrices,²²2In practice, besides the covariance matrices, we can also leverage other methods to define the feature-level and sample-level similarities, e.g., the cosine similarity and other kernel matrices. respectively, i.e., $\bm{\Sigma}_{1}=\frac{1}{N}(\bm{X}-\bm{\mu}_{1}\bm{1}_{N}^{T})(\bm{X}-\bm{\mu}_{1}\bm{1}_{N}^{T})^{T}$ and $\bm{\Sigma}_{2}=\frac{1}{D}(\bm{X}-\bm{1}_{D}\bm{\mu}_{2}^{T})^{T}(\bm{X}-\bm{1}_{D}\bm{\mu}_{2}^{T})$, where $\bm{\mu}_{1}=\frac{1}{N}\bm{X1}_{N}$ and $\bm{\mu}_{2}=\frac{1}{D}\bm{X}^{T}\bm{1}_{D}$. Following the work in [75, 74], we would like to make the feature-level covariance highly correlated with the sample-level covariance such that the features can preserve the structural relations among the samples. Therefore, we construct a structural cost as follows:

and $\langle C(\bm{X},\bm{P}),\bm{P}\rangle=-\text{tr}(\bm{\Sigma}_{1}\bm{P}\bm{\Sigma}_{2}^{T}\bm{P}^{T})$, whose significance is controlled by $\alpha_{0}$. It is easy to find that this structural regularization is the same as the objective of the Gromov-Wasserstein discrepancy problem in [76, 17], penalizing the difference between the sample-level covariance and the feature-level covariance. As shown in (3), combining the original optimal transport term with the structural regularizer leads to the well-known fused Gromov-Wasserstein (FGW) discrepancy [53], which is an optimal transport-based metric for structured data (like graphs and sets) [54, 77].

Compared with (2), (3) is an unconstrained optimization problem, and thus we can apply differentiable algorithms to optimize it.³³3When $\alpha_{0}=0$ and $\text{R}(\cdot)$ is a strictly-convex function, the objective function in (3) is strictly-convex. The optimal transport matrix $\bm{P}_{\text{rot}}^{*}$ can be viewed as a function of $\bm{X}$, whose parameters are the weights of the regularizers and the prior distributions, i.e., $\bm{P}_{\text{rot}}^{*}(\bm{X};\bm{\theta})$, where $\bm{\theta}=\{\alpha_{0},\alpha_{1},\alpha_{2},\alpha_{3},\bm{p}_{0},\bm{q}_{0}\}$ represents the model parameters for convenience. Plugging it into (1), we obtain the proposed regularized optimal transport pooling (ROTP) operation:

Our ROTP is a feasible global pooling operation. In particular, it satisfies the requirement of permutation-invariance under mild conditions.

Theorem 1 provides us with sufficient conditions to ensure the permutation-invariance of the proposed ROTP framework. Note that the two conditions in Theorem 1 are common in practice. When $\text{R}(\bm{P})$ is an entropic or a quadratic function of $\bm{P}$, Condition 1 is always held. The $\bm{q}_{0}$ is a permutation-equivariant function of $\bm{X}$ for the max-pooling and the attention-pooling, and it is uniform for the mean-pooling.

Moreover, our ROTP operation provides a generalized global pooling framework that unifies many representative pooling operations. In particular, the mean-pooling, the max-pooling, and the attention-pooling can be formulated as the specializations of (5) under different parameter configurations.

Proposition 1 demonstrates that a single ROTP operation can imitate various global pooling methods. Moreover, the hierarchical combination of multiple ROTP operations can reproduce more complicated pooling mechanisms, including mixed pooling operation and set fusion.

We first consider leveraging the proximal point method [17, 78, 48] to solve (3) iteratively. In particular, in the $t$-th iteration, given the current OT plan $\bm{P}^{(t)}$, we solve the following sub-problem:

Here, $\text{KL}(\bm{P}|\bm{P}^{(t)})$ is the proximal term based on the current variable, which is implemented as a KL-divergence. The weight $\tau$ controls its significance. As a result, our ROTP layer is built by stacking $T$ feed-forward modules, and each module corresponds to the optimization of (15).

When the smoothness regularizer is entropic, i.e., $\text{R}(\bm{P})=\langle\bm{P},\log\bm{P}-\bm{1}\rangle$, (15) becomes the following entropic unbalanced optimal transport (EUOT) problem:

where the matrix $\bm{C}^{(t)}=[c^{(t)}_{dn}]=-\bm{X}-\alpha_{0}\bm{\Sigma}_{1}\bm{P}^{(t)}\bm{\Sigma}_{2}^{T}-\tau\log\bm{P}^{(t)}$ is determined by the input data and the current variable $\bm{P}^{(t)}$. According to [46, 24], we consider the Fenchel’s dual form of the EUOT problem:

This problem can be solved by the Sinkhorn-scaling algorithm [44, 23]. In particular, the Sinkhorn-scaling algorithm solves this dual problem by the following iterative steps: $i)$ Initialize dual variables as $\bm{a}^{(0)}=\bm{0}_{D}$ and $\bm{b}^{(0)}=\bm{0}_{N}$. $ii)$ In the $k$-th iteration, the current dual variables $\bm{a}^{(k)}$ and $\bm{b}^{(k)}$ are updated by

where $\alpha_{1}^{\prime}=\alpha_{1}+\tau$. $iii)$ After $K$ steps, the variables converges and the optimal transport plan is updated as

The convergence of the algorithm has been proven in [79] — with the increase of $t$, $\bm{P}^{(t)}$ converges to a stationary point. Therefore, after repeating the above process $T$ times, we set $\bm{P}^{*}:=\bm{P}^{(T)}$.

The algorithm above leads to a Sinkhorn-based ROTP layer $f_{\text{rot}}(\bm{X};\bm{\theta})$. As illustrated in Fig. 2, this layer is implemented by unrolling the above iterative scheme by stacking $T$ proximal point modules, and each module is implemented by $K$ Sinkhorn-scaling steps. Furthermore, we can leverage the logarithmic stabilization strategy [46, 47] to improve the numerical stability of the Sinkhorn-scaling algorithm — instead of updating $\bm{T}(\bm{a}^{(k)},\bm{b}^{(k)})$ directly, we can update $\log\bm{T}(\bm{a}^{(k)},\bm{b}^{(k)})$. Accordingly, the exponential and scaling operations in (18) are integrated as “LogSumExp” operations, which helps to avoid numerical instability issues. Algorithm 1 shows the Sinkhorn-based ROTP layer implemented based on the stabilized Sinkhorn-scaling algorithm. Here, LogSumExp${}_{\text{col}}$ and LogSumExp${}_{\text{row}}$ apply column-wise and row-wise summation, respectively.

The Sinkhorn-based ROTP layer requires the smoothness regularizer $\text{R}(\bm{P})$ to be entropic, which limits its generalizability. Additionally, even if applying the logarithmic stabilization strategy, the Sinkhorn-scaling algorithm still has a high risk of numerical instability because the parameters $\alpha_{0}$ and $\alpha_{1}$ may change in a wide range during training. To overcome the challenges, we apply a Bregman alternating direction method of multipliers (Bregman ADMM or BADMM) to build another ROTP layer. In particular, we first rewrite (3) in an equivalent format by introducing three auxiliary variables $\bm{S}$, $\bm{\mu}$ and $\bm{\eta}$:

These three auxiliary variables correspond to the joint distribution $\bm{P}$ and its marginals. This problem can be further rewritten in a Bregman augmented Lagrangian form by introducing three dual variables $\bm{Z}$, $\bm{z}_{1}$, $\bm{z}_{2}$ for the three constraints in (20), respectively. For the ROT problem with auxiliary variables, we can write its Bregman augmented Lagrangian form as

Here, $\text{Div}(\cdot,\cdot)$ represents the Bregman divergence term, which is implemented as the KL-divergence as the work in [26, 27] did. Its significance is controlled by $\rho>0$. The last three lines of (21) contain the Bregman augmented Lagrangian terms, which correspond to the three constraints in (20). Here, $\text{R}(\bm{S},\bm{P})$ corresponds to the smoothness regularizer. When applying the entropic regularizer, we set $\text{R}(\bm{S},\bm{P})=\langle\bm{S},\log\bm{S}-\bm{1}\rangle$. When applying the quadratic regularizer, we set $\text{R}(\bm{S},\bm{P})=\langle\bm{S},\bm{P}\rangle$.

In particular, we solve the ROT problem by alternating optimization: At the $t$-th iteration, we first update $\bm{P}$ while fix other variables. We can ignore Constraint 3 and the three regularizers (because they are irrelevant to $\bm{P}$) and write Constraint 2 explicitly. The problem becomes:

where $\Pi(\bm{\mu}^{(t)},\cdot)=\{\bm{P}>\bm{0}|\bm{P}\bm{1}_{N}=\bm{\mu}^{(t)}\}$ is the one-side constraint of $\bm{P}$. We can derive the closed-form solution of this problem based on the first-order optimality condition:

where $\bm{Y}=$