Variational Search Distributions

Venue: NeurIPS_cc_2024_Workshop_BDU
Published: 2024-08-30 02:31:21

Abstract: We develop variational search distributions (VSD), a method for finding designs of a rare desired class in a batch sequential manner with a set experimental budget. We formalise the requirements and desiderata for this problem and formulate a solution via variational inference that fulfil these. In particular, VSD uses off-the-shelf gradient based optimisation routines, and can take advantage of scalable predictive models. We show that VSD can outperform existing baseline methods on a set of real sequence-design problems in various biological systems.

Main paper: 
 introduction: 
 We consider a variant of the active search problem [10, 13, 35] , where we wish to find as many members (designs) of a rare desired class in a batch sequential manner with a set experimental budget. Examples of this are compounds that could be useful pharmaceutical drugs, or highly active enzymes for catalysing chemical reactions. We assume the design space is discrete or partially discrete, high-dimensional, and practically innumerable. For example, the number possible configurations of a single protein is 20 O(100) [see, e.g. , 25] . We are interested in this objective for a variety of reasons. For example, we may wish to study the properties of the "fitness landscape" [21] to gain a better scientific understanding of a phenomenon such as natural evolution. Additionally, in other problems, we may not be able to completely specify the constraints and objectives of a task, but we would like to characterise the space of feasible designs. For example, we want enzymes that are able to degrade plastics in an industrial setting, but we still do not know the exact conditions (e.g. temperature, pH), some of which may be anti-correlated with enzyme catalytic activity.Assuming we can take advantage of a prior distribution over designs, we formulate the search problem as inferring the posterior distribution over rare, desirable designs. Importantly, this posterior can be used for generating new designs. Specifically, we use (black-box) variational inference (VI) [22] , and so refer to our method as variational search distributions (VSD). Our major contributions are; (1) we formulate the batch active search objective over an innumerable discrete design space as an instance of variational inference, (2) we present a relatively simple and modular algorithm, VSD, which solves this objective, and (3) we show that VSD satisfies well-defined requirements and desiderata specific to our problem. For example, it uses off-the-shelf gradient based optimisation routines, and can take advantage of scalable predictive models. In our experiments we show that VSD can outperform existing baseline methods on a set of real applications. 
problem formulation: 
 We are given a design space X , which can be discrete or mixed discrete-continuous and high dimensional, and where for each instance that we choose x ∈ X , we measure some corresponding property of interest (so-called fitness) y ∈ R. For example, in our motivating application of DNA/RNA or protein sequences (henceforth referred to as just sequences), X = A m where A is the sequence vocabulary (e.g., amino acid labels, |A| = 20) and m is the length of the sequence. However, we do not limit the application of our method to sequences. Using this framing, a real world experiment (for example, measuring the activity of an enzyme) can be modelled as an unknown relationship, y = f (x) + ϵ, for some black-box function (e.g. the experiment), f , and measurement error ϵ ∈ Rand E p(ϵ) [ϵ] = 0.Instead of wanting to model the whole space, we are only interested in a set of events which we choose based on fitness. For instance, we may be interested in the fittest measurable design; all designs which above a minimum level of feasibility (e.g. a wild-type sequence), τ ∈ R; the probability distribution of these feasible designs; or the shape of the black-box function for these feasible designs,EQUATIONrespectively. These tasks are depicted in Figure 1 . Our primary focus in this work is to estimate the density p(x|y > τ ), in a sequential manner. We assume that we have access to an existing dataset, Estimating the conditional density above is computationally and statistically challenging and, therefore, we cast this as a variational inference problem. As we shall see later, our solution allows us to satisfy the requirements and additional desiderata for our problem, as given in Table 1 .D N = {(y n , x n ) : n ∈ I N }Requirements (R) Desiderata (D) (R1) Rare: Identify feasible solutions, y > τ , which are rare events in X (D1) Guaranteed: guaranteed to estimate the true posterior p(x|y > τ ) given enough data. (R2) Sequential: generate candidates, sequentially and non-myopically (D2) Gradient: use established gradient-based optimisation strategies for candidate search (R3) Discrete: search over large discrete input spaces, x ∈ X = A m (D3) Generative: learn generative models x (s) ∼ q(x) for fit sequences (R4) Batch: generate diverse batches of O(1000) candidates per round (D4) Scale: Use scalable predictive models, critical in high-throughput experiments. Table 1 : Requirements (R) and additional desiderata (D) in developing our method. 
