On Learning to Solve Cardinality Constrained Combinatorial Optimization in One-Shot: A Re-parameterization Approach via Gumbel-Sinkhorn-TopK

1 INTRODUCTION . In this paper , we aim to solve a family of cardinality constrained combinatorial optimization : given the cardinality budget k , our aim is to find the optimal set with size k which minimizes the objective function f ( x ) . Denote ||x||0 ≤ k as the cardinality constraint , we have : min x f ( x ) s.t . ||x||0 ≤ k ( 1 ) 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 influential 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 efficiency-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 inefficient . 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 Gumbel-Sinkhorn-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 . 2 ) 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 . 3 ) Experiments on max covering and discrete clustering problems show that under a comparative or smaller inference time budget , 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 . 2 RELATED WORK . 2.1 LEARNING OF COMBINATORIAL OPTIMIZATION . 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 . ( 2021 ) 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 they can handle combinatorial constraints that are non-trivial to be encoded in a differentiable manner , by adding constraints on the feasible actions in the multi-step roll-out procedure . However , the training of reinforcement learning may be inefficient 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.2 OPTIMAL TRANSPORT AND GUMBEL RE-PARAMETERIZATION . Optimal Transport . The optimal transport ( OT ) is a linear programming problem that aims to find 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 . ( 2018 ) . 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 . 3 THE PROPOSED GUMBEL-SINKHORN-TOPK . 3.1 PRELIMINARY : SOFT-TOPK ALGORITHM . The TopK problems are ubiquitous in machine learning , yet the MAP inference approach ( performing sorting at the time cost of O ( m logm ) and selecting the first 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 . Specifically , for selecting the k largest items from the probability vector s of size m ( k < m ) , the OT problem refers to moving 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 : c = [ 1 1 ... 1 ] ︸ ︷︷ ︸ m items r = [ m− k k ] } 2 destinations ( 2 ) Algorithm 1 : Gumbel-Sinkhorn-TopK ( GS-TopK ) for Solving Cardinality Constrained Combinatorial Optimization Input : List s with m items ; TopK size k ; Sinkhorn factor τ ; noise factor σ ; sample size # G. 1 for a ∈ { 1 , 2 , ... , # G } do 2 for all si , s̃i = si − σ log ( − log ( ui ) ) , where ui is sampled from ( 0 , 1 ) uniform distribution ; 3 D̃ = [ s̃1 −min ( s ) ... s̃m −min ( s ) max ( s ) − s̃1 ... max ( s ) − s̃m ] ; c = [ 1 1 ... 1 ] ︸ ︷︷ ︸ m items ; r = [ m− k k ] ; 4 T̃a = exp ( −D̃/τ ) ; 5 while not converged do 6 D̃r = diag ( T̃a1 r ) ; T̃a = D̃−1r T̃a ; 7 D̃c = diag ( T̃ > a 1 c ) ; T̃a = T̃aD̃−1c ; Output : A list of transport matrices [ T̃1 , T̃2 , ... , T̃ # G ] . 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 : D = [ s1 −min ( s ) s2 −min ( s ) ... sm −min ( s ) max ( s ) − s1 max ( s ) − s2 ... max ( s ) − sm ] . ( 3 ) While the OT is formulated as an integer linear programming problem : min T tr ( T > D ) s.t . T ∈ { 0 , 1 } m×2 , T1 = r , T > 1 = c , ( 4 ) 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 firstly sorting all items and then selecting the TopK items . It is also equivalent to the MAP inference by regarding s1 , ... , sm 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 . ( 4 ) is modified with entropic regularization : min Tτ tr ( Tτ > D ) + τh ( Tτ ) s.t . Tτ ∈ [ 0 , 1 ] m×2 , Tτ1 = r , Tτ > 1 = c , ( 5 ) where h ( Tτ ) = ∑ i , j T τ ij log T τ ij is the entropic regularizer ( Cuturi , 2013 ) . Given any real-valued matrix D , Eq . ( 5 ) is solved by firstly enforcing the regularization factor τ : Tτ = exp ( −D/τ ) . Then Tτ is row- and column-wise normalized alternatively : Dr = diag ( T τ1 r ) , Tτ = D−1r Tτ , Dc = diag ( Tτ > 1 c ) , Tτ = TτD−1c , ( 6 ) where means element-wise division . We denote Tτ∗ as the converged solution , which is the optimal solution to Eq . ( 5 ) . Theorem 2 of ( Xie et al. , 2020 ) . Denote xk , xk+1 as the k-th ( k + 1 ) -th largest items , we have ||T∗ −Tτ∗||F ≤ 2mτ log 2 |xk − xk+1| , ( 7 ) where the gap is controlled by τ if |xk − xk+1| > 0 . However , the above bound diverges if xk = xk+1 , which is unavoidable because xk , xk+1 are outputs by a neural network . It motivates us to develop our improved version Gumbel-Sinkhorn-TopK by introducing the Gumbel noise .

Summary: This submission studies the learning-based formulation of cardinality constrained combinatorial optimization using probabilistic methods, specifically using Gumble-Sinkhorn method. By introducing Gumble noise into the existing Soft-TopK (Xie et al., 2020) formulation, the authors claimed that the bound of the difference between approximate objective function value and the ground-truth value is tightened, compared to the results in Xie et al., 2020. The authors presented Gumbel-Sinkhorn-TopK (GS-TopK) and its application in TopK combinatorial optimization problems. Experimental results for max covering and k-means clustering were presented to demonstrate the claim.