Variational Search Distributions

Authors: Daniel M. Steinberg, Rafael Oliveira, Cheng Soon Ong, Edwin V. Bonilla
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
Keywords: black box optimization, Bayesian optimization, variational inference, generative models, level set estimation, biological sequence design, protein engineering, fine tuning

Abstract:
We develop VSD, a method for conditioning a generative model of discrete, combinatorial designs on a rare desired class by efficiently evaluating a black-box (e.g. experiment, simulation) in a batch sequential manner. We call this task active generation; we formalize active generation's requirements and desiderata, and formulate a solution via variational inference. VSD uses off-the-shelf gradient based optimization routines, can learn powerful generative models for desirable designs, and can take advantage of scalable predictive models. We derive asymptotic convergence rates for learning the true conditional generative distribution of designs with certain configurations of our method. After illustrating the generative model on images, we empirically demonstrate that VSD can outperform existing baseline methods on a set of real sequence-design problems in various protein and DNA/RNA engineering tasks.

Decision: accept (poster)
Reviewer scores: 8, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730293193054

thread_type:
official_review

<thread_object>
creator: reviewer_9cf6

<content>
summary:
The authors present a novel variational method for learning to sample from rarely observed events, aiming to minimize a distance between the distribution of interest, namely $p(x∣y>t)$, and its parametric variational counterpart $q(x|\phi)$. The problem is reformulated to leverage the “whole dataset,” not just rarely observed events, and is expressed as Equation (5), which comprises two terms: $log p (y>t∣x)$ and the negative KL divergence between$q(x|\phi)$ and $p(x)$. The authors' final proposal is to estimate $p(y>t∣x)$ using a parametric function instead of a simple PI estimate. The variational distribution is optimized by a REINFORCE gradient estimator.

strengths:
The paper is clear, well-written, and aligns with well-established benchmarks in the field, such as CBAS (Brookes et al.). 
The model is supported by convergence analysis and an extensive set of well-handled experiments.

weaknesses:
While the model description is clear, the model comprise a parametric distribution $p(x|D_0)$ which might be the biggest model shortcoming originating from the model own formulation. 

Its major impact is that it reweights the gradient estimates of $q(x|\phi)$. Intuitively, how would that compare simply to the iterative strategy of Cbas ?

questions:
1.Since your algorithm heavily relies on another model ($p(x | D)$), I would be highly interested in better understanding the influence of a good prior on your variational distribution.
2. Regarding the GFP experiments, do you sample already existing sequences ? What is the influence of the relative poor performance of the oracle on ood data on the interpretation of the results  ?
3. How can you explain that only a very simple prior such as a mean field performs on average better ?  It seems quite logical for GFP for instance where a wild type exists, however it is less intuitive for datasets without wild type.

Typo: the recall and precision have the same expression.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730406142982

thread_type:
official_review

<thread_object>
creator: reviewer_z76b

<content>
summary:
The paper develops the variational search distribution method to solve the active search problem in biological design. VSD estimates the super level-set distribution in a sequential manner by generating batches of data points for sequential experiments. Empirical results on optimizing protein fitness in several datasets showcase the effectiveness of VSD.

strengths:
- The paper formulates the batch active search problem in the variational inference framework and provides theoretical guarantees to the learned distribution based on the sequentially attained data.
- Experimental results on real-world biological datasets demonstrate the practical use of the algorithm and its effectiveness to solve the problem.

weaknesses:
- The precision of VSD and most other methods is decreasing with more rounds in TrpB and TFBIND8 datasets while the recall values are in general low. However, an ideal method should achieve a better estimation of the ground truth super level-set distribution as more samples are collected. This may be due to the initial training set size being too large or the fitness landscape being easy to model. How do the models perform with a smaller initial training set size?
- How is VSD compared with the simple and commonly used directed evolution method?

questions:
- How robust are the results to the selection of the threshold $\tau$ and the batch size $B$?
- While the reviewer is not familiar with the field, could the authors give some intuitions about the difference between VSD and active learning approaches like Bayesian optimization, and why VSD is better?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730889577962

thread_type:
official_review

<thread_object>
creator: reviewer_1vhr

<content>
summary:
The authors propose a black box varoiational inference approachfor discrete designs generation. The authors derive asymptotic convergence rates for learning the true conditional generative distribution of designs. Compelling results on high dimensional sequence-design problems are demonstrated.