variational search distributions: 
 We cast the estimation of p(x|y > τ ) as a sequential optimisation problem using variational inference. To do this we optimise the well-known evidence lower bound (ELBO), L ELBO (ϕ) = E q(x|ϕ) [log Pr(y > τ |x)] -D KL [q(x|ϕ)∥p(x)] , where we assume access to a prior distribution over the space of designs, p(x). Furthermore, we estimate log Pr(y > τ |x, D N ) using a surrogate model by recognising an equivalence between this distribution and the probability of improvement (PI) acquisition function from Bayesian optimisation (BO) [16] ,EQUATIONHere 1 : R × R → {0, 1} is the indicator function and p(y|x, D N ) is typically estimated using the posterior predictive distribution of a Gaussian process (GP). So Pr(y > τ |x, D N ) = Ψ((µ(x)τ )/σ(x)), where Ψ(•) is a cumulative standard normal distribution, and µ(x), σ(x) are the posterior predictive mean and standard deviation of the GP. We can now rewrite the ELBO as,L ELBO (ϕ, τ, D N ) = E q(x|ϕ) [log α P I (x, D N , τ )] -D KL [q(x|ϕ)∥p(x)] .(3) We refer to our method that optimises the objective in Equation 3 as VSD, as we are using the variational posterior distribution as a means of searching the space of fit sequences, satisfying (R1), (R2) and (R4). Concretely, we draw a set of sample candidates from our search distribution, (D3), each round,EQUATIONIn general, because of the discrete combinatorial nature of our problem, we cannot readily use the re-parameterization trick to estimate the gradients of the ELBO above. Instead, we use of the score function gradient estimator [19] with standard gradient descent methods (D2),EQUATIONwhere we use Monte-Carlo sampling to approximate this expectation with a suitable variance reduction scheme, such as using a control variate or baseline. We find that the exponentially smoothed average of the ELBO works well in practice, and is the same strategy employed in Daulton et al. [8] . Effectively, VSD implements black-box variational inference [22] for parameter estimation, and despite the high-dimensional nature of X , we find we only need O(1000) samples to estimate the required expectations for ELBO optimisation on problems with m = O(100), satisfying (R3).Class probability estimation: So far our method indirectly computes the PI acquisition function by transforming the predictions of a GP surrogate model, p(y|x, D N ), as in Equation 2. Instead we may choose to follow the reasoning used by Bayesian optimisation by density-ratio estimation (BORE)in [32, 20, 27] and directly estimate the quantity we care about, Pr(y > τ |x, D N ). We do this with class probability estimation (CPE) using p(z=1|x, D N ) ≈ π θ (x), where 1[y > τ ] =: z ∈ {0, 1}, and π θ : X → [0, 1]. We can recover the class probability estimates using a proper scoring rule [11] such as Brier score or log-loss on training data,EQUATIONThe VSD objective using CPE becomes,EQUATIONinto which we plug θ * t = argmax θ L CPE (θ, D z N ). Using a CPE also opens up the choice of estimators that are more scalable than a GP surrogate, satisfying our last desiderata (D4). This may be crucial if we choose to run more than a few rounds of experiments with B = O(1000). Additionally, since VSD is a black box method, we can choose to use CPEs that are non-differentiable, such as decision tree ensembles.The complete VSD algorithm is given in Algorithm 1 (Appendix D), in which we have allowed for a threshold function, τ t = f τ ({y : y ∈ D N }, γ t ). This function can be used to modify the threshold each round, e.g. following [32] , an empirical quantile function τ t = Qy (γ t ) where γ t ∈ (0, 1), or a constant τ in the case of estimating the density of the super-level set.Theoretical analysis and related work: We show in Appendix A that the VSD objective, in fact, generalises the BO objective, providing a lower bound that is tight iff the posterior collapses to a Dirac delta distribution. In Appendix B we provide a formulation that generalises several related optimisation algorithms (and VSD) including Bayesian optimisation with probabilistic reparameterisation (BOPR) [8] , design by adaptive sampling (DbAS) [6] , conditioning by adaptive sampling (CbAS) [5] , BORE [32] and amortised BO [31] . The key takeaway is that, as seen in Table 3 , VSD satisfies all the requirements and desiderata for our problem.  
