SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Dinghuai Zhang*  
Mila
Hanjun Dai Google DeepMind
Nikolay Malkin, Aaron Courville, Yoshua Bengio, Ling Pan Mila
Combinatorial optimization (CO) problems are often NP-hard and thus out of reach for exact algorithms, making them a tempting domain to apply machine learning methods. The highly structured constraints in these problems can hinder either optimization or sampling directly in the solution space. On the other hand, GFlowNets have recently emerged as a powerful machinery to efficiently sample from composite unnormalized densities sequentially and have the potential to amortize such solution-searching processes in CO, as well as generate diverse solution candidates. In this paper, we design Markov decision processes (MDPs) for different combinatorial problems and propose to train conditional GFlowNets to sample from the solution space. Efficient training techniques are also developed to benefit long-range credit assignment. Through extensive experiments on a variety of different CO tasks with synthetic and realistic data, we demonstrate that GFlowNet policies can efficiently find high-quality solutions. Our implementation is open-sourced at https://github.com/zdhNarsil/GFlowNet-CombOpt.
Combinatorial optimization (CO) is a branch of optimization that studies problems of minimizing or maximizing some cost over a finite feasible set. CO problems usually involve discrete structures, such as graphs, networks, and permutations, and require optimizing an objective function subject to discrete constraints in the solution space, which is often NP-hard. CO problems have broad applications, including in medicine, engineering, operations research, and management (Paschos, 2010), and have spurred the development of discrete mathematics and theoretical computer science for a century (Kuhn, 1955; Kruskal, 1956; Ford & Fulkerson, 1956).
During the past few decades, researchers have developed numerical solvers such as GUROBI (Gurobi Optimization, 2023) to give approximate solutions via integer programming. In recent years, interest in
Figure 1: Illustration of GFlowNet for a toy MIS problem where all the states form a DAG. Every trajectory starts from the same initial state (whose vertices are all gray). Each transition denotes adding one vertex to the solution set, i.e., changing one vertex to black. See Section 3.2 for details.
learning-based methods for solving CO problems has grown significantly. These approaches leverage the power of deep networks to learn the inherent structure of CO problems and provide efficient
*Correspondence to dinghuai.zhang@mila.quebec.
and effective solutions. One family of machine learning methods utilizes solver-found solutions to provide a supervised learning (SL) signal for training neural networks (Selsam et al., 2018; Gasse et al., 2019; Nair et al., 2020). This line of algorithms requires expensive precomputation by the numerical solver to produce supervision labels. On the other hand, unsupervised learning (UL) methods search for the solution without the help of a solution oracle. One branch of UL methods, called probabilistic methods, decodes the heatmap generated from one pass of a neural network to get solutions (Toenshoff et al., 2019; Karalias & Loukas, 2020). These methods can achieve fast inference at the cost of a large optimality gap. Another branch of unsupervised methods is reinforcement learning (RL), which iteratively refines or constructs the problem solution with practitioner-specified MDPs (Bello et al., 2016; Deudon et al., 2018; Wu et al., 2019).
Despite many recent efforts to apply deep RL to CO problems, such approaches have fundamental limitations. For example, due to the symmetry in problem configurations, there could be multiple optimal solutions to the same CO problem (Li et al., 2018). Standard RL algorithms such as Fujimoto et al. (2018) are grounded in the nature of cumulative reward maximization and fail to promote diversity in the solutions. Although entropy-regularized RL (Haarnoja et al., 2017, 2018; Zhang et al., 2023b) converges to a stochastic policy instead of a deterministic one, these methods target trajectory-level entropy rather than solution-level entropy. Consequently, the agent may get trapped in solutions that can be reached by many trajectories and lack the ability to generate diverse candidate solutions. In addition, the performance of RL methods largely depends on the designed reward function and relies on a dense per-step reward for learning the value functions and policies. As a result, it is challenging to apply RL in our problems if only a terminal reward is provided.
Although there are attempts to fix these issues (Kwon et al., 2020; Ahn et al., 2020), they are mostly problem-specific and only achieve marginal improvement. In this work, we turn to a more principled framework, namely generative flow networks (Bengio et al., 2021, GFlowNets), to search for high-quality diverse candidates in CO problems. GFlowNet is a novel decision-making framework for learning stochastic policies to sample composite objects with probability proportional to a given terminal reward, which is suitable for problems where the solution is only related to the terminal state of generative trajectories. We design MDPs for a variety of NP-hard CO problems, where the intermediate states form a flow network in the latent space, and GFlowNet learns an agent to sequentially make decisions in this environment (Figure 1). Another challenge of applying GFlowNets is learning from long trajectories. In large-scale graph CO problems, the GFlowNet agent will encounter a very long trajectory before termination, rendering the task of credit assignment challenging. To this end, we develop efficient learning algorithms to train GFlowNets from transitions rather than complete trajectories, which greatly helps the learning process, especially in large-scale setups. Through extensive experiments on different CO tasks, we demonstrate the advantage of our proposed GFlowNet approach. In summary, our contributions are as follows:
- We design a problem-specific MDP for GFlowNet training on four different CO tasks.
- We propose an efficient GFlowNet learning algorithm to enable fast credit assignment for the GFlowNet agents with long trajectories that emerge in our graph CO problems.
- The empirical advantage of GFlowNets is validated through experiments on different CO problems.
2.1 GFlowNets
Generative flow networks, or GFlowNets, are variational inference algorithms that treat sampling from a target probability distribution as a sequential decision-making process (Bengio et al., 2021, 2023). We briefly summarize the formulation and the main training algorithms for GFlowNets.
We assume that a fully observed, deterministic MDP with set of states  $S$  and set of actions  $\mathcal{A} \subseteq S \times S$  is given. The MDP has a designated initial state, denoted  $\mathbf{s}_0$ . Certain states are designated as terminal and have no outgoing actions; the set of terminal states is denoted  $\mathcal{X}$ . All states in  $S$  are assumed to be reachable from  $\mathbf{s}_0$  by a (not necessarily unique) sequence of actions (see Figure 2). A complete trajectory is a sequence of states  $\tau = (\mathbf{s}_0 \to \mathbf{s}_1 \to \dots \to \mathbf{s}_n)$ , where  $\mathbf{s}_n \in \mathcal{X}$  and each pair of consecutive states is related by an action, i.e.,  $\forall i (\mathbf{s}_i, \mathbf{s}_{i+1}) \in \mathcal{A}$ .
A policy on the MDP is a choice of distribution  $P_F(\mathbf{s}'|\mathbf{s})$  for each  $\mathbf{s} \in S \setminus \mathcal{X}$  over the states  $\mathbf{s}'$  reachable from  $\mathbf{s}$  in a single action. A policy induces a distribution over complete trajectories via
$$
P _ {F} (\mathbf {s} _ {0} \to \mathbf {s} _ {1} \to \dots \to \mathbf {s} _ {n}) = \prod_ {i = 0} ^ {n - 1} P _ {F} (\mathbf {s} _ {i + 1} \mid \mathbf {s} _ {i}).
$$
The marginal distribution over the final states of complete trajectories is denoted  $P_F^\top$ , a distribution on  $\mathcal{X}$  that may in general be intractable to compute exactly, as  $P_F^\top(\mathbf{x}) = \sum_{\tau \to \mathbf{x}} P_F(\tau)$ , with the sum taken over all complete trajectories that end in  $\mathbf{x}$ .
A reward function is a mapping  $\mathcal{X}\rightarrow \mathbb{R}_{>0}$ , which is understood as an unnormalized probability mass on the set of terminal states (we will typically make the identification  $R(\mathbf{x}) = \exp (-\mathcal{E}(\mathbf{x}) / T)$ , where  $\mathcal{E}:\mathcal{X}\to \mathbb{R}$  is an energy function and  $T > 0$  is a temperature parameter). The learning problem approximately solved by a GFlowNet is to fit a policy  $P_F(\mathbf{s}'|\mathbf{s})$  such that the induced distribution  $P_F^\top (\mathbf{x})$  is proportional to the reward function, i.e.,
$$
P _ {F} ^ {\top} (\mathbf {x}) \propto R (\mathbf {x}) = \exp (- \mathcal {E} (\mathbf {x}) / T). \tag {1}
$$
The policy  $P_F(\mathbf{s}'|\mathbf{s})$  is parametrized as a neural network with parameters  $\theta$  taking  $\mathbf{s}$  as input and producing the logits of transitioning to each possible subsequent states  $\mathbf{s}'$ . This problem is made difficult both by the intractability of computing  $P_F^\top$  given  $P_F$  and by the unknown normalization constant (partition function) on the right side of (1). Learning algorithms overcome these difficulties by introducing auxiliary objects into the optimization. Next, we review two relevant objectives.
Detailed balance (DB) The DB objective (Bengio et al., 2023), requires learning two objects in addition to a parametric forward policy  $P_{F}(\mathbf{s}'|\mathbf{s};\boldsymbol{\theta})$  (we omit  $\boldsymbol{\theta}$  when it causes no ambiguity):
- A backward policy, which is a distribution  $P_B(\mathbf{s}|\mathbf{s}';\pmb{\theta})$  over the parents (predecessors) of any noninitial state in the MDP;
- A state flow function  $F(\cdot; \pmb{\theta}): \mathcal{S} \rightarrow \mathbb{R}_{>0}$ .
The detailed balance loss for a single transition  $\mathbf{s} \rightarrow \mathbf{s}'$  is defined as
$$
\ell_ {\mathrm {D B}} (\mathbf {s}, \mathbf {s} ^ {\prime}; \boldsymbol {\theta}) = \left(\log \frac {F (\mathbf {s} ; \boldsymbol {\theta}) P _ {F} (\mathbf {s} ^ {\prime} | \mathbf {s} ; \boldsymbol {\theta})}{F (\mathbf {s} ^ {\prime} ; \boldsymbol {\theta}) P _ {B} (\mathbf {s} | \mathbf {s} ^ {\prime} ; \boldsymbol {\theta})}\right) ^ {2}. \tag {2}
$$
The DB training theorem states that if  $\ell_{\mathrm{DB}}(\mathbf{s},\mathbf{s}';\pmb {\theta}) = 0$  for all transitions  $\mathbf{s}\rightarrow \mathbf{s}'$ , then the policy  $P_F$  satisfies (1), i.e., samples proportionally to the reward. The manner of selecting transitions  $\mathbf{s}\rightarrow \mathbf{s}'$  on which to minimize (2) is discussed below.
The loss that performs best for the problems in this paper (see (6) in §3.3) is equivalent to DB with a particular parametrization of  $\log F(s)$  that bootstraps learning by expressing the log-state flow as an additive correction to a partially accumulated negative energy.
Trajectory balance (TB) The TB objective (Malkin et al., 2022) features a simpler parametrization: in addition to the action policy  $P_F$ , one learns a backward policy  $P_B$  and only a single scalar  $Z_{\theta}$ , an estimator of the partition function corresponding to the initial state flow  $F(\mathbf{s}_0)$  in the DB parametrization. The TB loss for a complete trajectory  $\tau = (\mathbf{s}_0 \to \mathbf{s}_1 \to \dots \to \mathbf{s}_n = \mathbf{x})$  is
$$
\ell_ {\mathrm {T B}} (\tau ; \boldsymbol {\theta}) = \left(\log \frac {Z _ {\boldsymbol {\theta}} \prod_ {i = 0} ^ {n - 1} P _ {F} \left(\mathbf {s} _ {i + 1} \mid \mathbf {s} _ {i} ; \boldsymbol {\theta}\right)}{R (\mathbf {x}) \prod_ {i = 0} ^ {n - 1} P _ {B} \left(\mathbf {s} _ {i} \mid \mathbf {s} _ {i + 1} ; \boldsymbol {\theta}\right)}\right) ^ {2}. \tag {3}
$$
The TB training theorem states that if  $\ell_{\mathrm{TB}}(\tau; \boldsymbol{\theta}) = 0$  for all complete trajectories  $\tau$ , then the policy  $P_F$  satisfies (1). Furthermore,  $Z$  then equals the normalization constant of the reward,  $\hat{Z} = \sum_{\mathbf{x} \in \mathcal{X}} R(\mathbf{x})$ . In practice, policies and flows are typically output in the log domain, i.e., a neural network predicts logits of the distributions  $P_F(\cdot | \mathbf{s})$ ,  $P_B(\cdot | \mathbf{s}'')$  and the log-flows  $\log F(\mathbf{s})$  and  $\log Z$ .
Note that one can unambiguously write the policy as a distribution over subsequent states and not over actions because the MDP is deterministic. Most GFlowNet work uses the equivalent language of directed acyclic graphs, see analysis in Pan et al. (2023b); we use the MDP formalism to be consistent with RL terminology and avoid clashes of notation and language with the graphs that are the inputs to CO problems.
Training policy and exploration The DB and TB losses depend on individual transitions or trajectories, but leave open the question of how to choose the transitions or trajectories on which they are minimized. A common choice is to train in an on-policy manner, i.e., rollout trajectories  $\tau \sim P_F(\tau; \pmb{\theta})$  and perform gradient descent steps on  $\ell_{\mathrm{TB}}(\tau; \pmb{\theta})$  or on  $\ell_{\mathrm{DB}}(\mathbf{s}_i, \mathbf{s}_{i+1}; \pmb{\theta})$  for transitions  $\mathbf{s}_i \rightarrow \mathbf{s}_{i+1}$  in  $\tau$ . In this case, DB and TB have close connections to variational (ELBO maximization) objectives (Malkin et al., 2022).
However, an exploratory behaviour policy can also be used, for example, by sampling  $\tau$  from a version of  $P_F$  that is tempered or mixed with a uniform distribution (resembles  $\epsilon$ -greedy exploration in RL). Note that, unlike policy gradient methods in RL, GFlowNet objectives require no differentiation through the sampling procedure that yields  $\tau$ . The ability to stably learn from off-policy trajectories is a key advantage of GFlowNets over hierarchical variational models (Zimmermann et al., 2022; Malkin et al., 2023). See related study in Section C.
Conditional GFlowNets The MDP and the reward function in a GFlowNet can depend on some conditioning information. For example, in the tasks we study, a GFlowNet policy sequentially constructs the solution to a CO problem on a graph  $\mathbf{g}$ , and the set of permitted actions depends on  $\mathbf{g}$ . The conditional GFlowNets we train achieve amortization by sharing the policy model between different  $\mathbf{g}$ , enabling generalization to  $\mathbf{g}$  not seen in training.
2.2 Graph combinatorial optimization problems
We focus on the following four NP-hard CO problems on graphs: maximum independent set (MIS), maximum clique (MC), minimum dominating set (MDS), and maximum cut (MCut). A CO problem can be described with an undirected graph  $\mathbf{g} = (V, E)$ , where  $V$  is the set of vertices and  $E$  is the set of edges. Such problems typically require one to optimize over variables in a finite composite space to maximize or minimize some particular graph properties, as follows. Without loss of generality, we assume that all the graph weights equal one for simplicity, as our method can be easily extended to weighted graphs.
Maximum independent set In graph theory, an independent set is a set of vertices  $S$  in a given graph structure  $\mathbf{g}$  where any pair of vertices  $\{i,j\} \subseteq S$  are not neighbors i.e.,  $\forall i,j \in S, (i,j) \notin E$ . The MIS problem is to find such an independent set that has the largest possible size  $|S|$ .
Maximum clique A clique is a subset of the vertices  $S \subseteq V$  where all pairs of vertices are adjacent, i.e.,  $\forall i,j \in S, (i,j) \in E$ . The MC problem is to find a clique that has the largest size. We also remark that MC can be considered as a complementary problem of MIS, in the sense that the MC of any graph is actually the MIS of its complementary graph.
Minimum dominating set A dominating set  $S\subseteq V$  for a graph  $\mathbf{g}$  is a subset of vertices such that, for any vertex in this graph, it is either in  $S$ , or it has a neighbor in  $S$ . The MDS problem is to find the smallest dominating set for a given graph structure.
Maximum cut Given a subset of graph vertices  $S \subseteq V$ , the cut is defined as the number of edges between  $S$  and  $V \setminus S$ . The MCut problem is to find a set of vertices  $S$  that maximizes the cut.
3.1 Optimization as probabilistic inference
A CO problem can be seen as a constrained energy minimization problem in a discrete composite space  $\arg \min_{\mathbf{x} \in \mathcal{X}} \mathcal{E}(\mathbf{x})$ , where  $\mathcal{E}$  is the target energy function and  $\mathcal{X}$  is the solution space or feasible set. This optimization problem can be considered roughly as sampling from an energy-based model  $p_T^*(\mathbf{x}) \propto \exp\left\{-\mathcal{E}(\mathbf{x}) / T\right\}$ , where  $T$  is the temperature parameter that controls the smoothness of the density landscape. Sampling from such highly structured space is non-trivial, therefore we propose to use GFlowNets to amortize this inference process. According to the GFlowNet theory, a perfectly trained GFlowNet with the reward function  $R(\mathbf{x}) = \exp\left\{-\mathcal{E}(\mathbf{x}) / T\right\}$  will be able to accurately sample from  $p_T^*(\mathbf{x})$ . In this way, GFlowNets trained with a reasonably small temperature  $T$  can be used to search for solutions to given CO problems, as stated below.
Proposition 1. Assume that the GFlowNet is perfectly trained for a given temperature  $T$ , i.e., the training loss over a policy with full support equals zero. Then, if  $T \to \infty$ , the distribution sampled
(a) Independent set
(b) Clique
(c) Dominating set
(d) Cut
Figure 2: Our proposed MDP designs for different CO tasks. All the tasks aim at a set of vertices as the problem solution, thus we use  $0 / 1 / \emptyset$  to represent "not in the set" / "in the set" / "unspecified" for each vertex, which corresponds to white / black / gray in the figures. For each figure, the first arrow denotes conducting one action (marked with hammer), and the second arrow denotes the designed transition to guarantee any intermediate state represents a valid independent set / clique / dominating set / cut. The rightmost graph in each figure shows a feasible solution for the CO problem, where all vertices are specified (i.e., belongs to the GFlowNet terminal state space  $\mathcal{X}$ ).
by GFlowNet will converge to a uniform distribution on  $\mathcal{X}$ ; if  $T\to 0$ , the distribution sampled by GFlowNet will converge to the uniform distribution on optimal solutions to the optimization problem.
Notice that each graph  $\mathbf{g}$  corresponds to one unique sampling problem; thus, here we learn a graph conditional GFlowNet to amortize this condition distribution  $p(\mathbf{x}|\mathbf{g})$ , i.e., every learnable GFlowNet component (like the forward / backward policy or the flow function) is conditioned on the graph  $\mathbf{g}$ .
3.2 Designing Markov decision processes for GFlowNets
This section illustrates the design of appropriate Markov decision process (MDP) formulations for GFlowNet learning; see Figure 2 for a summary.
State For all the CO problems studied in this work, the solution  $\mathbf{x}$  is a subset of vertices for a given graph, represented as a binary vector  $\mathbf{x} = (\mathbf{x}^1,\dots ,\mathbf{x}^{|V|})\in \mathcal{X}\subseteq \{0,1\}^{|V|}$ , where  $\mathbf{x}^i = 1$  indicates that the  $i$ -th vertex belongs to the set and  $\mathbf{x}^i = 0$  indicates that it does not. Note that the feasible solution space  $\mathcal{X}$  is a subset of the full space  $\{0,1\}^{|V|}$ , as some binary vectors encode a vertex set outside the feasible set (e.g., a non-independent set in the MIS problem). The solution space is also the GFlowNet's terminal state space. The design of the GFlowNet state space is inspired by that in Zhang et al. (2022b). We begin by defining
$$
\bar {S} \triangleq \left\{\left(\mathbf {s} ^ {1}, \dots , \mathbf {s} ^ {| V |}\right): \mathbf {s} ^ {d} \in \{0, 1, \emptyset \}, d = 1, \dots , | V | \right\}, \tag {4}
$$
where  $\varnothing$  represents an unspecified "void" or "yet unspecified" situation for a particular vertex. Notice that  $\overline{S}$  is a superset of the terminal state space  $\mathcal{X}$ . The initial state is the all-void vector  $\mathbf{s}_0 = (\emptyset, \emptyset, \dots, \emptyset)$ . In our MDP design, we only allow transforming vertex values from void  $(\emptyset)$  to non-void (0 or 1) according to problem-specific rules. The trajectory terminates when all the entries are non-void, i.e., every entry is either 0 or 1. Notice that we do not need to have an explicit stop action in the agent's action space.
Notice that not all vectors in  $\overline{S}$  that have no void entries lie in  $\mathcal{X}$ . Therefore, we must restrict  $S$  so that the set of terminal states is exactly the set of vectors encoding feasible solutions  $\mathcal{X}$ . To be precise, we define the state space  $S$  to be the set of states  $\mathbf{s} \in \overline{S}$  such that there exists at least one  $\mathbf{x} \in \mathcal{X}$  that can be obtained from  $\mathbf{s}$  by valid transitions.
Care must be taken to modify the set of permitted actions from each state accordingly, so that actions always produce states in  $S$ . This turns out to be doable in the problems we study (as, for example, any subset of an independent set is independent). As an example, we next describe our action / transition / reward design for the MIS problem. Details for other tasks are deferred to the Appendix B.
Action The initial state is all-void, meaning that the partially constructed independent set is initialized with the empty set. The action of the MIS MDP is simply to choose one void vertex and add it to the current solution set by turning the entry value of the chosen vertex from  $\varnothing$  to 1.
Figure 3: Illustration of transition-based GFlowNet training. We break complete trajectories into transitions, and then randomly choose some of the them to form a buffer to train the GFlowNet agent. We use  $\mathbf{x}$  to denote the terminal state  $\mathbf{s}_n$ .  $\tilde{\mathcal{E}} (\cdot)$  denotes designed intermediate learning signals.
Algorithm 1 GFlowNet training algorithm
Require: Graph dataset, GFlowNet model with parameters  $\theta$ , reward function.
Sample graph g from dataset;
Sample  $\tau$  with the GFlowNet forward policy  $P_F(\cdot |\cdot ;\mathbf{g},\pmb {\theta})$
Randomly choose transitions from  $\tau$  to create buffer  $\mathcal{B}$
$$
\triangle \boldsymbol {\theta} \leftarrow \sum_ {(\mathbf {s} \rightarrow \mathbf {s} ^ {\prime}) \in \mathcal {B}} \nabla_ {\boldsymbol {\theta}} \ell_ {\mathrm {F L}} (\mathbf {s}, \mathbf {s} ^ {\prime}; \mathbf {g}, \boldsymbol {\theta})
$$
(as per Equation 6);
Update  $\theta$  with some optimizer;
until some convergence condition
Transition To ensure that the state can always be completed to an independent set, it is essential to carefully handle the actions taken. When a void vertex is chosen and its entry value is modified to 1, we also update the entry values of all its neighboring vertices to 0, which excludes the possibility of getting two adjacent vertices in the following steps. This proactive approach ensures that the independent set constraint is not violated in subsequent steps.
We remark that the feasible set  $\mathcal{X}$  in this problem consists of all independent sets, but of order-maximal independent sets, i.e., those to which no vertex can be added while keeping them independent. Non-order-maximal independent sets cannot be constructed with such transitions.
Reward We set the log reward to be the resulting independent set size, i.e.,  $\mathcal{E}(\mathbf{x}) = -|\mathbf{x}|_1$ , where  $|\mathbf{x}|_1$  denotes the  $\ell_1$  norm, i.e., the number of 1's in the binary vector  $\mathbf{x}$ .
3.3 Factors affecting training efficiency