strengths:
* The problem is important as it has applications in pharmaceutical drugs/enzyme design.
* The paper paper is well written and the method is sound
* Experimental results on high dimensional datasets demonstrate superiority of the approach

weaknesses:
* The method lacks novelty, it's based on putting together blocks that have already been proposed in the litterature
* The paper clarity can be improved with an overview plot of the method

questions:
* What's 'x' in the title of Figure 1?
* What are the limitations of this approach?
* How is diversity within a batch enforced? 
* The reverse KLD is known to result in mode collapse. Why wasn't this an issue?
* Which variation reduction method did you use for the gradient estimator?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730959428074

thread_type:
official_review

<thread_object>
creator: reviewer_xamp

<content>
summary:
This paper casts sequential black-box optimization as a variational inference (i.e. amortized optimization) problem, and uses this perspective to unify a collection of different black-box optimization algorithms under a common theoretical framework and presents some proof of concept results on easy sequence optimization tasks.

#### UPDATE - 12/02/2024 ####
After extended discussion and an exceptionally thorough response from the authors, my concerns have been addressed and I am recommending acceptance. I believe this paper presents a very nice conceptual view of the topic of active generation and the empirical results raise interesting questions. I encourage my fellow reviewers to review the updated manuscript and re-evaluate their scores. I have left my original review unaltered for any interested readers.

strengths:
This paper demonstrates a clarity of thought and composition that is commendable, I particularly enjoyed the related work section.

Likewise I do not have any major concerns regarding the technical soundness of the results presented.

As a good conceptual introduction to the topic, I think this draft could be useful to researchers new to the topic with some revisions.

weaknesses:
I have two general impressions of this paper. 

First, it seems like the authors have not really chosen a direction for the paper. There are at least three different directions here, A) a unifying view of sequential black box optimization algorithms, B) a practical algorithm for sequential BBO, and C) theoretical analysis of convergence rates of a particular sequential BBO algorithm under strong assumptions. I would suggest you pick no more than two directions, preferably one. I actually think this particular subfield could really benefit from a more holistic perspective of the work that has been done, as I constantly see minor variations of these algorithms in my social media feed and review stack with no apparent awareness of the relationships between them. From what I can tell from this draft, it seems that A and C likely play more to your strengths.


