SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Sampling. At each round  $i$ , we use the trained HN-GFN to sample a diverse batch of  $b$  candidates. Let  $\Lambda^i$  be the set of  $k$  target preference vectors  $\lambda_{\mathrm{target}}^i$ . We sample  $\frac{b}{k}$  molecules for each  $\lambda_{\mathrm{target}}^i \in \Lambda^i$  and evaluate them on the oracle in parallel. We simply construct  $\Lambda^i$  by sampling from  $\mathrm{Dir}(\alpha)$ , but it is worth noting that this prior distribution can also be defined adaptively based on the trade-off of interest. As the number of objectives increases, we choose a larger  $k$  to cover the objective space.
4.3 Hindsight-like off-policy strategy
Resorting to the conditioning mechanism, HN-GFN can learn a family of policies to achieve various goals, i.e., one can treat sampling high-reward molecules for a particular preference vector as a separate goal. As verified empirically in Jain et al. [28], since the FM objective is off-policy and offline, we can use offline trajectories to train the target policy for better exploration, so long as the assumption of full support holds. Our key insight is that each policy can learn from the valuable experience (high-reward molecules) of other similar policies.
Inspired by hindsight experience replay [2] in RL, we propose to share high-performing molecules among policies by re-examining them with different preference vectors. As there are infinite possible preference vectors, we focus on target preference vectors  $\Lambda^i$ , which are based on to sample candidates, and build a replay buffer for each  $\lambda_{\mathrm{target}}^i \in \Lambda^i$ . After sampling some trajectories during training, we store in the replay buffers the complete objects  $x$  with the corresponding reward  $R_{\lambda_{\mathrm{target}}^i}(x)$ .
To be specific, at each iteration, we first sample a preference vector from a mixture between  $\mathrm{Dir}(\alpha)$  and a uniform distribution over  $\Lambda^i$ :  $(1 - \gamma)\mathrm{Dir}(\alpha) + \gamma$  Uniform. If  $\Lambda^i$  is chosen, we construct half of the training batch with offline trajectories from the corresponding replay buffer of molecules
encountered with the highest rewards. Otherwise, we incorporate offline trajectories from the current observed dataset  $\mathcal{D}_i$  instead to ensure that HN-GFN samples correctly in the vicinity of the observed Pareto set. Our strategy allows for flexible trade-offs between generalization and specialization. As we vary  $\gamma$  from 0 to 1, the training distribution of preference vectors moves from  $\mathrm{Dir}(\alpha)$  to  $\Lambda^i$ . Exclusively training the HN-GFN with the finite target preference vectors  $\Lambda^i$  can lead to poor generalization. In practice, although we only sample candidates based on  $\Lambda^i$ , we argue that it is vital to keep the generalization to leverage the trained HN-GFN to explore various preference vectors adaptively. The detailed algorithm can be found in Appendix A.
4.4 Surrogate model and acquisition function
While GPs are well-established in continuous spaces, they scale poorly with the number of observations and do not perform well in discrete spaces [57]. There has been significant work in efficiently training non-Bayesian neural networks to estimate epistemic uncertainty [21, 39]. We benchmark widely used approaches (in Appendix C.3), and use a flexible and effective one: evidential deep learning [1]. As for the acquisition function, we use Upper Confidence Bound [55] to incorporate the uncertainty. To be precise, the objectives are modeled with a single multi-task MPNN, and the acquisition function is applied to the scalarization.
We first verify that HN-GFN has adequate capacity to generalize over preference vectors in a single-round synthetic scenario. Next, we evaluate the effectiveness of the proposed MOBO algorithm based on HN-GFN in multi-round practical scenarios, which are more in line with real-world molecular optimization. Besides, we conduct several ablation studies to empirically justify the design choices.
5.1 Single-round synthetic scenario
Here, our goal is to demonstrate that we can leverage the HN-GFN to sample molecules with preference-conditioned property distributions. The HN-GFN is used as a stand-alone optimizer outside of MOBO to directly optimize the scalarizations of oracle scores. As the oracle cannot be called as many times as necessary practically, we refer to this scenario as a synthetic scenario. To better visualize the trend of the property distribution of the sampled molecules as a function of the preference weight, we only consider two objectives: inhibition scores against glycogen synthase kinase-3 beta  $(\mathrm{GNK}3\beta)$  and c-Jun N-terminal kinase-3 (JNK3) [40, 33].
Compared methods. We compare HN-GFN against the following methods. Preference-specific GFlowNets (PS-GFN) is a set of vanilla unconditional GFlowNets, each trained separately for a specific preference vector. Note that PS-GFN is treated as "gold standard" rather than a baseline, as it is trained and evaluated using the same preference vector. Concat-GFN and FiLM-GFN are two alternative conditional variations of GFlowNet based on the concatenation and FiLM [51], respectively. MOEA/D [64] and NSGA-III [14] are two multi-objective evolutionary algorithms that also incorporate preference information. We perform evolutionary algorithms over the 32-dim latent space learned by HierVAE [32], which gives better optimization performance than JT-VAE [31].
Metrics. All the above methods are evaluated over the same set of 5 evenly spaced preference vectors. For each GFlowNet-based method, we sample 1000 molecules per preference vector as solutions. We compare the aforementioned methods on the following metrics: Hypervolume indicator (HV) measures the volume of the space dominated by the Pareto front of the solutions and bounded from below by the preference point  $(0,0)$ . Diversity (Div) is the average pairwise Tanimoto distance over Morgan fingerprints. Correlation (Cor) is the Spearman's rank correlation coefficient between the probability of sampling molecules from an external test set under the GFlowNet and their respective rewards in the logarithmic domain [48]. See more details in Appendix B.1.2. In a nutshell, HV and Div measure the quality of the solutions, while Cor measures how well the trained model is aligned with the given preference vector. Each experiment is repeated with 3 random seeds.
Experimental results. As shown in Table 1, HN-GFN outperforms the baselines and achieves competitive performance to PS-GFN (gold standard) on all the metrics. Compared to GFlowNet-based
Table 1: Evaluation of different methods on the synthetic scenario.
Method	HV (↑)	Div (↑)	Cor (↑)
MOEA/D	0.182 ± 0.045	n/a	n/a
NSGA-III	0.364 ± 0.041	n/a	n/a
PS-GFN	0.545 ± 0.055	0.786 ± 0.013	0.653 ± 0.003
Concat-GFN	0.534 ± 0.069	0.786 ± 0.004	0.646 ± 0.008
FiLM-GFN	0.431 ± 0.045	0.795 ± 0.014	0.633 ± 0.009
HN-GFN	0.550 ± 0.074	0.797 ± 0.015	0.666 ± 0.010
Figure 2: Left: The distribution of Top-100 JNK3 scores. Right: The progression of the average Top-20 rewards over the course of training of the HN-GFN in optimizing GSK3 $\beta$  and JNK3.
methods, evolutionary algorithms (MOEA/D and NSGA-III) fail to find high-scoring molecules, especially the MOEA/D. HN-GFN outperforms Concat-GFN and FiLM-GFN, and is the only conditional variant that can match the performance of PS-GFN, implying the superiority of the well-designed hypernetwork-based conditioning mechanism. In Figure 2 (left), we visualize the empirical property (JNK3) distribution of the molecules sampled by HN-GFN and PS-GFN conditioned on the set of evaluation preference vectors. We observe that the distributions are similar and the trends are consistent: the larger the preference weight, the higher the average score. The comparable performance and consistent sampling distributions illustrate that HN-GFN has adequate capacity to generalize over preference vectors. Since the runtime and storage space of PS-GFN scale linearly with the number of preference vectors, our unified HN-GFN provides a significantly efficient way to explore various trade-offs between objectives.
5.2 Multi-objective Bayesian optimization
Next, we evaluate the effectiveness of HN-GFN as an acquisition function optimizer within MOBO in practical scenarios. We consider the following objective combinations of varying sizes:
- GNK3 $\beta$  + JNK3: Jointly inhibiting Alzheimer-related targets GNK3 $\beta$  and JNK3.
-  $\mathrm{GNK3\beta + JNK3 + QED + SA}$ : Jointly inhibiting  $\mathrm{GNK3\beta}$  and JNK3 while being drug-like and easy-to-synthesize.
We rescale the SA score (initially between 10 and 1) so that all the above properties have a range of [0,1] and higher is better. For both combinations, we consider starting with  $|\mathcal{D}_0| = 200$  random molecules and further querying the oracle  $N = 8$  rounds with batch size  $b = 100$ . Each experiment is repeated with 3 random seeds.
Baselines. We compare HN-GFN with three representative LSO methods (qParEGO [37], qEHVI [12], and LaMOO [65]), as well as a variety of state-of-the-art combinatorial optimization
(a)  $\mathrm{GSK3}\beta +\mathrm{JNK3}$
(b)  $\mathrm{GSK3}\beta +\mathrm{JNK3} + \mathrm{QED} + \mathrm{SA}$
Figure 3: Optimization performance (hypervolume) over MOBO loops.
Table 2: Diversity for different methods in MOBO scenarios.
Div (↑)
GSK3β + JNK3	GSK3β + JNK3 + QED + SA
Graph GA	0.347 ± 0.059	0.562 ± 0.031
MARS	0.653 ± 0.072	0.754 ± 0.027
P-MOCO	0.646 ± 0.008	0.350 ± 0.130
HN-GFN	0.810 ± 0.003	0.744 ± 0.008
HN-GFN w/ hindsight	0.793 ± 0.007	0.738 ± 0.009
methods: Graph GA [30] is a genetic algorithm, MARS [61] is a MCMC sampling approach, and P-MOCO [41] is a multi-objective RL method. We provide more details in Appendix B.2.
Experimental results. Figure 3 illustrates that HN-GFN achieves significantly superior performance (HV) to baselines, especially when trained with the proposed hindsight-like off-policy strategy. HN-GFN outperforms the best baselines P-MOCO and Graph GA of the two objective combinations by a large margin (0.67 vs. 0.50 and 0.42 vs. 0.34), respectively. Besides, HN-GFN is more sample-efficient, matching the performance of baselines with only half the number of oracle evaluations. All combinatorial optimization methods result in more performance gains compared to the LSO methods, implying that it is promising to optimize directly over the discrete space. In Table 2, we report the Div (computed among the batch of 100 candidates and averaged over rounds). For a fair comparison, we omit the LSO methods as they only support 160 rounds with batch size 5 due to memory constraints. Compared with Graph GA and P-MOCO, which sometimes propose quite similar candidates, the superior optimization performance of HN-GFN can be attributed to the ability to sample a diverse batch of candidates. Another interesting observation, in the more challenging setting  $(\mathrm{GNK}3\beta + \mathrm{JNK}3 + \mathrm{QED} + \mathrm{SA})$ , is that MARS generates diverse candidates via MCMC sampling but fails to find a high-quality Pareto front, indicating that HN-GFN can find high-reward modes better than MARS. The computational costs are discussed in Appendix B.4.
Effect of the hindsight-like off-policy strategy. In the first round of MOBO, for each  $\lambda_{\mathrm{target}} \in \Lambda$  we sample 100 molecules every 500 training steps and compute the average Top-20 reward over  $\Lambda$ . In Figure 2 (right), we find that the hindsight-like off-policy strategy significantly boosts average rewards, demonstrating that sharing high-performing molecules among policies effectively speeds up the training of HN-GFN. On the other hand, further increasing  $\gamma$  leads to slight improvement. Hence, we choose  $\gamma = 0.2$  for the desired trade-off between generalization and specialization.
Effect of  $\mathrm{Dir}(\alpha)$ . Here we consider the more challenging  $\mathrm{GNK3}\beta +\mathrm{JNK3} + \mathrm{QED} + \mathrm{SA}$  combination, where the difficulty of optimization varies widely for various properties. Table 3 shows that the distribution skewed toward harder properties results in better optimization performance. In our early experiments, we found that if the distribution of  $\Lambda^i$  is fixed, HN-GFN is quite robust to changes in  $\alpha$ .
Effect of scalarization functions  $s_{\lambda}$ . In Table 3, we observe that WS leads to a better Pareto front than Tchebycheff. Although the simple WS is generally inappropriate when a non-convex Pareto front is encountered [10], we find that it is effective when optimized with HN-GFN, which can sample diverse high-reward candidates and may reach the non-convex part of the Pareto front. In addition, we conjecture that the non-smooth reward landscapes induced by Tchebycheff are harder to optimize.
Table 3: Ablation study of the  $\alpha$  and scalarization functions on  $\mathrm{{GNK}}3\beta  + \mathrm{{JNK}}3 + \mathrm{{QED}} + \mathrm{{SA}}$  .
α	scalarization function
(1,1,1,1)	(3,3,1,1)	(3,4,2,1)	WS	Tchebycheff
HV	0.312 ± 0.039	0.385 ± 0.018	0.416 ± 0.023	0.416 ± 0.023	0.304 ± 0.075
Div	0.815 ± 0.015	0.758 ± 0.018	0.738 ± 0.009	0.738 ± 0.009	0.732 ± 0.014
We have introduced a MOBO algorithm for sample-efficient multi-objective molecular optimization. This algorithm leverages a hypernetwork-based GFlowNet (HN-GFN) to sample a diverse batch of candidates from the approximate Pareto front. In addition, we present a hindsight-like off-policy strategy to improve optimization performance. Our algorithm outperforms existing approaches in synthetic and practical scenarios. Future work includes extending this algorithm to other discrete optimization problems such as biological sequence design and neural architecture search. One limitation of the proposed HN-GFN is the higher computational costs than other training-free optimization methods. However, the costs resulting from model training are negligible compared to the costs of evaluating the candidates in real-world applications. We argue that the higher quality of the candidates is much more essential than lower costs.
This research was partially supported by National Key R&D Program of China under grant No.2018AAA0102102, National Natural Science Foundation of China under grants No.62176231, No.62106218, No.82202984, No.92259202, No.62132017 and U22B2034, Zhejiang Key R&D Program of China under grant No.2023C03053. Tingjun Hou's research was supported in part by National Natural Science Foundation of China under grant No.22220102001.
[1] Alexander Amini, Wilko Schwarting, Ava Soleimany, and Daniela Rus. Deep evidential regression. Advances in Neural Information Processing Systems, 33:14927-14937, 2020.
[2] Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, OpenAI Pieter Abbeel, and Wojciech Zaremba. Hindsight experience replay. Advances in neural information processing systems, 30, 2017.
[3] Christof Angermueller, David Dohan, David Belanger, Ramya Deshpande, Kevin Murphy, and Lucy Colwell. Model-based reinforcement learning for biological sequence design. In International Conference on Learning Representations, 2020. URL https://openreview.net/forum?id=Hk1xbgBKvr.
[4] Maximilian Balandat, Brian Karrer, Daniel Jiang, Samuel Daulton, Ben Letham, Andrew G Wilson, and Eytan Bakshy. Botorch: a framework for efficient monte-carlo Bayesian optimization. Advances in neural information processing systems, 33:21524-21538, 2020.
[57] Kevin Swersky, Yulia Rubanova, David Dohan, and Kevin Murphy. Amortized Bayesian optimization over discrete spaces. In Conference on Uncertainty in Artificial Intelligence, pages 769-778. PMLR, 2020.
[58] Austin Tripp, Erik Daxberger, and José Miguel Hernández-Lobato. Sample-efficient optimization in the latent space of deep generative models via weighted retraining. Advances in Neural Information Processing Systems, 33:11259-11272, 2020.
[59] Austin Tripp, Gregor NC Simm, and José Miguel Hernández-Lobato. A fresh look at de novo molecular design benchmarks. In NeurIPS 2021 AI for Science Workshop, 2021.
[60] Robin Winter, Floriane Montanari, Andreas Steffen, Hans Briem, Frank Noé, and Djork-Arné Clevert. Efficient multi-objective molecular optimization in a continuous latent space. Chemical science, 10(34):8016-8024, 2019.
[61] Yutong Xie, Chence Shi, Hao Zhou, Yuwei Yang, Weinan Zhang, Yong Yu, and Lei Li. Mars: Markov molecular sampling for multi-objective drug discovery. In International Conference on Learning Representations, 2021.
[62] Jiaxuan You, Bowen Liu, Zhitao Ying, Vijay Pande, and Jure Leskovec. Graph convolutional policy network for goal-directed molecular graph generation. Advances in neural information processing systems, 31, 2018.
[63] Dinghuai Zhang, Nikolay Malkin, Zhen Liu, Alexandra Volokhova, Aaron Courville, and Yoshua Bengio. Generative flow networks for discrete probabilistic modeling. arXiv preprint arXiv:2202.01361, 2022.
[64] Qingfu Zhang and Hui Li. Moea/d: A multiobjective evolutionary algorithm based on decomposition. IEEE Transactions on evolutionary computation, 11(6):712-731, 2007.
[65] Yiyang Zhao, Linnan Wang, Kevin Yang, Tianjun Zhang, Tian Guo, and Yuandong Tian. Multi-objective optimization by learning space partition. In International Conference on Learning Representations, 2022. URL https://openreview.net/forum?id=FlwzVjfMryn.
[66] Yiheng Zhu, Zhenqiu Ouyang, Ben Liao, Jialu Wu, Yixuan Wu, Chang-Yu Hsieh, Tingjun Hou, and Jian Wu. Molhf: A hierarchical normalizing flow for molecular graph generation. In Proceedings of the Thirty-Second International Joint Conference on Artificial Intelligence, pages 5002–5010, 8 2023.
[67] Eckart Zitzler and Lothar Thiele. Multiobjective evolutionary algorithms: a comparative case study and the strength pareto approach. IEEE transactions on Evolutionary Computation, 3(4): 257-271, 1999.
Algorithm 1 describes the overall framework of the proposed MOBO algorithm, where HN-GFN is leveraged as the acquisition function optimizer. Algorithm 2 describes the training procedure for HN-GFN within MOBO.
Algorithm 1 MOBO based on HN-GFN
Input: oracle  $f = (f_{1},\dots ,f_{M})$  , initial dataset  $\mathcal{D}_0 = \{(x_i^0,f(x_i^0))\}_{i = 1}^n$  acquisition function a, parameter of Dirichlet distribution  $\alpha$  , number of rounds  $N$  , batch size b   
Initialization: surrogate model  $\mathcal{M}$  , parameters of HN-GFN  $\pi_{\theta}$    
for  $i = 1$  to  $N$  do Fit surrogate model  $\mathcal{M}$  on dataset  $\mathcal{D}_{i - 1}$  Sample the set of target preference weights  $\Lambda \sim \mathrm{Dir}(\alpha)$  Train  $\pi_{\theta}$  with reward function  $R_{\lambda}(x) = a(\mu (s_{\lambda}(x)),\sigma (s_{\lambda}(x));\mathcal{M})$ $\triangleright$  Algorithm 2 Sample query batch  $\mathcal{B}_i = \{x_j^i\}_{j = 1}^b$  based on  $\lambda_{target}\in \Lambda$  Evaluate batch  $\mathcal{B}_i$  with  $f$  and augment the dataset  $\mathcal{D}_{i + 1} = \mathcal{D}_i\cup \{(x_j^i,f(x_j^i))\}_{j = 1}^b$    
end for
Algorithm 2 Training procedure for HN-GFN with the hindsight-like off-policy strategy
Input: available dataset  $\mathcal{D}_i$  , reward function  $R$  , minibatch size  $m$  , set of target preference vectors   
 $\Lambda$  , proportion of hindsight-like strategy  $\gamma$  , replay buffers  $\{\mathcal{R}_{\lambda}\}_{\lambda \in \Lambda}$    
