AdaLead: A simple and robust adaptive greedy search algorithm for sequence design

1 INTRODUCTION . An important problem across many domains in biology is the challenge of finding DNA , RNA , or protein sequences which perform a function of interest at a desired level . This task is challenging for two reasons : ( i ) the map φ between sequences X = { x1 , · · · , xn } and their biological function y = { y1 , · · · , yn } is non-convex and ( ii ) has sparse support in the space of possible sequences AL , which also grows exponentially in the length of the sequence L for alphabet A . This map φ is otherwise known as a “ fitness landscape ” ( de Visser et al. , 2018 ) . Currently , the most widely used practical approach in sequence design is “ directed evolution '' ( Arnold , 1998 ) , where populations of biological entities are selected through an assay according to their function y , with each iteration becoming more stringent in the selection criteria . However , this model-free approach relies on evolutionary random walks through the sequence space , and most attempted optimization steps ( mutations ) are discarded due to their negative impact on y . Recent advances in DNA sequencing and synthesis technologies allow large assays which query y for specific sequences x with up to 105 physical samples per batch ( Barrera et al. , 2016 ) . This development presents an opening for machine learning to contribute in building better surrogate models φ′ : X → y which approximate the oracle φ that maps each sequence to its true function . We may use these models φ′ as proxies of φ , in order to generate and screen sequences in silico before they are sent to synthesis ( Yang et al. , 2019 ; Fox et al. , 2007 ) . While a large body of work has focused on building better local approximate models φ′ on already published data ( Otwinowski et al. , 2018 ; Alipanahi et al. , 2015 ; Riesselman et al. , 2017 ; Sinai et al. , 2017 ) , the more recent work is being done on optimization in this setting ( Biswas et al. , 2018 ; Angermueller et al. , 2020 ; Brookes & Listgarten , 2018 ; Brookes et al. , 2019 ) . Although synthesizing many sequences within a batch is now possible , because of the labor-intensive nature of the process , only a handful of iterations of learning can be performed . Hence data is often collected in serial batches bi , comprising data Dt = { b0 , · · · , bt } and the problem of sequence design is generally cast as that of proposing batches so that we may find the optimal sequence x∗t = argmaxx∈Dt φ ( x ) over the course of these experiments . In this paper , we focus our attention on ML-augmented exploration algorithms which use ( possibly non-differentiable ) surrogate models φ′ to improve the process of sequence design . While the work is under an active learning setting , in which an algorithm may select samples to be labelled , with data arriving in batches bi , our primary objective is black-box optimization , rather than improving the accuracy of surrogate model . We define Eθ ( D , φ′ ) to denote an exploration algorithm with parameters θ , which relies on dataset D and surrogate model φ′ . When the context is clear , we will simply use E as shorthand . In most contexts , the sequence space is large enough that even computational evaluation is limited to a very small subset of possible options . For this reason , we consider the optimization as samplerestricted , not only in the number of queries to the ground truth oracle , but also the number of queries to the surrogate model ( Among other reasons , this allows us to thoroughly study the algorithms on landscapes that can be brute-forced , simulating a similar situation when the sequence space is very large compared to computation power , a very common setting . ) The algorithm E may perform v sequence evaluations in silico for every sequence proposed . For example , v × B samples may be evaluated by the model before B samples are proposed for measurement . Ideally , E should propose strong sequences even when v is small ; that is , the algorithm should not need to evaluate many sequences to arrive at a strong one . 2 CONTRIBUTIONS . In this study , we make three main contributions towards improving algorithms for sequence design : 1 . To build on recent progress in biological sequence design , the research community needs good benchmarks and reference algorithm implementations against which to compare new methods . We implement an open-source simulation environment FLEXS that can emulate complex biological landscapes and can be readily used for training and evaluating sequence design algorithms . We hope that FLEXS will help ensure meaningful and reproducible research results and accelerate the process of algorithm development for ML-guided biological sequence design . 2 . We introduce an abstracted oracle to allow the empirical study of exploration strategies , independent of the underlying models . This helps us understand relevant properties , such as robustness and consistency of the algorithms . 3 . Inspired by evolutionary and Follow the Perturbed Leader approaches in combinatorial optimization , we propose a simple model-guided greedy approach , termed Adapt-with-the-Leader ( ADALEAD ) . ADALEAD is simple to implement and is competitive with previous state-of-the-art algorithms . We propose ADALEAD as a strong , accessible baseline for testing sequence design algorithms . We show that in general , simple evolutionary algorithms are strong benchmarks to compete against and should be included in future analyses of new methods . 3 EVALUATION . We evaluate the algorithms on a set of criteria designed to be relevant to both the biological applicability as well as the soundness of the algorithms considered ( Purohit et al. , 2018 ) . We run the algorithms using FLEXS , where all of these algorithms and criteria evaluators are implemented . • Optimization : We let maximization be the objective . Most optimization algorithms operate under the assumption that critical information such as the best possible y∗ or the set of all local maxima M is unknown . While it is reasonable to assume that the best sequence is the one with the highest fitness , this is not necessarily the case in reality . For instance , we might wish to bind a particular target , but binding it too strongly may be less desirable than binding it at a moderate level . As measurements of this criterion , we consider the maximum y = maxx φ ( x ) over all sequences considered , as well as the cardinality |S| , where S = { xi | φ ( xi ) > yτ } and yτ > 0 is a minimum threshold value . It is noteworthy that we often do not know if any solutions y > yτ exists , hence finding many solutions by an algorithm is a sign of its strength . • Robustness : A major challenge for input design in model-guided algorithms is that optimizing directly on the surrogate φ′ can result in proposing a sequence x with large error , instead of approximating x∗ ( e.g . if the proposed input x is far outside D ) . Additionally , while biological models φ may contain regularities that are universally shared , those regularities are not known . Hence , a desired property of the algorithm is that it is robust in the face of a poor model . • Consistency : The algorithm E should produce better performing sequences if it has access to a higher quality model φ′ . Additionally , we desire that the high-performing sequences proposed by the algorithm are diverse . Because a sequence may be disqualified for reasons which are unknown during the design phase , we would like to find distinct solutions which meet the optimality criteria . However , metrics for measuring diversity can be ambiguous and we only focus on measuring diversity in a narrow sense . We measure diversity using |S| . When the ground truth model can be fully enumerated to find its local optima ( peaks ) by brute force ( i.e . we know the maximaM ) , we can use the number of found maxima |M′yτ | as a measure of diversity , whereM′yτ ⊂ M represents the maxima found by the algorithm above fitness yτ . 4 RELATED WORK . Bayesian Optimization ( BO ) . BO algorithms are designed to optimize black-box functions which are expensive to query . Importantly , these algorithms make use of the uncertainty of model estimates to negotiate exploration versus exploitation . In a pioneering study , Romero et al . ( 2013 ) demonstrate the use of BO for protein engineering . Many successful attempts have followed since ( e.g . Gonzalez et al . ( 2015 ) and Yang et al . ( 2019 ) ) , however in each of these cases a specific model of the design space is assumed to shrink the search space significantly . Otherwise BO is known to perform poorly in higher dimensional space ( Frazier , 2018 ) , and to our knowledge , no general purpose sequence design algorithm using BO has performed better than the models considered below . For this reason , while we implement variations of BO as benchmarks ( see similar interpretations in Belanger et al . ( 2019 ) ) , we do not consider these implementations as competitive standards . In our figures , we use the EI ( Expected Improvement ) acquisition function with an evolutionary sequence generator as our BO algorithm , and show comparisons with alternatives ( on TF landscapes ) in the supplement . Generative models . Another class of algorithms approach the task of sequence design by using regularized generative models . At a high level , these approaches pair a generative model Gϕ with an oracle φ′ , and construct a feedback loop that updates Gϕ ( and sometimes φ′ ) to produce high-performing sequences . In Feedback-GAN ( FBGAN ) , Gupta & Zou ( 2018 ) pair a generative adversarial network ( Goodfellow et al. , 2014 ) that is trained to predict whether sequences belong to the functional set using a frozen oracle φ′ which filters a subset of sequences at each training step . They bias the training of the generator and discriminator towards high-performing sequences . Killoran et al . ( 2017 ) pursue the sequence optimization by regularizing an “ inverted model '' , φ′θ −1 ( yi ) = xi with a Wasserstein GAN ( Arjovsky et al. , 2017 ) which is trained to produce realistic samples . In this case , both φ′θ and the generator are trained jointly . Brookes & Listgarten ( 2018 ) propose an algorithm , Design by Adaptive Sampling ( DbAS ) , that works by training a generative model Gϕ on a set of sequences X0 , and generating a set of proposal sequences X̂ ∼ Gϕ . They then use φ′θ to filter X̂ for high-performing sequences , retrain Gϕ , and iterate this process until convergence . This scheme is identical to the cross-entropy method with a VAE as the generative model , an important optimization scheme ( Rubinstein , 1999 ) . In follow-up work , termed Conditioning by Adaptive Sampling ( CbAS ) ( Brookes et al. , 2019 ) , the authors aim to address the pitfall in which the oracle is biased , and gives poor estimates outside its training domain D. The authors enforce a soft pessimism penalty for samples that are very distinct from those that the oracle could have possibly learned from . Specifically , they modify the paradigm so that as the generator updates its parameters ϕt → ϕ′t while training on samples in the tail of the distribution , it discounts the weight of the samples xi by Pr ( xi|G ; ϕ0 ) Pr ( xi|G ; ϕt ) . In other words , if the generative model that was trained on the original data was more enthusiastic about a sample than the one that has updated according to the oracle ’ s recommendations , that sample is up-weighted in the next training round ( and vice versa ) . Notably , as the oracle is not updated during the process , there are two rounds of experiments where they maximize the potential gains from their oracle given what it already knows : a round to create the oracle , and a round to improve the generative model . While it is trivial to repeat the process for multiple rounds , the process can be improved by incorporating information about how many rounds it will be used for . We use DbAS and CbAS as state of the art representatives in this class of regularized generative algorithms . Reinforcement Learning ( RL ) . RL algorithms are another method used to approach this problem . As these algorithms learn to perform tasks by experience , their success is often dependent on whether interactions with the environment are cheap or if there is a good simulator of the environment which they can practice on ( Mnih et al. , 2015 ) . However , in our setting , good simulators often do not exist , and sampling the environment directly can be very expensive . As a result , recent approaches to this problem have built locally accurate simulators and used them to train an RL agent . With DyNA-PPO , Angermueller et al . ( 2020 ) achieve state-of-the-art results in sequence design tasks following this approach . They train a policy network based on Proximal Policy Optimization ( PPO ) ( Schulman et al. , 2017 ) by simulating φ through an ensemble of models φ′′ . The models are trained on all measured data so far , and those that achieve a high R2 score in cross-validation are selected as part of the ensemble . Additionally , to increase the diversity of proposed sequences , they add a penalty for proposing samples that are close to previously proposed samples . They compare their results against some of the methods mentioned above , as well as other RL training schemes , showing superior results . While DyNA-PPO reaches state of the art performance , a major drawback of policy gradient algorithms is that they are complex to implement , and their performance is highly sensitive to implementation ( Engstrom et al. , 2019 ) . We use DyNA-PPO as a representative state of the art for all sequence design algorithms for our study . Most recent sequence design algorithms do not use evolution-inspired methods . Despite their lack of popularity , as we show here , they are fairly strong baselines that should be considered . Here we investigate classic algorithms like standard Wright-Fisher process , without and without models , Covariance Matrix Adaptation evolutionary strategy , ( CMA-ES ) ( Hansen & Ostermeier , 2001 ) and ADALEAD . It ’ s important to note that the focus of our study is distinct from those focused on better local approximate models ( e.g . ( Rao et al. , 2019 ; Riesselman et al. , 2017 ) ) , or semi-supervised and transfer learning approaches ( e.g . ( Madani et al. , 2020 ; Biswas et al. , 2020 ) ) .

Summary: In this work, the authors propose a greedy search approach for designing biological sequences in an active learning, batch setting. The algorithm is a fairly standard evolutionary algorithm which identifies the set of candidates at each iteration by adapting the best candidates from the previous iteration, and then evaluating those adaptations using a surrogate model. A set of experiments (using an evaluation sandbox also proposed by the authors for this work) suggest the proposed approach finds more high-quality solutions than competing approaches; however, especially in the experiments using the realistic/empirical surrogate model (an ensemble of 3 CNNs), the quality of the best solutions found by several approaches are statistically similar.