SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yiheng Zhu $^{1*}$ , Jialu Wu $^{2*}$ , Chaowen Hu $^{3}$ , Jiahuan Yan $^{1}$ , Chang-Yu Hsieh $^{2\dagger}$ , Tingjun Hou $^{2\dagger}$ , Jian Wu $^{4,5,6\dagger}$
$^{2}$ College of Pharmaceutical Sciences, Zhejiang University
$^{3}$ Polytechnic Institute, Zhejiang University
$^{4}$ Second Affiliated Hospital School of Medicine, Zhejiang University
$^{5}$ School of Public Health, Zhejiang University
$^{6}$ Institute of Wenzhou, Zhejiang University
{zhuyiheng2020, jialuwu, ChaowenHu, jyansir, kimhsieh, tingjunhou, wujian2000} @zju.edu.cn
Many crucial scientific problems involve designing novel molecules with desired properties, which can be formulated as a black-box optimization problem over the discrete chemical space. In practice, multiple conflicting objectives and costly evaluations (e.g., wet-lab experiments) make the diversity of candidates paramount. Computational methods have achieved initial success but still struggle with considering diversity in both objective and search space. To fill this gap, we propose a multi-objective Bayesian optimization (MOBO) algorithm leveraging the hypernetwork-based GFlowNets (HN-GFN) as an acquisition function optimizer, with the purpose of sampling a diverse batch of candidate molecular graphs from an approximate Pareto front. Using a single preference-conditioned hypernetwork, HN-GFN learns to explore various trade-offs between objectives. We further propose a hindsight-like off-policy strategy to share high-performing molecules among different preferences in order to speed up learning for HN-GFN. We empirically illustrate that HN-GFN has adequate capacity to generalize over preferences. Moreover, experiments in various real-world MOBO settings demonstrate that our framework predominantly outperforms existing methods in terms of candidate quality and sample efficiency. The code is available at https://github.com/violet-sto/HN-GFN.
Designing novel molecular structures with desired properties, also referred to as molecular optimization, is a crucial task with great application potential in scientific fields ranging from drug discovery to material design. Molecular optimization can be naturally formulated as a black-box optimization problem over the discrete chemical space, which is combinatorially large [52]. Recent years have witnessed the trend to leverage computational methods, such as deep generative models [31] and combinatorial optimization algorithms [62, 30], to facilitate optimization. However, the applicability of most prior approaches in real-world scenarios is hindered by two practical constraints: (i) Chemists commonly seek to optimize multiple properties of interest simultaneously [60, 33]. For example, in addition to effectively inhibiting a disease-associated target, an ideal drug is desired to be easily synthesizable and non-toxic. Unfortunately, as objectives often conflict, in most cases, there is no single
Figure 1: MOBO loop for molecular optimization using a surrogate model  $\mathcal{M}$  for uncertainty estimation and HN-GFN for acquisition function optimization. In each round, the policy  $\pi_{\theta}$  is trained with reward function  $R_{\lambda}$ , where  $\lambda$  is sampled from  $\mathrm{Dir}(\alpha)$  per iteration. A new batch of candidates is sampled from the approximate Pareto front according to target preference vectors  $\lambda_{\mathrm{target}} \in \Lambda$ .
optimal solution, but rather a set of Pareto optimal solutions defined with various trade-offs [18, 45]. (ii) Realistic oracles (e.g., wet-lab experiments and high-fidelity simulations) require substantial costs to synthesize and evaluate molecules [22]. Hence, the number of oracle evaluations is notoriously limited. In such scenarios, the diversity of candidates is a critical consideration.
Bayesian optimization (BO) [34, 54] provides a sample-efficient framework for globally optimizing expensive black-box functions. The core idea is to construct a cheap-to-evaluate surrogate model approximating the true objective function (also known as the oracle) on the observed dataset and optimize an acquisition function (built upon the surrogate model) to obtain informative candidates with high utility for the next round of evaluations. Due to the common large-batch and low-round settings in biochemical experiments [3], batch BO is prioritized to shorten the entire cycle of optimization [25]. MOBO also received broad attention and achieved promising performance in continuous problems by effectively optimizing differentiable acquisition functions [12, 13]. Nevertheless, it is less prominent in discrete problems, where no gradients can be leveraged to navigate the search space.
Although existing discrete molecular optimization methods can be adopted as the acquisition function optimizer to alleviate this issue, they can hardly consider diversity in both search and objective space: 1) many neglect the diversity of proposed candidates in search space [30, 33]; 2) many multi-objective methods simply rely on a predefined scalarization function (e.g., mean) to turn the multi-objective optimization (MOO) problem into a single-objective one [61, 20]. As the surrogate model cannot exactly reproduce the oracle's full behaviors and the optimal trade-off is unclear before optimization (even with domain knowledge), it is desired to not only propose candidates that bring additional information about the search space but also explore more potential trade-offs of interest. To achieve this goal, we explore how to extend the recently proposed generative flow networks (GFlowNets) [6, 7], a class of generative models that aim to learn a stochastic policy for sequentially constructing objects with a probability proportional to a reward function, to facilitate multi-objective molecular optimization. Compared with traditional combinatorial optimization methods, GFlowNets possess merit in generating diverse and high-reward objects, which has been verified in single-objective problems [6, 28].
In this work, we present a MOBO algorithm based on GFlowNets for sample-efficient multi-objective molecular optimization. We propose a hypernetwork-based GFlowNet (HN-GFN) as the acquisition function optimizer within MOBO to sample a diverse batch of candidates from an approximate Pareto front. Instead of defining a fixed reward function as usual in past work [6], we train a unified GFlowNet on the distribution of reward functions (random scalarizations parameterized by preference vectors) and control the policy using a single preference-conditioned hypernetwork. While sampling candidates, HN-GFN explores various trade-offs between competing objectives flexibly by varying the input preference vector. Inspired by hindsight experience replay [2], we further introduce a hindsight-like off-policy strategy to share high-performing molecules among different preferences
and speed up learning for HN-GFN. As detailed in our reported experiments, we first empirically verify that HN-GFN is capable of generalizing over preference vectors, then apply the proposed framework to real-world scenarios and demonstrate its effectiveness and efficiency over existing methods. Our key contributions are summarized below:
- We introduce a GFlowNet-based MOBO algorithm to facilitate real-world molecular optimization. We propose HN-GFN, a conditional variant of GFlowNet that can efficiently sample candidates from an approximate Pareto front.
- We introduce a hindsight-like off-policy strategy to speed up learning in HN-GFN.
- Experiments verify that our MOBO algorithm based on HN-GFN can find a high-quality Pareto front more efficiently compared to state-of-the-art baselines.
Molecular optimization. Molecular optimization has been approached with a wide variety of computational methods, which can be mainly grouped into three categories: 1) Latent space optimization (LSO) methods [24, 31, 58, 66] perform the optimization over the low-dimensional continuous latent space learned by generative models such as variational autoencoders (VAEs) [36]. These methods require the latent representations to be discriminative, but the training of the generative model is decoupled from the optimization objectives, imposing challenges for optimization [58]. Instead of navigating the latent space, combinatorial optimization methods search for the desired molecular structures directly in the explicit discrete space with 2) evolutionary algorithms [30] and 3) reinforcement learning (RL) [62] guiding the search. However, most prior methods only focus on a single property, from non-biological properties such as drug-likelihood (QED) [8] and synthetic accessibility (SA) [19], to biological properties that measure the binding energy to a protein target [6]. Multi-objective molecular optimization has recently received wide attention [33, 61, 20]. For example, MARS [61] employs Markov chain Monte Carlo (MCMC) sampling to find novel molecules satisfying several properties. However, most approaches require a notoriously large number of oracle calls to evaluate molecules on the fly [33, 61]. In contrast, we tackle this problem in a sample-efficient manner.
GFlowNet. GFlowNets [6, 7] aim to sample composite objects proportionally to a reward function, instead of maximizing it as usual in RL [56]. GFlowNets are related to the MCMC methods due to the same objective, while amortizing the high cost of sampling (mixing between modes) over training a generative model. GFlowNets have made impressive progress in various applications, such as biological sequence design [28], discrete probabilistic modeling [63], and Bayesian structure learning [15]. While the concept of conditional GFlowNet was originally discussed in Bengio et al. [7], we are the first to study and instantiate this concept for MOO, in parallel with Jain et al. [29]. Compared with them, we delicately design the conditioning mechanism and propose a hindsight-like off-policy strategy that is rarely studied for MOO (in both the RL and GFlowNet literature). Appendix D provides a detailed comparison.
Bayesian optimization for discrete spaces. The application of BO in discrete domains has proliferated in recent years. It is much more challenging to construct surrogate models and optimize acquisition functions in discrete spaces, compared to continuous spaces. One common approach is to convert discrete space into continuous space with generative models [24, 16, 44]. Besides, one can directly define a Gaussian Process (GP) with discrete kernels [46, 53] and solve the acquisition function optimization problem with evolutionary algorithms [35, 57].
Multi-objective Bayesian optimization. BO has been widely used in MOO problems for efficiently optimizing multiple competing expensive black-box functions. Most popular approaches are based on hypervolume improvement [12], random scalarizations [37, 49], and entropy search [27, 5]. While there have been several approaches that take into account parallel evaluations [11, 38] and diversity [38], they are limited to continuous domains.
We address the problem of searching over a discrete chemical space  $\mathcal{X}$  to find molecular graphs  $x\in \mathcal{X}$  that maximize a vector-valued objective  $f(x) = (f_{1}(x),f_{2}(x),\ldots ,f_{M}(x)):\mathcal{X}\to \mathbb{R}^{M}$ , where  $f_{m}$  is a black-box function (also known as the oracle) evaluating a certain property of molecules. Practically, realistic oracles are extremely expensive to evaluate with either high-fidelity simulations or wet-lab experiments. We thus propose to perform optimization in as few oracle evaluations as possible, since the sample efficiency is paramount in such a scenario.
There is typically no single optimal solution to the MOO problem, as different objectives may contradict each other. Consequently, the goal is to recover the Pareto front – the set of Pareto optimal solutions which cannot be improved in any one objective without deteriorating another [18, 45]. In the context of maximization, a solution  $f(x)$  is said to Pareto dominate another solution  $f(x')$  iff  $f_{m}(x) \geq f_{m}(x') \forall m = 1, \dots, M$  and  $\exists m'$  such that  $f_{m'}(x) > f_{m'}(x')$ , and we denote  $f(x) \succ f(x')$ . A solution  $f(x^{*})$  is Pareto optimal if not Pareto dominated by any solution. The Pareto front can be written as  $\mathcal{P}^* = \{f(x^*): \{f(x): f(x) \succ f(x^*)\} = \emptyset\}$ .
The quality of a finite approximate Pareto front  $\mathcal{P}$  is commonly measured by hypervolume [67] - the M-dim Lebesgue measure  $\lambda_{M}$  of the space dominated by  $\mathcal{P}$  and bounded from below by a given reference point  $\boldsymbol{r} \in \mathbb{R}^{M}$ :  $HV(\mathcal{P}, \boldsymbol{r}) = \lambda_{M}(\cup_{i=1}^{| \mathcal{P}|} [\boldsymbol{r}, y_{i}])$ , where  $[\boldsymbol{r}, y_{i}]$  denotes the hyper-rectangle bounded by  $\boldsymbol{r}$  and  $y_{i} = f(x_{i})$ .
3.2 Batch Bayesian optimization
Bayesian optimization (BO) [54] provides a model-based iterative framework for sample-efficient black-box optimization. Given an observed dataset  $\mathcal{D}$ , BO relies on a Bayesian surrogate model  $\mathcal{M}$  to estimate a posterior distribution over the true oracle evaluations. Equipped with the surrogate model, an acquisition function  $a: \mathcal{X} \to \mathbb{R}$  is induced to assign the utility values to candidate objects for deciding which to be evaluated next on the oracle. Compared with the costly oracle, the cheap-to-evaluate acquisition function can be efficiently optimized. We consider the scenario where the oracle is given an evaluation budget of  $N$  rounds with fixed batches of size  $b$ .
To be precise, we have access to a random initial dataset  $\mathcal{D}_0 = \{(x_i^0,y_i^0)\}_{i = 1}^n$ , where  $y_{i}^{0} = f(x_{i}^{0})$  is a true oracle evaluation. In each round  $i\in \{1,\dots ,N\}$ , the acquisition function is maximized to yield a batch of candidates  $\mathcal{B}_i = \{x_j^i\}_{j = 1}^b$  to be evaluated in parallel on the oracle  $y_{j}^{i} = f(x_{j}^{i})$ . The observed dataset  $\mathcal{D}_{i - 1}$  is then augmented for the next round:  $\mathcal{D}_i = \mathcal{D}_{i - 1}\cup \{(x_j^i,y_j^i)\}_{j = 1}^b$ .
In this section, we present the proposed MOBO algorithm based on hypernetwork-based GFlowNet (HN-GFN), shown in Figure 1. Due to the space limitation, we present the detailed algorithm in Appendix A. Our key idea is to extend GFlowNets as the acquisition function optimizer for MOBO, with the objective of sampling a diverse batch of candidates from the approximate Pareto front. To begin, we introduce GFlowNets in the context of molecule design, then describe how GFlowNet can be biased by a preference-conditioned hypernetwork to sample molecules according to various trade-offs between objectives. Next, we propose a hindsight-like off-policy strategy to speed up learning in HN-GFN. Finally, we introduce the surrogate model and acquisition function.
GFlowNets [6] seek to learn a stochastic policy  $\pi$  for sequentially constructing discrete objects  $x\in \mathcal{X}$  with a probability  $\pi (x)\propto R(x)$ , where  $\mathcal{X}$  is a compositional space and  $R:\mathcal{X}\to \mathbb{R}_{\geq 0}$  is a non-negative reward function. The generation process of object  $x\in \mathcal{X}$  can be represented by a sequence of discrete actions  $a\in \mathcal{A}$  that incrementally modify a partially constructed object, which is denoted as state  $s\in S$ . Let the generation process begin at a special initial state  $s_0$  and terminate with a special action indicating that the object is complete ( $s = x\in \mathcal{X}$ ), the construction of an object  $x$  can be defined as a complete trajectory  $\tau = (s_0\rightarrow s_1\rightarrow \dots \rightarrow s_n = x)$ .
Following fragment-based molecule design [6, 61], the molecular graphs 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}$ . There are multiple action sequences leading to the same state, and no fragment removal actions, the space of possible action sequences can thus be denoted by a directed acyclic graph (DAG)  $\mathcal{G} = (\mathcal{S},\mathcal{E})$ , where the edges in  $\mathcal{E}$  are transitions  $s\rightarrow s^{\prime}$  from one state to another. To learn the aforementioned desired policy, Bengio et al. [6] propose to see the DAG structure as a flow network.
Markovian flows. Bengio et al. [7] first define a trajectory flow  $F: \mathcal{T} \to \mathbb{R}_{\geq 0}$  on the set of all complete trajectories  $\mathcal{T}$  to measure the unnormalized density. The edge flow and state flow can then be defined as  $F(s \to s') = \sum_{s \to s' \in \tau} F(\tau)$  and  $F(s) = \sum_{s \in \tau} F(\tau)$ , respectively. The trajectory flow  $F$  determines a probability measure  $P(\tau) = \frac{F(\tau)}{\sum_{\tau \in \mathcal{T}} F(\tau)}$ . If flow  $F$  is Markovian, the forward transition probabilities  $P_F(s'|s) = \frac{F(s \to s')}{F(s)}$ .
Flow matching. A flow is consistent if the following flow consistency equation is satisfied  $\forall s\in S$  ..
$$
F (s) = \sum_ {s ^ {\prime} \in P a _ {\mathcal {G}} (s)} F \left(s ^ {\prime} \rightarrow s\right) = R (s) + \sum_ {s ^ {\prime \prime}: s \in P a _ {\mathcal {G}} \left(s ^ {\prime \prime}\right)} F \left(s \rightarrow s ^ {\prime \prime}\right) \tag {1}
$$
where  $Pa_{\mathcal{G}}(s)$  is a set of parents of  $s$  in  $\mathcal{G}$ . As proved in Bengio et al. [6], if the flow consistency equation is satisfied with  $R(s) = 0$  for non-terminal state  $s$  and  $F(x) = R(x) \geq 0$  for terminal state  $x$ , a policy  $\pi$  defined by the forward transition probability  $\pi(s'|s) = P_F(s'|s)$  samples object  $x$  with a probability  $\pi(x) \propto R(x)$ . GFlowNets propose to approximate the edge flow  $F(s \to s')$  using a neural network  $F_{\theta}(s, s')$  with enough capacity, such that the flow consistency equation is respected at convergence. To achieve this, Bengio et al. [6] define a temporal difference-like [56] learning objective, called flow-matching (FM):
$$
\mathcal {L} _ {\theta} (s, R) = \left(\log \frac {\sum_ {s ^ {\prime} \in P a _ {\mathcal {G}} (s)} F _ {\theta} \left(s ^ {\prime} , s\right)}{R (s) + \sum_ {s ^ {\prime \prime} : s \in P a _ {\mathcal {G}} \left(s ^ {\prime \prime}\right)} F _ {\theta} \left(s , s ^ {\prime \prime}\right)}\right) ^ {2} \tag {2}
$$
One can use any exploratory policy  $\widetilde{\pi}$  with full support to sample training trajectories and obtain the consistent flow  $F_{\theta}(s,s^{\prime})$  by minimizing the FM objective [6]. Consequently, a policy defined by this approximate flow  $\pi_{\theta}(s'|s) = P_{F_{\theta}}(s'|s) = \frac{F_{\theta}(s\rightarrow s')}{F_{\theta}(s)}$  can also sample objects  $x$  with a probability  $\pi_{\theta}(\mathbf{x})$  proportionally to reward  $R(x)$ . Practically, the training trajectories are sampled from a mixture between the current policy  $P_{F_{\theta}}$  and a uniform distribution over allowed actions [6]. We adopt the FM objective in this work because the alternative trajectory balance [43] was also examined but gave a worse performance in early experiments. Note that more advanced objectives such as subtrajectory balance [42] can be employed in future work.
4.2 Hypernetwork-based GFlowNets
Our proposed HN-GFN aims at sampling a diverse batch of candidates from the approximate Pareto front with a unified model. A common approach to MOO is to decompose it into a set of scalar optimization problems with scalarization functions and apply standard single-objective optimization methods to gradually approximate the Pareto front [37, 64]. Here we consider the weighted sum (WS):  $s_{\lambda}(x) = \sum_{i} \lambda_{i} f^{i}(x)$  and Tchebycheff [45]:  $s_{\lambda}(x) = \max_{i} \lambda_{i} f^{i}(x)$ , where  $\lambda = (\lambda_{i}, \dots, \lambda_{M})$  is a preference vector that defines the trade-off between the competing properties.
To support parallel evaluations, one can simultaneously obtain candidates according to different scalarizations. Practically, this approach hardly scales efficiently with the number of objectives for discrete problems. Taking GFlowNet as an example, we need to train multiple GFlowNets separately for each choice of the reward function  $R_{\lambda}(x) = s_{\lambda}(x)$  to cover the objective space:
$$
\theta_ {\lambda} ^ {*} = \underset {\theta} {\arg \min } \mathbb {E} _ {s \in \mathcal {S}} \mathcal {L} _ {\theta} (s, R _ {\lambda}) \tag {3}
$$
Our key motivation is to design a unified GFlowNet to sample candidates according to different reward functions, even ones not seen during training. Instead of defining the reward function with a
fixed preference vector  $\lambda$ , we propose to train a preference-conditioned GFlowNet on the distribution of reward functions  $R_{\lambda}$ , where  $\lambda$  is sampled from a simplex  $S_M = \{\lambda \in \mathbb{R}^m | \sum_i \lambda_i = 1, \lambda_i \geq 0\}$ :
$$
\theta^ {*} = \underset {\theta} {\arg \min } \mathbb {E} _ {\lambda \in S _ {M}} \mathbb {E} _ {s \in \mathcal {S}} \mathcal {L} _ {\theta} (s, R _ {\lambda}) \tag {4}
$$
Remarks. Assuming an infinite model capacity, it is easy to prove that the proposed optimization scheme (Equation (4)) is as powerful as the original one (Equation (3)), since the solutions to both loss functions coincide [17]. Nevertheless, the assumption of infinite capacity is extremely strict and hardly holds, so how to design the conditioning mechanism in practice becomes crucial.
4.2.1 Hypernetwork-based conditioning mechanism
We propose to condition the GFlowNets on the preference vectors via hypernetworks [26]. Hypernetworks are networks that generate the weights of a target network based on inputs. In vanilla GFlowNets, the flow predictor  $F_{\theta}$  is parameterized with the message passing neural network (MPNN) [23] over the graph of molecular fragments, with two prediction heads approximating  $F(s,s')$  and  $F(s)$  based on the node and graph representations, respectively. These two heads are parameterized with multi-layer perceptrons (MLPs). For brevity, we write  $\theta = (\theta_{\mathrm{mpnn}},\theta_{\mathrm{pred}})$ .
One can view the training of HN-GFN as learning an agent to carry out multiple policies that correspond to different goals (reward functions  $R_{\lambda}$ ) defined in the same environment (state space  $S$  and action space  $\mathcal{A}$ ). We thus propose only to condition the weights of prediction heads  $\theta_{\mathrm{pred}}$  with hypernetworks, while sharing the weights of MPNN  $\theta_{\mathrm{mpnn}}$ , leading to more generalizable state representations. More precisely, the hypernetwork  $h(\cdot; \phi)$  takes the preference vector  $\lambda$  as inputs and outputs the weights  $\theta_{\mathrm{pred}} = h(\lambda; \phi)$  of prediction heads. The parameters  $\phi$  of the hypernetwork are optimized like normal parameters. Following Navon et al. [47], we parametrize  $h$  using an MLP with multiple heads, each generating weights for different layers of the target network.
4.2.2 As the acquisition function optimizer
Training. At each iteration, we first sample a preference vector  $\lambda$  from a Dirichlet distribution  $\mathrm{Dir}(\alpha)$ . Then the HN-GFN is trained with the reward function  $R_{\lambda}(x) = a(\mu(s_{\lambda}(x)), \sigma(s_{\lambda}(x)); \mathcal{M})$ , where  $\mu$  and  $\sigma$  are the posterior mean and standard deviation estimated by  $\mathcal{M}$ . In principle, we retrain HN-GFN after every round. We also tried only initializing the parameters of the hypernetwork and found it brings similar performance and is more efficient.