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: The paper proposes an unsupervised method for cardinality-constrained combinatorial problems. By building on previous work on Soft TopK through optimal transport (OT), the authors provide an enhanced Soft TopK framework that includes Gumbel noise in the OT problem. This provides a more complete theoretical bound on the original OT problem but also strong experimental performance on two combinatorial optimization problems.