experiments & conclusion: 
 We evaluate our method, VSD, on a number of real-world sequence design tasks involving various Figure 2 shows the results for the fitness landscape experiments (and the BBO experiments can be found in Appendix E). VSD is clearly the best performing method for the DHFR and TrpB fitness landscape tasks, with the related method CbAS also performing well. For the TFBIND8 task, AdaLead and proximal exploration (PEX) seem to outperform all other methods until the last round. We have consistently found these evolutionary-search based methods to be highly effective on lower-dimensional problems, however we consistently observe their performance degrading as the dimension of the problem increases. We suspect this is a direct consequence of their random mutation strategies being suited to exploration in low dimensions, but less efficient in higher dimensions compared to the learned generative models employed by VSD, CbAS, and DbAS. Our version of BORE (which is just the expected log-likelihood component of Equation 7) performs badly in all cases, and this is a direct consequence of its proposal distribution collapsing to a Kronecker delta centred on x * t . In a non-batch setting, this behaviour is not problematic, but shows how crucial the Kullback-Liebler (KL) divergence regularisation of VSD is in this batch setting. 
 
 Appendix: We have presented variational search distributions (VSD), a method for finding as many designs of a rare class sequentially under some experimental constraints. VSD is underpinned by variational inference, wich allows it to satisfy several critical requirements and important desiderata specific to this problem, including batch generation and learning generative models for fit sequences. We have showcased the benefits of this method empirically on a set of sequential-design biological problems and intend to explore more flexible variational distributions in future work.[ A natural question is how VSD relates to the BO objective [9, Ch.7],x * t = argmaxEQUATIONFirstly, we can see that the expected log-likelihood of term of Equation 3 lower-bounds this quantity, Proposition 1. For a parametric model, q(x|ϕ), given ϕ ∈ Φ ⊆ R m and q ∈ P :EQUATIONand the bound becomes tight as q(x|ϕ * t ) → δ(x * t ), a Kronecker delta function at the maximiser x * t .Taking the argmax of the RHS will result in the proposal distribution collapsing to a delta distribution at x * t for appropriate choice of q(x|ϕ). The intuition for Equation 9 is that the expected value of a random variable is less than or equal to its maximum. The proof of this is in [8, 29] . Extending this lower bound, Proposition 2. For a divergence D : P(X ) × P(X ) → [0, ∞), and prior p(x|D N ) where p ∈ P :EQUATIONWe can see that this bound is trivially true given the range of divergences, and this covers VSD as a special case. However, this bound is tight iff p(x|D N ) → δ(x * t ) with an appropriate choice of q(x|ϕ).VSD is one of many methods that makes use of a generalisation of the bound on which Proposition 1 is based,EQUATION)This bound is useful for black-box optimisation of f -this may be a physical experiment directly, or a surrogate for which gradients with respect to x are not defined (as in the case of VSD). Other well known methods that make use of this bound are Evolution strategies (ES) and natural evolution strategies (NES) [36], variational optimisation (VO) [29, 4] , estimation of distribution algorithms (EDA) [17] , and BOPR [8] . They variously make use of maximum likelihood estimation of q(x|ϕ),or the score function gradient estimator -REINFORCE [37] as it is known in the reinforcement learning literature.Algorithms that make an explicit sequential modification to Equation 11 include DbAS [6] ,CbAS [5] and amortised BO [31] . Where in the case of DbAS and CbAS the expectation is taken with respect to q(x|ϕ t-1 ) in order to stop the collapse of this distribution to a point mass. Amortised BO parameterises the proposal distribution as q(x t |x t-1 , ϕ).Method w(x) ϕ ′ ζ Fixed x (s) ∼ q(x|ϕ ′ )? VSD log π θ * (x) + log p(x) -log q(x|ϕ) ϕ - No CbAS π θ * (x)p(x|DN )/q(x|ϕ * t-1 ) ϕ * t-1 - Yes DbAS π θ * (x) ϕ * t-1 - Yes BORE * π θ * (x) ϕ - No BOPR α(x, DN ) ϕ - No Amortised BO α(x, DN ) -α(xt-1, DN ) ϕ xt-1 NoTable 2 : How related methods can be adapted from Equation 12. ζ = -means ζ is an optional input or parameter for this model.We can take a unifying view of these algorithms, as well as many others by recognising the general gradient estimator,EQUATIONwhere we give each component in Table 2 .In Table 2 BORE * has been adapted to discrete X by using the score function gradient estimator.CbAS and DbAS have been adapted here to use a CPE -though were originally derived as using the equivalent of a PI acquisition function. CbAS and DbAS use fixed samples x (s) from q(x|ϕ * t-1 ) for approximating the expectation, and then optimise ϕ using a maximum-likelihood style procedure.Whereas the other methods adapt their samples of this expectation during optimisation by using the score function gradient estimator.Conceptually one of the closest methods to VSD is CbAS, which can be viewed as optimising the forward KL divergence, D KL [p(x|y > τ )∥q(x|ϕ)] in a sequential fashion using cross entropy estimation [24] with importance weights. Batch-BORE [20] also optimises the reverse KL divergence and uses CPE, but uses Stein variational inference [18] for direct optimisation of continuous batch candidates, {x b : b ∈ I B }. Other methods that have been applied to batch biological sequence optimisation tasks include finite horizon methods such as GFlowNets [12] , and the reinforcement learning based DynaPPO [2] . Heuristic stochastic search methods such as AdaLead [26] and PEX [23] have also demonstrated strong empirical performance on these tasks. We compare the properties of some of the most relevant methods to our problem in Table 3 . m o r t i s a t i o n BOPR [8] ✓ ✗ ✓ ✓ - ✗ ✗ ✗ ✓ - BORE [32] - ✗ ✓ ✓ - ✗ ✗ ✓ ✗ - Batch BORE [20] ✗ ✓ ✓ ✓ ✓ - ✓ ✓ ✗ ✓ DbAS [6] ✓ ✓ ✓ ✓ ✓ ✗ ✗ ✓ ✗ ✓ CbAS [5] ✓ ✓ ✓ ✓ ✓ ✓ ✓ ✓ ✗ ✓ Amortized BO [31] ✓ ✓ ✓ ✓ ✓ ✗ ✗ ✗ ✓ ✓ GFlowNets [12] ✓ ✓ ✓ ✓ ✓ ✗ ✗ ✗ ✓ ✓ DynaPPO [2] ✓ ✓ ✓ ✓ ✓ ✗ ✗ ✗ ✓ ✓ AdaLead [26] ✓ ✓ ✓ ✗ ✗ ✗ ✗ - ✓ ✗ PEX [23] ✓ ✓ ✓ ✗ ✗ ✗ ✗ - ✓ ✗ GGS [15] ✓ ✓ ✗ ✗ ✓ ✗ - ✗ ✓ ✗ VSD (ours) ✓ ✓ ✓ ✓ ✓ ✓ ✓ ✓ ✗ ✓Table 3 : Feature table of competing methods: ✓ has feature, ✗ does not have feature, -partially has feature, or requires simple modification. We follow [31] in their definition of amortisation referring to the ability to use q(x|ϕ * t-1 ) for warm-starting the optimisation of ϕ t .We compare our method, VSD, on a number of sequence design tasks and compare to existing baseline methods.We use three well established datasets; a green fluorescent protein (GFP) from Aequorea victoria [25] , an adeno-associated virus (AAV) [7] ; and DNA binding activity to a human transcription factor (TFBIND8) [34, 3] . These datasets have been used variously by [6, 5, 1, 15, 12] among others. The GFP task is to maximize fluorescence, this protein consists of 238 amino acids, of which 237 can mutate. The AAV task us to maximise the genetic payload that can be delivered, and this protein has 28 amino acids, all of which can mutate. A complete combinatorial assessment is infeasible for these tasks, and so we use the convolution neural network oracle presented in [15] as in-silico ground truth. TFBIND8 contains a complete combinatorial enumeration of the effect of changing 8 nucleotides on binding to human transcription factor SIX6 REF R1 [3] . The dataset we use contains all 65536 sequences, prepared by [34] . We also use two datasets from recent works that enumerate the (near) complete combinatorial space of short sequences. The first dataset measures the antibiotic resistance of Escherichia coli metabolic gene folA, which encodes dihydrofolate reductase (DHFR) [21] . Only a sub-sequence of this gene is varied (9 nucleic acids which encode 3 amino acids), and so a near-complete (99.7%) combinatorial scan is available. For variants that have no fitness (resistance) data available, we give a score of -1. The next dataset is near-complete combinatorial scan of four interacting amino acid residues near the active site of the enzyme tryptophan synthase (TrpB) [14] , with 159,129 unique sequences and fitness values, we use -0.2 for the missing fitness values. These residues are explicitly shown to exhibit epsistasis -or non-additive effects on catalytic function -which makes navigating this landscape a more interesting challenge from an optimisation perspective. The properties of these datasets are summarised in Table 4 . For all experiments we run a predetermined number of experimental rounds, T = 10, and we set the batch size to B = 128. In the first set of experiments, we use a fixed threshold, τ , with the aim of estimating p(x|y > τ ) (or S for non probabilistic models). For the next set of experiments, we set the threshold, τ , adaptively for testing VSD's ability to find the fittest sequence.We compare against DbAS [6] , CbAS [5] , AdaLead [26] , PEX [23] , BORE [32] adapted with the reinforce gradient estimator and random samples from the prior, p(x). All methods share the same surrogate model, and other components where appropriate (acquisition functions, priors and proposal distributions). For Adalead, we set κ = 0.5 since the CPE using τ is already performing the same thresholding.In this setting we fix τ over all rounds, for all competing methods, and we only consider the combinatorially (near) complete datasets. The primary measures by which we compare methods are where Qt y is the empirical quantile function of targets at round t, p t-1 is a percentile from the previous round, and η ∈ [0, 1] is a parameter that controls an 'annealing'-like schedule for τ t that prioritises exploration of the fitness landscape in earlier rounds and exploitation of known fit regions in later rounds. This is a strategy loosely-similar to [28] . The main measure of interest for these experiments is simple/instantaneous regret r t which quantifies how close the methods get to obtaining the globally fittest sequence,EQUATIONwhere y * is the fitness value of the maximum fitness sequence x * .For these experiments we use the higher dimension AAV (y * =19.54) and GFP (y * =4.12) datasets toshow that VSD can scale to higher dimensions. However, the X of these experiments is completely intractable to fully explore experimentally, and so we use a predictive oracle trained on all of the original experimental data as the ground-truth black-box function. This is the same strategy used in [5, 12, 33, 15] amongst others, and we use the exact CNN-based oracles from [15] for these experiments. Unfortunately, it has been shown that some of the oracles used in these experiments do not predict well out-of-distribution [30] , limiting their applicability to real-world problems.We follow [15] in the experimental settings for the AAV and GFP datasets, but we modify the max- (same as the previous MLP), followed by a convolutional layer (with a kernel width of 5 residues), followed by two hidden layers with max-pooling.Again we find the simple proposal distribution in Equation 16works as well as other more complex auto-regressive and transition, q(x t |x t-1 , ϕ), proposal distributions. However, this time we make use of a prior distribution of the same form as Equation 16, fit using maximum likelihood on the same training sequences used for the CPE, regardless of fitness. We find that a uniform prior leads to far inferior convergence results for all methods (BORE and DbAS use this as the initial proposal distribution) apart from AdaLead and PEX, which use alternative generative heuristics. The Random method draws sequences randomly from this prior.The results are summarised in Figure 4 . VSD is amongst the leading methods, CbAS and DbAS for both experiments, but it is never significantly better. We can se that AdaLead, PEX and BORE all perform worse than random for reasons previously mentioned. Simple regret can drop below zero for these experiments since an oracle is used as the black box function, but the global maximiser is taken from the experimental data. This potentially highlights some of the overconfidence issues inherent in these oracles outlined in [30] .The VSD algorithm is summarised in Algorithm 1.This section presents additional results.A full set of results for the fitness landscape experiments can be seen in Figure 3 .Results for the BBO experiments are shown in Figure 4 .  