while not converged do Flag  $\sim$  Bernoulli(  $\gamma$  ) if Flag  $= 1$  then  $\lambda \sim \Lambda$  Sample  $\frac{m}{2}$  trajectories from replay buffer  $\mathcal{R}_{\lambda}$  else  $\lambda \sim \mathrm{Dir}(\alpha)$  Sample  $\frac{m}{2}$  trajectories from the available dataset  $D_{i}$    
end if   
 $\theta = (\theta_{\mathrm{mpnn}},h(\lambda ;\phi))$    
Sample  $\frac{m}{2}$  trajectories from policy  $\widetilde{\pi}$  and store terminal states  $x$  in  $\mathcal{R}_{\lambda}$  for all  $\lambda \in \Lambda$  Compute reward  $R_{\lambda}(x)$  on terminal states  $x$  from each trajectory in the minibatch Update parameters  $\theta_{\mathrm{mpnn}}$  and  $\phi$  with a stochastic gradient descent step w.r.t Equation (4)   
end while
B.1 Experimental settings
B.1.1 Molecule domain
Following Bengio et al. [6], the molecules are generated by sequentially attaching a fragment, which is chosen from a predefined vocabulary of building blocks, to an atom of the partially constructed molecules. The maximum trajectory length is 8, with the number of actions varying between 100 and 2000 depending on the state, making  $|\mathcal{X}|$  up to  $10^{16}$ . We adopt the property prediction models released by Xie et al. [61] to evaluate the inhibition ability of generated molecules against GSK3 $\beta$  and JNK3.
B.1.2 Metrics
Diversity. Diversity (Div) is the average pairwise Tanimoto distance over Morgan fingerprints. In the synthetic scenario, for each preference vector, we sample 1000 molecules, calculate the Div among the Top-100 molecules, and report the averages over preferences. In MOBO, the DiV is computed among the batch of 100 candidates per round, as Graph GA and MARS are not preference-conditioned. And we believe this metric possibly is more aligned with how these methods might be used in a biology or chemistry experiment.
Correlation. Correlation (Cor) is the Spearman's rank correlation coefficient between the probability of sampling molecules from an external test set under the GFlowNet and their respective rewards in the logarithmic domain:  $\mathrm{Cor} =$  Spearman's  $\rho_{\log (\pi (x)),\log (R(x))}$ . The external test set is obtained in two steps: First, we generate a random dataset containing 300K molecules uniformly based on the number of building blocks; Next, we sample the test sets with uniform property distribution corresponding to  $\mathrm{GSK3}\beta$  and JNK3, respectively, from the 300K molecules.
All baselines are implemented using the publicly released source codes with adaptations for our MOBO scenarios. The evolutionary algorithms (MOEA/D and NSGA-III) are implemented in PyMOO [9], and the LSO methods (qParEGO, qEHVI, and LaMOO) are implemented in BoTorch [4]. In MOBO scenarios, LaMOO and Graph GA utilize EHVI as the acquisition function. For all GP-based methods, each objective is modeled by an independent GP.
B.3 HN-GFN
We implement the proposed HN-GFN in PyTorch [50]. The values of key hyper-parameters are illustrated in Table 4.
Surrogate model: We use the 12-layer MPNN as the base architecture of the surrogate model in our experiments. In MOBO, a single multi-task MPNN is trained with a batch size of 64 using the Adam optimizer with a dropout rate of 0.1 and a weight decay rate of 1e-6. We apply early stopping to improve generalization.
HN-GFN: HN-GFN contains a vanilla GFlowNet and a preference-conditioned hypernetwork. The architecture of GFlowNet is a 10-layer MPNN, and the hypernetwork is a 3-layer MLP with multiple heads, each generating weights for different layers of the target network. The HN-GFN is trained with Adam optimizer to optimize the Flow Matching objective.
Table 4: Hyper-parameters used in the real-world MOBO experiments.
Hyper-parameter	GSK3β + JNK3	GSK3β + JNK3 + QED + SA
Surrogate model
Hidden size	64	64
Learning rate	2.5e-4	1e-3
λ for evidential regression	0.1	0.1
Number of iterations	10000	10000
Early stop patience	500	500
Dropout	0.1	0.1
Weight decay	1e-6	1e-6
Acquisition function (UCB)
β	0.1	0.1
HN-GFN
Learning rate	5e-4	5e-4
Reward exponent	8	8
Reward norm	1.0	1.0
Trajectories minibatch size	8	8
Offline minibatch size	8	8
hindsight γ	0.2	0.2
Uniform policy coefficient	0.05	0.05
Hidden size for GFlowNet	256	256
Hidden size for hypernetwork	100	100
Training steps	5000	5000
α	(1,1)	(3,4,2,1)
B.4 Empirical running time
The efficiency is compared on the same computing facilities using 1 Tesla V100 GPU. In the context of MOBO, the running time of three LSO methods (i.e., qParEGO, qEHVI, and LaMOO) is around 3 hours, while Graph GA optimizes much faster and costs only 13 minutes. In contrast, the time complexity of deep-learning-based discrete optimization methods is much larger. MARS costs 32 hours, while our proposed HN-GFN costs 10 hours. With the hindsight-like training strategy, the running time of HN-GFN will increase roughly by  $33\%$ .