A Pareto-optimal compositional energy-based model for sampling and optimization of protein sequences
Anonymous Author(s)
AffiliationAddress email
Abstract
Generative models are one of the most 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 both real world (antibody-design)empirical analysis.
Introduction
Generative models have recently gained increased interest within the machine learning community.
Part of the excitement stems from the realistic-looking images created by generative adversarialnetworks [13], variational autoencoders [18], energy-based models, (EBMs) [9] or more recentlyand diffusion models [28] that made headlines [25, 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 challenging tocontrol a generative model; and finally, we are still not able to generate or imagine samples that are(very) different from what we have seen in training.
Overall, the lack of constraints and control ingenerative 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.
Prominent generative models have already been proposed for such purposes,[10, 12, 24] to name a few.
However, generating new samples without guidance or control, doesnot guarantee usability or success of the proposed molecules.
Each molecule has to satisfy multipleproperties such as affinity, aggregation, thermal stability [1, 30], before it even qualifies as a potentialcandidate for a clinical trail.
Not taking such fundamental properties into consideration, leads topotential complications later on during clinical trials.
Motivated by this real-world challenge, we set out to propose a new energy-based generative methodtargeted at protein design and discovery that simultaneously takes into account multiple properties thatan antibody (protein) has to satisfy.
The concurrent adherence to multiple properties is a challenge initself, as it often involves optimizing multiple conflicting objectives.
In antibodies, for instance, naiveoptimization of affinity can be achieved at the expense of other desirable developability properties.
Unfortunately, when one cares about the optimization of multiple conflicting properties, it is oftenimpossible to find a single sample that satisfies all the objectives simultaneously.
Instead, we wish topropose a set of diverse data points from the Pareto front [3] that correspond to different choices forthe different objective functions.
By doing so, one gains a global perspective the optimal trade-offbetween objectives and can select a molecule according to their preference.
To achieve this, we relyon recent advances in multi-objective optimization (MOO)[27, 22] as well as compositional samplingwith EBMs [9], to build a Pareto-optimal compositional EBM, pcEBM .
In what follows, we first present the related work and background on multiple gradient descent (MGD)[7, 22] and EBMs in section 2, which compose the pcEBM described in section 3.
In section 4 weinclude empirical evaluation and discussion, before concluding in section 5.
Background2.1 Problem setup
Antibodies (as any protein) are composed of 2 chains of amino acids (AA), which can be representedas sequences of characters (22 characters corresponding to each of the 22 possible amino acids),typically of total length L ∼ 250 .
We will denote the sequences as x = ( x 1 , ...x L ) where each x lcorresponds to the AA at position l .
For each of those sequences, we measure m properties, such thatf i : R L → R , i ∈ 1 ..m .
Our goal is to generate new sequences x ∗ with prefered values for each ofthe m properties.
Since we can not fully satisfy all of the properties/functions simultaneously (theymay be conflicting with each other), we are interested in finding samples which can not be furtherimproved simultaneously for all the objectives, yielding the notion of Pareto optimality [4].
2.2 Multi objective optimization
A straightforward approach to solving MOO is Linear Scalarization [26] based on a preferencevector λ
= [ λ 1 , ...λ m ] from the probability simplex on [ m ] , i.e., S = { λ : (cid:80) i = 1 m λ i = 1 , λ i
≥0 , i ∈
[ m ] } .
Each λ ∈ S leads to a weighted objective function f λ ( x ) =
(cid:80) i = 1 m λ i f i ( x ) andits minimizer x ∗ λ = argmin x f λ ( x ) .
As we evaluate on λ from a grid of S , we hope that thecorresponding x ∗ λ approximates the Pareto front.
This simple to implement strategy suffers fromproblems such for example solutions that can only lie on the convex envelope of the Pareto front,consequently making this plausible approach only in cases when the Pareto front is convex [22].
A recent alternative for non-convex Pareto fronts is Multiple Gradient Descent, (MGD) , [7], basedon the idea to iteratively update the variable x along a direction that ensures that all the objectivesare decreased simultaneously, i.e. guids the optimization towards a Pareto improvement.
We denotethe gradient of the i-th objective as g i ( x ) =
∇
i
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.
Following
Taylor approximation ⟨ g i , g ∗ ⟩ ≈ − ( f i ( x k )
− f i ( x k − 1 )) /ϵ represents the decreasing rate of f i whenwe update x along direction g ∗ ( x ) .
In MGD, g ∗ is chosen to maximize the slowest decreasing rateamong all the objectives, that is:
[Equation]
By doing so, g ∗ ( x ) is pushed to have positive inner products with all g i ( x ) .
If that is impossible,{ g i ( x ) } mi =1 will contain conflicting directions and we will get g ∗ ( x ) =
which will terminate theprocedure.
Using Lagrangian duality, [7] has shown that the above objective has the optimal value
[Equation] ( x ) where { λ i g i ( x ) } i =1 m is the solution of
[Equation]
The above optimisation problem has a closed form solution when m = 2 , and a fast algorithm form > 2 [27, 17].
By construction, when the step size is small, MGD decreases monotonically allobjectives simultaneously and will terminate at local Pareto point.
2.3 Compositional energy-based models
Although EBMs have been around since many years ago [20, 19], they have recently been rediscoveredin the quest of generative modeling.
More concretly EBMs learn to represent data by approximatingan unnormalized probability distribution across.
They do so by fitting an energy function E θ ( x ) foreach data point, parameterized by a neural network, which outputs a scalar real energy such that themodel distribution is:
[Equation]
Training an EBM on a given data distribution is done using contrastive divergence [16].
Besidesapproximating data distributions, EBMs have the inherent generative property and so we can use themto sample new data points x - from p θ during both training and generation.
For this we use MCMCbased on Langevin dynamics [29].
In this framework samples are initialized from gaussian/uniformrandom noise followed by iterative refinement:
[Equation]
where k is the k − th iteration step and η is the step size.
This sampling procedure allows us to use thegradient of a combined distribution to generate samples from distributions composed of p θ and otherdistributions.
We use this ability to generate from composed distributions, consisted of m concepts c ,as proposed by [8], an approach naturally named compositional EBM (cEBM).
[8] propose multiplevariations for executing different logical operations over properties of interest such as conjunction,negation and disjunction.
In the rest of this paper we only consider conjunction, which correspondsto combining individual EBMs as product of experts:
[Equation]
Furthermore, the sampling from a composed models is correspondingly adapted:
[Equation]
Pareto compositional EBM
[22] improves MGD multi objective optimization via sampling, encouraged by the opportunity toaccess a diverse set of solutions i.e. Pareto front (not only a single point local Pareto estimate), andimportantly, compared to existing methods, without relaying on predefined preference functions (asin Linear scaling).
The main idea is to incorporate the MGD minimum norm gradient updates withinsampling optimization procedures as Stein variational gradient descent [21] or Langevin dynamics[29]:
[Equation]
Where α is positive constant scaling the Gaussian noise ω .
Remembering that g ∗ ( x ) ≈ (cid:80) mi
=1 λ ∗ i g i ( x ) and that g i ( x ) = ∇ i
f i ( x ) , and replacing the gradientsin Equation 2.3 with minimum norm gradients, leads us to the pareto optimal compositional EBM,pcEBM .
When generating new sequence x ∗ , weather starting form a random noise or improving onan existing one, at each iteration, if x ∗ is far away from the Pareto front with large multiple gradients|| g ∗ ( x ∗ )
|| , Equation 3 will mainly drive x ∗ towards the Pareto front.
On the other hand, when x ∗ isclose to the Pareto front with small || g ∗ ( x ) || , the noise will dominate and let x ∗ randomly diffuse the set.
Differently from [22] that uses MGD with sampling for multi property optimization, here, weleverage the joint framework of MGD and cEBM for multi-property sampling.
Empirical evaluation
Next, we empirically evaluate the four considered baselines: multiple gradient descent, ( MGD ), compositional
EBM, ( cEBM ), linearly scaled cEBM ( ls-cEBM ) andpareto compositional EBM, ( pcEBM ).
Experimental setup We wish to generate plausible antibody sequences.
Antibodies are proteins constructedfrom a sequence of amino-acids whose functional properties depend on the contents and ordering of that sequence.
Therefore, in order to sample new, or improve properties ofexisting antibodies, we choose to implement our baselineswith sequence based convolution networks [11].
Moredetails on the architecture can be found in Appendix.
The datasets we use are from both public and proprietarysources.
We consider three significant properties for thegenerated sequences: (i) similarity to known human antibodies, Ab-like , (ii) affinity, or binding to antigen targets of interest Aff and (iii) stickiness ornon-specificity of a the molecule Stc .
To compare the performance of the baselines in terms of satisfying multiple-properties, we include ametric from the multi-objective optimization literature - Hyper Volume ( HV ), [14], computed overthe energy scores per property model.
Additionally, to validate the similarity of the new sampledsequences (a proxy for realness) to the ones seen during the training of each property energy basedmodel; we include edit distances ( e dist ) [23].
The edit-distance represents the score of the bestpossible alignment between the two genetic sequences over all possible alignments.
The computationof both metrics is included in appendix.
Next, we empirically evaluate the four considered baselines: multiple gradient descent, ( MGD ),compositional EBM, ( cEBM ), linearly scaled cEBM ( ls-cEBM ) and pareto compositional EBM,( pcEBM ).
First, we train three separate EBMs one for each property of interest.
Once they have been trainedto convergence, we use the same models for sampling new sequences with each of the baselinesusing Langevin dynamics.
The differences in implementation are as follows: MGD does not useany noise, at each iteration, simply optimizes the gradients, cEBM uses LD but has uniform weightsfor all properties, ls-cEBM is same as cEBM but including domain-informed weights, and finallypcEBM learns the optimal weights for each property per sequence.
Each baseline has the same set ofimportant hyper parameters, namely step size - η , number of steps - k .
Results and discussion To quantitatively evaluate the effectiveness of the proposed methods, weaddress the following research questions:
• Q.1 Can we use EBMs to generate valid/reasonable 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?
To tackle Q.1 , we include Table 4.
We notice that multi property EBMs have smaller distancesto all properties compared to the single objective ones as expected have sequences most similarto the property they correspond for, but are less like the other objectives.
Additionally from the [Table] multi-objective baselines, pcEBM generates sequences closest to the training datasets.
The differencein the results regarding different step sizes in the gradient descent, hint a faster convergence of thepcEBM variant.
We confirm this in Figure 2, which presents the energy scores at each iteration stepfor both cEBM and pcEBM for the same step size η = 1 e − 2 .
For the same starting sequence, theslope pcEBM converges much faster.
This result also answers partially Q.3 .
To further investigate the motivation to use multiple gradients toachieve better Pareto optimally, we calculate the hyper volume acrossthe three properties we wish to minimize for.
Here we consider theenergy of each properly EBM as proxy for the objective {+of interest 
+} The advantage of pcEBM is noticeable only at smaller steps, butalso the variance in both edit distance and hypervolume for differentη is much smaller for pcEBM than others.
However, if we turn to Figure 1, where we plot the pareto front obtained from candidatesequences from each baseline, we see that the pcEBM front hasspread than competitors.
Moreover, cEBM and pcEBM fonts arecloser to the origin, which in out case of minimizing objectivesshows the advantage of multi properly optimization with Langevindynamics, i.e. our Q.2 .
Finally, for Q.4 we wish to explore the possibility for improvingproperties of already existing antibodies by using the variants ofcompositional EBMs.
Due to time and space limitations, in the current version we present results forthe property of stickiness Stc .
We start by screening the existing sequences to select only those withbad Stickiness scores.
We have an external SeqCNN predictor (classifier) that acts as surrogate oraclefor stickiness.
We use that oracle to evaluate the probability for having a good Stc score, that is, themolecule is likel to not stick to anything in it’s environment.
We pass the selected sequences withmeasured (as well as predicted) bad scores, through cEBM and pcEBM respectively.
Next, we usethe surrogate oracle to evaluate the stickiness of the newly proposed designs.
In Table 2, we see theimprovements per baseline, i.e. The probability for a good score is increased by at least 30% with the EBMs.
Figure 4 (a) and (b) offer an additional visual confirmation - they represents a 2D tSNE plotof all sequences with measured Stickiness, colored by their score, blue means good and red a bad [Table] [Figure] score.
On the same plot, we overlay an example sequence (seed) with particularly bad stickiness andits corresponding proposed designs.
We see that both cEBM and pcEBM improve the probability fora good stickiness score, and move the seed towards the section of the manifold covered by datapointswith good scores.
Conclusion
We propose a new approach to sampling data points that simultaneously satisfy multiple desiredproperties to the best possible extent, i.e. they are Pareto-optimal.
We do so by integrating multiplegradients within compositional energy based models.
In a series of experiments for generating andimproving existing sequences, we validate the performance of the compositional EBMs as well astheir Pareto optimality.
The generality and modularity of pcEBM allow for multiple expansionswe are excited to explore: (i) other data types, in particular graph structures suitable for molecules,(ii) incorporating other sampling frameworks such as Stein variational gradient descent [21], (iv)incorporating uncertainty estimates and (iv) pareto-optimal training for EBMs, to name a few.
Appendix5.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 3 Conv1D layers with kernel size 9 and padding 1, with ReLUnon-linearities and penultimate layer of size 256.
All EBMs were trained with contrastive training,
Adam optimizer and early stopping criterion.
Hyperparameters For each multipropery sampling baseline, a hyper parameter grid was exploredwith 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} 5.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
5.3 Additional results from experiments
HV(Ab-Like, Aff, Stc)
HV(Aff, Stc)
HV(Ab-Like, Stc)
HV(Ab-Like, Aff) η =
[Equation]
MGD
[Table]