SlowGuess/NeurIPS_2023_OCR huggingface dataset:
T2T: From Distribution Learning in Training to Gradient Search in Testing for Combinatorial Optimization
Yang Li, Jinpei Guo, Runzhong Wang, Junchi Yan*
Dept. of Computer Science and Engineering & MoE Key Lab of AI
{yanglily, jinpei, runzhong.wang, yanjunchi}@sjtu.edu.cn
Extensive experiments have gradually revealed the potential performance bottleneck of modeling Combinatorial Optimization (CO) solving as neural solution prediction tasks. The neural networks, in their pursuit of minimizing the average objective score across the distribution of historical problem instances, diverge from the core target of CO of seeking optimal solutions for every test instance. This calls for an effective search on each problem instance, while the model should serve to provide supporting knowledge that benefits the search. To this end, we propose T2T (Training to Testing) framework that first leverages the generative modeling to estimate the high-quality solution distribution for each instance during training, and then conducts a gradient-based search within the solution space during testing. The proposed neural search paradigm consistently leverages generative modeling, specifically diffusion, for graduated solution improvement. It disrupts the local structure of the given solution by introducing noise and reconstructs a lower-cost solution guided by the optimization objective. Experimental results on Traveling Salesman Problem (TSP) and Maximal Independent Set (MIS) show the significant superiority of T2T, demonstrating an average performance gain of  $49.15\%$  for TSP solving and  $17.27\%$  for MIS solving compared to the previous state-of-the-art.
Machine Learning (ML) for Combinatorial Optimization  $(\mathrm{CO})^2$ , abbreviated as ML4CO [1; 2; 3; 4], is a rapidly growing area, lying at the intersection between the two well-established communities. CO problems are essential in operational research as well as computer science, in both theory and practice. However, due to the inherent computational difficulty, e.g. NP-hardness, (approximately) solving these problems with effectiveness and efficiency poses significant challenges. Compared with traditional heuristic-based solvers, ML has recently shown its promising potential to automatically learn heuristics in a data-driven scheme [1; 2; 3; 4], or meanwhile with the aid of human knowledge which would normally improve the sample-efficiency [5; 6; 7; 8; 9]. Moreover, the introduction of neural networks also allows for the speedup by parallel matrix computation in forward inference.
A popular framework in ML4CO is developing neural networks to directly predict a solution or generate a solution with the aid of neural predictions, such that the corresponding objective score is minimized [10; 11; 12; 13]. However, the neural network can only fit limited mappings from problem
*Correspondence author. This work was partly supported by National Key Research and Development Program of China (2020AAA0107600), NSFC (62222607) and STCSM (22511105100).
In this paper, CO problems mainly refer to optimization problems on graphs [1] with an objective function.
Figure 1: Diagram of proposed framework.
Figure 2: Diffusion modeling for CO solving where the model learns how to gradually denoise from the random noise to predict each step's data distribution, eventually modeling the solution distribution.
parameters to (near) optimal solutions in the training set, resulting in subpar performance on unseen test instances [14; 15]. Therefore, it calls for neural solvers to efficiently explore solutions tailored to each specific problem instance. In the traditional framework, the network prediction typically converges to a local optimum with little diversity, struggling to further explore superior solutions. In this paper, we resort to a more effective framework that first learns to model the distribution of high-quality solutions conditioned on problem instances, then conducts a gradient-based search to obtain a superior solution specific to the given instance, as shown in Fig. 1. Existing alternative methods that conduct the tailored search in the testing phase, e.g. active search [16; 17] and meta-learning [18], require updating the model weights during testing, hindering their applications in computational intensive scenarios. In contrast, we develop a more efficient gradient search framework to obtain a superior solution specific to the given instance, without the need to update model weights.
We exploit the powerful generative neural CO solver [19] that learns to model the distribution of high-quality solutions conditioned on specific problem instances. The generative models are welcomed in their ability to output high-quality images [20], texts [21] and graphs [22] with great diversity, whereby these advantages are also welcomed in the ML4CO community. Some recent successes in generative learning of CO are demonstrated in classic CO problems e.g. Traveling Salesman Problem (TSP) and Maximal Independent Set (MIS) (via Diffusion [19] and Variational Auto-Encoder [23]), Satisfiability Problem (SAT) and Mixed-Integer Linear Programming (MILP) (via graph generators [24; 25; 26]), also in more realistic applications e.g. routing for chip design (via Diffusion [27] and Generative Adversarial Network [28]). Notably, the very recent adaptation of the flourishing diffusion models in probabilistic generation [19] has resulted in state-of-the-art performance for solving CO problems e.g. TSP. The diffusion modeling for CO solving is illustrated in Fig. 2, where the model is designed to estimate the solution distribution given a certain instance  $G$ , i.e.  $p_{\theta}(\mathbf{x}_0, G)$ . Despite its high-quality solutions, a major drawback is the unawareness of optimizing the CO objective as the generative model simply learns to fit a distribution of solutions. To further improve the efficiency of the current costly from-scratch sampling, it stands crucial to introduce a more efficient search-based paradigm to further explore high-quality solutions in terms of the final objective of the problem.
To this end, we propose a gradient-based search paradigm to further leverage the learned solution distribution through training, bridging the solution distribution learning during generative modeling training to gradient search during testing for CO, dubbed T2T. The proposed search paradigm well aligns with the behavior of a progressive generative model especially for the SOTA embodiment i.e. diffusion. The diffusion model provides a probabilistic framework that enables a smooth and flexible transition between random noises and feasible solutions, allowing for varying degrees of structure disruption, with a reference denoising process to restore a solution. Leveraging this framework, it is natural to model the local search heuristic as two stages, i.e., disrupting the local structure and reconstructing it with a better objective score. Specifically, T2T first adds certain noise to the given solution and then denoises it to reconstruct a lower-cost solution in iterations, which is achieved by incorporating the objective optimization guidance into the denoising steps through gradient feedback in the loop, steering the recovered sampling towards the optimization direction, as illustrated in Fig. 3. T2T enhances the solving paradigm by incorporating a search procedure with explicit objective guidance, which enables more efficient sampling and allows for deeper exploitation of the estimated solution distribution.
For the discrete nature of CO, T2T builds itself upon discrete diffusion models [29; 30; 31], which in training learns the transition between categorical noises and feasible discrete solutions that are
represented as probability maps of nodes/edges on graphs indicating whether each node/edge is included in the solution. The transition is tailored for each instance as it is conditioned on the corresponding instance's statistics. The training of T2T follows the standard training paradigm of diffusion models, which primarily equips the model with feasibility awareness and estimates the distribution of high-quality solutions. While during testing, we introduce the objective optimization guidance which serves as the complementary optimization and gradient search utility.
We showcase the efficacy of T2T on two typical NP-hard combinatorial optimization problems for edge-selecting and node-selecting types respectively, i.e., TSP and MIS, which have been attracting most attention in the community of ML4CO. Empirically, T2T exhibits strong performance superiority over existing state-of-the-art neural solvers on benchmark datasets across various scales. Notably, a significant performance gain is observed when compared to the purely constructive diffusion solver [19], thereby highlighting the effectiveness of the designed search-based paradigm.
The highlights of this paper include: 1) We introduce the framework of T2T which first estimates the high-quality solution distribution for each instance during training, and then leverages the diffusion process to perform an effective local search in the solution space during testing; 2) We incorporate the optimization objective guidance into the denoising model with theoretical analysis, which equips the model with gradient feedback directly from the instance-wise objective score, enabling it to perform optimization while ensuring general feasibility; 3) Extensive experiments show the promising performance of the proposed T2T, even when trained without high-quality solution labels.
Machine Learning for CO Solving. ML solvers can be categorized into construction solvers and improvement solvers. Constructive solvers include autoregressive solvers [32; 12; 6; 23; 8] which step-by-step extend a partial solution to produce a complete feasible solution, and non-autoregressive solvers [33; 34; 14; 35; 19] which predict a solution (typically soft constrained) in one shot. Improvement solvers [36; 37; 5] often define local operators such as 2OPT [36; 37], node swapping [5; 37], sub-problem resolving [7; 38] to perform local search towards minimizing the optimization objective. The algorithm pipelines generally involve selecting specific regions and locally optimizing sub-solutions. While T2T can be a rather different improvement solver, which allows for parallel optimization across multiple sub-regions and equips the generative modeling with the search procedure.
Generative modeling for CO leverages the powerful expressive and modeling capabilities of generative models for direct CO solving (e.g. VAE [23; 39], GAN [28; 40; 41], Diffusion [19; 27; 42], as well as GFlowNets [43; 44], which, while not a typical generative model itself, is closely related to generative modeling), and for supporting CO solving [24; 25]. DIFUSCO [19] has become state-of-the-art for the TSP problem by leveraging powerful diffusion models. However, no instance-specific search paradigms are introduced in existing methods to fully utilize the estimated solution distribution. Instead, they rely on multiple from-scratch samplings to leverage the distribution, which is computationally expensive. Though [23] attempts to utilize the evolution algorithm to conduct multiple iterations of sampling, it yet incurs even more significant costs without explicit objective guidance in the loop. In contrast, T2T serves as an improvement solver that locally modifies the given solution with objective guidance in the loop. It offers the flexibility to introduce any degree of disruption and rewrite without the necessity to construct the solution from scratch.
Diffusion Probabilistic Models. Diffusion models (score-based models) typically comprise a noising process and a learnable denoising process, where neural networks predict the noise introduced in each step. Diffusion's inference can be seen as a score (gradient) guided search, which is a general principle across different settings [45; 46; 47]. Diffusion in continuous space (e.g., for image generation) [29; 48; 42; 49; 50; 51; 52] have established and consolidated the theoretical foundations of the general diffusion framework. Several landmark works have advanced the diffusion model through various aspects. DDIM [49] generalizes DDPM via a class of non-Markovian processes which gives rise to implicit models that produce high-quality samples much faster. For conditional generation, classifier guidance [29; 52] and classifier-free guidance [53] are introduced with the score estimate of the diffusion models to produce conditional results. Similar paradigms are adopted in discrete diffusion models for generating discrete data using binomial noises [29] or multinomial/categorical noises [30; 31]. Diffusion also demonstrates its capability to learn offline policies in reinforcement learning tasks [54; 55].
This section presents our proposed T2T framework, including the preliminaries, offline training pipeline based on generative modeling, and online testing pipeline with objective guided search.
3.1 Preliminaries: Combinatorial Optimization on Graphs
Following the notations adopted in [13; 56; 57] we define  $\mathcal{G}$  as the universe of CO problem instances represented by graphs  $G(V,E)\in \mathcal{G}$ , where  $V$  and  $E$  denote the node set and edge set respectively. CO problems can be broadly classified into two types based on the solution composition: edge-selecting problems which involve selecting a subset of edges, and node-selecting problems which select a subset of nodes. Let  $\pmb{x}\in \{0,1\}^N$  denote the optimization variable. For edge-selecting problems,  $N = n^2$  and  $\pmb{x}_{i\cdot n + j}$  indicates whether  $E_{ij}$  is included in  $\pmb{x}$ . For node-selecting problems,  $N = n$  and  $\pmb{x}_i$  indicates whether  $V_{i}$  is included in  $\pmb{x}$ . The feasible set  $\Omega$  consists of  $\pmb{x}$  satisfying specific constraints as feasible solutions. A CO problem on  $G$  aims to find a feasible  $\pmb{x}$  that minimizes the given objective function  $f(\cdot ;G):\{0,1\} ^N\to \mathbb{R}_{\geq 0}$ :
$$
\min  _ {\boldsymbol {x} \in \{0, 1 \} ^ {N}} f (\boldsymbol {x}; G) \quad \text {s . t .} \quad \boldsymbol {x} \in \Omega \tag {1}
$$
In this paper, we study two primary and representative CO problems: TSP and MIS. TSP defines on an undirected complete graph  $G = (V,E)$ , where  $V$  represents  $n$  cities and each edge  $E_{ij}$  has a non-negative weight  $w_{ij}$  representing the distance between cities  $i$  and  $j$ . The problem can then be formulated as finding a Hamiltonian cycle of minimum weight in  $G$ . For MIS, given an undirected graph  $G = (V,E)$ , an independent set is a subset of vertices  $S \subseteq V$  such that no two vertices in  $S$  are adjacent in  $G$ . MIS is the problem of finding an independent set of maximum cardinality in  $G$ .
3.2 Offline Solution Distribution Learning on Training Set based on Diffusion
The basic diffusion modeling for solution generation in training is primarily based on [19]'s design while this paper explicitly formulate the solving task as a conditional generation task. As mentioned in Sec. 3.1, the solutions of CO problems can be represented as  $\mathbf{x} \in \{0,1\}^N$  with  $\mathbf{x} \in \Omega$ . For each entry, the model estimates a Bernoulli distribution indicating whether this entry should be selected. In implementation, each entry of solution  $\mathbf{x}$  is represented by a one-hot vector such that  $\mathbf{x} \in \{0,1\}^{N \times 2}$ . Generative modeling aims to model the distribution of high-quality solutions conditioned on the given instance, while feasibility constraints can be broadly captured through learning and eventually hard-guaranteed by post-processing. For the purpose of problem-solving, the generative model is desired to estimate the solution distribution given a certain instance  $G$ , i.e.,  $p_{\theta}(\mathbf{x}|G)$ . In this section, we model this distribution through the discrete diffusion models [42; 30].
Following the notations of [42; 30], the general framework of diffusion includes a forward noising and a reverse denoising Markov process. The noising process, also denoted as the diffusion process, takes the initial solution  $\mathbf{x}_0$  sampled from the distribution  $q(\mathbf{x}_0|G)$  and progressively introduces noise to generate a sequence of latent variables  $\mathbf{x}_{1:T} = \mathbf{x}_1,\mathbf{x}_2,\dots ,\mathbf{x}_T$ . Specifically, the noising process is formulated as  $q(\mathbf{x}_{1:T}|\mathbf{x}_0) = \prod_{t = 1}^{T}q(\mathbf{x}_t|\mathbf{x}_{t - 1})$ . The denoising process is learned by the model, which starts from the final latent variable  $\mathbf{x}_T$  and denoises  $\mathbf{x}_t$  at each time step to generate the preceding variables  $\mathbf{x}_{t - 1}$  based on the instance  $G$ , eventually recovering the target data distribution. The formulation of the denoising process is expressed as  $p_{\theta}(\mathbf{x}_{0:T}|G) = p(\mathbf{x}_T)\prod_{t = 1}^{T}p_{\theta}(\mathbf{x}_{t - 1}|\mathbf{x}_t,G)$ , where  $\theta$  denotes the parameters of the model. The training optimization aims to align  $p_{\theta}(\mathbf{x}_0|G)$  with the data distribution  $q(\mathbf{x}_0|G)$  and the objective adopts the variational upper bound of the negative log-likelihood:
$$
\mathcal {L} = \mathbb {E} _ {q} \left[ - \log p _ {\theta} \left(\mathbf {x} _ {0} \mid \mathbf {x} _ {1}, G\right) + \sum_ {t > 1} D _ {K L} \left[ q \left(\mathbf {x} _ {t - 1} \mid \mathbf {x} _ {t}, \mathbf {x} _ {0}\right) \| p _ {\theta} \left(\mathbf {x} _ {t - 1} \mid \mathbf {x} _ {t}, G\right) \right] \right] + C \tag {2}
$$
Specifically, the forward noising process is achieved by multiplying  $\mathbf{x}_t \in [0,1]^{N \times 2}$  at step  $t$  with a forward transition probability matrix  $\mathbf{Q}_t \in [0,1]^{2 \times 2}$  where  $[\mathbf{Q}_t]_{ij}$  indicates the probability of transforming  $\mathbf{e}_i$  in each entry of  $\mathbf{x}_t$  to  $\mathbf{e}_j$ . We set  $\mathbf{Q}_t = \left[ \begin{array}{cc} \beta_t & 1 - \beta_t \\ 1 - \beta_t & \beta_t \end{array} \right]$  [30], where  $\beta_t \in [0,1]$
3Each entry with [0, 1] indicates that it is included in  $\mathbf{x}$  and [1, 0] indicates the opposite.
such that the transition matrix is doubly stochastic with strictly positive entries, ensuring that the stationary distribution is uniform which is an unbiased prior for sampling. The noising process for each step and the  $t$ -step marginal are formulated as:
$$
q \left(\mathbf {x} _ {t} \mid \mathbf {x} _ {t - 1}\right) = \operatorname {C a t} \left(\mathbf {x} _ {t}; \mathbf {p} = \mathbf {x} _ {t - 1} \mathbf {Q} _ {t}\right) \quad \text {a n d} \quad q \left(\mathbf {x} _ {t} \mid \mathbf {x} _ {0}\right) = \operatorname {C a t} \left(\mathbf {x} _ {t}; \mathbf {p} = \mathbf {x} _ {0} \overline {{\mathbf {Q}}} _ {t}\right) \tag {3}
$$
where  $\operatorname{Cat}(\mathbf{x};\mathbf{p})$  is a categorical distribution over  $N$  one-hot variables with probabilities given by vector  $\mathbf{p}$  and  $\overline{\mathbf{Q}}_t = \mathbf{Q}_1\mathbf{Q}_2\dots \mathbf{Q}_t$ . Through Bayes' theorem, the posterior can be achieved as:
$$
q \left(\mathbf {x} _ {t - 1} \mid \mathbf {x} _ {t}, \mathbf {x} _ {0}\right) = \frac {q \left(\mathbf {x} _ {t} \mid \mathbf {x} _ {t - 1} , \mathbf {x} _ {0}\right) q \left(\mathbf {x} _ {t - 1} \mid \mathbf {x} _ {0}\right)}{q \left(\mathbf {x} _ {t} \mid \mathbf {x} _ {0}\right)} = \operatorname {C a t} \left(\mathbf {x} _ {t - 1}; \mathbf {p} = \frac {\mathbf {x} _ {t} \mathbf {Q} _ {t} ^ {\top} \odot \mathbf {x} _ {0} \overline {{\mathbf {Q}}} _ {t - 1}}{\mathbf {x} _ {0} \overline {{\mathbf {Q}}} _ {t} \mathbf {x} _ {t} ^ {\top}}\right) \tag {4}
$$
The neural network is trained to predict the logits of the distribution  $\tilde{p}_{\theta}(\tilde{\mathbf{x}}_0|\mathbf{x}_t,G)$ , such that the denoising process can be parameterized through  $q(\mathbf{x}_{t - 1}|\mathbf{x}_t,\tilde{\mathbf{x}}_0)$ :
$$
p _ {\theta} \left(\mathbf {x} _ {t - 1} \mid \mathbf {x} _ {t}\right) \propto \sum_ {\tilde {\mathbf {x}} _ {0}} q \left(\mathbf {x} _ {t - 1} \mid \mathbf {x} _ {t}, \tilde {\mathbf {x}} _ {0}\right) \tilde {p} _ {\theta} \left(\tilde {\mathbf {x}} _ {0} \mid \mathbf {x} _ {t}, G\right) \tag {5}
$$
Specifically for implementation, the condition  $G$  is enforced as an input to the network  $\theta$ , which is embodied as an anisotropic graph neural network with edge gating mechanisms [33] following [19]. For TSP, the instance condition consists of 2D coordinates of the vertices. The network input includes node embeddings from the 2D coordinates, along with edge embeddings from  $\mathbf{x}_t$ . For MIS, the edges  $E$  serve as the instance condition, complemented by the inclusion of node embeddings from  $\mathbf{x}_t$  to collectively form the input. The output of the network is  $\tilde{p}_{\theta}(\tilde{\mathbf{x}}_0|\mathbf{x}_t,G)\in [0,1]^{N\times 2}$  parameterizing  $N$  Bernoulli distributions for  $N$  entries in  $\tilde{\mathbf{x}}_0$ .
3.3 Test-stage Gradient-based Search with Instance-wise Objective Feedback
In this section, we first present the technical solution for the objective gradient incorporation in denoising steps, and then illustrate the neural search paradigm that leverages the guided denoising steps.
3.3.1 Objective-aware Denoising Process with Gradient Feedback
In the context of combinatorial optimization, the incorporation of objective optimization is necessary and important, which enables the direct involvement of objective and effective search over the solution space towards minimizing the score. As discussed in Sec. 3.2, at step  $t$  the model denoise  $\mathbf{x}_t$  by  $p_{\theta}(\mathbf{x}_{t - 1}|\mathbf{x}_t,G)$ . For the purpose of objective optimization, we aim to estimate  $p_{\theta}(\mathbf{x}_{t - 1}|\mathbf{x}_t,G,y^*)$  where  $y^{*}$  is the optimal objective score given the instance  $G$ :
$$
y ^ {*} = \min  _ {\mathbf {x}} l (\mathbf {x}; G) \quad \text {s . t .} \quad \mathbf {x} \in \Omega \tag {6}
$$
to guide the denoising process to  $\mathbf{x}^{*} = \arg \min_{\mathbf{x}}l(\mathbf{x};G)$ . We will show later that this ability does not necessarily require the access of  $y^{*}$  or  $\mathbf{x}^*$ , and can readily be achieved by the model learned in Sec. 3.2 without relying on any additionally trained networks. The adopted objective is defined below.
**Optimization Objective.** Since the generated samples are not guaranteed to satisfy the hard constraints, we incorporate constraint satisfaction into the objective and utilize a relaxed objective with constraint penalty, as adopted in [13; 56; 57]. Consider a relaxed cost function  $f_{r}(\cdot ;G): \{0,1\}^{N} \to \mathbb{R}_{\geq 0}$  which defined on the constraint relaxed space  $\{0,1\}^{N}$  satisfying  $f_{r}(\pmb{x};G) = f(\pmb{x};G), \forall \pmb{x} \in \Omega$ , and a constraint penalty  $g_{r}(\pmb{x};G): \{0,1\}^{N} \to \mathbb{R}_{\geq 0}$  where  $g_{r}(\pmb{x};G) = 0$  if  $\pmb{x} \in \Omega$  and  $g_{r}(\pmb{x};G) \geq 1$  if  $\pmb{x} \notin \Omega$ , the objective can be formulated as:
$$
l (\boldsymbol {x}; G) \triangleq f _ {r} (\boldsymbol {x}; G) + \beta g _ {r} (\boldsymbol {x}; G) \tag {7}
$$
where  $\beta >\max_{x\in \Omega}f(x;G)$ . For MIS, the constraint can be straightforwardly modeled as  $x_{i}x_{j} = 0$  if  $(i,j)\in E$ , thus the objective can be formulated as  $l_{\mathrm{MIS}}(\pmb {x};G)\triangleq -\sum_{1\leq i\leq N}\pmb{x}_i + \beta \sum_{(i,j)\in E}\pmb{x}_i\pmb{x}_j$ . While for TSP, since the constraints are difficult to explicitly express, we directly select the objective as  $l_{\mathrm{TSP}} = x\odot D$  where  $D\in \mathbb{R}_{+}^{n\times n}$  denotes the distance matrix.
Objective Guidance in Denoising Steps. The estimation of  $p_{\theta}(\mathbf{x}_{t - 1}|\mathbf{x}_t,G,y^*)$  can be modeled through the following proposition, adapted from the classifier guidance technique [52]:
Proposition 1. The optimization-enforced denoising probability estimation  $p_{\theta}(\mathbf{x}_{t - 1}|\mathbf{x}_t,G,y^*)$  equals to  $Zp_{\theta}(\mathbf{x}_{t - 1}|\mathbf{x}_t,G)p(y^* |\mathbf{x}_{t - 1},G)$ , where  $Z$  is a normalizing constant.