Sample-efficient Multi-objective Molecular Optimization with GFlowNets

Summary

Molecule generation involves the optimization of multiple potentially competing objectives simultaneously. As evaluating these objectives can be a time-consuming and costly task, sample efficiency is paramount. This work proposes a multi-objective Bayesian optimization approach leveraging GFlowNets to tackle this problem. The authors propose a preference-based decomposition and a hyper-network-based parameterization for incorporating preferences and conditioning on those. They also employ a hindsight-experience replay strategy to utilize offline data during learning. The hyper-network-based GFlowNet is then used as an acquisition function optimizer to find a diverse set of molecules on the empirical Pareto front.

Strengths

- The paper covers an important and impactful topic of multi-objective optimization in molecule design. - The experimental results, in principle, are strong. - It is an original combination of existing techniques. - Many benchmarks are compared. - Ablation study

Weaknesses

- It is not directly evident to me how this work is different from Pareto GFlowNet that was proposed in [GFlowNet Foundations](https://arxiv.org/pdf/2111.09266.pdf) - Some parts remain somewhat unclear to me. How is the hyper-network trained? Is it stable, does it do what it is expected to do? - I am somewhat missing the link between the surrogate model and HN-GFN from what is described in the main paper. Why do we actually need the surrogate model? Can't the GFlownet propose diverse candidates by itself? - Perhaps the work would benefit from a clearer description and/or toy tasks to clear this up.

Questions

- It would be nice if the authors could dedicate some more information about the test problem. What is the search space of the molecules, what atom types, how are impossible molecular structures avoided, etc? Are the building blocks created in such a way that they always can be connected to each other without violations? - Why a batch size of 100 and 8 rounds? In what setting would a practitioner be able to perform 100 wet lab experiments simultaneously? Perhaps with a simulator, this would be possible. But wouldn't the surrogate models (at least from the BO qEHVI and qParEGO) benefit from lower batch sizes and more iterations? This could change the benchmarks significantly.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

The authors have addressed limitations.

Summary

A GFlowNet for molecular optimization conditioned on the preference weights of multiple objectives is proposed. To be precise, the model is trained to sample molecules from a target space with probability proportional to some combination of reward functions (weighted sum or Chebyshev scalarization) and trained with varying weights of the rewards, with the GFlowNet conditioning achieved by a hypernetwork parametrization. The system is tested in both single-step and Bayesian optimization / active learning settings.

Strengths

- The writing is generally clear, problem is introduced well and placed in context. - Good results on the tasks studied relative to the chosen baselines - Two interesting GFlowNet contributions not specific to molecules: - Study of conditioning algorithms (FiLM and Concat variants), with the hypernetwork-based approach found to perform best. - The use of a conditional replay buffer, which is, as far as I know, new in GFlowNet literature (a replay buffer was used with GFlowNets in [arXiv:2202.13903, UAI 2022] but in a vanilla RL manner). - Both are illustrated well with ablation studies. - The code is a helpful addition that aids in understanding the details of what was done.

Weaknesses

Meta-weakness: Due to my closely following the literature, I was aware of an earlier submission of this paper to ICLR 2023. The present submission is almost identical to ICLR version (post-rebuttal) and therefore cannot incorporate the feedback of the reviewers. That said, to avoid biasing myself here, I did not read the ICLR reviews and this assessment is entirely my own. Comparison with prior work: - There is the paper "Multi-objective GFlowNets" [MO; arXiv:2210.12765, ICML 2023], which also studies GFlowNets for multiobjective Bayesian optimization by conditioning on scalarization weights. Can you comment on the differences in problem setting or method with that paper? - The attached code includes the implementations of two baselines, in `code/generator/gfn.py`, but they are not included in the paper: - MOReinforce, which is a natural candidate for an RL-based baseline that can readily share the parametrization of the agent with the GFlowNet. - GFlowNet with trajectory balance [TB; arXiv:2201.13259, NeurIPS 2022]. Trajectory balance is the objective used in nearly all work on GFlowNets for molecular optimization after [5], but for some reason it is not evaluated or mentioned in the text. In particular, it is evaluated on molecule synthesis in both [TB] and [MO]. Small issue in presentation/math: There are some inconsistencies in definitions related to GFlowNets. Equation (1) writes $R(s)$ for the reward received when terminating at $s$, as in [5]. However, the text immediately following it talks about terminal and nonterminal states (with nonterminals having reward 0), and the text above says there is a special termination action that leads to terminal states. This is more consistent with the convention from [TB], where terminal states have no children and are the only states with nonzero reward. This is not a major bug as the two conventions are equivalent, but it would be good to be careful.

Questions

Please see above.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

yes

Summary

This paper proposes the use of GFlotNets to address the problem of sample-efficient multi-objective molecular optimization, an important problem in various scientific discovery application - such as materials design and drug discovery. The key idea proposed in this work is to leverage hypernetwork-based GFlowNets - referred to as HN-GFN - to optimize the acquisition function for multi-objective Bayesian optimization (MOBO). The goal is to enable efficient sampling of a diverse high-quality batch of molecular candidates from an approximate Pareto front.

Strengths

The proposed hyper network-based GFlowNets as a MOBO acquisition function optimizer provides a novel and intuitive way of using GFlowNets for sampling novel molecular candidates, where a "unified" GFlowNet is trained that considers the distribution of different reward functions that correspond to different preference vectors instead of relying on a single GFlowNet for a fixed preference vector. This may allow the resulting model, HN-GFN, to naturally explore the various trade-offs between multiple objectives that may compete with one another by adapting to varying the input preference vector. The proposed method builds on the recently proposed and widely popular GFlowNets and extends its flexibility for multi-objective molecular optimization by incorporating a hypernetwork-based approach. Additionally, the adoption of a hindsight-like off-policy strategy is proposed to improve the learning efficiency, and ultimately, the multi-objective molecular optimization performance. Overall, the paper is well-organized and written in a clear manner, the proposed method is well-motivated and novel, and the performance evaluation results demonstrate the potential advantages of the proposed HN-GFN.

Weaknesses

Although the batch size may significantly affect the overall computational cost as well as the optimization performance, there is no discussion on the impact of selecting a specific batch size nor any empirical evaluation based on different batch sizes. While the evaluation results provide some preliminary evidence of the potential advantages of HN-GFN, the evaluations in the current study are limited to (virtually) a single problem: i.e., inhibition of GSK3β + JNK3 (with potential additional considerations for synthesizability and drug-likeness). Additional examples are needed to more convincingly demonstrate the general applicability (and merits) of the proposed HN-GFN. I suggest providing further evaluation results based on other benchmark problems often used for evaluating other generative models (e.g., JT-VAE, HierVAE, etc.) In Figure 2, only trends for JNK3 are shown, but it would be helpful to show the optimization trends for GNK3β as well for completeness.

Questions

Please see the questions and suggestions in the weaknesses section.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

The conclusion section includes a very brief discussion of a limitation of the current method and suggests directions for future work. The broader implication of the work is not explicitly discussed in the manuscript.

Summary

This work addresses molecular design with GFLOWNET, an algorithm learning a sampling policy $\pi$ proportional to the reward function, i.e. where $\pi(x) \propto R(x)$. In particular, this work tackles the essential yet under-adressed multiobjective optimization setting. They do so using a hypernetwork (conditioned on a preference vector) providing weights for the sampler.

Strengths

In my opinion the presentation is clear and the motivation behind this work are nicely presented. In addition, the paper is well written and easy to follow. As a side note, its presentation of the GFN is interesting and worth the reminder to ease the understanding of the paper. The approach is both original and sound, boosting its potential use in practical application. Finally, the presented results are conclusive and seem to validate the intuition behind the method.

Weaknesses

The main weaknesses in the paper is in my opinion in the lack of diverse experiments or explanation about them. See questions.

Questions

1- A part of the protocole is unclear to me: are the generated molecules evaluated on the objective using an actual oracle or using the learned surrogate function ? 2 - How are the objectives computed ? Is it “physics” based or does it rely on learning a surrogate model ? Then, how are the generated molecules evaluated ? 3 - What is the length of the generated molecules ? Is it a hyperparameter of the model ? 4 - How sensitive are the results to the original dataset $D_0$? In my personal experience, Gflownets tend to be rather unstable or difficult to train due to the sparse feedback (and depends on the sequence size) and 200 initial molecules seem low to be able to learn a robust policy. 5 - Does the proposed Hindsight-like off policy strategy amounts to the proposition of [1], i.e. to hybrid the learning between generated molecules vs already generated molecules ? 6 - What are the variable predicted by the MPNN ? 7 - Can the authors discuss their implementation choice regarding the gflow network. For instance [1] use a MLP as sampler. [1]: Jain et al, Biological sequence designs with Gflownets.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

A limitation of the paper is the lack of diversity in the experimental settings. For instance, the paper Gflownet for biological sequence design presents optimization results for several benchmark datasets such as GFP. Can the authors comment on the ability of their method to adapt to such an ‘off-line’ setting (which is an essential in biological sequence design) ?

Response to rebuttal

Understood about baselines. I hope that this discussion will be added to the revised version of the paper. I am also aware of the NaN-in-logits issue with TB. I believe that past work reported the performance at the last state before NaN, but expect this would happen earlier with a higher reward exponent. Overall, I am satisfied with the response; having read the other reviews and responses, I maintain my score. Thank you again for the clarifications!

Response to the authors

I thank the authors for their valuable response and additional technical details that significantly improve my understanding of their work. I believe the discussion regarding the choice of the datasets could be included in the main paper. I raise my score to 5.