Second, the authors seem blissfully unaware of a substantial body of work on this topic. To be quite candid, the paper reads like it was written circa September 2021. This is not mere rhetoric. The most recent baseline the authors consider was published at ICML 2021. It is also odd that two of the baselines you did include, DbAS and CbAS, are not even designed for the sequential setting. As a very active researcher in this exact area, I struggle to understand who this paper is for and how the authors pictured their place in the broader dialogue on this topic. I am sure you worked very hard on this paper and I commend your effort, but I honestly believe the best advice I can give you is to talk to more people working on this topic, preferably from outside your immediate academic circle. While it is difficult to hear this feedback, one of the functions of peer review is to reveal "unknown unknowns". I want to be sure this review is constructive, so I will provide some key references if you are serious about diving into this topic. You should also consider making use of tools like [Connected Papers](https://www.connectedpapers.com/) to improve your literature review process and avoid this situation in the future. 

You can start with [A survey and benchmark of high-dimensional Bayesian optimization of discrete sequences](https://arxiv.org/abs/2406.04739). This work is the most up-to-date complete survey on the topic I have seen, and the benchmarking rigor is notably good. This paper is associated with two repositories, [poli](https://github.com/MachineLearningLifeScience/poli) and [poli-baselines](https://github.com/MachineLearningLifeScience/poli-baselines). The former contains a suite of test functions that are much more up to date than the combinatorially complete landscapes considered in this paper, and the latter contains a suite of baseline solvers. You may even want to consider contributing your method as a solver to poli-baselines at some point.

Some key axes of variation to consider: 

How is the optimization problem solved? Most fall into one of three categories, directed evolution (which you seem to be familiar with based on your inclusion of AdaLead and PEX), generative search with explicit guidance, e.g. [2, 3, 4, 5, 6], and generative search with implicit guidance [7, 8], which can also be seen as a kind of amortized search. I could cite more papers but I believe I have made my point. Algorithms also differ in their handling of constraints, and their approach to managing the feedback covariate shift induced by online active data collection by an agent. 

In particular I will draw your attention to [a tutorial for LaMBO-2](https://github.com/prescient-design/cortex/blob/main/tutorials/4_guided_diffusion.ipynb) if you want to start considering more up to date baselines, however I would recommend using the solver interface provided in poli-baselines for actual experiments. You may also be interested in Ehrlich functions if you would like a convenient test function that is much more difficult to solve than small combinatorially complete landscapes but still easy to work with [9]. Ehrlich functions are available in [a small standalone package](https://github.com/prescient-design/holo-bench) or [as part of the poli package](https://machinelearninglifescience.github.io/poli-docs/using_poli/objective_repository/ehrlich_functions.html).

While I'm sure this is not the outcome you hoped for, science is a dialogue, and good science requires awareness of what is happening outside your academic niche. Hopefully my feedback is clear and actionable enough to benefit this work and your progression as a scientist.

References

- [1] González-Duque, M., Michael, R., Bartels, S., Zainchkovskyy, Y., Hauberg, S., & Boomsma, W. (2024). A survey and benchmark of high-dimensional Bayesian optimization of discrete sequences. arXiv preprint arXiv:2406.04739.
- [2] Tripp, A., Daxberger, E., & Hernández-Lobato, J. M. (2020). Sample-efficient optimization in the latent space of deep generative models via weighted retraining. Advances in Neural Information Processing Systems, 33, 11259-11272.
- [3] Stanton, S., Maddox, W., Gruver, N., Maffettone, P., Delaney, E., Greenside, P., & Wilson, A. G. (2022, June). Accelerating bayesian optimization for biological sequence design with denoising autoencoders. In International Conference on Machine Learning (pp. 20459-20478). PMLR.
- [4] Gruver, N., Stanton, S., Frey, N., Rudner, T. G., Hotzel, I., Lafrance-Vanasse, J., ... & Wilson, A. G. (2023). Protein design with guided discrete diffusion. Advances in neural information processing systems, 36.
- [5] Maus, N., Jones, H., Moore, J., Kusner, M. J., Bradshaw, J., & Gardner, J. (2022). Local latent space bayesian optimization over structured inputs. Advances in neural information processing systems, 35, 34505-34518.
- [6] Maus, N., Wu, K., Eriksson, D., & Gardner, J. (2022). Discovering many diverse solutions with bayesian optimization. arXiv preprint arXiv:2210.10953.
- [7] Tagasovska, N., Gligorijević, V., Cho, K., & Loukas, A. (2024). Implicitly Guided Design with PropEn: Match your Data to Follow the Gradient. arXiv preprint arXiv:2405.18075.
- [8] Chen, A., Stanton, S. D., Alberstein, R. G., Watkins, A. M., Bonneau, R., Gligorijevi, V., ... & Frey, N. C. (2024). LLMs are Highly-Constrained Biophysical Sequence Optimizers. arXiv preprint arXiv:2410.22296.
- [9] Stanton, S., Alberstein, R., Frey, N., Watkins, A., & Cho, K. (2024). Closed-Form Test Functions for Biophysical Sequence Optimization Algorithms. arXiv preprint arXiv:2407.00236.

questions:
The following questions are sincere:

- Who is the audience for this paper? 

- What questions is this paper answering?

- What does the variational inference framing get us in the end? Access to a set of tools for theoretical analysis?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732258940708

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We thank the reviewers for their constructive comments and criticisms of our work, and for the feedback that has made this paper stronger.

The major concerns and criticisms of our work come from reviewer xamp:

1. We have neglected to cite major works in the related setting of latent space optimization (LSO), and where VSD sits in relation to this work.
2. We have neglected to compare to LSO on difficult BBO baselines.

We discuss these major criticisms below, but we have also updated our manuscript to (1) include LSO in our related work section and Table 2 and (2) we have benchmarked VSD against LaMBO-2 on the suggested Ehrlich function BBO task, in which VSD gets the best performance. The latter experiment is currently in appendix C.4, but we are working on improving and expanding it, and we aim to move it to the main paper by the end of this discussion period.

Update: also based on xamp's feedback -- we have slightly reframed the problem we are solving as "active generation" (as distinct from, but related to active learning, active search and BBO), and updated the introduction and problem formulations to clarify this.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper considers an active search problem and presents a method to generate new designs of rare desired class under budget constraints. The target problem is of discrete and combinatorial nature. The proposed algorithm is based on variational inference. Theoretical analysis is conducted to understand its performance. All reviewers agree the paper contains interesting and potentially impactful results. One major concern is the lacking of comparison with strong baselines. One reviewer also points out the paper misses some important related work. Overall, this is a nice contribution to active search algorithms.