Antigen-Specific Antibody Design and Optimization with Diffusion-Based Generative Models for Protein Structures
Abstract
Antibodies are immune system proteins that protect the host by binding to specific antigens such as viruses and bacteria.
The binding between antibodies and antigens is mainly determined by the complementarity-determining regions (CDR) of the antibodies.
In this work, we develop a deep generative model that jointly models sequences and structures of CDRs based on diffusion probabilistic models and equivariant neural networks.
Our method is the first deep learning-based method that generates antibodies explicitly targeting specific antigen structures and is one of the earliest diffusion probabilistic models for protein structures.
The model is a “Swiss Army Knife” capable of sequence-structure co-design, sequence design for given backbone structures, and antibody optimization.
We conduct extensive experiments to evaluate the quality of both sequences and structures of designed antibodies.
We find that our model could yield competitive results in binding affinity measured by biophysical energy functions and other protein design metrics.
Introduction
Antibodies are important immune proteins generated during an immune response to recognize and neutralize the pathogen [Janeway et al., 2001].
As illustrated in Figure 1a, an antibody contains two heavy chains and two light chains, and their overall structure is similar.
Six variable regions determine the specificity of an antibody to the antigens.
They are called the Complementarity Determining
Regions (CDRs), denoted as H1, H2, H3, L1, L and L3.
Therefore, the most important step for developing effective therapeutic antibodies is to design CDRs that bind to the specific antigen [Presta, 1992, Akbar et al., 2022a]
Similar to other protein design tasks, the search space of CDRs is vast.
A CDR sequence with L amino acids has up to 20 L possible protein sequences.
It is not feasible to test all the possible sequences using experimental approaches, so computational methods are needed.
Traditional computational approaches rely on sampling protein sequences and structures from complex biophysical energy functions [Pantazes and Maranas, 2010, Lapidoth et al., 2015, Adolf-Bryfogle et al., 2018, Warszawski et al., 2019].
They are generally time-consuming and are prone to get trapped in local optima.
Recently, various deep generative models have been developed to design antibodies [Saka et al., 2021,
Akbar et al., 2022b, Jin et al., 2022].
Compared to conventional algorithms, deep generative models
Antibody
Antigen
Heavy Chain
Light Chain
CDR
CDR-H
CDR-L 2 1 3
H2 H1H3 L3 L1 L (a)
N C α R C N
C α R C N C α R C N C α R C N C α R C N C α
R
C
NC α R C
Antigen
CDR
Interaction
N C α R C N C α R C
N C α R C N C α R C (b)
Non-interacting
Interacting
Same C α Distance Between Amino Acids could directly capture higher-order interactions among amino acids on antibodies and antigens and generate antibodies more efficiently [Akbar et al., 2022a].
Recently, Jin et al. proposed a generative model for antibody structure-sequence co-design.
Their model addresses two important computational challenges: First is how to model the intrinsic relation between CDR sequences and 3D structures, and second is how to model the distribution of CDRs conditional on the rest of the antibody sequence.
However, there is still a large gap to fill before generative models become practical for antibody design.
Here, we identify another three challenges for antibody sequence-structure co-design.
First , the model should be explicitly conditional on the 3D structures of the antigen and generate CDRs that fit the antigen structure in the 3D space.
This is indispensable for the model to generalize to new antigens.
Second , the interactions between amino acids are mainly determined by side-chains which are groups of atoms stretching out from the protein backbone (Figure 1b)
[Liljas et al., 2016].
Therefore, the model should be able to consider both the position and orientation of amino acids.
Third , in drug discovery, pharmacologists collect multiple initial antibodies either from humanized mice or patients [Presta, 1992, Barlow et al., 2018, Warszawski et al., 2019].
Therefore, instead of de novo design, the model should be applicable to another realistic scenario: optimizing a particular antibody to increase the binding affinity to the antigen.
To the best of our knowledge, no previous machine learning model satisfies all of the above design principles.
AntigenStructure
AntibodyFramework
CDR Designs
ARALYYYDSSGYDAYYFDYARGRRVSAPNAYGAREGLQ ARGHTLYGAVLMENYHYGMARGLDPLYCNNSTTCYRVG
[Figure]
a t i v e
M o d e l
To address these challenges, we propose a diffusion-based generative model [Sohl-Dickstein et al., 2015, Song and
Ermon, 2019, Ho et al., 2020, Yang et al., 2022] capable of jointly sampling antibody CDR sequences and structures.
More importantly, the joint distribution of a CDR sequence and its structure is directly conditional on antigen structures .
Given a protein complex consisting of an antigen and an antibody framework as input 2 (as illustrated in Figure 2), we first initialize the CDR with an arbitrary sequence, positions, and orientations.
The diffusion model first aggregates information from the antigen and the antibody framework.
Then, it iteratively updates the amino acid type, position, and orientation of each amino acid on CDRs.
In the last step, we reconstruct the CDR structure at the atom level using side-chain packing algorithms based on the predicted orientations [Alford et al., 2017].
From the perspective of model capability, one of the most important reasons for us to choose the diffusion-based model over other generative models such as generative adversarial networks [Goodfellow et al., 2014] and variational auto-encoders [Kingma and Welling, 2013] is that it generates CDR candidates iteratively in the sequence-structure space so that we can interfere and impose constraints on the sampling process to support a broader range of design tasks.
We summarize our contributions as follows:
• We propose the first deep learning models to perform antibody sequence-structure design by considering the 3D structures of the antigen.
• In our model, we not only design protein sequences and coordinates but also side
-chain orientations (represented as SO(3) element) of each amino acid.
It is the first deep learning
model that could achieve atomic-resolution antibody design and is equivariant to rotation and translation.
• We show that our model can be applied to a wide range of antibody design tasks, including sequence-structure co-design, fix-backbone CDR design, and antibody optimization.
Related Work
Computational Antibody Design Conventional computational approaches are mainly based on sampling algorithms over hand-crafted and statistical energy functions and iteratively modify protein sequences and structures [Adolf-Bryfogle et al., 2018, Lapidoth et al., 2015, Warszawski et al., 2019,
Pantazes and Maranas, 2010, Ruffolo et al., 2021].
These methods are inefficient and prone to getting stuck at local optima due to the rough energy landscape.
In recent years, deep learning methods have shown potential in antibody design by using language models to generate protein sequences [Alley et al., 2019, Shin et al., 2021, Saka et al., 2021, Akbar et al., 2022b].
Although much more efficient, the sequence-based methods can only generate new antibodies based on previously observed antibodies but can hardly generate antibodies for specific antigen structures.
Jin et al.
proposed the first CDR sequence-structure co-design deep generative model which focuses on designing antibodies to neutralize SARS-CoV-2.
It relies on an additional antigen-specific predictor to predict the neutralization of the designed antibodies, which is not generalizable to arbitrary antigens .
In comparison to their model, we explicitly model the 3D structure of an antigen, opening the door to generalizing the prediction to unseen antigens with solved 3D structures.
Another advantage of our model is that we consider not only backbone atom coordinates but also the orientation of amino acids.
The orientation is critical to protein-protein interactions as most of the atoms interacting between antibodies and antigens are in the side-chain
[Liljas et al., 2016] (as illustrated in Figure 1b).
Lastly, the model proposed by Jin et al.
is not equivariant by construction, which is fundamental in molecular modeling.
Protein Structure Prediction Protein structure prediction algorithms take protein sequences and
Multiple Sequence Alignments (MSAs) as input and translate them to 3D structures [Jumper et al., 2021, Baek et al., 2021, Yang et al., 2020].
Accurate protein structure prediction models predict not only the position of amino acids but also their orientation [Jumper et al., 2021, Yang et al., 2020].
The orientation of amino acids determines the direction in which its side chain stretches, so it is indispensable for reconstructing full-atom structures.
AlphaFold2 [Jumper et al., 2021] predicts per-amino-acid orientations in an iterative fashion, similar to our proposed model.
However, it is not generative, unable to efficiently sample diverse structures for protein design.
Recently, based on prior protein structure prediction algorithms, methods for predicting antibody CDR structures have emerged [Ruffolo et al., 2022b,a], but they are not able to design CDR sequences.
Diffusion-Based Generative Models Diffusion probabilistic models learn to generate data via denoising samples from a prior distribution [Sohl-Dickstein et al., 2015, Song and Ermon, 2019,
Ho et al., 2020].
Recently, progress has been made in developing equivariant diffusion models for molecular 3D structures [Shi et al., 2021, Hoogeboom et al., 2022, Jing et al., 2022].
Atoms in a molecule do not have natural orientations, so the generation process differs from generating protein structures.
Diffusion models have also been extended to non-Euclidean data, such as data in the
Riemannian manifolds [Leach et al., 2022, De Bortoli et al., 2022].
These models are relevant to modeling orientations which are represented by elements in SO(3) .
In addition, diffusion models can also be used to generate discrete categorical data [Hoogeboom et al., 2021, Austin et al., 2021].
Methods
This section is organized as follows: Section 3.1 introduces notations used throughout the paper and formally states the problem.
Section 3.2 formulates the diffusion process for modeling antibodies.
Section 3.3 introduces details about the neural network parameterization for the diffusion processes.
Section 3.4 presents sampling algorithms for various antibody design tasks.
3.1 Definitions and Notations
An amino acid in a protein complex can be represented by its type, C α atom coordinate, and the orientation, denoted as s i ∈ { ACDEFGHIKLMNPQRSTVWY } , x i ∈ R 3 , O
i ∈ SO(3) , respectively.
Here i = 1 . . .
N , and N is the number of amino acids in the protein complex 3 .
In this work, we assume the antigen structure and the antibody framework is given (Figure 2), and we focus on designing CDRs on the antibody framework.
Assume the CDR to be generated has m amino acids with index from l +1 to l + m .
They are denoted as R = { ( s j , x j , O j ) | j = l + 1 , . . .
, l + m } .
Formally, our goal is to jointly model the distribution of
R given the structure of the antibody-antigen complex C =
{ ( s i , x i , O i ) | i ∈ { 1 . . .
N }\{ l + 1 , . . .
, l + m }} .
3.2 Diffusion Processes
A diffusion probabilistic model defines two Markov chains of diffusion processes.
The forward diffusion process gradually adds noise to the data until the data distribution approximately reaches the prior distribution.
The generative diffusion process starts from the prior distribution and iteratively transforms it to the desired distribution.
Training the model relies on the forward diffusion process to simulate the noisy data.
Let (cid:0) s tj , x tj , O tj (cid:1) denote the intermediate state of amino acid j at time step t .
R t = { s tj , x tj , O tj } l + m j = l +1 represents the sequence and structure sampled at step t .
t = 0 represents the state of real data (observed sequences and structures of CDRs) and t =
T represents samples from the prior distribution.
Forward diffusion goes from t = 0 to T , and generative diffusion proceeds in the opposite way.
The diffusion processes for amino acid types s tj , coordinates x tj , and orientations
O tj are defined as follows:
Multinomial Diffusion for Amino Acid Types
The forward diffusion process for amino acid types is based on the multinomial distribution defined as follows [Hoogeboom et al., 2021]:
[Equation]
where onehot represents a function that converts amino acid type to a 20-dimensional one-hot vector and 1 is an all-one vector.
β t type is the probability of resampling another amino acid over 20 types uniformly.
When t → T , β t type is set close to 1 and the distribution is closer to the uniform distribution.
The following probability density provides an efficient way to perturb s 0 j for timestep t during training [Hoogeboom et al., 2021]:
[Equation]
where ¯
α t type = (cid:81)
[Equation]
The generative diffusion process is defined as:
[Equation]
where F ( · )[ j ] is a neural network model taking the structure context (antigen and antibody framework) and the CDR state from the previous step as input and predicts the probability of the amino acid type for the j -th amino acid on the CDR.
Note that, different from the forward diffusion process, the generative diffusion process must rely on the structure context C and the CDR state of the previous step including positions and orientations.
The main difference between these two processes is that the forward diffusion process adds noise to data so it is irrelevant to data or contexts but the generative diffusion process depends on the given condition and full observation of the previous step.
The generative diffusion process needs to approximate the posterior q ( s t − 1 j | s tj , s 0 j ) derived from Eq.1 and
Eq.2 to denoise.
Therefore, the objective of training the generative diffusion process for amino acid types is to minimize the expected KL divergence between Eq.3 and the posterior distribution:
[Equation]
N C α C
[Table]
α C N C α
[Table]
N C α C N C α
[Table]
[Figure]
of the CDR’s sequences, positions, and orientations for the next step.
In the end, full-atom structures are constructed by the side-chain packing algorithm.
Diffusion for C α Coordinates As the coordinate of an atom could be an arbitrary value, we scale and shift the coordinates of the whole structure such that the distribution of atom coordinates roughly match the standard normal distribution.
We define the forward diffusion for the normalized C α coordinate x j as follows:
[Equation]
where β t pos controls the rate of diffusion and its value increases from 0 to 1 as time step goes from 0 to t , and ¯ α t pos = (cid:81) tτ
=1 (1 − β τ pos ) .
Using the reparameterization trick proposed by Ho et al.
, the generative diffusion process is defined as:
[Equation]
Here, G ( · )[ j ] is a neural network that predicts the standard Gaussian noise ϵ j ∼ N ( 0 , I ) added to (cid:113) ¯ α 0 pos x 0 j (scaled coordinate of amino acid j ) based on the reparameterization of Eq.6: x tj = (cid:113) ¯ α 0 pos x
[Equation]
+ (cid:113) 1 − ¯ α 0 pos ϵ j .
The objective function of training the generative process is the expected
MSE between G and ϵ j , which is simplified from aligning distribution p to the posterior q ( x t − 1 j | x tj , x 0 j )
[Ho et al., 2020]:
[Equation]
SO(3)
Denoising for Amino Acid Orientations We empirically formulate an iterative perturbdenoise scheme for learning and generating amino acid orientations represented by SO(3) elements [Leach et al., 2022].
Note that we do not use the term diffusion because the formulation does not strictly follow the framework of diffusion probabilistic models though the overall principle is the same.
Similar to the typical diffusion process, the distribution of orientations perturbed for t steps is defined as, according to Leach et al.
[Equation]
IG SO(3) denotes the isotropic Gaussian distribution on SO(3) parameterized by a mean rotation and a scalar variance [Leach et al., 2022, Matthies et al., 1970, Nikolayev and Savyolov, 1970].
ScaleRot modifies the rotation matrix by scaling its rotation angle with the rotation axis fixed [Gallier and Xu,
¯ α t ori =
[Equation]
tτ =
1 (1 − β τ ori ) , where β t ori is the variance increases with the step t .
The conditional distribution used for the generation process of orientations is defined as:
[Equation]
where H ( · )[ j ] is a neural network that denoises the orientation and outputs the denoised orientation matrix of amino acid j .
Training the conditional distribution requires aligning the predicted orientation from H ( · ) to the real orientation.
Hence, we formulate the training object that minimizes the expected discrepancy measured by the inner product between the real and the predicted orientation matrices:
[Equation]
where (cid:98)
O t − 1 j = H ( · )[ j ] is the predicted orientation for amino acid j .
The Overall Training Objective By summing Eq.4, 9, and 12 and taking the expectation w.r.t.
t , we obtain the final training objective function:
[Equation] [Equation]
To train the model, we first sample a time step t and then sample noisy states { s tj , x tj , O tj } l + m j = l +1 ∼ p by adding noise to the training sample using the diffusion process defined by Eq.2, 6, and 10.
We compute the loss using the noisy data and backpropagate the loss to update model parameters.
3.3 Parameterization with Neural Networks
In this section, we briefly introduce the neural network architectures used in different components of the diffusion process.
The purpose of the networks is to encode the CDR state at a time step t along with the context structure: { s tj , x tj , O tj } l + m j = l +1 ∪ { s ti , x ti , O ti } i
= { 1 ...N }\{ l +1 ...
l + m } , and then denoises the CDR amino acid types ( F ), positions ( G ), and orientations ( H ).
First, we adopt Multiple Layer Perceptrons (MLPs) to generate embeddings for single and pairs of amino acids.
The single amino-acid embedding MLP creates vector e i for amino acid i , which encodes the information of amino acid types, torsional angles, and 3D coordinates of all the heavy atoms.
The pairwise embedding MLP encodes the Euclidean distances and dihedral angles between amino acid i and j to feature vectors z ij .
We adopt IPA [Jumper et al., 2021], an orientation-aware roto-translation invariant network to transform e i and z ij into hidden representations h i , which aims to represent the amino acid itself and its environment.
Next, the representations are fed to three different MLPs to denoise the amino acid types, 3D positions, and orientations of the CDR, respectively.
In particular, the MLP for denoising amino acid types outputs a 20-dimensional vector representing the posterior probabilities.
The MLP for denoising C α coordinates predicts the scaled change of the coordinate in terms of the current orientation of the amino acid.
As the coordinate deviation is calculated in the local frame, we left-multiply it by the orientation matrix and transform it back to the global frame.
Formally, this can be expressed as ˆ ϵ j =
O t
j MLP G ( h j ) .
Predicting coordinate deviations in the local frame and projecting it to the global frame ensures the equivariance of the prediction, as when the entire 3D structure rotates by a particular angle, the coordinate deviations also rotate by the same angle.
The MLP for denoising orientations first predicts a so (3) vector
[Gallier and Xu, 2003].
The vector is converted to a rotation matrix M j ∈ SO(3) right-multiplied to the orientation to produce a new mean orientation for the next generative step:
[Equation]
The proposed networks are equivariant to the rotation and translation of the overall structure:
Proposition 1.
For any proper rotation matrix R ∈ SO(3) and any 3D vector r ∈ R 3 (rigid transformation ( R , r ) ∈ SE(3) ), F , G and H satisfy the following equivariance properties:
[Equation]
[Equation]
where R R t + r := { s tj , x tj + r , R O tj } l + m j = l +1 and R C
+ r :=
{ s i , x i + r , RO i } i ∈{ 1 ...N }\{ l +1 ,...,l + m } denote the rotated and translated structure.
Note that F , G , and H are not single MLPs.
Each of them includes the shared encoder and a specific MLP.
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
antigen worse.
This could explain their difference in the binding energy ( ∆ G ).
3.4 Sampling Algorithms
The sampling algorithm first samples amino acid types from the uniform distribution over 20 classes: s Tj ∼ Uniform(20) , C α positions from the standard normal distribution: x Tj ∼ N ( 0 , I 3 ) , and orientations from the uniform distribution over SO(3) :
O Tj ∼ Uniform(SO(3)) .
Note that we normalize the coordinates of the structure in the same way as training such that C α positions in the
CDR roughly follow the standard normal distribution.
Next, we iteratively sample sequences and structures from the generative diffusion kernel by denoising amino acid types, C α coordinates, and orientations until t = 0 .
To build a full atom 3D structure, we construct the coordinates of N , C α , C ,
O , and side-chain C β (except glycine that does not have C β ) according to their ideal local coordinates relative to the C
α position and orientation of each amino acid
[Engh and Huber, 2012].
Based on the five reconstructed atoms, the rest of the side-chain atoms are constructed using the side-chain packing function implemented in Rosetta [Alford et al., 2017].
In the end, we adopt the AMBER99 force field [Lindorff-Larsen et al., 2010] in OpenMM
[Eastman et al., 2017] to refine the full atom structure.
In addition to the joint design of sequences and structures, we can constrain partial states for other design tasks.
For example, by fixing the backbone structure (positions and orientations) and sampling only sequences, we can do fix-backbone sequence design .
Another usage is to optimize an existing antibody .
Specifically, we first add noise to the existing antibody for t steps and denoise the perturbed antibody sequence starting from the t -th step of the generative diffusion process.
Experiments
We present the application of our model, named DiffAb , in three antibody design tasks: sequencestructure co-design (Section 4.1), antibody sequence design based on antibody backbones (Section 4.2), and antibody optimization (Section 4.3).
In Section 4.4, we show how to use our model without known antibody frameworks bound to the antigen.
4.1 Sequence-Structure Co-design
The dataset for training the model is derived from the SAbDab database[Dunbar et al., 2014].
We first remove structures whose resolution is worse than 4Å and discard antibodies targeting non-protein
[Table] [Table]
antigens.
We cluster antibodies in the database according to CDR-H3 sequences at 50% sequence identity.
We manually select five clusters as the test set, containing 19 antibody-antigen complexes in total.
The test set includes antigens from several well-known pathogens including SARS-CoV-2,
MERS, influenza, and so on.
Structures in the remaining clusters are used for training.
To evaluate the performance, we remove the original CDR from the antibody-antigen complex in the test set and sample both the sequence and structure of the removed region.
We set the length of the CDR to be identical to the length of the original CDR for simplicity.
In practice, one can enumerate different lengths of CDRs.
We compare our model to RosettaAntibodyDesign (RAbD) [Adolf-Bryfogle et al., 2018], an antibody design software based on Rosetta energy functions.
For each model, we draw 100 samples for each CDR.
Both the original structures and designed structures from different methods are refined by OpenMM and Rosetta.
We use the following metrics to evaluate designed antibodies: (1) IMP: is the percentage of designed
CDRs with lower (better) binding energy ( ∆ G ) than the original CDR.
The binding energy is calculated by InterfaceAnalyzer in the Rosetta software package [Alford et al., 2017].
(2) RMSD: is the C
α root-mean-square deviation (RMSD) between the generated structure and the original structure with only antibody frameworks aligned .
(3) AAR: is the amino acid recovery rate measured by the sequence identity between the reference CDR sequences and the generated sequences [AdolfBryfogle et al., 2018].
Note that different from Jin et al.
[2022], we do not use neutralization prediction models because they are sequence-based and are specified to a limited class of antigens, which deviates from our goal of developing a general antibody design model.
Table 1 shows that our model (DiffAb) recovers CDR sequences more accurately than RAbD (higher
AAR).
The RMSDs of CDRs generated by DiffAb are higher in CDR-H3, which indicates that our generated samples are more diverse structurally.
The IMP score of DiffAb is on par with RAbD in
CDR-H3, and lower in other CDRs.
However, it should be noted that RAbD optimizes the Rosetta energy function, which is also used for evaluation.
Our model achieves reasonably good binding energy without explicit supervision signal from Rosetta energy functions.
Figure 4 presents three generated examples of CDR-H3 targeting SARS-CoV-2 RBD.
Sample 1 has the lowest binding energy and it can be observed that it has better complementarity to the antigen.
The binding energy of Sample 3 is higher than the original one and visually, the shape of the CDR does not fit the antigen well.
4.2 Fix-Backbone Sequence Design and Structure Prediction
In this setting, the backbone structure of CDRs is given and we only need to design the CDR sequence, which transforms the task into a constrained sampling problem.
Fix-backbone design is a common setting in the area of protein design [Ingraham et al., 2019, Hsu et al., 2022, Anishchenko et al., 2021, Strokach et al., 2020, Tischer et al., 2020].
For this task, we consider FixBB , a Rosetta-based sequence design software given CDR backbone structure, as the baseline.
We use the AAR metric introduced in Section 4.1 to evaluate the designed CDRs.
As shown in Table 2, our model achieves better AAR in all the CDRs.
This shows that our model is also powerful in modeling the conditional probability of sequences given backbone structures.
Admittedly, the training data is clustered only by CDR-H3 sequences, so the model might have seen
[Figure]
Diffusion
[Figure]
other CDRs in the test set during training, leading to even higher AAR.
However, we believe this is not an issue as CDRs other than H3 are generally conserved and contribute less to the specificity Xu and Davis [2000].
Our model can predict CDR structures by fixing the sequence.
Table 3 shows that it accurately predicts the structure of CDR H1, H2, L1, L2, and L3 (RMSD ≤ 1.5Å).
The accuracy of CDR-H prediction is lower due to the high variability.
Figure 5a separately shows the accuracy of different
CDR-H3 lengths.
The prediction is generally more accurate for shorter ones.
When the CDR-H contains more than 10 amino acids, the prediction accuracy drops.
4.3 Antibody Optimization
In contrast to redesigning the CDR, the optimization method can improve binding energy while keeping the optimized CDR similar to the original one .
Figure 5c shows the line plot of the results.
t
IMP
RMSD
SeqID 18.52% 1.194Å 92.42% 16.67% 1.252Å 91.61%
[Table] [Table]
290Å
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
We use our model to optimize existing antibodies which is another common pharmaceutical application.
To optimize an antibody, we first perturb the
CDR sequence and structure for t steps using the forward diffusion process.
Then, we denoise the sequences starting from the ( T − t ) -th step ( t steps remaining) of the generative diffusion process and obtain a set of optimized antibodies.
This process is illustrated in Figure 5b.
We optimize CDR-H of the antibodies in the test set with various t values.
For each antibody and t , we perturb the CDR independently 100 times and collect 100 optimized
CDRs different from the original CDR.
We report the percentage of optimized antibodies with improved binding energy (IMP), RMSD, and sequence identity (SeqID) of the optimized CDR in comparison to the original antibody.
We also compare the optimized antibodies with the de novo ( t = T = 100 ) designed antibodies introduced in Section 4.1.
As shown in
Table 4 and Figure 5c, the optimization method could produce antibodies with improved binding energy measured by the Rosetta energy function.
In contrast to redesigning CDRs, optimization improves binding energy while keeping the optimized
CDR similar to the original one, which is desired in many practical applications.
4.4 Design Without Bound Antibody Frameworks
In the last experiment, we consider designing antibodies without a known binding pose against the antigen, a more general and challenging setting.
We show that this challenging task could be achieved with docking software.
Specifically, we create an antibody template from an existing antibody structure by removing its CDR-H3.
This is because CDR-H3 is the most variable one and accounts
[Figure]
for most of the specificity, while other CDRs are much more conserved [Xu and Davis, 2000].
Next, we use HDOCK [Yan et al., 2017] to dock the antibody template to the target antigen to produce the antibody-antigen complex.
In this way, the problem reduces to the original problem so we can adapt our model to design the CDR-H3 sequence and structure and re-design other CDRs.
We demonstrate using this method to design antibodies for the SARS-CoV-2 Omicron RDB structure (PDB: 7wvn , residue A322-A590, the structure is not bound to any antibodies).
The antibody template is derived from a human antibody against influenza (PDB: 3qhf ).
Figure 6 shows the docking structure, five designed CDR-H3s, and the binding energy distribution.
It is hard to confidently conclude that the generated antibodies are effective without a reference antibody.
However, according to the binding energy distribution, we can still say the generated antibodies are at least reasonable.
Conclusions and Limitations
In this work, we propose a diffusion-based generative model for antibody design.
Our model is capable of a wide range of antibody design tasks and can achieve competitive performance.
One main limitation of this work is that it relies on an antibody framework bound to the target antigen.
Therefore, we leave it for future work to design an effective model for generating antibodies without bound structures.
Another limitation is that it remains unclear whether the generated antibodies can be produced in the wet lab and actually binds to the target.
More efforts are needed to design a biologically effective antibody.
Acknowledgments and Disclosure of Funding
Supported by National Key R&D Program of China No. 2021YFF1201600.
A Diffusion Processes A.1 Posteriors
Posterior of Amino Acid Types
The generative diffusion kernel for amino acid types p (
[Equation]
1 j
[Equation]
(Eq.3) should align to the posterior q ( s t − 1 j | s tj , s 0 j ) .
It can be derived from Eq.1 and Eq.2 [Hoogeboom et al., 2021]:
[Equation]
The vector inside Multinomial( · ) might not sum to one.
In this case, the probability of a class is the ratio of the value in the sum of the vector.
Posterior of C α Coordinates The generative diffusion kernel
p
[Equation]
(Eq.7) should align to the posterior obtained from Eq.5 and Eq.
[Ho et al., 2020]:
[Equation]
x
[Equation]
− 1 j
[Equation]
where µ q ( · · · ) = (cid:113)
[Equation]
A.2 Amino Acid C α Position Normalization
As amino acid C
α positions could be arbitrary in the 3D space.
We need to normalize them to use the standard normal distribution with zero mean and unit variance as the prior distribution.
First, we need to derive the statistics of CDR positions.
For each CDR in the SAbDab dataset, we shift the overall structure such that the center point of the two CDR anchors is located in origin.
Then, we aggregate C α positions in the shifted CDRs.
Finally, we calculate the mean and standard deviation of them.
Before training and inference, we shift the whole structure according to their CDR anchors and further shift and scale the structure according to the pre-calculated mean and standard deviation to obtain the normalized coordinates.
B Distributions on SO(3) B.1 Preliminary: Axis-Angle Representation of Rotations
Conventionally, a rotation is usually represented by 3 Euler angles ( α, β, γ ) , which can be interpreted as the composition of counter-clockwise rotations by α, β, γ about x, y, z axes.
However, the
Euler representation is unsuitable for defining useful operations and distributions w.r.t.
rotations considered in this work.
Alternatively, we introduce another rotation representation called axis-angle representations .
This representation parameterized a rotation with an rotational axis u ( ∥ u ∥ 2 = 1 ) and an angle θ ( θ ∈ R ).
B.2 Logarithm of Rotation Matrices and Exponential of Skew-Symmetric Matrices
Logarithm of Rotation Matrices Derived from the definition of matrix logarithm, the logarithm of a rotation matrix R is a skew-symmetric matrix
[Gallier and Xu, 2003], which can be represented as:
[Equation] [Equation]
It can be proven that v =
[ v x , v y , v z ] is the rotational axis of R , and ∥ v ∥ 2 is the rotational angle.
For brevity, we can use the vector notation v to
To efficiently compute the logarithm of a rotation matrix without computing matrix logarithm or solving rotational axis-angle, we can use the following formula [Gallier and Xu, 2003]:
[Equation]
(21) where θ can be obtained from θ = cos −(cid:0)
Tr R
[Equation]
(cid:1) by the fact that Tr( R ) = 1 + 2 cos θ .
Specially, when θ =
(or R = I ), log R = [ 0 , 0 , 0 ] .
Exponential of Skew-Symmetric Matrices
The inversion of the rotation matrix logarithm is the exponential of skew-symmetric matrices.
Derived from the definition of matrix exponential, the conversion formula is [Gallier and Xu, 2003]:
[Equation]
where S is a skew-symmetric matrix parameterized by three values v =
[ v x , v y , v z ] , identical to the definition in Eq.20.
Remarks The logarithm and exponential defined above provide an easy way to create and manipulate rotations in the axis-angle parameterization space.
For example, when we would like to create a rotation matrix with an axis and an angle, we can first create a vector v whose direction is the same as the given axis and whose length equals the angle.
Then, we rewrite the vector v into a skew-symmetric matrix S , and finally convert it to a rotation matrix by Eq.22.
We can also manipulate a rotation matrix, for example, changing its rotational angle, by mapping it to the logarithm space, modifying the skew-symmetric matrix, and finally converting it back to a rotation matrix using the exponential formula.
B.3 ScaleRot : Rotation Scaling Function
When we parameterize a rotation matrix with an axis and an angle, it is natural to define the rotation scaling function ScaleRot as scaling the rotational angle.
Formally, the definition is:
[Equation] [Equation]
where k is the scaling factor and R is a rotation matrix.
Specially, ScaleRot(0 , R )
=
I for all rotation matrix R .
Intuitively, scaling a rotation matrix by 0 cancels its effect, leading to identity transformation.
B.4 IG SO(3) : Isotropic Gaussian Distribution on SO(3)
The isotropic Gaussian distribution on SO(3) , denoted as IG SO(3) , is defined on the axis-angle space of rotation:
S 2 × [0 , π ] , where S 2 = {∥ x ∥ 2 = 1 | x ∈ R 3 } is the unit sphere in R 3 .
IG SO(3) is parameterized by a mean rotation M and a scalar variance σ 2 .
Let u ∈ S 2 and θ denote the rotational axis and angle random variables respectively.
We first consider IG SO(3) with the identity matrix as its mean: IG SO(3) ( u , θ | I , σ 2 ) .
Its p.d.f.
is defined by the product of the uniform distribution on Sand a special angular distribution [Matthies et al., 1970, Nikolayev and Savyolov, 1970, Leach et al., 2022]:
[Equation]
(24) where p
[Equation]
( S 2 ) ( u ) = 1 4 π δ ( ∥ u ∥ 2 − 1) , ( u ∈
[Equation]
(25) and p angular ( θ | σ 2 ) = 1 − cos θ π
[Equation]
[Equation]
When the mean is other than I , to sample from the distribution, we can first sample an rotation E from IG SO(3) ( u , θ | I , σ 2 ) .
Then, we left-multiply R to E to obtain the desired random value RE .
Sampling The algorithm for drawing samples from IG SO(3) ( I , σ 2 ) (here, the mean rotation is identity) can be broken down into two steps.
The first step is to draw a unit vector u from the uniform distribution on S 2 , p uniform ( S 2 ) ( u ) .
This can be done efficiently by sampling from the 3D standard Gaussian distribution and then normalizing the sampled vector to unit length.
The second part is drawing samples from p angular ( θ | σ 2 ) , which could be more tricky.
We empirically use two different proximate sampling strategies depending on the variance σ 2 .
When σ is larger than 0.1, the series (Eq.26) converges fast.
In such cases, we use histograms to approximate the distribution.
In specific, we evenly partition [0 , π ] into 8192 bins, and use the probability density p angular ( θ | σ 2 ) at the center of each bin as the bin weight.
We randomly select a bin according to the weights to draw samples from the discretized distribution.
Then, we sample from the uniform distribution spanning from the lower bound to the upper bound of the bin.
The discretization process is time-consuming.
However, since the variances in the diffusion processes are predetermined, we pre-compute and cache the bins and weights to draw samples efficiently.
When σ is smaller than 0.1, we approximate the distribution using the truncated Gaussian distribution whose mean is 2 σ and the standard deviation is σ .
Empirically, we find that the above proximate sampling algorithm is sufficient for training and sampling from our diffusion model.
To sample from IG SO(3) with an arbitrary mean rotation R , we first draw a rotation from
IG SO(3)
( I , σ 2 ) , denoted as E .
Finally, we left-multiply R to E to get the desired sample.
B.5 Uniform Distribution on SO(3)
The uniform distribution on SO(3) is equivalent to the uniform distribution of normalized quaternions on S 3 [Shoemake, 1992].
To sample a random rotation uniformly, we first sample a random vector from the 4D standard normal distribution.
Next, we normalize the vector and treat it as a quaternion.
Finally, we convert the quaternion to a rotation matrix which can be regarded as a sample from the uniform distribution on SO(3) .
C Neural Network Parameterization C.1 Computing Residue Orientations
The orientation of a residue is determined by the coordinate of its three backbone atoms: C α , C , and
N .
Let x αi , x C i , and x N i denote the 3D coordinates of the three backbone atoms of the i -th residue respectively.
The orientation of the residue, denoted by O i , can be constructed using the following
Gram-Schmidt-based algorithm:
[Equation]
C.2 Architectures Amino Acid Embedding
Layer The embedding layer for each amino acid takes into account the following information:
• Amino acid type : Each of the 20 amino acid types is represented by an embedding vector denoted by e type i .
• Heavy atom local coordinates : The coordinate of each heavy atom in an amino acid is projected to the local coordinate frame using the rule x local i
[Equation]
⊺ i
[Equation]
− x αi ) .
All of the local coordinates are concatenated into a single vector denoted by e coord i .
If some heavy atoms are missing, their local coordinates are filled with zeros.
Note that the local coordinates are invariant to global rotation and translation thanks to the projection rule.
• Backbone dihedral angles : The backbone dihedrals of amino acid, including ϕ , ψ , and ω [Liljas et al., 2016, Ingraham et al., 2019], are transformed using a series of sine and cosine functions with different frequencies, which are then concatenated into a single vector e dihed i .
• CDR flags and anchor flags : Amino acids on the CDR or by the two ends of the CDR (anchors) are differentiated from other amino acids by special 0-1 flags denoted as e flag i .
All of the vectors above are concatenated and fed to an MLP to produce the final embedding vector for each residue.
Pairwise Embedding Layer Pairwise embeddings include information about the relationship between two residues.
The pairwise embedding for residue i and j involves the following information:
• Amino acid types of both amino acids : There are 20 × 20 = 400 combinations of two amino acid types.
We represent each of them using an embedding vector denoted by z type ij .
• Sequential relative position : If two residues are on the same chain and their distance on the sequence is less than or equal to 32 ( d seq ij ∈ {− 32 . . .
, the distance is represented by an embedding vector z seq ij .
Otherwise, the distance embedding is filled with zeros.
• Pairwise distances :
The distances between all pairs of atoms are flattened into a vector and transformed by e − c d ij ( c is a learnable coefficient) into the spatial distance embedding z dist ij .
Missing pairs are filled with zeros.
• Pairwise backbone dihedrals : The backbone dihedrals between any two amino acids i and j are defined as ϕ ij =
Dihedral( x C i , x N j , x αj , x C j ) and ψ ij =
Dihedral( x N i , x αi , x C i , x N j ) .
These two dihedrals are transformed by a series of sine and cosine functions into pairwise dihedral embeddings z dihed ij .
We concatenate the above vectors and feed them into an MLP to get the final pairwise embeddings for each pair of amino acids z ij .
Encoder
The encoder for encoding the current diffusion state consists of a stack of orientationaware invariant 3D attention layers.
Its aim is to capture relationships between amino acids and provide high-level representations for each residue to denoise.
Let h ℓi denote the hidden representation output from the last layer (when ℓ = 0 , the representation is the initial residue embedding).
The formulas for computing the logit of attention weight between residue i (query) and j (key) is:
[Equation]
where q ( · ) , k ( · ) , f ( · ) , and g ( · ) are MLP subnetworks.
The attention weights can be obtained by taking softmax: w ij = softmax Nj =1 ( a ij ) .
Note that, for simplicity, we do not consider attention heads in the formula, but in practice we use multiple attention heads and different heads can be combined easily via concatenation.
The formula for computing the value passed from residue j to i is:
[Equation]
where v ( · ) is a network consisting of MLPs.
Finally, the values along with attention weights are used to update the amino acid representations with residual connection and layer normalization, same as the standard transformer [Vaswani et al., 2017].
C.3 Notes on the Notations of the Denoising Networks F , G , and H
The notations F , G , and H do not only denote the MLPs following the encoder that outputs denoising results.
It refers to the embedding layers, the encoder, and the specific output MLP (for example,
F includes the MLP for denoising amino acid types).
Therefore, the input to F , G , and H is the diffusion state (sequence and structure) rather than hidden representations.
Treating the three sections as a whole allows us to neatly express the equivariance property of the model.
C.4 Proof of Equivariance
Lemma 1.
The Euclidean distance function between two points is invariant to rotations and translations, i.e. d ( Rx 1 + r , Rx 2 + r )
=
d ( x 1 , x 2 ) , ∀ R ∈ SO(3) , r ∈ R 3 .
Proof.
[Equation]
Lemma 2.
The dihedral function for four points is invariant to rotations and translations, i.e.
Dihedral( Rx 1 + r , Rx 2 + r , Rx 3 + r , Rx 4 + r ) =
Dihedral( x 1 , x 2 , x 3 , x 4 ) , ∀ R ∈ SO(3) , r ∈ R 3 .
Here, Dihedral( · · · ) is defined as:
[Equation]
where v i
= x i +1 − x i ( i = 1 , 2 , 3 ).
Proof.
First, we note that:
[Equation]
By the equivariance of cross product ( Ra × Rb = R ( a × b ) ) and the invariance of inner product ( Ra · Rb = a · b ), we have:
[Equation]
Lemma 3.
The per-amino-acid orientation O i is equivariant to rotations and translations, i.e.,
O ( Rx αi + r , Rx C
i + r , Rx N i + r ) =
RO ( x αi , x C i , x N i )
Proof.
First, we show that the first two basis vectors e
and e 2 are equivariant:
[Equation]
Let v 2 =
x N i − x αi .
We have ( Rx N i + r ) − ( Rx αi + r ) =
Rv 2 .
Then, we can prove the equivariance of e 2 :
[Equation]
By the equivariance of cross product, it is straightforward to show that e 3 is also equivariant.
Finally, combining the equivariance of e 1 , e 1 , and e 3 , we prove the equivariance of the orientation matrix:
[Equation]
Lemma 4.
The per-amino-acid and pairwise embedding layers are invariant to rotations and translations of the input structure.
i.e. e ( s i , { x atom i } atom , ϕ i , ψ i , ω
i , e flag i ) =
e ( s i , { Rx atom i + r } atom , ϕ i , ψ i , ω
i , e flag i ) , and z ( { d ( x atom1 i , x atom2 j ) } atom1, atom2 , · · · ) =
z ( { d ( Rx atom1 i + r , Rx atom2 j + r ) } atom1, atom2 , · · · ) .
Proof.
Before embedding atom positions for an amino acid, the network first projects the positions using the orientation by the rule:
[Equation]
The projection operation is invariant to rotations and translations, using Lemma 3:
[Equation]
The formulas for computing dihedral angles ( ϕ i , ψ i , ω i ) are also invariant by Lemma 2 Other variables (amino acid types and CDR flags) are independent of the 3D structure and hence they are invariant.
So far, we have shown that all the components of embedding layers are invariant to rotations and translations of the overall 3D structure.
Therefore, the embedding layer is invariant.
Pairwise embedding layers involve distances between residues, which are invariant by Lemma 2.
Other variables are irrelevant to 3D structures.
Hence, the pairwise embedding layer is invariant.
Lemma 5.
The orientation-aware attention layer is invariant to rotations and translations if the input hidden representations
h i , z ij ( i, j = 1 . . .
N ) come from invariant functions.
Proof.
First, we show that projecting atoms on the j -th amino acid to the orientation of the i -th amino acid is invariant to rotations and translations by Lemma 3:
[Equation]
As other inputs to the attention layer ( h i , z ij ( i, j = 1 . . .
N ) ) are invariant to rigid transforms on the structure, the networks for computing attention weights and values are invariant.
Hence, the attention layer is invariant.
In the case where we stack multiple attention layers, each layer outputs invariant representations for its next layer.
Therefore, the network consisting of multiple attention layers is invariant.
Proposition 2.
For any proper rotation matrix R ∈ SO(3) and any 3D vector r ∈ R 3 (rigid transformation ( R , r ) ∈ SE(3) ), F , G and H satisfy the following equivariance properties:
[Equation]
where R R t + r := { s tj , x tj + r , R O tj } l + m j = l +1 and R C
+ r :=
{ s i , x i + r , RO i } i ∈{ 1 ...N }\{ l +1 ,...,l + m } denote the transformed structure.
Proof.
By Lemma 5, we know that the encoder network produces invariant representations.
Therefore, the MLP for predicting amino acid types that transforms the invariant representations into a probability over 20 categories is invariant, so F is invariant.
The MLP for predicting local coordinate changes MLP G ( h i ) is invariant.
The local coordinate change is converted to the global coordinate change using the following rule:
[Equation]
By Lemma 3, the above rule is equivariant to rotations, and hence G is equivariant to rotations.
Similarly, the MLP for predicting changes in orientation MLP H
( h i ) is invariant.
The changes is applied to the original orientation by:
[Equation]
which is equivariant to rotations according to Lemma 3.
Therefore, H is equivariant to rotations.
D Sampling Algorithms D.1 Backbone Atoms and Sidechain C β Construction
The coordinates of backbone atoms ( N , C α , C , O ) and sidechain C β can be determined by the orientation and the C α position of an amino acid because the geometry of these atoms is inflexible [Liljas et al., 2016].
To construct the position of N , C α , C , and C β for the i -th amino acid, we use the following formula:
[Equation]
where O i and x i is the model-predicted amino acid orientation and C α position.
c atom is the local coordinate derived from experimental data relative to the orientation and the C α position, as shown in the following table.
Atom c x c y c
z
N -0.
C α 0.
C 1.
C β -0.
The position of O depends on the ψ angle of the amino acid, which relies on the next amino acid in the sequence.
Therefore, after constructing backbone atoms, we calculate the ψ angle for each amino acid ( ψ i = Dihedral(N
i , C iα , C i , N i +1 ) ), and use the following rule to construct O coordinates:
[Equation] [Equation]
where
[Equation]
D.2 Sidechain Packing and Full Atom Refinement
We use PackRotamersMover in PyRosetta
[Chaudhury et al., 2010] to pack sidechains only for amino acids on the generated CDR.
The packing program is based on the Dunbrack 2010 rotamer library [Shapovalov and Dunbrack Jr, 2011] and the REF2015 energy function [Alford et al., 2017].
After packing sidechains, we refine the structure with OpenMM [Eastman et al., 2017].
Specifically, we first use PDBFixer to prepare the structure for refinement.
We minimize the potential energy of the structure.
The potential energy is AMBER99SB force field plus quadratic constraint terms that restrain the position of atoms outside the generated CDR.
E Source Code
Code and data are available at https://github.com/luost26/diffab