A Pareto-optimal compositional energy-based model for sampling and optimization of protein sequences
Anonymous Author(s)
AffiliationAddress email
Abstract
Generative models constitute one of the prominent machine learning methodsleveraged by engineering and life sciences.
Nevertheless, besides resemblance tothe data we wish to sample, oftentimes our aim is to propose new data points thatwill satisfy, not one, but multiple properties of interest.
Such aspiration becomeseven more challenging when those properties are independent or orthogonal toeach other.
In this work, we propose a Pareto-compositional energy-based model(pcEBM) , a generative framework for sampling new datapoints that simultaneouslyadhere to various constraints and optimize distinct properties.
pcEBM emergesnaturally from an integrated multi-gradient energy-based sampling framework andis shown to lead to practical benefits in a series of real world (antibody-design)empirical analysis.
Introduction
Generative models have recently attracted significant attention within the machine learning community.
Part of the excitement stems from the convincing images created by generative adversarial networks(GAN) [13], variational autoencoders [20], energy-based models, (EBMs) [9] or more recentlydiffusion models [31] that made headlines [29, 5].
However, there are far less success stories inindustry applications [6, 15].
There are multiple reasons for this fallback: not all datasets are in imageformat; there is a lack of evaluation protocols and metrics for synthetic data; it is still challengingto control and train some of the generative models (i.e., GANs); and finally, we are still not able togenerate or imagine samples that are (very) different from what we have seen in training.
Overall, thelack of constraints and control in generative modeling limits the current use cases.
This manuscript is motivated by the need to accelerate the development and discovery of new drugs,namely, therapeutic antibodies by incorporating machine learning methods in classical computationalbiology pipelines.
Though multiple generative models have already been proposed for such purposes [10, 12, 28], generating new samples without guidance or control does not guarantee usabilityor success of the proposed molecules.
In practice, each molecule has to satisfy multiple properties,such as affinity, aggregation, and thermal stability [1, 33], before it qualifies as a potential candidatefor a clinical trail.
Not taking such fundamental properties into consideration leads to potentialcomplications later on during clinical trials and the optimization of straightforward objectives doesnot necessarily translate into progress in the laboratory [25].
Motivated by this real-world challenge, we here propose a new energy-based generative methodtargeted at antibody design and discovery that simultaneously takes into account multiple propertiesthat an antibody (protein) has to satisfy.
The concurrent adherence to multiple properties is achallenge in itself, as it often involves optimizing multiple conflicting objectives.
In antibodies, for instance, naive optimization of affinity can be achieved at the expense of other desirable developabilityproperties.
In general, when one cares about the optimization of multiple conflicting properties, it is oftenimpossible to find a single sample that satisfies all of the objectives simultaneously.
We argue that,instead, one should aim to propose a set of diverse data points from the Pareto front [3] that correspondto different choices for the different objective functions.
By doing so, one gains a global perspectivethe optimal trade-off between objectives and can select a molecule according to their preference.
Toachieve Pareto optimality, we rely on recent advances in multi-objective optimization (MOO)[30, 24]as well as on compositional sampling with EBMs [9] to build a Pareto-optimal compositional EBM( pcEBM ).
In what follows, we first present the the necessary background on multi-objective optimization insubsection 2.1 and related work on compositional sampling with EBM in subsection 2.2 leading topcEBM described in subsection 2.3.
In section 3 we include empirical evaluation and discussion,before concluding in section 4.
Background and method
Problem setup.
Antibodies are composed of two chains of amino acids (AA), which can berepresented as sequences of characters (20 characters corresponding to each of the possible aminoacids found in humans), typically of combined length L ∼ 250 .
We will denote the sequences asx = ( x 1 , . . . , x L ) , where x l ∈ { 1 , . . , 20 } corresponds to the AA type at position l .
For each ofthose sequences, we measure m properties, such that f i : R L → R for all i ∈ 1 , . . . , m .
Our goal is to generate new sequences x ∗ with preferred values for each of the m properties.
Sincewe cannot fully satisfy all of the properties/objectives simultaneously (they may be conflicting witheach other), we are interested in finding samples which can not be further improved simultaneouslyfor all the objectives, yielding the notion of Pareto optimality [4].
2.1 Multi-objective optimization
A straightforward approach to solving MOO is Linear Scalarization [19] based on a preference vectorλ =
[ λ 1 , . . . , λ m ] from the probability simplex on [ m ] , i.e., S = { λ : (cid:80) mi
{+=
 λ i = 1 , λ i ≥ 0+} , i ∈ [ m ] } .
