O N
L EARNING
TO
S OLVE
C ARDINALITY
C ONSTRAINED C OMBINATORIAL
O
PTIMIZATION IN O NE -
S HOT :
A
R E - PARAMETERIZATION
A PPROACH
VIA
G UMBEL -S INKHORN -T OP K
Anonymous authors
Paper under double-blind review
A BSTRACT
Cardinality constrained combinatorial optimization requires selecting an optimal subset of k elements, and it will be appealing to design data-driven algorithms that perform TopK selection over a probability distribution predicted by a neural network.
However, the existing differentiable TopK operator suffers from an unbounded gap between the soft prediction and the discrete solution, leading to inaccurate estimation of the combinatorial objective score.
In this paper, we present a self-supervised learning pipeline for cardinality constrained combinatorial optimization, which incorporates with Gumbel-Sinkhorn-TopK (GS-TopK) for neardiscrete TopK predictions and the re-parameterization trick resolving the nondifferentiable challenge.
Theoretically, we characterize a bounded gap between the Maximum-A-Posteriori (MAP) inference and our proposed method, resolving the divergence issue in the previous differentiable TopK operator and also providing a more accurate estimation of the objective score given a provable tightened bound to the discrete decision variables.
Experiments on max covering and discrete clustering problems show that under comparative time budget, our method outperforms state-of-the-art Gurobi solver and the novel one-shot learning method Erdos Goes Neural.
I NTRODUCTION
In this paper, we aim to solve a family of cardinality constrained combinatorial optimization: given the cardinality budget k , our aim is to ﬁnd the optimal set with size k whose covered objects correspond to the maximum sum of values.
Extensive applications of these problems can be found in machine learning and operations research, for example, discovering coresets for dataset summarization and continual learning (Borsos et al., 2020), discovering the most inﬂuential seed users in social networks (Chen et al., 2021), and planning facility locations in operation research (Liu, 2009).
Although there has been a success in discovering the important submodularity (Fujishige, 1991) for a family of these problems with a bounded approximation ratio of greedy algorithms, the ubiquitous NP-hard challenge is always motivating us to develop improved algorithms with better efﬁciency-accuracy trade-off.
The recent success of machine learning has inspired researchers to develop learning-based algorithms for combinatorial optimizations (Vinyals et al., 2015; Dai et al., 2016; Karalias & Loukas, 2020; Wang et al., 2021b), and this paper falls in line with these works.
According to the taxonomy by a recent survey (Peng et al., 2021), existing machine learning methods for combinatorial optimization can be mainly categorized into two groups: multi-step methods (Vinyals et al., 2015;
Dai et al., 2016; Lu et al., 2019; Wang et al., 2021a) and one-shot methods (Wang et al., 2019; Karalias & Loukas, 2020; Li et al., 2019).
The former line of works usually formulates the combinatorial optimization solving procedure as a Markov Decision Process (MDP) and refers to novel reinforcement learning methods (Mnih et al., 2013; Schulman et al., 2017) to tackle the MDP.
Perhaps the largest drawback of such a multi-step pipeline is that the training and inference procedures can be inefﬁcient.
Based on some previous efforts (Wang et al., 2019; Li et al., 2019), there is a recent seminal work (Karalias & Loukas, 2020) aiming to develop a one-shot pipeline for general combinatorial problems.
Although being a fast and general pipeline, Karalias & Loukas (2020) do not encode constraint in the model prediction and thus the estimation of the combinatorial objective score may be inaccurate.
In this paper, we aim to facilitate the one-shot learning pipeline with cardinality constraints, such that the constraints can be softly encoded by the output of a neural network.
Regarding the input to
TopK as a probability distribution, the maximum-a-posterior (MAP) inference is done by sorting all probabilities and selecting the K -largest ones, however, this process is non-differentiable.
Though appealing, it is non-trivial to encode the constraint in a differentiable manner with bounded gap w.r.t.
the MAP inference and the bound of the existing SOFT-TopK algorithm (Xie et al., 2020) diverges if the k -th and ( k +1) -th probabilities are equal.
Unfortunately, such a dilemma is possible because the probabilities are outputs from a neural network.
Besides, concerning the discrete decision variables in combinatorial optimization, it is also non-trivial to build an accurate estimation of the objective score with neural network output (which has to be continuous in order to preserve the gradient).
To tackle the non-differentiable challenge and the ill-estimated objective score, we refer to the seminal reparameterization technique with Gumbel noise, which can mimic the sampling procedure from the continuous probability distributions to near-discrete decision variables, while still preserving the gradient.
Such a technique has been found successful for Top-1 (Jang et al., 2017), permutation (Mena et al., 2018), and sorting (Grover et al., 2019) distributions, and in this paper, we improve the SOFT-TopK operator (Xie et al., 2020) by introducing the Gumbel re-parameterization trick, and we theoretically prove that its gap can be bounded with the existence of the Gumbel noise, in contrast to the gap of SOFT-TopK (Xie et al., 2020) that may turn diverged.
To this end, we present G umbel- S inkhorn- TopK (GS-TopK) to facilitate the one-shot selfsupervised learning pipeline (Karalias & Loukas, 2020) with constrained neural network output.
The proposed GS-TopK operator improves the SOFT-TopK (Xie et al., 2020), where we theoretically characterize a bounded gap by introducing non-zero noise factors.
Besides, we show how to integrate the proposed GS-TopK method into the combinatorial optimization pipeline, with an accurate estimation of the objective score, formulating a self-supervised learning pipeline following (Karalias & Loukas, 2020).
The contributions of this paper are summarized as follows
: 1) We propose a self-supervised learning approach to solve cardinality constrained combinatorial optimizations in one-shot.
The cardinality constraints are encoded in the model output via our proposed GS-TopK, and our pipeline is differentiable via the re-parameterization trick.
We theoretically characterize a bounded gap between the MAP inference and the approximate solution by Gumbel-Sinkhorn-TopK without requiring the large gap of k -th and ( k +1) -th probabilities, which improves the result of the state-of-the-art SOFT-TopK algorithm (Xie et al., 2020).
This property is crucial for an accurate estimation of the combinatorial objective score in learning.
Experimental results on max covering and discrete clustering problems show that our method surpasses the state-of-the-art one-shot learning methods e
g. (Karalias & Loukas, 2020), and also surpasses the state-of-the-art Gurobi and SCIP solvers with even smaller inference time budgets.
R ELATED W ORK 2. L EARNING OF C OMBINATORIAL O PTIMIZATION
Recently it is a trending topic of developing machine learning based combinatorial optimization methods to achieve a better trade-off between time and accuracy by modern data-driven approaches.
Here we follow the taxonomy by Peng et al. and discuss some representative works.
Multi-step methods predict the solution by a multi-step decision process.
Since the fundamental work (Vinyals et al., 2015) that proposes an encoding-decoding architecture, researchers have explored this general pipeline to various real-world combinatorial optimization problems.
The most popular framework follows the reinforcement learning pipeline introduced by Khalil et al.
(2017), where the state is encoded by a graph neural network, and multi-step actions are taken to sequentially construct the solution.
Applications have been explored on job scheduling (Zhang et al., 2020), computer resource allocation (Mao et al., 2019), quadratic assignment programming (Liu et al., 2020), bin-packing (Duan et al., 2019), just to name a few.
The advantage of multi-step methods is that it can handle combinatorial constraints that are non-trivial to be encoded in a differentiable manner, by adding constraint on the feasible actions in the multi-step roll-out procedure.
However, the training of reinforcement learning may be inefﬁcient and vulnerable, and there are also some recent efforts on developing one-shot combinatorial optimization learning methods.
One-shot methods predict the solution by a single forward pass.
Despite the challenges in encoding general constraints, one-shot learning methods have been successful on certain problems where the constraint can be encoded by the neural network.
For example, in graph matching (Wang et al., 2019; 2021b) the constraint can be encoded by Sinkhorn algorithm (Sinkhorn & Rangarajan, 1964), and also for graph similarity regression problems (Bai et al., 2019; Li et al., 2019) that are unconstrained.
The seminal work (Karalias & Loukas, 2020) aims to develop a general pipeline for one-shot learning of combinatorial optimization, by encoding the violation of the constraint as part of its objective score.
However, our experimental results show that compared to Karalias & Loukas (2020), it will be more appealing if we can explicitly encode the constraint in the neural network output.
2. O PTIMAL T RANSPORT AND G UMBEL R E - PARAMETERIZATION
Optimal Transport.
The optimal transport (OT) is a linear programming problem that aims to ﬁnd the optimal transportation plan with minimal cost.
Though the development of the Sinkhorn algorithm can be date back to 1960s (Sinkhorn & Rangarajan, 1964), it has drawn increasing attention from the machine learning community since the pioneering work (Cuturi, 2013) which introduces the entropic-regularization for the optimal transport problem.
Since the Sinkhorn algorithm (Sinkhorn & Rangarajan, 1964) incorporates row- and column-wise normalization which is differentiable, OT can be incorporated within any end-to-end deep learning pipelines.
Thus, OT has been found successful for deep generative models (Arjovsky et al., 2017; Patrini et al., 2019), solving jigsaw puzzles (Santa Cruz et al., 2018), and deep graph matching (Fey et al., 2020; Yu et al., 2020).
In complement to the recent line of differentiable sorting papers (Petersen et al., 2021; Cuturi et al., 2019), Xie et al.
(2020) develop the differentiable TopK algorithm namely SOFT-TopK, by formulating the TopK problem as an OT problem.
However, the theoretical gap between the prediction of SOFT-TopK and the result obtained by sorting diverges if the k -th and ( k + 1) -th largest items are equal.
In this paper, we prove that this gap can be bounded by introducing Gumbel noise.
Gumbel Re-parameterization.
The Gumbel distribution with the re-parameterization trick is shown effective to mimic the sampling procedure from a given distribution in a differentiable manner.
Jang et al.
(2017); Maddison et al. (2017) take the initiative to develop Gumbel-Softmax, where one-hot vectors are sampled from a softmax probability distribution.
The non-differentiable sampling procedure turned to be differentiable by the namely re-parameterization trick (Kingma & Welling, 2013).
Inspired by Jang et al. (2017), Mena et al.
(2018) propose a Gumbel-Sinkhorn method to learn the latent permutations, which is relevant to the GS-TopK method developed in this paper.
Besides, Paulus et al.
(2020) develop a general framework for generating combinatorial distributions via random distributions.
However, these existing methods do not study the application in combinatorial optimization learning, and the gap between the MAP inference and the result after Gumbel re-parameterization is not well characterized in Mena et al.
In this paper, we explore the application of the re-parameterization trick with Gumbel noise in combinatorial optimization, also we theoretically characterize the gap between the MAP inference and the proposed Gumbel method.
T HE P ROPOSED G UMBEL -S INKHORN -T OP K 3. P RELIMINARY : SOFT-T OP K A LGORITHM
The TopK problems are ubiquitous in machine learning, yet the MAP inference approach (performing sorting at the time cost of O ( m log m ) and selecting the ﬁrst k items) is unfriendly for current gradient-based deep neural networks.
To introduce differentiable TopK for machine learning, Xie et al.
(2020) formulate the TopK problem as an optimal transport (OT) problem.
Speciﬁcally, for selecting the k largest items from the probability vector s of size m ( k < m ), the OT problem refers to move k items to one destination (selected), and the other m − k elements to another destination (not selected).
Thus, the marginal distributions of the OT are:
[Equation]
We can design the destinations to be the min/max values of s , such that the TopK items are moved to max( s ) , and the other items are moved to min( s ) .
Then, the distance matrix of OT is given as:
[Equation]
Combinatorial Optimization
While the OT is formulated as an integer linear programming problem:
[Equation]
where T is the transportation plan which also corresponds to the TopK solution of the problem, and 1 is a column vector whose all elements are 1s.
The optimal solution T ∗ should be equivalent to the solution by ﬁrstly sorting all items and then select the TopK items.
It is also equivalent to the MAP inference by regarding s 1 , ..., s m as probabilities.
Following the differentiable Sinkhorn algorithm for OT (Cuturi, 2013; Sinkhorn & Rangarajan, 1964), the binary constraint on T is relaxed to continuous values [0 , 1] , and Eq.
(3) is modiﬁed with entropic regularization:
[Equation]
where h ( T τ ) =
(cid:80) i,j T τij log T τij is the entropic regularizer (Cuturi, 2013).
Given any real-valued matrix D , Eq.
(4) is solved by ﬁrstly enforcing the regularization factor τ : T τ = exp( − D /τ ) .
Then
T τ is row- and column-wise normalized alternatively:
[Equation]
where (cid:11) means element-wise division.
We denote T τ ∗ as the converged solution, which is the optimal solution to Eq.
Theorem 2 of (Xie et al., 2020) .
Denote x k , x k +1 as the k -th
( k + 1) -th largest items, we have
[Equation]
where the gap is controlled by τ if | x k − x k +1 | > 0 .
However, the above bound diverges if x k = x k +1 , which is unavoidable because x k , x k +1 are outputs by a neural network.
It motivates us to develop our improved version Gumbel-Sinkhorn-TopK by introducing the Gumbel noise.
3. G UMBEL -S INKHORN -T OP K (GS-T OP K) A LGORITHM
In this section, we present our proposed G umbel- S inkhorn- TopK (GS-TopK) as summarized in
Alg.
and we theoretically characterize the gap between the MAP inference and the prediction by
GS-TopK.
As the re-parameterization trick (Jang et al., 2017), instead of sampling from a distribution which is non-differentiable, we add random variables to probabilities predicted by neural networks.
The Gumbel distribution is characterized as:
[Equation]
where σ controls the variance and u is from (0 , 1) uniform distribution.
We can update s and D as:
[Equation]
Again we formulate the integer linear programming version of the OT with Gumbel noise:
[Equation]
where the optimal solution to Eq.
(10) is denoted as T σ ∗ .
Without loss of generality, in the following we sort s 1 , s 2 , s 3 , ..., s m in descending order, into X = x 1 , x 2 , x 3 , ..., x m , and x k , x k +1 are the k -th and ( k + 1) -th largest items, respectively.
We characterize the gap between T ∗ and T σ ∗ under the expectation over Gumbel distribution.
Now we present the following two lemmas to help establish our theoretical results, in terms of a bounded gap between the MAP inference and the approximate solution by Gumbel-Sinkhorn-TopK,
Lemma 1 ( Bias by the Gumbel Noise ) .
For the integer linear programming solutions with and without Gumbel noises, we have
[Equation]
Proof sketch:
The gap between T ∗ , T σ ∗ is derived by counting the expected number of times that, after adding the Gumbel noise, an element may become larger than x k or smaller than x k +1 .
For the top k elements, the probability that it will become smaller than x k +1 is upper bounded by 1 / (1 + exp | x k − x k +1 | σ ) .
For the last ( m − k ) elements, the probability that it will become larger than x k is also upper bounded by 1 / (1 + exp
| x k − x k +1 | σ ) .
Since || T ∗ − T σ ∗ || 2 F changes at most by 4 if an item becomes larger/smaller than x k / x k +1 , and we have m items, then we can proof the upper bound of this gap.
The detailed proof is referred to Appendix A.1.
To make the integer linear programming problem feasible for gradient-based deep learning methods, we also relax the integer constraint and add the entropic regularization term:
[Equation]
which is solved by Sinkhorn algorithm on (cid:101) D : ﬁrstly (cid:101) T = exp( − (cid:101) D /τ )
, then
[Equation]
Here we denote the optimal solution to Eq.
(12) as (cid:101) T ∗ .
We theoretically characterize the gap between
T σ ∗ and (cid:101) T ∗ , which corresponds to the optimal solutions of Eq.
(10) and Eq.
(12), respectively.
Lemma 2 ( Gap to a Discrete Solution ) .
Under probability (1 − (cid:15) ) , we have
[Equation]
where
[Equation]
Proof sketch:
The gap between T σ ∗ and (cid:101) T ∗ is a generalization from the Theorem 2 of
Xie et al.
We denote x π k , x π k +1 as the k -th and ( k + 1) -th largest items after disturbed by the Gumbel noise, and our aim becomes to prove the upper bound of E u (cid:2) 1 / ( | x π k + g σ ( u π k ) − x π k +1 − g σ ( u π k +1 ) | ) (cid:3) , where the probability density function of g σ ( u π k ) − g σ ( u π k +1 ) can be bounded by f ( y )
=
1 / ( y 2 + 4) .
Thus we can compute the upper bound under probability (1 − (cid:15) ) by integration.
The detailed proof is referred to Appendix A.2.
Theorem 1.
For the proposed Gumbel TopK solver, denote T ∗ as the solution by deterministic TopK algorithm, (cid:101) T ∗ as the solution of Gumbel-Sinkhorn-TopK. With probability (1 − (cid:15) ) , we have
[Equation]
where σ is the noise factor and τ is the temperature of Sinkhorn algorithm.
m is the number of candidates to be selected.
Ω( x i , x j , σ, (cid:15) ) is deﬁned in Eq. [Figure]
Proof: Based on Lemma 1 and Lemma [Figure] Theorem 1 is proved by triangle inequality:
[Figure]
and we compare GS-TopK and SOFT-TopK concerning the gap to a discrete solution w.r.t.
different τ, σ conﬁgurations. Here the gap to discrete solutions is tighten by a larger σ and a smaller τ for GS-TopK, compared to SOFT-TopK whose gap is larger and can only be controlled by τ .
Remarks .
As derived above, the gap between
T ∗ and (cid:101) T is composed of two parts: the bias by the Gumbel noise (in Lemma 1) and the gap to a discrete solution (in Lemma 2).
Interestingly, the ﬁrst term becomes smaller given a smaller σ , and the second term is smaller given a larger σ or a smaller τ .
See the toy example in Fig.
1, a larger σ can tighten the gap of GSTopK to a discrete solution, which is welcomed in our combinatorial optimization learning scenario where (cid:101) T ∗ is used to compute the objective score in a self-supervised learning pipeline.
If | x k − x k +1 | > 0 , with σ → 0 + , Eq.
(16) degenerates to the bound derived by (Xie et al., 2020) and only differs by a constant factor (see
Appendix A.3 for details):
[Equation]
It makes a strong assumption that | x k − x k +1 | > 0 , and the bound diverges if x k = x k +1 .
Unfortunately, x k , x k +1 are predictions by a neural network, whereby such an assumption may not be satisﬁed.
In comparison, given σ > 0 , our conclusion in Eq.
(16) is bounded for any x k , x k +1 .
3. A PPLICATION TO S ELF -S UPERVISED C OMBINATORIAL O PTIMIZATION L EARNING
In the following, we discuss the application of GS-TopK to self-supervised combinatorial optimization learning, by taking the max covering problem as an example.
We consider the max covering problem with m sets and n objects.
Each set may cover any number of objects, and each object is associated with a value.
Our aim is to ﬁnd k sets (where k < m ) such that the covered objects have the maximum sum of values.
The general illustration of our pipeline can be found in Fig.
Firstly, we build a bipartite graph whose two disjoint sets of vertices are the sets and the objects.
An edge is deﬁned if an object is covered by a set.
Denote v ∈ R n as the value of each object, and
A ∈ { 0 , 1 } m × n as the adjacency matrix, the problem is formulated as
[Equation]
where I ( x ) is an indicator I ( x ) i
= [-
-]1 if x i ≥ 1 else I ( x ) i = 0 .
encode the bipartite graph we exploit three layers of GraphSage (Hamilton et al., 2017) followed by a fully-connected layer with sigmoid to predict the probability of selecting each set, which is denoted as s ∈ [0 , 1] m .
The probabilities s are fed into the GS-TopK algorithm, which outputs a batch of near-discrete transportation matrices [ (cid:101) T 1 , (cid:101) T 2 , ..., (cid:101)
T #G ] , whose second row is regarded as a continuous approximation of the decision variables.
The objective value is estimated as:
[Equation]
For the discrete clustering problem, our aim is to select k objects from the full set of m objects, by minimizing the sum of distance for each object between itself and its nearest selected object.
It differs from the k-means clustering by adding the ‘discrete’ constraint that cluster centers has to be a subset of objects, which is also known as facility location problem in operations research (Liu, 2009).
Denote ∆ ∈ R + m × m as the distance matrix for each pair of objects, we formulate the discrete clustering problem as follows (please refer to Appendix B for details):
[Equation]
As illustrated in Alg. 2, during training, the Adam optimizer (Kingma & Ba, 2014) is applied for learning the neural network weights.
During inference, the neural network prediction is regarded as initialization, and we again optimize the distribution w.r.t.
the combinatorial objective score by
Adam, during which a search procedure is also performed among all Gumbel samples.
Homotopy GS-TopK .
Based on our Theorem 1, the gap between T ∗ and (cid:101) T ∗ can be tightened by shrinking τ and σ , which motivates us to develop a homotopy variant (Xiao & Zhang, 2012; Xu et al., 2016) of GS-TopK with gradually decreased τ, σ , such that (cid:101) T ∗ and (cid:101) J can provide more accurate estimation during inference.
See line 13 in Alg. 2 for details.
E XPERIMENTS AND D ISCUSSION 4. P ROTOCOLS AND B ASELINES
We apply the proposed GS-TopK algorithm to learning two representative problems of cardinality constrained NP-hard combinatorial optimization: i) max covering and ii) discrete clustering problem.
For each problem, both a smaller-sized case and a larger-sized case are considered, and we build separate training/testing datasets for learning methods.
In this paper, we referred to randomly generated datasets following previous machine learning for combinatorial optimization papers (Khalil
[Figure]
We compare the following representative and diverse baselines with our proposed method:
Classic Algorithms .
Classic algorithms are well-developed strong baselines with theoretical guarantees.
In our experiment, we consider the greedy algorithm as the baseline for the max covering problem, which has the worst-case approximation ratio of (1 − 1 /e )
due to the submodularity of the max covering problem (Fujishige, 1991).
For the discrete clustering problem, we resort to the classic kmeans (Macqueen, 1967) and kmeans++ (Arthur & Vassilvitskii, 2007) with improved initialization techniques.
These algorithms are fast but effective, even surpassing commercial solvers with only 1 / 1000 inference time on certain test cases.
General Purpose Solvers .
We also compare two general-purpose solvers Gurobi 9.0 (under educational license)
(Gurobi Optimization, LLC, 2021) and SCIP 7.0 (Gamrath et al., 2020).
Gurobi is the state-of-the-art commercial solver for general combinatorial optimization tasks, and it is known heavily engineered by various speedup techniques.
SCIP is the state-of-the-art open-source solver, which is also regarded as a baseline.
We formulate the combinatorial problems as integer programming and then call the solver to solve the problems, with a limited time budget.
Our implementation of the solvers is with the Google ORTools API 2 .
One-shot Learning Baselines .
We compare with the general Erdos Goes Neural (EGN) (Karalias &
Loukas, 2020) framework that models each decision variable as a Bernoulli distribution and encodes the constraint as part of the learning objective.
Since the authors of (Karalias & Loukas, 2020) do not consider cardinality constrained problems in their experiment, we re-implement and tune the EGN model based on their ofﬁcial implementation 3 .
We also compare with SOFT-TopK (Xie et al., 2020), which is theoretically characterized as a special case of our proposed GS-TopK when σ = 0 .
Also to ensure the effectiveness of the re-parameterization trick, we facilitate SOFT-TopK with Gumbel sampling during inference.
All one-shot learning methods are with the same model structure for a fair comparison and are trained and tested on separate datasets.
Besides, we empirically ﬁnd the learning process of one-shot methods converges within tens of minutes.
Since the reinforcement learning methods (Khalil et al., 2017; Chen & Tian, 2019) require signiﬁcantly more training time, we only include one-shot methods for fair comparison.
4. E XPERIMENTAL R ESULTS
Table 1 and Table 2 report results on max covering and discrete clustering problems, respectively.
An interesting ﬁnding is that the classic algorithms are efﬁcient but also very effective.
However, if we can afford a higher time budget, there seems no universal approach to improve the results of the classic algorithms.
In comparison, Gurobi and SCIP solvers can trade-off time for better accuracy, but they are less effective than our learning methods when considering the objective scores under the same time budget.
Among one-shot learning methods, Erdos Goes Neural (Karalias & Loukas, 2020) is the only method that does not encode the constraint in the model output, and the experimen-
[Table] [Table]
tal results suggest that it is appealing if we can encode the constraint in the model output. Directly applying Gumbel sampling to SOFT-TopK seems either an effective improvement, and satisfying results are achieved when exploiting the Gumbel re-parameterization trick and enabling gradientbased optimization over the combinatorial objective score.
We further validate the effectiveness of
Homotopy GS-TopK, where the objective scores are improved at the cost of more inference time.
Further Discussions .
Our insight for the effectiveness of GS-TopK over other state-of-the-art oneshot learning methods (Karalias & Loukas, 2020; Xie et al., 2020) is that it can provide a more accurate estimation of the combinatorial objective score with discrete decision variables.
Fig.shows that our GS-TopK estimates a lower bound of the true objective score at the beginningiterations, while providing an accurate estimation afterwards.
Besides, the Homotopy version of our algorithm can gradually reduce the variance and help converge to a better solution.
C ONCLUSION
We have presented Gumbel-Sinkhorn-TopK, which performs differentiable TopK operation over the input distribution and is facilitated with the Gumbel re-parameterization trick.
We theoretically characterize the gap between the MAP inference and the Gumbel-Sinkhorn-TopK, which improves the state-of-the-art SOFT-TopK operator (Xie et al., 2020).
We further show the successful application in learning to solve the important cardinality constrained combinatorial optimization problems.
R EPRODUCIBILITY S TATEMENT
The following efforts are made to ensure the reproducibility of this paper:
• We provide a complete proof about our Lemmas and Theorems in the Appendix, see Appendix A; • We discuss the hyper-parameter conﬁgurations and the experiment setup in Sec 4.1; • We provide the implementation details on max covering problem in Sec. 3.3, and the implementation details on discrete problem problem in Appendix B; •
We discuss the way of generating random dataset in Sec 4.1; • The code will be made publicly available once this paper is accepted.
A P ROOF OF T HEOREMS
Before starting the detailed proof of the lemmas and theorems, ﬁrstly we recall the notations used in this paper:
• T ∗ =
TopK ( D ) is the optimal solution of the integer linear programming form of the OT problem Eq.
(3), which is equivalent to the solution by ﬁrstly sorting all items and then select the TopK items; • T τ ∗ =
Sinkhorn ( D ) is the optimal solution of the entropic regularized form of the OT problem Eq.
(4) solved by Sinkhorn algorithm.
It is also the output by SOFT-TopK (Xie et al., 2020); • T σ ∗ =
TopK ( (cid:101) D ) is the optimal solution of the integer linear programming form of the OT problem after disturbed by the Gumbel noise Eq.
(10), which is equivalent to the solution by ﬁrstly adding the Gumbel noise, then sorting all items and ﬁnally select the TopK items; • (cid:101) T ∗ =
Sinkhorn ( (cid:101) D ) is the optimal solution of the entropic regularized form of the OT problem after disturbed by the Gumbel noise Eq.
(12) solved by Sinkhorn algorithm.
It is also the output by our proposed GS-TopK.
Lemma 3.
Given real numbers x i , x j , and u i , u j are from i.i.d.
(0 , 1) uniform distribution.
After
Gumbel perturbation, the probability that x i + g σ ( u i ) >
x j + g σ ( u j ) is:
[Equation]
Proof.
Since g σ ( u i ) = − σ log( − log( u i )) , P ( x i + g σ ( u i ) >
x j + g σ ( u j )) is equivalent to the probability that the following inequality holds:
[Equation]
And we have
[Equation] [Equation]
Since u j ∈ (0 , 1) , log( u j ) < 0 .
Then we have
[Equation] [Equation]
Since u i , u j are i.i.d.
uniform distributions, the probability when the above formula holds is
[Equation]
[Equation]
Thus the probability that x i + g σ ( u i ) > x j + g σ ( u j ) after Gumbel perturbation is:
[Equation]
A. P ROOF OF L EMMA
Proof of Lemma 1.
If one of the top- k items becomes smaller than x k +1 after distorted by the Gumbel noise, then || T ∗ − T σ ∗ ||
F increases at most by 2.
Without loss of generality, here we assume x i as one of the top- k items i.e. i ≤ k . According to Lemma 3 , the probability that x i becomes smaller than x k +1 after perturbation is
[Equation]
Similarly, if one of the last- ( m − k ) items becomes larger than x k after distorted by the Gumbel noise, then || T ∗ − T σ ∗ || F also increases at most by 2. Without loss of generality, here we assume x j as one of the last- ( m − k ) items i.e. j ≥ k
+ 1 . According to Lemma 3 , the probability that x j becomes larger than x k after perturbation is [Equation] To compute the expectation of || T ∗ − T σ ∗ || , we sum all probabilities and we have
[Equation]
where the constant 4 appears because || T ∗ − T σ ∗
|| 2 F will change at most by 4 if one more item crosses the boundary.
Since for all i ≤ k and j ≥ k + 1 , we have
[Equation]
Thus we have
Eq.
(32) ≤
[Equation] [Equation]
Since E u [ || T ∗ − T σ ∗ || F ] ≤
[Equation]
|| T ∗ − T σ ∗
|| 2 F (cid:105) , we have
[Equation]
A. P ROOF OF L EMMA
Proof of Lemma 2.
By disturbing X with i.i.d.
Gumbel noise, we have
[Equation]
[Equation] ) (39) where g σ ( u ) = − σ log( − log( u )) is the Gumbel noise modulated by noise factor σ , and u 1 , u 2 , u 3 , ..., u m are i.i.d.
uniform distribution.
We deﬁne π as the permutation of sorting X (cid:48) in descending order, i.e. x π 1 + g σ ( u π 1 ) , x π 2 + g σ ( u π 2 ) , x π 3 + g σ ( u π 3 ) , ..., x π m + g σ ( u π m ) are in descending order.
Recall Theorem 2 in (Xie et al., 2020) , for x 1 , x 2 , x 3 , ..., x m we have
[Equation]
By substituting X with X (cid:48) and taking the expected value, we have
[Equation]
Based on Lemma 3 , the probability that π k
= i, π k +
=
j is
[Equation]
where the ﬁrst term denotes x
i + g σ ( u i ) >
x j + g σ ( u j ) , the second term denotes all conditions that there are ( k − 1) items larger than x i + g σ ( u i ) and the rest items are smaller than x j + g σ ( u j ) .
In the following we derive the upper bound of E u
[Equation]
We denote
A i,j as
[Equation]
where (cid:15) is a sufﬁciently small number.
Then we have
[Equation]
=
[Equation]
i
[Equation]
[Equation] [Equation] [Equation] [Equation]
[Equation]
[Equation]
[Equation]
We denote f ( δ, σ, z ) as:
[Equation]
For the probability terms in Eq. (47), for all permutations π , there must exist π a , π b , such that
[Equation]
Thus we have
[Equation]
By transforming Eq.
(20) in Lemma 3 , and substituting x j − x i by y , we have
[Equation]
where the right side is the form of the cumulative distribution function (CDF) of standard Logistic distribution by setting σ = 1 : [Equation]
[Equation]
Thus − log( − log( u i ))
+log( − log( u j )) is equivalent to the Logistic distribution whose probability density function (PDF) is [Equation]
[Equation]
and in this proof we exploit an upper bound of PDF ( y ) :
[Equation]
Based on the Logistic distribution, we can replace − σ log( − log( u i ))
+ σ log( − log( u j )) by σy where y is from the Logistic distribution.
Thus we can derive the upper bound of f ( δ, σ, z ) as follows
[Equation]
where Eq.
(65) is because
[Equation]
4 , and Eq.
(73) is because π − 2 arctan(
[Equation]
With probability (1 − (cid:15) )
[Equation]
we have
[Equation]
Thus
[Equation]
Thus we have
[Equation]
In conclusion
[Equation]
with the probability (1 − (cid:15) )
[Equation]
we have
[Equation]
And we denote Ω( x i , x j , σ, (cid:15) ) as
[Equation]
Finally, Theorem 1 is proved by triangle inequality by jointly considering Lemma 1 and Lemma 2 .
A. R EMARKS W . R . T . SOFT-T OP K
In the following we add some remarks about the relationship between our conclusion of GS-TopK and the conclusion derived by the authors of SOFT-TopK (Xie et al., 2020): the SOFT-TopK is a special case of our proposed algorithm, when we set σ = 0 .
We have the following conclusion:
Theorem 2.
Assume the values of x k , x k +1 are unique 4 , under probability (1 − (cid:15) ) , we have
[Equation]
x k x k+1 x i x j x k x
k+1 x j x i
Condition
Conditionx k x k+1 x i x j x k x k+ x j x i
[Figure]
Proof.
Since σ → 0+ , the ﬁrst term becomes 0.
For the second term, we discuss four conditions as shown in Fig.
4, except for the following condition: x i = x k , x j = x k +1 .
Condition 1 .
If x i ≥ x k , x j ≤ x k +1 (equalities do not hold at the same time), we have at least x i − x k > 0 or x k +1 − x j > 0 .
Then we have
[Equation]
Condition 2 .
For any case that x i < x j , we have x i − x j < 0 , thus
[Equation] [Equation]
Condition 3 .
If x i > x j ≥ x k , we have x
i − x k > 0 .
Then we have
[Equation] [Equation]
Condition 4 .
If x k +1 ≥ x i > x j , we have x k +1 − x j > 0 .
Then we have
[Equation]
Therefore [Equation] if x i (cid:54) = x k and x j (cid:54) =
x k +
[Equation]
the second term Ω( x i , x j , σ [Equation] (cid:15) ) degenerates to 0 when σ → 0 + .
Thus we have the following conclusion by only considering x i = x k , x j = x k +1 :
[Equation]
B I MPLEMENTATION D ETAILS ON D ISCRETE C LUSTERING P ROBLEM
For the discrete clustering problem, recall that our aim is to select a subset of objects as cluster centers.
It differs from the k-means clustering problem whose cluster centers may not have to be a subset of the objects, and the constraint for discrete clustering is more welcomed for real-world facility location planning applications.
Denote ∆ ∈ R + m × m as the distance matrix for each pair of objects.
Following our implementation of the max covering problem, we can also derive a GPU-friendly formulation of the discrete clustering problem using PyTorch syntax:
[Equation]
We also model the discrete clustering problem as a graph and encode the problem structure by a graph neural network.
Speciﬁcally, for points lies on the 2D plane, we deﬁne there is an edge between two points if their distance are smaller than a predeﬁned threshold.
Then the geometric information is encoded by 3-layer SplineCNN (Fey et al., 2018) which was found successful for geometric feature learning.
Similar to the max covering problem, we rewrite the GPU-friendly form of the objective score in a differentiable manner, being estimated by the output of GS-TopK algorithm:
[Equation]
where we use a softmax operator with temperature T to approximate the min operator in a differentiable way.
(cid:12) means the tensor-product operator in PyTorch.
The GS-TopK learning algorithm is summarized in Alg. 3.
Algorithm 3: Gumbel-Sinkhorn-TopK Learning for Solving the Discrete Clustering ProblemInput: the distance matrix ∆
for all
update
respect to the
and learning
descend;
if Inference then update s with respect to the gradient ∂J∂ s and learning rate α by gradient descend; for all i ,
J i TopK J best ; 14 if Homotopy Inference then Decline the values of τ, σ and jump to line
C E XPERIMENT S ETUP
Our algorithms are implemented by PyTorch and the graph neural network modules are based on (Fey & Lenssen, 2019).
In our paper, we optimize the hyperparameters by greedy search on a small subset of problem instances ( ∼ 5) and set the best conﬁguration of hyperparameters for both GSTopK,
SOFT-TopK. We empirically set learning rate = 0 . 1 , τ = 0 . 05 , σ = 0 .
25 as the default conﬁguration, and we set ( τ, σ ) =
for Homotopy GS-TopK.
We set # G = 1000 for max covering, # G = 500 for discrete clustering.
All experiments are done on a workstation with i7-9700K@3.60GHz, 16GB memory, and 2080Ti GPU.
D A BLATION S TUDY ON H YPERPARAMETERS
Firstly, we want to add some remarks about the selection of hyperparameters:
• # G (number of Gumbel samples): # G affects how many samples are taken during training and inference for GS-TopK.
A larger # G (i
e. more samples) will be more appealing, because GS-TopK will have a more accurate estimation of the objective score, and it will have a higher probability of discovering better solutions.
However, # G cannot be arbitrarily large because we are running GS-TopK in parall on GPU, and the GPU has limited memory.
Therefore, in experiments, we set a large enough # G (e.g. # G = 1000 ) and ensure that it can ﬁt into the GPU memory of our workstation (2080Ti, 11G).
• τ (entrophic regularization factor of Sinkhorn): Theoretically, τ controls the gap of the continuous Sinkhorn solution to the discrete solution, and a smaller τ will lead to a tightened gap.
This property is validated by our theoretical ﬁndings in Lemma 2 .
Unfortunately, τ cannot be arbitrarily small, because a smaller τ requires more Sinkhorn iterations to converge.
Therefore, given a ﬁxed number of Sinkhorn iterations (100), we need trialand-error to discover the suitable τ for both SOFT-TopK and
GS-TopK.
The grid search results below show that our selection of τ fairly balances the performances of both SOFTTopK and GS-TopK. • σ (Gumbel noise factor): As derived in Theorem 1 , we need to balance between the two gaps ( Lemma 1 and Lemma 2 ) by a suitable σ .
Since | x k − x k +1 | is unknown, we cannot directly ﬁnd the theoretically optimal σ given τ .
In the experiments, we ﬁrstly determine a τ , and then ﬁnd a suitable σ by greedy search on a small subset ( ∼ 5) of problem instances.
We conduct an ablation study about the sensitivity of hyperparameters by performing an extensive grid search near the conﬁguration used in our max covering experiments ( τ = 0 .
05 , σ = 0 .
15 , # G = 1000 ).
We choose the k=50, m=500, n=1000 max covering problem, and we have the following results for GS-TopK (ours) and
SOFT-TopK (Xie et al., 2020) (higher is better):
[Table]
[Table]
Under the conﬁguration used in our paper, both SOFT-TopK and GS-TopK have relatively good results, thus we do not ﬁnd a set of hyperparameters to intentionally poison the peroformance of
[Table]
SOFT-TopK.
By grid search, we also discover better hyperparameters for both GS-TopK and SOFTTopK, but it is worth noting that SOFT-TopK is still inferior to GS-TopK with the best discovered hyperparameters.
Our grid search result shows that our GS-TopK is not very sensitive to σ if we have τ = 0 .
05 or 0 . 1 , and the result of τ = 0 .
01 is inferior because the Sinkhorn algorithm may not converge.
The resutls of # G = 1000 are all better than # G = 800 , suggesting that a larger # G is appealing if we have enough GPU memory.
It is also discovered that SOFT-TopK seem to be able to accept a smaller value of τ comapred to GS-TopK, probably because adding the Gumbel noise will increase the divergence of elements thus performs in a sense similar of decreasing τ when considering the convergence of Sinkhorn.