SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Transition-based GFlowNet training Most successful GFlowNet implementations (e.g., those in https://github.com/GFNOrg/) require a complete trajectory to compute the training loss and its gradient. Existing implementations with DB also specify the training loss at the trajectory level:
$$
\mathcal {L} (\tau ; \boldsymbol {\theta}) = \frac {1}{n} \sum_ {t = 0} ^ {n - 1} \ell_ {\mathrm {D B}} \left(\mathbf {s} _ {t}, \mathbf {s} _ {t + 1}; \boldsymbol {\theta}\right), \quad \tau = \left(\mathbf {s} _ {0}, \mathbf {s} _ {1}, \dots , \mathbf {s} _ {n}\right), \tag {5}
$$
and calculate the parameter gradient update with  $\mathbb{E}_{\tau \sim \pi (\tau)}[\nabla_{\pmb{\theta}}\mathcal{L}(\tau ;\pmb {\theta})]$  with some potentially off-policy distribution  $\pi (\tau)$ . This is also the case for other GFlowNet algorithms, such as flow matching (Bengio et al., 2021) and subtrajectory balance (Madan et al., 2023).
These implementations work well for moderate-scale trajectories, as previous GFlowNet works have shown (Malkin et al., 2022; Madan et al., 2023). However, for very long trajectories, such a design hinders efficient training: for a single complete trajectory which contains  $n$  transitions, one needs  $n$  calls of the neural network forward passes and storing all the intermediate feature maps and parameters in the GPU memory. Each forward pass contains multiple message passing operations on given graph structures, which is computationally expensive (linearly increasing with  $n$ ) in terms of speed and memory storage. In our experiments, we encounter large-scale problems where the trajectory is as long as  $\sim 400$  in length; nonetheless, with such graphs, our adopted graph neural networks can only support
Figure 4: Comparison between different GFlowNet variants.
forward and backward passes with batch size approximately 250 (which is much smaller than 400) on a 40 GB GPU memory device. This prohibits the usage of these trajectory-based GFlowNet training objectives on large-scale graph applications. In addition, the correlation between consecutive samples may incur stability issues (Mnih et al., 2015).
To this end, we turn to use a transition-based GFlowNet training approach without the knowledge of complete trajectories, which is first proposed in Deleu et al. (2022) and has shown effectiveness in Nishikawa-Toomey et al. (2022); Deleu et al. (2023). We randomly sample  $B$  transitions  $\mathcal{B} = \{(s^b, s'^b)\}_{b=1}^B$  from a complete trajectory and construct detailed balance-based loss:  $\mathcal{L}(\boldsymbol{\theta}) = \frac{1}{B}\sum_{b=1}^{B}\ell(s^b, s'^b; \boldsymbol{\theta})$ . This enables GFlowNets to be efficiently trained with long trajectories and limited GPU memory, which also converges faster. A corresponding schematic illustration of the algorithm can be found in Figure 3. In Figure 4, we show a comparison of learning efficiency between transition-based and trajectory-based GFlowNet implementations, where “-traj” denotes the latter variant. Different methods here share similar speed for one epoch training, thus from the figure we can see that transition-based approaches learn more efficiently. Besides the improved computation efficiency, our transition-based approach is also more memory efficient – as the memory occupation is proportional to the batch size instead of trajectory lengths (as for trajectory-based implementation), demonstrating its applicability to long-horizon problems.
Improving credit assignment with intermediate learning signals For normal GFlowNet training methods, the only learning signal in a trajectory comes from the terminal states and their associated reward values. This results in a relatively slow credit assignment process, and is inefficient to propagate information from near-terminal states to the early states due to the difficulty of attributing credits of each action in a long trajectory. The ability to learn from incomplete trajectories is especially important in our proposed transition-based training, since we are not using all the transitions from the complete trajectories. Therefore, we incorporate intermediate learning signals via the forward-looking (Pan et al., 2023a, FL) technique:
$$
\ell_ {\mathrm {F L}} (\mathbf {s}, \mathbf {s} ^ {\prime}; \boldsymbol {\theta}) = \left(- \tilde {\mathcal {E}} (\mathbf {s}) + \log \tilde {F} (\mathbf {s}; \boldsymbol {\theta}) + \log P _ {F} (\mathbf {s} ^ {\prime} | \mathbf {s}; \boldsymbol {\theta}) + \tilde {\mathcal {E}} (\mathbf {s} ^ {\prime}) - \log \tilde {F} (\mathbf {s} ^ {\prime}; \boldsymbol {\theta}) - \log P _ {B} (\mathbf {s} | \mathbf {s} ^ {\prime}; \boldsymbol {\theta})\right) ^ {2}, \tag {6}
$$
where  $\tilde{\mathcal{E}} (\cdot):S\to \mathbb{R}$  is a continuation of the reward energy  $\mathcal{E}(\cdot):\mathcal{X}\rightarrow \mathbb{R}$  which is only defined in the terminal state space  $\mathcal{X}$ . For simplicity, here we ignore the conditioning on graph structure  $\mathbf{g}$ . This FL method enables dense supervision signals to GFlowNet training, resulting in faster credit assignment as can be seen in Figure 3. Notice here that we need to design a handcrafted reward  $\tilde{\mathcal{E}} (\mathbf{s})$  for all possible latent states to reflect our estimation on intermediate signals. For MIS problems, we naturally define  $-\tilde{\mathcal{E}} (\mathbf{s})$  to be the number of vertices in the current set, i.e., the number of 1 entries in state s. This semantically coincides with the definition of an MIS terminal reward. We defer the intermediate reward design for other tasks to the Appendix. The effectiveness of FL against the DB or TB algorithms can be seen in Figure 4. We summarize the resulting algorithm in Algorithm 1.
GFlowNets GFlowNets were intended as diversity-seeking samplers for biological sequence and molecule design, an application area that continues to motivate research (Bengio et al., 2021; Jain et al., 2022, 2023b; Shen et al., 2023; Jain et al., 2023a). However, much recent work has used GFlowNets as samplers for Bayesian posterior distributions, e.g., causal discovery (Deleu et al., 2022; Nishikawa-Toomey et al., 2022; Atanackovic et al., 2023), amortized variational EM with discrete latents (Hu et al., 2023), neurosymbolic inference (van Krieken et al., 2022), and feature attribution in classifiers (Li et al., 2023). The theory and optimization techniques for GFlowNets have also evolved, with improved training objectives and exploration techniques (Malkin et al., 2022; Madan et al., 2023; Pan et al., 2022, 2023a; Shen et al., 2023), better understanding of their connections to variational methods (Zhang et al., 2022a; Malkin et al., 2023; Zimmermann et al., 2022), and extensions to stochastic (Zhang et al., 2023c; Pan et al., 2023b) and continuous (Lahlou et al., 2023) MDPs.
ML for combinatorial optimization The surge of machine learning for CO problems alleviates the reliance on hand-crafted heuristics while enabling generalization to new instances (Bengio et al., 2018; Cappart et al., 2021). Some methods (Li et al., 2018; Gasse et al., 2019; Gupta et al., 2020; Sun & Yang, 2023) rely on supervised information from expert solvers, which can be hard to obtain. Alternative approaches that leverage reinforcement learning (Dai et al., 2017; Kool et al., 2019; Chen & Tian, 2019; Yolcu & Poczos, 2019; Ahn et al., 2020; Delarue et al., 2020; Drori et al., 2020) or other unsupervised learning objectives (Karalias & Loukas, 2020; Sun et al., 2022; Wang et al., 2022) broaden the applicability of learning for CO problems. However, the mode-collapse issue might hinder the diversity and thus the solution coverage. In this regard, GFlowNets are easy to train while also designed for discovering multiple modes. As the first example to show the advantage of
Table 1: Max independent set experimental results. We report the absolute performance, approximation ratio (relative to KAMIS), and inference time. All algorithms fall into three categories: OR (operations research), SL (supervised learning), and UL (unsupervised learning). "—" denotes no reasonable result is achieved by the corresponding algorithm in 10 hours. Time shown as H:M:S.
METHOD	TYPE	SMALL	LARGE	SATLIB
SIZE ↑	DROP ↓	TIME ↓	SIZE ↑	DROP ↓	TIME ↓	SIZE ↑	DROP ↓	TIME ↓
GUROBI	OR	19.98	0.01%	47:34	40.90	5.21%	2:10:26	425.95	0.00%	3:43:19
KAMIS	OR	20.10	0.00%	1:24:12	43.15	0.00%	2:03:36	425.96	0.00%	4:15:41
PPO	UL	19.01	5.42%	1:17	32.32	25.10%	7:33	421.49	1.05%	13:12
INTEL	SL	18.47	8.11%	13:04	34.47	20.12%	20:17	—	—	—
DGL	SL	17.36	13.61%	12:47	34.50	20.05%	23:54	—	—	—
OURS	UL	19.18	4.57%	0:32	37.48	13.14%	4:22	423.54	0.57%	23:13
GFlowNets in CO problems, Zhang et al. (2023a) tackles the robust job scheduling problem that is central to compiler optimization. In our paper, we formalize CO under the GFlowNet framework with principled MDP design and demonstrate its effectiveness in a wide range of graph CO tasks.
We conduct extensive experiments on various graph CO tasks to demonstrate the effectiveness of the proposed GFlowNet approach. For MIS problems, we follow the setup in the MIS benchmark from Böther et al. (2022), while for other tasks, we follow the experimental setup in Sun et al. (2022).
Datasets As Dai et al. (2021) have pointed out that problems in existing synthetic graph data are relatively easy for MIS and MC, we take the more complicated RB graphs (Xu & Li, 2000) following Karalias & Loukas (2020). For realistic data, we take the SATLIB dataset (Hoos et al., 2000), which is reduced from SAT instances in conjunctive normal form. For the other two tasks, namely MDS and MCut, we adopt BA graphs (Barabási & Albert, 1999) following Sun et al. (2022). For all types of synthetic graph data, we generate two scales of datasets (denoted SMALL and LARGE), which respectively contain around 200 to 300 vertices and 800 to 1200 vertices.
Baselines For MIS problems, we compare with the baselines in the MIS benchmark. For classical operation research (OR) methods, we include a general-purpose mixed-integer program solver (GUROBI) and a MIS-specific solver (Lamm et al., 2017, KAMIS). For learning-based methods, we compare with a reinforcement learning-based PPO method (Ahn et al., 2020), and supervised learning with tree search refinement, in two different implementations (Li et al. (2018, INTEL) and Bother et al. (2022, DGL)). For the non-MIS tasks, we compare with the GUROBI solver, two heuristic methods which are greedy and mean-field annealing (Bilbro et al., 1988, MFA), and two state-of-the-art probabilistic methods (Karalias & Loukas, 2020) and the annealed version (Sun et al., 2022). We use ERDOS and ANNEAL to denote these two learning-based methods. For max-cut problems, we also adopt a semi-definite programming baseline (coined SDP in the result tables) which aims at a relaxation of the MCut task. We set its maximal running time to 10 hours. For the methods already contained in the MIS benchmark, we train them from scratch and report performance with the same protocol; for other algorithms, we write our own implementations and test them in the same way.
Results & analysis We evaluate both the performance and the inference time and report the mean value of the objective (e.g., set size in MIS) and the approximation ratio relative to the best-performing non-ML solver, treated as an oracle. The time denotes the total latency of evaluating on the test set. The best results among non-OR methods are marked in bold. MIS results are demonstrated in Table 1, where the problem-specific solver KAMIS is served as the gold standard for calculating the drop. For MIS methods, the larger the size of the independent set found, the better the algorithm. The "UL" (unsupervised learning) refers to the algorithms that do not need labels (i.e., solutions found by solvers) in the training set. For SATLIB, none of the supervised learning baselines can
3The approximation ratio is computed on the total cost aggregated over the test set and is defined as the DROP  $(1 - \frac{\mathrm{algorithm}}{\mathrm{oracle}})$  for maximization problems and the GAP  $(1 - \frac{\mathrm{oracle}}{\mathrm{algorithm}})$  for minimization problems.
Table 2: Max clique, min dominating set, and max cut results on small graphs (  $|V|$  between 200 and 300). We report absolute performance, approximation ratio, and inference time. All algorithms fall into three categories: OR (operations research), H (heuristic), and UL (unsupervised learning). The time latency is shown in the form of hour:minute:second.
METHOD	TYPE	MC	MDS	MCUT
SIZE ↑	DROP ↓	TIME ↓	SIZE ↓	GAP ↓	TIME ↓	SIZE ↑	DROP ↓	TIME ↓
GUROBI	OR	19.05	0.00%	1:55	27.89	0.00%	1:47	732.47	0.00%	13:04
SDP	OR	—	—	—	—	—	—	700.36	4.38%	35:47
GREEDY	H	13.53	28.98%	0:25	37.39	25.41%	2:13	688.31	6.03%	0:13
MFA	H	14.82	22.15%	0:27	36.36	23.29%	2:56	704.03	3.88%	1:36
ERDOS	UL	12.02	36.90%	0:41	30.68	9.09%	1:00	693.45	5.33%	0:46
ANNEAL	UL	14.10	25.98%	0:41	29.24	4.62%	1:01	696.73	4.88%	0:45
OURS	UL	16.24	14.75%	0:42	28.61	2.52%	2:20	704.30	3.85%	2:57
Table 3: Max clique, min dominating set, and max cut results on large graphs (whose  $|V|$  is between 800 and 1200). We use the same format as Table 2.
METHOD	TYPE	MC	MDS	MCUT
SIZE ↑	DROP ↓	TIME ↓	SIZE ↓	GAP ↓	TIME ↓	SIZE ↑	DROP ↓	TIME ↓
GUROBI	OR	33.89	0.00%	16:40	103.80	0.00%	13:48	2915.29	0.00%	1:05:29
SDP	OR	—	—	—	—	—	—	2786.00	4.43%	10:00:00
GREEDY	H	26.71	21.17%	0:25	140.52	26.13%	35:01	2761.06	5.29%	3:07
MFA	H	27.94	17.56%	2:19	126.56	17.98%	36:31	2833.86	2.79%	7:16
ERDOS	UL	25.43	24.96%	2:16	116.76	11.10%	3:56	2870.34	1.54%	2:49
ANNEAL	UL	27.46	18.97%	2:16	111.50	6.91%	3:55	2863.23	1.79%	2:48
OURS	UL	31.42	7.29%	4:50	110.28	5.88%	32:12	2864.61	1.74%	21:20
achieve meaningful results (average MIS size  $< 400$ ) within 10 hours, thus we use “—” to mark performance. One can see that GFlowNets surpass all the learning-based baselines in the sense of achieving the largest independent set solution.
For the MC & MDS & MCut problems, we exhibit the experimental results on small graphs in Table 2 and results on large graphs in Table 3, where GUROBI serves as the gold standard to calculate the performance drop ratio. Notice that for MDS problems, the smaller the vertex size result, the better the performance, which is opposite to other tasks. We observe that GFlowNet outperforms other baselines across different tasks and problem scales, with the only exception being the large-scale max-cut problem. This reflects the fairly universal effectiveness of the proposed method. We also notice that GFlowNets require a longer inference time compared to other methods. This is because the GFlowNet only adds one vertex into the set at each step, thus it needs multiple steps to output a vertex set solution. In contrast, probabilistic methods such as "Erdős goes neural" (Karalias & Loukas, 2020) only require one neural network forward pass during inference time to give a coarse estimate of the solution. We can see how this compares with a sequential generation approach as in GFlowNets, where each decision is taken in the context of the previously already taken decisions, ensuring a better coordinated set of decisions.
Ablation study We now conduct in-depth ablation studies to evaluate the key design choices in our method. We first study the difference between a series of GFlowNet variants in Figure 4, based on the MIS task with small scale graph data. We compare transition-based FL, trajectory-based FL, transition-based DB, trajectory-based DB, and TB (which only has trajectory-based implementation). Our result indicates that GFlowNet's transition-based FL implementation learns the fastest for CO tasks, which supports our modeling choice in Section 3.3. Figure 5 in Section C summarizes additional ablation studies including the temperature annealing, off-policy exploration strategy, and network architectures. We vary the GFlowNet's temperature coefficient which scales the temperature hyperparameter; we ablate the off-policy exploration during the rollout stage of GFlowNet training; we also ablate the GFlowNet architecture. Our results indicate that the GFlowNet is robust to a
wide range of hyperparameters, making it appealing to be applied to different CO problems, which validates the effectiveness of our proposed methodology from another perspective.
Our work focuses on the challenges of solving CO problems using unsupervised learning approaches. This work contributes to this growing field by proposing to apply the principled GFlowNet decision-making framework to CO tasks. By combining the power of probabilistic inference and sequential decision-making, GFlowNets offer a promising direction for finding a diverse set of high-quality candidate solutions in CO problems. Technically, we have developed problem-specific MDPs and efficient learning algorithms to address the challenges associated with learning from long trajectories, making our approach practical and scalable in large-scale setups. Our extensive numerical results showcase the effectiveness and efficiency of GFlowNets in solving NP-hard level CO problems, highlighting their ability to generate a set of diverse high-quality solutions. We believe that our work opens up new possibilities for addressing the limitations of existing approaches and paves the way for future research at the intersection of machine learning and combinatorial optimization.
The authors are thankful to Yuandong Tian, David Wei Zhang, Haoran Sun, Zhiqing Sun, Taoan Huang, Sophie Xhonneux, and Zhaoyu Li. Yoshua acknowledges funding from CIFAR, NSERC, Intel and Samsung. Aaron is supported by Hitachi, Samsung, Canadian Research Chair and CIFAR. Dinghuai acknowledges Daozhen Lin for being his tour guide during their visit to beautiful Italy.
Ahn, Sungsoo, Seo, Younggyo, and Shin, Jinwoo. Learning what to defer for maximum independent sets. International Conference on Machine Learning (ICML), 2020.
Atanackovic, Lazar, Tong, Alexander, Hartford, Jason, Lee, Leo J., Wang, Bo, and Bengio, Yoshua. DynGFN: Bayesian dynamic causal discovery using generative flow networks. arXiv preprint 2302.04178, 2023.
Barabási, Albert-Laszlo and Albert, Reka. Emergence of scaling in random networks. Science, 286 5439:509-12, 1999.
Bello, Irwan, Pham, Hieu, Le, Quoc V., Norouzi, Mohammad, and Bengio, Samy. Neural combinatorial optimization with reinforcement learning. arXiv preprint 1611.09940, 2016.
Bengio, Emmanuel, Jain, Moksh, Korablyov, Maksym, Precup, Doina, and Bengio, Yoshua. Flow network based generative models for non-iterative diverse candidate generation. Neural Information Processing Systems (NeurIPS), 2021.
Bengio, Yoshua, Lodi, Andrea, and Prouvost, Antoine. Machine learning for combinatorial optimization: a methodological tour d'horizon. Eur. J. Oper. Res., 290:405-421, 2018.
Bengio, Yoshua, Deleu, Tristan, Hu, Edward J., Lahlou, Salem, Tiwari, Mo, and Bengio, Emmanuel. Gflownet foundations. Journal of Machine Learning Research (JMLR), 2023.
Bilbro, Griff L., Mann, Reinhold, Miller, Thomas K., Snyder, Wesley E., van den Bout, David E., and White, Mark W. Optimization by mean field annealing. Neural Information Processing Systems (NIPS), 1988.
Böther, Maximilian, Küssig, Otto, Taraz, Martin, Cohen, Sarel, Seidel, Karen, and Friedrich, Tobias. What's wrong with deep learning in tree search for combinatorial optimization. International Conference on Learning Representations (ICLR), 2022.
Cappart, Quentin, Chételat, Didier, Khalil, Elias, Lodi, Andrea, Morris, Christopher, and Velicković, Petar. Combinatorial optimization and reasoning with graph neural networks. International Joint Conference on Artificial Intelligence (IJCAI), 2021.
Chen, Xinyun and Tian, Yuandong. Learning to perform local rewriting for combinatorial optimization. Neural Information Processing Systems (NeurIPS), 2019.
Zhang, Dinghuai, Courville, Aaron C., Bengio, Yoshua, Zheng, Qinqing, Zhang, Amy, and Chen, Ricky T. Q. Latent state marginalization as a low-cost approach for improving exploration. International Conference on Learning Representations (ICLR), 2023b.
Zhang, Dinghuai, Pan, L., Chen, Ricky T. Q., Courville, Aaron C., and Bengio, Yoshua. Distributional GFlowNets with quantile flows. arXiv preprint 2302.05793, 2023c.
Zimmermann, Heiko, Lindsten, Fredrik, van de Meent, J.-W., and Naesseth, Christian A. A variational perspective on generative flow networks. Transactions of Machine Learning Research (TMLR), 2022.
Symbol	Description
S	state space
X	object (terminal state) space, subset of S
A	action / transition space (edges s → s')
s	state in S
s0	initial state, element of S
x	terminal state in X
τ	complete trajectory
θ	learnable parameter of GFlowNets
F: S → R	state flow
PF	forward policy (distribution over children)
PB	backward policy (distribution over parents)
R: X → R>0	reward function (unnormized target density)
Z	scalar, equal to ∑x∈X R(x)
E: X → R	energy function, equal to -T * log R(x)
E: S → R	energy in state space, to provide intermediate learning signals
T	scalar, temperature
g	graph configuration
V	vertex set of graph g
S	subset of V, solution to some CO problem
E	edge set of graph g
B MDP designs
In all MDPs, the agent maintain a set of vertices as the state. This is achieved by assign a value among  $\{0,1,\emptyset\}$  for each of the vertex, which corresponds to "not in the set" / "in the set" / "unspecified". At each step, the agent performs an action, which is essentially changing the entry value of one vertex from  $\emptyset$  to a specified binary value. After that, the designed transition may modify the value of other unspecified vertices to ensure the state will represent a feasible solution. The agent will receive a terminal reward when no action can be taken (i.e., termination). An illustration of the following descriptions can be found in Figure 2.
B.1 Maximum clique MDP
Here we think of the all-void initial state means that we have an empty set of vertex at the beginning.
Action The action is to choose an unspecified vertex and include it in the current set. This is changing the entry value of that vertex in the state vector representation from  $\varnothing$  to 1.
Transition We need to ensure that the current set always represents a clique, thus it is crucial to enforce the connectivity between vertices in the set. If a vertex could potentially join the clique in the future, then it must connect to all the vertices in the current set. Hence, after performing an action, we identify the vertices that fail to meet this requirement and mark them as 0. This ensures that the clique constraint is preserved throughout the generation process.
Reward We set the log reward to be the resulting clique size, i.e.,  $\mathcal{E}(\mathbf{x}) = -|\mathbf{x}|_1$ . For intermediate state  $\mathbf{s}$ , we define its intermediate signal by setting  $-\tilde{\mathcal{E}}(\mathbf{s})$  to be the number of 1 in  $\mathbf{s}$ .
B.2 Minimum dominating set MDP
In MDS, we start with an all-void initial state, where all vertices (which are in unspecified state) are considered to in the target dominating set. This is different from other MDP designs, where we think of unspecified void state as not in the vertex set. In this representation, the entire vertex set  $V$  forms a trivial dominating set. Starting from this, we perform actions to remove vertices from the current set but keep it to be a dominating set all the time.
Figure 5: Ablation study on GFlowNet training. These experiments are conduct on the small scale MIS task. From left to right: ablation on whether to use temperature annealing, inverse temperature, how much percentage of uniform noise used in off-policy exploration, and number of message passing layers in the graph neural network. These results demonstrate the robustness of GFlowNets against different hyperparameter setups.
Action Each action corresponds to selecting a vertex and removing it from the current dominating set. This is done by changing the entry value of the chosen vertex from  $\varnothing$  to 0.