SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Sebastian Sanokowski $^{1,2}$
Wilhelm Berghammer
Sepp Hochreiter $^{1,2}$
Sebastian Lehner $^{1,2}$
1 ELLIS Unit Linz and LIT AI Lab
$^{2}$  Institute for Machine Learning, Johannes Kepler University Linz, Austria sanokowski@ml.jku.at
Several recent unsupervised learning methods use probabilistic approaches to solve combinatorial optimization (CO) problems based on the assumption of statistically independent solution variables. We demonstrate that this assumption imposes performance limitations in particular on difficult problem instances. Our results corroborate that an autoregressive approach which captures statistical dependencies among solution variables yields superior performance on many popular CO problems. We introduce subgraph tokenization in which the configuration of a set of solution variables is represented by a single token. This tokenization technique alleviates the drawback of the long sequential sampling procedure which is inherent to autoregressive methods without sacrificing expressivity. Importantly, we theoretically motivate an annealed entropy regularization and show empirically that it is essential for efficient and stable learning.<sup>1</sup>
Combinatorial optimization (CO) problems are of central interest to a wide range of fields, including operations research, physics, and computational complexity [Papadimitriou and Steiglitz, 1998]. Since CO problems are typically NP-hard one would not expect that one can find the solutions to all arbitrarily large CO problem instances in polynomial time. However, these restrictions are concerned with worst case scenarios over entire problem families. For instance finding the Minimum Independent Set of any graph. Consequently, it is not surprising that a large body of work is focused on the design of solution strategies that yield a particularly good performance on a restricted subset of all possible instances. There is growing interest in exploring deep learning approaches to this restricted CO setting [Bello et al., 2017]. These methods aim to learn how to efficiently generate high quality solutions for CO problems and rely primarily on learned algorithmic components rather than on problem-specific, hand-crafted ones [Bengio et al., 2021]. Importantly, practicable methods should not rely on supervised training with solutions from other solvers since this limits the achievable performance of the learned model to that of the solver used to generate the training data [Yehuda et al., 2020]. As a result the development of unsupervised methods [Karalias and Loukas, 2020, Wang et al., 2022, Qiu et al., 2022, Karalias et al., 2022, Min et al., 2022, Wang and Li, 2023] and of Reinforcement Learning (RL) methods [Bello et al., 2017, Khalil et al., 2017, Kool et al., 2019, Ahn et al., 2020, Bother et al., 2022, Mazyavkina et al., 2021] for CO is an active field of research. In this work we are interested in training models that generate solutions at inference in contrast to approaches
Our code is available at: https://github.com/ml-jku/VAG-C0.
that require problem instance specific training. A popular approach is to take a representation of the CO problem instance as input to a deep learning model and to output the parameters of a distribution over solutions that has its probability mass concentrated on regions with high solution quality. This probabilistic optimization approach to CO is used in the recent works of e.g. [Karalias and Loukas, 2020, Min et al., 2022, Qiu et al., 2022, Sun et al., 2022]. In these works the distribution over solutions is built on the assumption of mutual statistical independence of the individual solution parameters and can, consequently, be represented by a product of Bernoulli distributions for the individual solution variables. This simplification is typically refereed to as a mean-field approximation (MFA) and is frequently used in various fields including e.g. statistical mechanics [Parisi, 1988], Bayesian statistics [Wainwright and Jordan, 2008], and the analysis of neural networks [Mei et al., 2019]. However, the simplifying assumption of the MFA restricts the expressivity of the corresponding distributions which limits its applicability when the target distribution represents strong correlations [Jaakkola and Jordan, 1998]. However, replacing MFA with more expressive approaches is computationally expensive and requires careful regularization to ensure efficient and stable training. Based on these considerations our contributions can be summarized as follows.
We demonstrate that the frequently used MFA imposes limits on the attainable solution quality for CO problems. Importantly, we introduce Variational Annealing on Graphs for Combinatorial Optimization (VAG-CO) a method that achieves new state-of-the-art performances on popular CO problems by combining expressive autoregressive models with annealed entropy regularization. We provide a theoretical motivation for this entropy regularization via considerations of the sample complexity of related density estimation problems. By introducing sub-graph tokenization VAG-CO drastically reduces the number of necessary steps to generate solutions and therein significantly improves the training and inference efficiency.
2 Problem Description
Ising Formulation of CO. To introduce VAG-CO it is convenient to adopt a point of view on CO that is motivated by statistical mechanics. As shown in Lucas [2014] many frequently encountered NP-complete CO problems can be reformulated in terms of a particular class of models known as Ising models which therefore allow the representation of different CO problem types in a unified manner. In this context the CO problem is equivalent to finding the minimum of an energy function  $E: \{-1, 1\}^N \mapsto \mathbb{R}$ . This function assigns an energy  $E(\sigma)$  to an N-tuple of discrete solution variables  $\sigma = (\sigma_1 \ldots \sigma_N) \in \{-1, 1\}^N$  which are often referred to as spins. The family of aforementioned Ising models is characterized by the following form of the energy function:
$$
E (\boldsymbol {\sigma}) = \sum_ {i <   j} J _ {i j} \sigma_ {i} \sigma_ {j} + \sum_ {i} ^ {N} B _ {i} \sigma_ {i}, \tag {1}
$$
where the first term represents the interaction between spins  $\sigma_{i}\in \{-1,1\}$  through couplings  $J_{ij}\in \mathbb{R}$ , while the second term determines the magnitude  $B_{i}\in \mathbb{R}$  of the contribution of an individual spin to the energy of the system. For brevity we will denote the parameters of the energy function  $(J_{ij},B_i)$  simply as  $E$ . Hence,  $E$  will represent a CO problem instance while  $E(\pmb {\sigma})$  will denote the energy of a solution  $\pmb{\sigma}$  under an energy function with parameters  $E$ . The Ising formulation for the four CO problem types studied in this work is given in Tab. 1.
Variational Learning. Next, we specify how we approach the learning problem of approximating the optimal solution for a given problem instance. We follow the frequently taken variational learning approach of using a neural network with parameters  $\theta$  to obtain for a given CO problem instance  $E$  an associated probability distribution  $p_{\theta}(\sigma |E)$  over the space of possible solutions  $\{-1,1\}^N$ . The problem instances  $E$  are assumed to be independently sampled from a probability distribution  $q(E)$  with  $\Omega = \mathrm{supp}(q)$ . In this setting the learning problem can be formulated as finding  $\operatorname{argmin}_{\theta}\sum_{\Omega}q(E)\sum_{\sigma}p_{\theta}(\sigma |E)E(\sigma)$ . Here  $\sum_{\sigma}$  is a shorthand for the sum over all  $\sigma \in \{-1,1\}^N$ . In practice both sums are approximated by empirical averages over samples from the corresponding distributions  $q$  and  $p_{\theta}$ . In CO  $E(\sigma)$  is typically characterized by many local minima. Therefore, directly approaching the learning problem described above via gradient descent methods is prone to getting stuck at sub-optimal parameters. As we demonstrate in Fig. 1 (Right) this problem can be alleviated by adding a regularization term to the optimization objective Eq. 1 which encourages  $p_{\theta}$  to avoid a collapse to local minima. The resulting loss function  $L:\{-1,1\}^N\to \mathbb{R}$  associated to the
learning problem is given by:
$$
L (\theta ; \beta) = \frac {1}{M} \sum_ {m = 1} ^ {M} \hat {F} (\theta ; \beta , E _ {m}), \tag {2}
$$
where  $M$  is the number of CO problem instances  $E_{m}$  in one batch and  $\beta$  plays the role of a regularization coefficient which is frequently referred to as the inverse temperature  $1 / T$ . The loss for an individual problem instance  $E_{m}$  is based on the so-called free-energy  $F(\theta ;\beta ,E_m)$ :
$$
F (\theta ; \beta , E _ {m}) = \sum_ {\boldsymbol {\sigma}} p _ {\theta} (\boldsymbol {\sigma} | E _ {m}) \left(E _ {m} (\boldsymbol {\sigma}) + 1 / \beta \log p _ {\theta} (\boldsymbol {\sigma} | E _ {m})\right). \tag {3}
$$
In practice, we approximate the expectation value  $F(\theta; \beta, E_m)$  with an empirical sample average  $\hat{F}(\theta; \beta, E_m)$  that is based on  $n_S$  samples  $\sigma \sim p_{\theta}(\sigma | E_m)$ . The term  $E_m(\sigma)$  on the right-hand side of Eq. 3 encourages the concentration of  $p_{\theta}(\sigma | E_m)$  on low-energy solutions. Small values of the regularization parameter  $\beta \in \mathbb{R}_{>0}$  encourage  $p_{\theta}(\sigma | E_m)$  to have a large entropy. For a given value of  $\beta$  the free-energy in Eq. 3 is known to be minimized by the Boltzmann distribution associated to  $E_m$  (see e.g. [MacKay, 2003]):
$$
p _ {B} (\boldsymbol {\sigma} \mid E _ {m}, \beta) = \frac {\exp - \beta E _ {m} (\boldsymbol {\sigma})}{\sum_ {\boldsymbol {\sigma} ^ {\prime}} \exp - \beta E _ {m} \left(\boldsymbol {\sigma} ^ {\prime}\right)}. \tag {4}
$$
Thus by minimizing  $L$  the model learns to approximate  $p_B(\sigma | E_m, \beta)$  for a given  $E_m$  and  $\beta$ . For  $\beta \to \infty$  inf the  $p_B(\sigma | E_m, \beta)$  has its mass concentrated on the global minima of  $E_m$  and for  $\beta \to 0$  it approaches the uniform distribution [Mézard and Montanari, 2009]. It is, therefore, to be expected that the minimization of  $L$  becomes harder at low temperatures, i.e. as  $\beta \to \infty$ , which opens the opportunity to a principled curriculum learning approach (Sec. 4). Based on these considerations we reformulate CO problems as the variational problem of approximating  $p_B(\sigma | E_m, \beta)$  in the limit  $\beta \to \infty$  with a variational ansatz  $p_\theta$  that has the variational parameters  $\theta$ . This problem can be formalized as  $\operatorname{argmin}_{\theta} \lim_{\beta \to \infty} L(\theta; \beta)$ .
3 Variational Annealing on Graphs
Our method VAG-CO addresses CO as a variational learning problem on graphs. In particular, given a set of CO problem instances it learns to approximate the Boltzmann distribution of the corresponding Ising models (Sec. 2) with an autoregressive distribution. To obtain efficient training with this expressive model we apply an annealed entropy regularization. We formulate this learning problem in an RL setting and use PPO to train our model. To alleviate the lengthy sampling process of autoregressive approaches we introduce sub-graph tokenization which allows us to generate multiple solution variables at in a single step without loosing expressive power. We further improve the memory efficiency of our approach by dynamically pruning the graph which represents the CO problem instance.
Autoregressive Solution Generation. In the following we specify how we represent  $p_{\theta}(\sigma |E_i)$  and how to sample it, i.e. how to generate solutions  $\pmb{\sigma}$ .
1. Draw a problem instance  $E = (B_i, J_{ij}) \sim q(E)$  from the data distribution  $q$  and construct a graph  $G = (V, \mathcal{E})$  based on  $E$ . The nodes  $\nu_i \in V$  correspond to the spins  $\sigma_i$  and the set of edges  $\mathcal{E}$  represents the non-zero components of  $J_{ij}$ . The graph  $G$  is equipped with node features  $x_i = [B_i, t_i]$  that are associated to its nodes  $\nu_i$ . Here  $t_i$  is a four-dimensional one-hot encoding which indicates the four possible states (I-IV) of a node  $\nu_i$  namely whether the corresponding spin  $\sigma_i$  is set to the value  $+1$  (I) or  $-1$  (II) or whether it is to be generated in the current step (III) or afterwards (IV). The edges between nodes  $\nu_i$  and  $\nu_j$  are associated with the scalar edge features  $J_{ij}$ .
2. Order the graph nodes according to the breadth-first search (BFS) algorithm. The  $i$ -th node in this ordering is denoted as  $\sigma_{i}$ . Now  $i = 1$  and  $t_i$  is set to the state (III) and all  $t_{>i}$  are set to (IV).
3. A GNN is used to parameterize a Bernoulli distribution  $p_{\theta}(\sigma_i|G) = \mathrm{GNN}_{\theta}(G)$  from which the value  $\pm 1$  of  $\sigma_{i}$  is sampled. The state encoding variables of  $G$  are updated by setting
$t_i$  associated to the graph is now updated accordingly and  $t_{i + 1}$  is set to (III). Now  $i$  is incremented.
4. The previous step is repeated until the values of all  $\sigma_{i}$  are set.
We note that at each step  $i$  the graph  $G$  depends on the problem instance  $E$  and the already generated spins  $\sigma_{< i}$ . Therefore this procedure represents an autoregressive parametrization of a distribution over the space of possible solutions  $\{-1,1\}^N$ . By denoting the graph  $G$  at step  $i$  as  $G_i$  we get:
$$
p _ {\theta} (\boldsymbol {\sigma} | E) = \prod_ {i = 1} ^ {N} p _ {\theta} \left(\sigma_ {i} \mid \sigma_ {<   i}, E\right) = \prod_ {i = 1} ^ {N} p _ {\theta} \left(\sigma_ {i} \mid G _ {i}\right). \tag {5}
$$
This approach is more expressive than MFA and can therefore be expected to be better suited to approximate the typically complex Boltzmann distributions (Eq. 4) encountered in CO. Next we specify how we realize a stable training procedure by employing RL methods. The model architecture and hyperparameters are detailed in App. A.10.
Reinforcement Learning Setting. To explain how we train our model it is convenient to adopt an RL perspective where an episode corresponds to the solution generation procedure described above. For this we consider a Markov Decision Process (MDP) that is given by  $(S, A, P, R)$ . Here,  $S$  is the set of possible states and the state at step  $i$  is denoted by  $s_i \in S$ . At each step  $s_i$  is represented by the graph  $G_i$ . Given a state  $s_i$  an action  $a_i$  that represents the assignments of the spin value  $\sigma_i \in \{-1, 1\}$  is sampled from a probability distribution which is parameterised by the policy  $p_\theta(\sigma_i | G_i)$ . After sampling an action the reward  $R_i(G_i, \beta)$  is computed according to an objective that is based on Eq. 3. We use the relation  $F(\theta; \beta, E) = -\mathbb{E}_{\sigma \sim p_\theta} \left[ \sum_{i=1}^N R_i(G_i, \beta) \right]$  and define the reward at step  $i$  as
$$
R _ {i} \left(G _ {i}, \beta\right) = - \left[ \Delta E _ {i} + \frac {1}{\beta} \log p _ {\theta} \left(\sigma_ {i} \mid \sigma_ {<   i}, E\right) \right], \tag {6}
$$
where  $\Delta E_{i} = \sigma_{i}\left[\sum_{j < i}J_{ij}\sigma_{j} + B_{i}\right]$ . With this definition maximising the reward is equivalent to minimizing the free-energy in Eq. 3. We approximate  $F(\theta ;\beta ,E)$  with the empirical mean of  $-\sum_{i = 1}^{N}R_{i}(G_{i},\beta)$  based on  $n_S$  samples  $\sigma \sim p_{\theta}(\sigma_i|\sigma_{< i},E)$ . Finally, the state is changed deterministically, i.e. the spin  $\sigma_{i}$  is set according to the sampled action  $a_{i}$ . The state update corresponds to the updated of  $t_i$  and  $t_{i + 1}$  as described in step 3 of the solution generation procedure. Our model learns to solve this MDP via the popular PPO algorithm [Schulman et al., 2017] which is described in more detail in App. A.9.
Annealing. As discussed in Sec. 2 our training objective is based on Eq. 3 which contains an entropy regularization term. The strength of this regularization is determined by temperature  $T$ . At first, the temperature is kept constant at  $T > 0$  for  $N_{\mathrm{warmup}}$  steps. Then, the temperature is slowly adapted for  $N_{\mathrm{anneal}}$  steps by following a predefined annealing schedule (see App. A.14) that converges to  $T = 0$  as the end of training is reached. This reduction of the temperature throughout training is motivated from a theoretical point of view in Sec. 4.
Subgraph Tokenization. By introducing subgraph tokenization we decrease the number of forward passes per CO problem instance without sacrificing expressivity. Instead of modelling the probability for the two possible values  $\pm 1$  of a single spin, we let the policy represent a probability distribution  $p_{\theta}(\sigma_{i:i+k}|\sigma_{<i},E)$  over all  $2^k$  possible configurations of  $k$  consecutive spins in the BFS ordering (step 2 of the solution generation procedure). We represent  $p_{\theta}(\sigma_{i:i+k}|\sigma_{<i},E)$  with a softmax function of a  $2^k$  dimensional vector that is output by the GNN (App. A.10). Subgraph tokenization represents a modification of step 3 of the solution generation procedure and yields an improvement of the performance (Sec. 6).
Dynamic Graph Pruning. We note that once the spin value of  $\sigma_{i}$  is sampled its interaction with adjacent spins  $\sigma_{j}$  is fixed. Therefore, assuming that we have sampled the first spin  $\sigma_{1}$ , its interaction with an adjacent spin  $\sigma_{j}$ , that is yet to be generated, can be expressed as  $B(1\mapsto j)_j\sigma_j$ , where we introduce  $B(1\mapsto j)_j = J_{1j}\sigma_1$ . Therefore, we can immediately remove generated spins from the graph and update the node embeddings  $x_{j} = [B_{j},t_{j}]$  of adjacent nodes with  $B_{j}\gets B_{j} + B(i\mapsto j)_{j}$ . Reducing the graph size during the solution generation process has the benefit of reducing the memory requirements when processing graphs with a GNN as in step 3 of the solution generation procedure.
4 Theoretical Motivation of Annealing
In the following we show that the concept of an annealed entropy regularization in Eq. 3 can be motivated based on the sample complexity of density estimation for Boltzmann distributions.
Consider the problem of approximating the probability density function of a Boltzmann distribution  $p_B(s;E,\beta)$ , where  $E:\mathbb{R}^{\mathrm{N}}\to [0,1]$  is a suitably normalized energy function. To solve this task we are given a set of samples  $S = \{s_1,\dots ,s_m\}$  which are independently drawn from the corresponding target Boltzmann distribution  $p_B(s;E,\beta^*)$  at a fixed inverse temperature  $\beta^{*}$ . For brevity we will denote the distribution associated to  $p_B(s;E,\beta)$  from now on as  $p(\beta)$  where  $E$  is suppressed since it is a fixed function. The empirical distribution corresponding to  $S$  will be denoted as  $\hat{p}$ . Further assume that we can evaluate  $E(s_{i})$  for all  $s_i\in S$ .
In the present context a natural feasibility criterion for the estimated distribution is that the expectation value of the energy function for the estimated distribution should be compatible with the corresponding value  $\mathbb{E}_{p(\beta *)}E(s)$  of the target distribution. We approach this problem with the maximum entropy principle [Jaynes, 1957]. Informally, this principle prescribes that among all feasible distributions one should prefer the one that yields the highest entropy. Since  $S$  contains only a finite number of samples we cannot determine  $\mathbb{E}_{p(\beta *)}E(s)$  with arbitrary accuracy. As already stated in a similar form in [Dudik et al., 2007] one obtains with Hoeffding's inequality that with a probability of  $1 - \delta$ :
$$
\left| \mathbb {E} _ {p (\beta *)} E (s) - \mathbb {E} _ {\hat {p}} E (s) \right| <   \sqrt {\ln (2 / \delta) / (2 m)}. \tag {7}
$$
As shown in [Kazama and Tsujii, 2003] the resulting maximum entropy optimization problem with the inequality constraint Eq. (7) is equivalent to the following regularized maximum likelihood problem over the family of Boltzmann distributions:  $\min_{\beta} \mathcal{L}_{\hat{p}}(\beta) + \lambda \beta$ , where  $\mathcal{L}_{\hat{p}}(\beta) = -\mathbb{E}_{\hat{p}} \log p(\beta)$  is the cross-entropy between  $\hat{p}$  and  $p(\beta)$ . Based on a closely related result in [Dudík et al., 2007] we obtain the following bound for the approximation of maximum entropy distributions:
Remark 1. Assume a bounded energy function  $E: \mathbb{R}^{\mathrm{N}} \to [0,1]$  and let  $\hat{\beta}$  minimize  $\mathcal{L}_{\hat{p}}(\beta) + \lambda |\beta|$  where  $\lambda = \sqrt{\ln(2 / \delta) / (2m)}$ . Then with a probability at least  $1 - \delta$ :
$$
D _ {K L} \left(p \left(\beta^ {*}\right) | | p (\hat {\beta})\right) \leqslant \frac {| \beta^ {*} |}{\sqrt {m}} \sqrt {2 \ln (2 / \delta)}. \tag {8}
$$
See App. A.15.2 for further details. According to Eq. 8 the sample complexity of approximating a maximum entropy distribution at an inverse temperature  $\beta$  is in  $O(\beta^2)$ .
Curriculum learning is a machine learning paradigm in which the difficulty of training tasks is gradually increased as training proceeds [Bengio et al., 2009]. In view of Theorem 1 the entropy annealing in VAG-CO (Sec. 3) can be regarded as a principled curriculum learning approach if an increased sample complexity is regarded as an indication of a more difficult learning problem. In supervised learning tasks Wu et al. [2021] find that the application of curriculum learning results in more resource efficient training but not necessarily in better performance. As we demonstrate in Sec. 6 our entropy annealing does actually yield better performing models.
As pointed out in [Yehuda et al., 2020] supervised training of CO solvers faces the problem of expensive data generation and that data augmentation cannot circumvent this problem. Consequently, there is a growing interest in RL and unsupervised methods. In the following we focus on methods that attempt to learn how to generate solutions rather than how to improve existing ones.
Unsupervised Learning and Reinforcement Learning. The work of [Bello et al., 2017] proposes to use an actor-critic approach to learn how to solve CO problems. They were the first to show that deep RL is a promising approach for CO. Their method represents an autoregressive distribution over solutions. However, in their setting rewards are sparse since they are only available for complete solutions. In our approach rewards are dense since they are available after every state transition. Another influential work is [Khalil et al., 2017] who were the first to exploit the graph structure of CO problems by learning problem representations with GNNs. Applying GNNs to CO problem has become a common choice since then [Cappart et al., 2021]. In [Tönshoff et al., 2020] the method