Each λ ∈ S leads to a weighted objective function f λ ( x ) = {+(cid:80) mi =
+}
λ i f i ( x ) and its minimizerx ∗ λ = argmin x f λ ( x ) .
As we evaluate on λ from a grid of S , we hope that the corresponding x ∗ λapproximates the Pareto front.
The solutions obtained by Linear Scalarization lie in the convexenvelope of the Pareto front, which renders the approach plausible for situations where the Paretofront is convex [24].
An alternative for non-convex Pareto fronts is Multiple Gradient Descent (MGD) [7], which is basedon the idea of iteratively updating the variable x along a direction that ensures that all the objectivesare decreased simultaneously, i.e., that guides the optimization towards a Pareto improvement.
Denote the gradient of the i -th objective as ∇ x f i ( x ) .
With gradient descent, we update the variablex k ← x k − 1 − ηg ( x ) , where g ( x ) is a vector to be determined and η is a small step size.
From a1st-order Taylor expansion, we can deduce that ⟨∇ x f i ( x ) , g ⟩ ≈ − ( f i ( x k ) − f i ( x k − 1 ))
/η whichrepresents the decreasing rate of f i when we update x along direction g ( x ) .
In MGD, g is chosen tomaximize the slowest decreasing rate among all the objectives, that is:
[Equation]
By doing so, g ( x ) is pushed to have positive inner products with all ∇ x f i ( x ) .
If the latter is impossible, {∇ x f i ( x ) } mi =1 will contain conflicting directions and we will get g ( x ) = [-
-]0 which terminates the procedure.
By using Lagrangian duality, Désidéri [7] has shown that the above objective has the optimal value g ( x ) ≈ (cid:80) mi =1 λ ∗ i ∇ x f i ( x ) where { λ i ∇ x f i ( x ) } i =1 m is the solution of min λ i
∥ m (cid:88) i =
λ i ∇ x f i ( x ) ∥ 2 subject to m {+(cid:88) i = λ i = 
+}
and λ i ≥ 0 for all i ∈ [ m ] .
The above convex optimisation problem has a closed form solution when m = 2 and a fast algorithmfor m > 2 [30, 18].
By construction, when the step size is small, MGD decreases monotonically allobjectives simultaneously and will terminate at a local Pareto point.
2.2 Compositional energy-based models
Although EBMs have been around since many years [22, 21], their utility was recently rediscoveredin the quest of generative modeling.
EBMs learn to represent data by approximating an unnormalizedprobability distribution across data.
They do so by learning an energy function E θ ( x ) parameterizedby a neural network that maps each input x to a scalar real value interpreted as the energy andapproximating the data distribution by the Boltzmann distribution under unit temperature:
[Equation]
Training an EBM on a given data distribution is usually done by contrastive divergence [16].
Besidespredicting the un-normalized data likelihood, EBMs can be used to sample new data points x fromp θ (during both training and generation). Sampling is typically performed by some Markov-Chain
Monte-Carlo (MCMC) procedure. We consider the common case of Langevin Diffusion (LD) wherethe samples are initialized from refinement/uniform random noise followed by iterative refinement:
[Equation]
with k being the sampling step and η the step size.
The above sampling procedure can be easily extended to produce samples from the composition of p θ with other distributions of interest. Along this direction, Du et al.
[8] considered the problem of sampling from m composed distributions, each modeling a different propertiey f i , an approach referred to as compositional EBM (cEBM). Du et al.
proposed multiple variations for executing different logical operations over properties of interest, such as conjunction, negation, and disjunction.
In the following, we focus on conjunction, which corresponds to combining individual EBMs as a product of experts:
[Equation]
We also note that cEBM utilized the following sampling procedure:
[Equation]
which amounts to applying LD w.r.t.
E ( x ) = (cid:80) i E ( x | f i ) . 2.3 Pareto compositional sampling Inspired by the previous work on multi-objective optimization [7, 24] discussed in subsection 2.1, wehere propose pcEBM, a method that samples from the Pareto front of m Boltzmann distributions ofinterest.
Our approach entails utilizing multiple gradient descent to select a locally optimal directionthat optimizes the slowest decreasing rate among all the objectives of the compositional EBMs:
[Equation]
where α is a positive constant and the internal optimization problem is the same utilized in MGD andthe external sampler is performing Langevin diffusion.
Suppose that we are generating a new sequence x , either by starting from random noise or byimproving an existing one. At each step k , if x k − 1 is far from the Pareto front and the gradients∇ x f i ( x k − 1 ) have non-negligible norm, Equation 4 will drive x k − 1 closer to the Pareto front.
Onthe other hand, when x ∗ is close to the Pareto front and the gradients have nearly vanished, thenoise dominates and x k − 1 performs Brownian motion.
As will be demonstrated experimentally,the addition of noise ω helps pcEBM to explore more efficiently the Pareto as compared to MGD,whereas the locally optimal determination of the optimization direction renders the optimization moreeffective than cEBM.
Empirical evaluation
In the following, we compare pareto compositional EBM ( pcEBM ) with three baselines: multiple gradient descent ( MGD ) [7], compositional EBM ( cEBM ) [8], and linearly scaled cEBM ( ls-cEBM ) [19].
The figure axes report the energy of EBMs that were separately trained on data with good affinity and BV scores, respectively, with smaller energies being more desirable.
Experimental setup. We wish to generate plausible antibody sequences.
We consider threekey properties for the generated antibodies: ( i )similarity to known human antibodies from thepublic database of observed antibodies [26] ( Ablike ), ( ii ) binding affinity to antigen targets ofinterest ( Aff ), and ( iii ) antibody clearance evaluated as off-target binding, presumed to accountfor the fast clearance [17], to which we refer asnon-specificity and shortly ( BV ) 1 . The datasetswe use are from both public and proprietarysources.
We train three separate EBMs, each pertaining to a property of interest. Though alternative choices are also possible, we instantiate allmodels as sequence-based convolution networks[11].
More details on the architecture can befound in the Appendix.
We use the same modelsfor sampling new sequences with each of thebaselines.
The differences in implementationare as follows: MGD does not add any noise, cEBM uses Lanvegin dynamics but has uniform weightsfor all properties; ls-cEBM differs from cEBM in that it includes domain-informed weights; andfinally pcEBM learns the optimal weights for each property per sequence.
Each baseline relies ontwo main hyper-parameters, namely the step size η and the number of steps k .
To compare different approaches in terms of their ability to satisfy multiple properties, we adopt the Hyper-Volume ( HV ) metric from the multi-objective optimization literature [14] computed over theenergy scores of each property model.
In all experiments, we aim for minimizing the energy scoresper property. Additionally, as a proxy for how well the sampled sequences capture the properties ofinterest we compute the minimal edit-distance of each sampled sequence to the ones seen during thetraining of each property EBM ( e dist ) [27]. More details about the computation of these metrics canbe found in the appendix.
Results and discussion. Our experiments are designed to provide quantitative answers to thefollowing research questions:
• Q.1 Can we use EBMs to generate valid/reasonable, multi-property compliant antibody sequences?
• Q.2 Does sampling with Langevin Dynamics propose more Pareto optimal sequences?
• Q.3 Does pcEBM have an advantage over cEBM?
• Q.
Starting from a seed sequence, can we improve a property of interest using (p)cEBM?
Table 1 and Table 2 address Q.1 .
As perhaps expected, EBMs trained on data possessing a singleproperty of interest propose sequences that are akin to those sharing the same property, but aredissimilar to the sequences preferred by other objectives. In contrary, sequences generated by amulti-property EBMs have smaller distances to all properties compared to the single-property EBMs.
Amongst the four considered multi-objective baselines, pcEBM generates sequences with higherhyper-volume and similarity to the training datasets, whereas MGD is a close second.
[Table]
Part of the difference between cEBM and pcEBM can be accounted by a difference in their speedof convergence.
We confirm this in Figure 2 by plotting the energy score at each sampling step forboth cEBM and pcEBM with a step size of η = 0 .
As observed, for the same starting sequence,pcEBM converges faster.
This result partially answers Q.3 .
To further test our approach, we calculate the hypervolume across the three properties we aim to accountfor.
Here, we consider the energy of each property EBM as proxy for the objective of interest
The advantage of pcEBM is noticeable at smaller step sizes(see Table 1), while the variance in both edit distance and hyper-volume for different η is significantlysmaller for pcEBM than others.
If we turn to Figure 1 that depicts the Pareto front for two propertiesobtained from candidate sequences from each baseline, we see that the pcEBM front is more spreadthan MGD approaches, agreeing with our suppositionthat employing LD leads to an improved coverageof the Pareto front as compared to pure gradient descent.
Further, pcEBM and MGD fronts are closer tothe origin that what is obtained by direct LD, whichdemonstrates the advantage of the multi-properly optimization aspect of pcEBM, addressing Q.2 and Q.3 .
Finally, for Q.4 we wish to explore the possibility for improving properties of already existingantibodies by using the variants of compositional EBMs.
We focus our analysis on non-specificity( BV ) as we have an external SeqCNN predictor (classifier) that can act as a surrogate pseudo-oracle. We start by screening the existing sequences to select only those with low BV scores, which werefer to as “seeds”.
We pass the selected seed sequences with poor measured (as well as predicted)scores, through our generative models, and then reuse the non-specificity pseudo-oracle to evaluate [Figure]
[Figure] the newly proposed designs. Figure 3 shows the BV score improvement achieved by each baseline.
We notice that the multi-objective methods manage to improve the worst BV scores from 0.27 toabove 0.9, with MGD and pcEBM giving a slightly lower average score than cEBM (though all scoresare comparably high).
Figures 3 (a) and (b) provide further insight into the non-specificity optimization process by examininga low-dimensional embedding of seed and optimized sequences.
The embedding is obtained byusing tSNE to project the last layer feature representations of the sequences obtained from thenon-specificity oracle, onto a 2D space.
We color each point (sequence) by their BV score with blueindicating a good score and red a poor one.
On the same plot, we overlay an example sequence (seed)with particularly bad BV score (i.e., sticky molecules) and its corresponding proposed designs.
Wesee that both cEBM and pcEBM improve the probability for a good non-specifisity score, and movethe seed towards the section of the manifold covered by datapoints with good scores.
Conclusion
We propose a new approach to sampling data points that simultaneously satisfy multiple desiredproperties.
We do so by integrating multiple gradients within compositional energy based models.
Ina series of experiments for generating and improving existing sequences, we validate the performance of the compositional EBMs as well as their Pareto optimality.
The generality and modularity ofpcEBM allow for multiple expansions we are excited to explore: ( i ) other data types, in particulargraph structures suitable for molecules, ( ii ) incorporating other sampling frameworks such as Steinvariational gradient descent [23], ( iii ) incorporating uncertainty estimates, and ( iv ) pareto-optimaltraining for EBMs, to name a few.
Appendix4.1 Additional details on experiments
SeqCNNs For all properties and baselines we use an identical NN architecture, that is one modelper protein chain, consisted of three Conv1D layers with kernel size 9 and padding 1, ReLU nonlinearities and penultimate layer of size 256.
All EBMs were trained with contrastive training, Adamoptimizer and early stopping criterion.
Hyperparameters For each multipropery sampling baseline, the following hyper parameter gridwas explored with random search:
• step size - η ∈ { 1 e − 4 , 1 e − 2 , 1 , 10 , 40 }
• number of steps - k ∈ { 100 , 200 , 300 , 400 }
• noise type - ω ∈ {N , U} .
4.2 Metrics
Hypervolume HV is an indicator for evaluating the quality of sets of solutions in multi objectiveoptimization.
For a given reference point of r =
[ r 1 , ..., r m ] we have an upper bound of the objectives,such that supxf i ( x ) ≤ r i , ∀ i ∈ [ m ] .
For a given set of solutions X = { x l } nl =1 , a hypervolumeindicator HV ( X ) is a measure of the region between all f i objectives and r :
[Equation]
where Λ( · ) denotes the Lebesgue measure.
In the context of our experiments, we are interested in minimizing the energy scores per each property,and we choose a reference point r
=
We compte HV from each Pareto front to thereference point with the implementation provided in PyMoo [2].
Edit distance Edit distance is the typical score for quantifying how dissimilar two strings (inour case, sequences) are to one another, measured by counting the minimum number of operationsrequired to transform one string into the other. Although different variants exist, here we rely on the Levenshtein distance which accounts for insertions, deletions and substitutions. In all experimentswe use the python library edist as implemented in [32].
4.3 Additional results from experiments
[Table]