Reviews: review: title: title: Clear paper with a novel algorithm for variational search
main_review: main_review: # Summary

This paper proposes a family of variational search problems, outlines their requirements and desiderata, and proposes an algorithm that satisfies them. By search problems, the authors refer to estimating densities of the sort $p(\boldsymbol{x} | y > \tau)$. Having access to said distributions is vital, and could conduce to generalizations of Bayesian optimization.

The authors estimate said probability of improvement using a variational approach, which they name Variational Search Distributions (VSD): approximate the posterior using a parametrized proposal distribution $q(\boldsymbol{x}|\phi)$, finding the parameters by maximizing the evidence lower bound. The authors derive a simple version of this objective, which can be readily optimized using out-of-the-box gradient methods. The method is presented as scalable, diverse, and generative.

# Review

The paper is a pleasure to read. The problem is well motivated, the proposed method is clear, and the derivations seem sound. Furthermore, the problem is relevant, and approached from a practitioner’s perspective. Sequence design indeed requires the left side of Table 1. The method is tested thoroughly against baselines.

With this in mind, I would like to recommend for clear acceptance.

Let me finish with some general comments and questions for improvement/discussion in future work:

- One detail in which I am not fully convinced is on whether (R4) is satisfied w. diversity: Could you/did you measure **the diversity** of the samples from the current variational approximation?
- There is no justification for the choice of batch size. One easy way to choose a batch size would be to approach practitioners and *ask*. Several of these analyses are done in plates of size ~98. Ideally, one would also see a sensitivity analysis of this parameter in the experimentation.
- When it comes to incorporating prior knowledge through $p(\boldsymbol{x})$. How difficult would it be to use log-likelihoods coming out of protein LLMs like the ESM family? My intuition says it should work out-of-the-box, but I might be mistaken. An experiment incorporating these likelihoods instead of a uniform prior would be pretty convincing, and I expect it to speed-up convergence significantly.

One final small comment: wich → which in the conclusions.
title: title: Variational inference for active search
main_review: main_review: This paper addresses the problem of active search (e.g., find compounds that would be useful as pharmaceuticals) by formulating the problem as finding an approximate posterior distribution over desirable designs given some prior over designs. The authors frame their method as variational inference and call it Variational Search Distributions (VSD). The exposition of the methods is clear, including the use of a score function gradient estimator and class probability estimation. VSD appears to generalize several related optimization algorithms and convincing experimental results are shown. 

Questions for the authors to consider: 

1.	How critical is the choice of prior distribution over designs? How might this depend on the dimension of the search space?
2.	Can you give an idea of the data necessary to either train a reasonable surrogate model or CPE and have the method succeed?