PropertyDAG: Multi-objective Bayesian optimization of partially ordered, mixed-variable properties for biological sequence design
Abstract
Bayesian optimization offers a sample-efficient framework for navigating the exploration-exploitation trade-off in the vast design space of biological sequences.
Whereas it is possible to optimize the various properties of interest jointly using a multi-objective acquisition function, such as the expected hypervolume improvement (EHVI), this approach does not account for objectives with a hierarchical dependency structure.
We consider a common use case where some regions of the Pareto frontier are prioritized over others according to a specified partial ordering in the objectives.
For instance, when designing antibodies, we maximize the binding affinity to a target antigen only if it can be expressed in live cell culture— modeling the experimental dependency in which affinity can only be measured for antibodies that can be expressed and thus produced in viable quantities.
In general, we may want to confer a partial ordering to the properties such that each property is optimized conditioned on its parent properties satisfying some feasibility condition.
To this end, we present PropertyDAG, a framework that operates on top of the traditional multi-objective BO to impose a desired partial ordering on the objectives, e.g. expression → affinity.
We demonstrate its performance over multiple simulated active learning iterations on a penicillin production task, toy numerical problem, and a real-world antibody design task.
Introduction
Designing biological sequences entails searching over vast combinatorial design spaces.
Recently, deep sequence generation models trained on large datasets of known, functional sequences have shown promise in generating physically and chemically plausible designs [e.g. 1–3].
Whereas these models accelerate the design process, limited resources place a cap on how many designs we can characterize in vitro for assessing their suitability.
Only once a design is validated in vitro and undergoes multiple rounds of optimization can it proceed down the drug development pipeline to preclinical development and clinical trials, where its performance is tested in vivo .
Because the wet lab cannot provide feedback on all of the candidate designs, we take an iterative, data-driven approach to select the the most informative subset to submit to the wet lab.
Many drug design applications call for such an active learning approach, as the initial datasets available to train predictive models on our desired properties of interest tends to be small or nonexistent.
The measurements returned by the lab in each iteration is appended to our training set
and we update our models using the augmented dataset for the next iteration.
The wet lab’s measurement process can be viewed as a black-box function that is expensive to evaluate.
In the context of identifying designs maximizing this function, Bayesian optimization (BO) emerges as a promising, sample-efficient framework that trades off exploration (evaluating highly uncertain designs) and exploitation (evaluating designs believed to carry the best properties) in a principled manner [4].
It relies on a probabilistic surrogate model that infers the posterior distribution over the objectives and an acquisition function that assigns an expected utility value to each candidate.
BO has been successfully applied to a variety of protein engineering applications [5–7].
In particular, we cast our problem as multi-objective BO, where multiple objectives are jointly optimized.
Our objectives originate from the molecular properties evaluated during in vitro validation.
This validation process involves producing the design, confirming its pharmacology, and evaluating whether it is active against a given drug target of interest.
If found to be potent, the design is then assayed for developability attributes—physicochemical properties that characterize the safety, delivery, and manufacturability [8].
The experimental process of in vitro validation signifies a hierarchy among the objectives.
Consider the property “expression” in the context of antibody design, for instance.
A designed antibody candidate must first be expressed in live cell culture.
If the level of expression does not meet a fixed threshold, the lab cannot produce it and it cannot be assayed for potency and developability downstream.
Supposing now that a design did express in viable amounts, if it does not bind to a target antigen with sufficient “affinity” (and is thus not potent), then the design fails as an antibody and there is little practical incentive in assaying it for developability (such as specificity and thermostability)— even if it is possible to do so.
The dependency between properties, whether experimental or biological in origin, motivates us prioritize some objectives before others when selecting the subset of designs to submit to the wet lab.
Our primary goal is to identify “joint positive” designs, designs that meet the chosen thresholds in all the parent objectives (expressing binders) according to the specified partial ordering and also perform well in the leaf-level objectives (high specificity, thermostability).
To this end, we propose PropertyDAG, a simple framework that operates on top of the traditional multi-objective BO to impose a desired partial ordering on the objectives, e.g. expression → affinity → { specificity, thermostability } .
Our framework modifies the posterior inference procedure within standard BO in two ways.
First, we treat the objectives as mixed-variable —in particular, each objective is modeled as a mixture of zeros and a wide dispersion of real-valued, non-zero values.
The surrogate model consists of a binary classifier and a regressor, which infer the zero mode and the non-zero mode, respectively.
We show that this modeling choice is well-suited for biological properties, which tend to carry excess zero, or null, values and fat tails [9].
Second, before samples from the posterior distribution inferred by the surrogate model enter the acquisition function, we transform the samples such that they conform to the specified partial ordering of properties.
We run multi-objective BO with PropertyDAG over multiple simulated active learning iterations to a penicillin production task, a toy numerical problem, and a real-world antibody design task.
In all three tasks, PropertyDAG-BO identifies significantly more joint positives compared to standard BO.
After the final iteration, the surrogate models trained under Property-BO also output more accurate predictions on the joint positives in a held-out test set than do the standard BO equivalents.
Background
Bayesian optimization (BO) is a popular technique for sample-efficient black-box optimization [see 10, 11, for a review].
Suppose our objective f : X → R is a black-box function of the design space
X that is expensive to evaluate.
Our goal is to efficiently identify a design x ⋆ ∈ X maximizing 1 f .
BO leverages two tools, a probabilistic surrogate model and a utility function, to trade off exploration (evaluating highly uncertain designs) and exploitation (evaluating designs believed to maximize f ) in a principled manner.
For each iteration t ∈ N , we have a dataset D t = { ( x (1) , y (1) ) , · · · , ( x ( N t ) , y ( N t ) ) } ∈ D t , where each y ( n ) is a noisy observation of f .
First, the probabilistic model ˆ f : X → R infers the posterior distribution p ( ˆ f |D t ) , quantifying the plausibility of surrogate objectives ˆ f ∈ F .
Next, we introduce a utility function u : X × F × D t : → R .
The acquisition function a ( x ) is simply the expected utility of x w.r.t.
our beliefs about f ,
[Equation]
For example, we obtain the expected improvement (EI) acquisition function if we take u EI ( x , ˆ f, D ) =
[ ˆ f ( x ) − max ( x ′ ,y ′ ) ∈D y ′ ]
+ , where [ · ] + = max( · , 0) [12, 4].
Generally the integral is approximated by Monte Carlo (MC) with posterior samples ˆ f ( j ) ∼ p
( ˆ f |D t ) .
We select a maximizer of a as the new design, measure its properties, and append the observation to the dataset.
The surrogate is then retrained on the augmented dataset and the procedure repeats.
Multi-objective optimization When there are multiple objectives of interest, a single best design may not exist.
Suppose there are K objectives, f : X → R K .
The goal of multi-objective optimization (MOO) is to identify the set of Pareto-optimal solutions such that improving one objective within the set leads to worsening another.
We say that x dominates x ′ , or f ( x ) ≻
f
( x ′ ) , if f k ( x ) ≥ f k ( x ′ ) for all k ∈ { 1 , . . .
, K } and f k ( x ) > f k ( x ′ ) for some k .
The set of non-dominated solutions X ∗ is defined in terms of the Pareto frontier (PF) P ∗ ,
X ⋆ = { x : f ( x ) ∈ P ⋆ } , where P ⋆ = { f ( x ) : x ∈ X , ∄ x ′ ∈
X s.t. f
( x ′ ) ≻
f ( x ) } .
[Equation]
MOO algorithms typically aim to identify a finite approximation to X ⋆ , which may be infinite, within a reasonable number of iterations.
One way to measure the quality of an approximate PF P is to compute the hypervolume HV( P| r ref ) of the polytope bounded by P ∪ { r ref } , where r ref ∈ R K is a user-specified reference point .
We obtain the expected hypervolume improvement (EHVI) acquisition function if we take
[Equation]
where P ′ =
P ∪ { ˆ f ( x ) } [13, 14].
Noisy observations In the noiseless setting, the observed baseline PF is the true baseline PF, i.e.
P t = { y :
y ∈ Y t , ∄
y ′ ∈
Y t s.t.
y ′ ≻
y } where Y t :=
{ y ( n ) } N t n =1 .
This does not, however, hold in many practical applications, where measurements carry noise.
For instance, given a zeromean Gaussian measurement process with noise covariance Σ , the feedback for a design x is y ∼ N ( f ( x ) , Σ) , not f ( x ) itself.
The noisy expected hypervolume improvement (NEHVI) acquisition function marginalizes over the surrogate posterior at the previously observed points X t = { x ( n )
[Equation]
where ˆ P t = { ˆ f ( x ) : x ∈ X t , ∄ x ′ ∈
X t s.t.
ˆ f ( x ′ ) ≻
ˆ f ( x ) } and ˆ P ′ =
ˆ P ∪ { ˆ f ( x ) } [15].
Related Work
Existing work on multi-objective BO does not account for objectives with a hierarchical dependency structure [16–19, 15].
We refer to [20] for a formulation of single-objective BO with a hierarchy in how the objective is computed.
A body of work focuses on constrained optimization, which optimizes a black-box function subject to a set of black-box constraints being satisfied [21–26].
For dealing with mixed-variable objectives, [27] propose reparameterizing the discrete random variables in terms of continuous parameters.
Our approach here is to model them explicitly using the zero-inflated formalism.
Method
Overview Figure 1 illustrates our proposed PropertyDAG-BO framework alongside standard BO.
The candidate generation step is identical in both cases; we first sample a large pool of design candidates from a proposal distribution, often implemented as a generative model.
The difference lies in the next step, in which the designs are scored by probabilistic surrogate models.
In PropertyDAGBO, we must first specify a partial ordering of our objectives (Section 4.1).
We then assign a zero-inflated generative model for each objective (Section 4.2) such that a probabilistic classifier models its “zero” mode and a probabilistic regressor models the remaining continuous-valued “nonzero” mode.
The raw posterior samples from the surrogates then undergo a “resampling” step (Section 4.3) that enforces the specified PropertyDAG.
Finally, the modified posterior samples enter the multi-objective acquisition function, which scores the design candidates just as in standard BO.
[Figure]
4.1 Defining a PropertyDAG
Many drug design applications motivate us to prescribe some hierarchy among our objective properties of interest.
The partial ordering may arise from an experimental dependency, e.g. a design candidate must pass a certain threshold in one property for its other properties to be measured.
In the context of antibody design, a design candidate is a sequence of amino acids representing an antibody that must first be expressed in cell culture.
If the level of expression does not exceed some threshold in mass per volume, the lab cannot produce it in viable amounts and it cannot be assayed for other properties, such as binding affinity to a target antigen.
Our PropertyDAG may then take the form: expression → affinity.
Experimental dependencies like this creates an asymmetry among the objectives; it reduces the information content of designs that do not express much more than that of designs that do not bind, because non-expressing designs cannot provide binding measurements.
Alternatively, the partial ordering may encode our preference for the types of designs we want to obtain.
We may prioritize a property, for instance, so that we reject designs that performs poorly in this property, no matter how well they perform in all the others.
If a designed antibody does not bind to the target antigen, it has failed in its primary function, so we have little interest in its developability properties, such as specificity to the target antigen and thermostability, even though, unlike for non-expressers, they often remain measurable.
We then impose the following PropertyDAG: expression → affinity → { specificity, thermostability } .
A PropertyDAG can be expressed as ordered sets of properties: [Equation] [Equation]
[Equation] → { y L, 0 , . . .
, y L,M L } , where y l,m is the property at level l ∈ { 0 , . . .
, L − 1 } of the hierarchy and m ∈ { 0 , . . .
, M l − 1 } is its index among the M l sibling properties at the same level l .
Figure 2 shows an example of a PropertyDAG with three properties and two levels.
4.2 Zero-inflated modeling
Biological properties tend to carry excess zeros, or null values.
Their zero-inflated nature motivates us to employ statistical models that account for large incidences of zeros [28].
The zero-inflated negative binomial (ZINB) distribution has been applied to model discrete counts in single-cell RNA-seq data [29].
For each objective y k , our zero-inflated model assigns a binary random variable b k ∈ { 0 , 1 } to generate the zeros and r k ∈ R to generate the remaining dispersion of continuous non-zero values.
Assume f is non-negative (equivalently, that it is bounded from below).
Given a training dataset D t available at time t , we decompose p ( ˆ f k | D t ) as follows:
[Equation]
where p ( r k | x , D t , θ r ) is a regressor parameterized by θ r trained to predict f k given b k = 1 .
For simplicity we have assumed p ( r k =
| x , D t , θ r ) = 0 .
Since Gaussian processes (GPs) are often used as surrogates for BO [see, e.g. 30, for a review on GPs] and common GP assumptions fail for sparse, multi-modal data, separating out the non-zero mode of the data can improve posterior inference.
To accommodate objective hierarchy, each p ( b k | x , D t ) decomposes further as
[Equation]
where p ( b k | x , D t , θ b ) is a classifier parameterized by θ b trained to predict 1 { f k ( x ) > 0 } and pred( k ) are the predecessors, or ancestral nodes, in the PropertyDAG corresponding to property k .
HV I
[Figure]
This general framework, presented in terms of a zero-inflated, continuous-valued objective (a mixture of a delta function at zero and a continuous distribution), applies to binary-valued objectives and continuous-valued objectives without zero inflation, which can be viewed as specific cases taking p ( r k | x , D t , θ r ) =
p ( r k ) =
N (0 , σ 2 ) with very small σ and p ( b k = 1 | x , D t , θ b ) = 1 , respectively.
4.3 Resampling
Using a simple resampling trick, we modify the posterior samples from the surrogate models to enforce the parent-child relationships specified in PropertyDAG.
As before, consider a property y k and refer to its predecessors, or ancestral nodes, as pred( k ) .
Suppose we have drawn single a posterior sample and obtained β k ′ ∼ p ( b k ′ | x ) ∈ { 0 , 1 } and ρ k ′ ∼ p ( r k ′ | x ) ∈ R for each k ′ ∈ { 1 , . .
, K } .
Without any modification, Equation 5 would yield the following sample γ k of y k :[-γ k = 
-]
[Equation]
Instead, we begin at the top level and proceed down the levels of PropertyDAG to impose dependencies between b k and its predecessor properties { b k ′ } k ′ ∈ pred( k ) .
If y k is a top-level property, then pred( k ) = ∅ and ˆ β k = β k .
Otherwise, y k has parent properties and we have
[Equation]
We can then use the modified binary samples { ˆ β k } Kk =1 to obtain our effective sample ˆ γ k of y k : [Equation] if [Equation]
or
[Equation] pred( [Equation] k else.
Let γ :=
[ γ 1 , . . . ,
γ K ] ∈ R K and ˆ γ :=
[ˆ γ 1 , . . .
, ˆ γ K ] ∈ R K and denote the transformation at the sample level described in Equations 8 and 9 as h : R K → R K
such that h ( γ ) =
ˆ γ .
We repeat this resampling procedure to other posterior samples.
The transformed samples are then used to evaluate NEHVI (Equation 4) via MC integration.
More precisely, suppose we draw S posterior samples in parallel for a design candidate x ∗ (reflecting both aleatoric and epistemic uncertainties) and the previously observed designs X t
= { x ( n ) } N t n =1 (reflecting the aleatoric uncertainty) and denote each draw as γ ∗ s and G s := { γ ( n ) s } N t n =1 , respectively, for s = 1 , . . . , S .
Then the MC approximation of NEHVI can be efficiently evaluated as:
[Equation]
where P [ s ] t = { γ s : γ s ∈ G s , ∄ γ ′ s ∈ G s s.t. γ ′ s ≻
γ s } and P ∗ [ s ] t = P [ s ] t
[Equation]
[Equation]
Figures 3a and 3b illustrate the effect of the PropertyDAG resampling on each posterior sample’s HVI contribution to EHVI.
Experiments
We perform simulated active learning experiments on two synthetic tasks and a real-world antibody design task.
We use NEHVI (Equation 4) as our acquisition function and evaluate it via MC integration (Equation 10).
In each experiment, we test three types of acquisitions: (1) batched, multi-objective
BO with PropertyDAG (“qNEHVI-DAG”), (2) without PropertyDAG (“qNEHVI”), and (3) random.
Our main metric is the number of acquired “joint positive” designs, designs that exceed the chosen thresholds in all objectives according to the specified PropertyDAG.
We refer to the batch size as q .
5.1 Penicillin production dataset
This task is based on the penicillin production simulator [31].
We defined the goal as minimizing the CO 2 byproduct emission while ensuring that the fermentation time is below a set threshold and the yield exceeds a set threshold ( K = 3 , X
= R 7 ).
We negated the latter two objectives to define a maximization problem and assume the PropertyDAG: { y 0 , 0 } → { y 1 , 0 } → { y 2 , 0 } , where y 0 , 0 =
Yield (“Objective 0”), y 1 , 0 =
Negative fermentation time (“Objective 1”), and y 2 , 0 =
Negative CObyproduct (“Objective 2”).
Zero-mean Gaussian noise was added to the input, following [19].
We fit an exact GP to model r k and an approximate GP with the variational evidence lower bound (ELBO) to model b k , separately for each Objective k .
We drew 512 posterior samples to evaluate qNEHVI.
We executed 10 rounds of simulated active learning by initializing the surrogates with 8 training points and selecting q = 4 out of 80 randomly-sampled pool of candidate points in iteration.
The three acquisition modes (qNEHVI-DAG, qNEHVI, and Random) were subject to the same initial training points and candidate pool each round.
The entire experiment was repeated 5 times.
Figure 4a shows that qNEHVI-DAG identifies significantly more joint positives than do qNEHVI and Random over active learning iterations. [Figure] [Table]
[Figure] [Table] [Figure] [Table]
[Figure] [Table] [Figure] [Table]
[Figure]
[Table]
[Figure] [Table] [Figure] [Table]
N u m b e r o f
j o i n t po s iti v e s f ound
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table] [Figure] 5.2 Branin-Currin toy problem This task is based on an analytic Branin-Currin test function from [19] with X = R 2 and K = 2 .
We reformulated this task to simulate the antibody design task (Section 5.3) in a controlled environment.
We defined the PropertyDAG, { y 0 , 0 } → { y 1 , 0 } , where y 0 , 0 =
Dimension 0 (“Objective 0”) and y 1 , [-
-]0 = Dimension 1 (“Objective 1”).
Objective 0 was transformed into binary values using a set threshold.
Objective 1 was zero-inflated and real-valued.
Posterior inference was performed following the same procedure described in Section 5.1.
We executed 20 rounds of simulated active learning by initializing the surrogates with 6 training points and selecting q = 4 out of 40 randomly-sampled pool of candidate points in iteration.
The entire experiment was repeated 10 times.
Figure 4b shows that qNEHVI-DAG identifies significantly more joint positives than do qNEHVI and Random over active learning iterations.
Figure 5 compares the qNEHVI-DAG, qNEHVI, and Random selections for Objective 1, for the final selections stacked across the 10 repeated trials.
Overall, qNEHVI-DAG identifies more examples to the right of the threshold (black dashed lines) than do qNEHVI and Random, and the improvement is more pronounced for the identification of joint positives (middle panel).
5.3 Antibody design
The antibody design task is derived from real-world dataset of antibody sequences and their measured in vitro properties, affinity and expression.
As in the toy problem (Section 5.2), we defined the
PropertyDAG, { y 0 , 0 } → { y 1 , 0 } , where y 0 ,
= Expression (“Objective 0”) and y 1 , 0 = Affinity [Figure] (“Objective 1”).
Objective 0 was binary-valued, i.e. expressing or not.
Objective 1 was zero-inflated and real-valued.
We executed 3 iterations of simulated active learning and repeated the entire procedure 5 times.
To simulate active learning, we split the entire dataset of 4,023 variable-length protein sequences, designed as antibodies for an anonymized target antigen A, into 5 groups of sizes 1230, 736, 746, 711, and 600.
The first group served as the initial training set for the surrogates, the following three groups as the “candidate pools” from which we selected 200 candidates in each iteration, and the last served as a held-out test set.
As shown in Figure 6a, qNEHVI-DAG once again outperforms qNEHVI and Random in the number of joint positives.
The log posterior density evaluated at the affinity measurements for the joint positives (expressing binders) in the test set is also highest for qNEHVI-DAG, which indicates that the surrogate models from qNEHVI-DAG had the most accurate beliefs about the joint positives after the final iteration.
[Figure]
Conclusion
Our proposed method, PropertyDAG, sits on top of the existing multi-objective BO framework to make it amenable to a common scenario in drug design, where a hierarchical structure, or partial ordering, exists among the objectives.
It modifies the surrogate posteriors so that each objective is modeled as zero-inflated (a mixture of excess zeros and a continuous distribution) and parent properties in PropertyDAG are prioritized before the children.
Empirical evaluations shows that
PropertyDAG-BO can identify significantly more designs that are jointly positive (i.e. exceeding a chosen threshold in all properties) than does standard BO.
By encapsulating our experimental and biological priors on the relationship between molecular properties, our method promises to accelerate computational drug discovery.
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information on how to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] , or [N/A] .
You are strongly encouraged to include a justification to your answer , either by referencing the appropriate section of your paper or providing a brief inline description.
For example:
• Did you include the license to the code and datasets?
[Yes] See Section ?? .
• Did you include the license to the code and datasets?
[No] The code and the data are proprietary.
• Did you include the license to the code and datasets?
[N/A]
Please do not modify the questions and only use the provided macros for your answers.
Note that the
Checklist section does not count towards the page limit.
In your paper, please delete this instructions block and only keep the Checklist section heading above along with the questions/answers below.
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] (c) Did you discuss any potential negative societal impacts of your work?
[Yes] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[N/A]
(b) Did you include complete proofs of all theoretical results?
[N/A] 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[No] We will release the material after the work is accepted by a peer-reviewed journal for publication.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[N/A] (b) Did you mention the license of the assets?
[N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
[N/A] (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
A Appendix A.1 Additional results for the penicillin production task
Figure 7 compares the qNEHVI and qNEHVI-DAG selections for every pair