I TERATIVE R EFINEMENT G RAPH N EURAL N ETWORK FOR A NTIBODY S EQUENCE -S TRUCTURE C O - DESIGN
Anonymous authors
Paper under double-blind review
A BSTRACT
Antibodies are versatile proteins that bind to pathogens like viruses and stimulate the adaptive immune system.
The speciﬁcity of antibody binding is determined by complementarity-determining regions (CDRs) at the tips of these Y-shaped proteins.
In this paper we propose a generative model to automatically design the CDRs of antibodies with enhanced binding speciﬁcity or neutralization capabilities
Previous generative approaches formulate protein design as a structureconditioned sequence generation task, assuming the desired 3D structure is given a priori.
In contrast, we propose to co-design the sequence and 3D structure of
CDRs as graphs.
Our model unravels a sequence autoregressively while iteratively reﬁning its predicted global structure.
The inferred structure in turn guides subsequent residue choices.
For efﬁciency, we model the conditional dependence between residues inside and outside of a CDR in a coarse-grained manner.
Our method achieves superior log-likelihood on the test set and outperforms previous baselines in designing antibodies capable of neutralizing the SARS
-CoV- virus
I NTRODUCTION
Monoclonal antibodies are increasingly adopted as therapeutics targeting a wide range of pathogens such as SARS-CoV-2 (Pinto et al., 2020).
Since the binding speciﬁcity of these Y-shaped proteins is largely determined by their complementarity-determining regions (CDRs), the main goal of computational antibody design is to automate the creation of CDR subsequences with desired properties.
This problem is particularly challenging due to the combinatorial search space of over 20 60 possible CDR sequences and the small solution space which satisﬁes the desired constraints of binding afﬁnity, stability, and synthesizability (Raybould et al., 2019).
There are three key modeling questions in CDR generation.
The ﬁrst is how to model the relation between a sequence and its underlying 3D structure.
Generating sequences without the corresponding structure (Alley et al., 2019; Shin et al., 2021) can lead to sub-optimal performance (Ingraham et al., 2019), while generating from a predeﬁned 3D structure (Ingraham et al., 2019) is not suitable for antibodies since the desired structure is rarely known a priori (Fischman & Ofran, 2018).
Therefore, it is crucial to develop models that co-design the sequence and structure.
The second question is how to model the conditional distribution of CDRs given the remainder of a sequence ( context ).
Attention-based methods only model the conditional dependence at the sequence level, but the structural interaction between the CDR and its context is crucial for generation.
The last question relates to the model’s ability to optimize for various properties.
Traditional physics-based methods (Lapidoth et al., 2015; Adolf-Bryfogle et al., 2018) focus on binding energy minimization, but in practice, our objective can be much more involved than binding energies (Liu et al., 2020).
In this paper, we represent a sequence-structure pair as a graph and formulate the co-design task as a graph generation problem.
The graph representation allows us to model the conditional dependence between a CDR and its context at both the sequence and structure levels.
Antibody graph generation poses unique challenges because the global structure is expected to change when new nodes are inserted.
Previous autoregressive models (You et al., 2018; Gebauer et al., 2019) cannot modify a generated structure because they are trained under teacher forcing.
Thus errors made in the previous steps can lead to a cascade of errors in subsequent generation steps.
To address these problems, we propose a novel architecture which interleaves the generation of amino acid nodes with the prediction of 3D structures.
The structure generation is based on an iterative reﬁnement of a global graph rather than a sequential expansion of a partial graph with teacher forcing.
Since the context sequence is long, we further introduce a coarsened graph representation by grouping nodes into blocks.
We apply graph convolution at a coarser level to efﬁciently propagate the contextual information to the CDR residues.
After pretraining our model on antibodies with known structures, we ﬁnetune it using a predeﬁned property predictor to generate antibodies with speciﬁc properties.
We evaluate our method on three generation tasks, ranging from language modeling to SARS-CoVneutralization optimization and antigen-binding antibody design.
Our method is compared with a standard sequence model (Saka et al., 2021; Akbar et al., 2021) and a state-of-the-art graph generation method (You et al., 2018) tailored to antibodies.
Our method not only achieves lower perplexity on test sequences but also outperforms previous baselines in property-guided antibody design tasks.
R ELATED WORK
Antibody/protein design.
Current methods for computational antibody design roughly fall into two categories.
The ﬁrst class is based on energy function optimization (Pantazes & Maranas, 2010; Li et al., 2014; Lapidoth et al., 2015; Adolf-Bryfogle et al., 2018), which use Monte Carlo simulation to iteratively modify a sequence and its structure until reaching a local energy minimum.
Similar approaches are used in protein design (Leaver-Fay et al., 2011; Tischer et al., 2020).
Nevertheless, these physics-based methods are computationally expensive (Ingraham et al., 2019) and our desired objective can be much more complicated than low binding energy (Liu et al., 2020).
The second class is based on generative models.
For antibodies, they are mostly sequence-based (Alley et al., 2019; Shin et al., 2021;
Saka et al., 2021; Akbar et al., 2021).
For proteins, O’Connell et al.
; Ingraham et al.
; Strokach et al.
Karimi et al.
(2020); Cao et al.
(2021) further developed models conditioned on a backbone structure or protein fold.
Our model also seeks to incorporate 3D structure information for antibody generation.
Since the best CDR structures are often unknown for new pathogens, we co-design sequences and structures for speciﬁc properties.
Generative models for graphs.
Our work is related to autoregressive models for graph generation (You et al., 2018;
Li et al., 2018; Liu et al., 2018; Liao et al., 2019; Jin et al., 2020a).
In particular,
Gebauer et al.
(2019) developed G-SchNet for molecular graph and conformation co-design.
Unlike our method, they generate edges sequentially and cannot modify a previously generated subgraph when new nodes arrive.
While Graphite (Grover et al., 2019) also uses iterative reﬁnement to predict the adjacency matrix of a graph, it assumes all the node labels are given and predicts edges only.
In contrast, our work combines autoregressive models with iterative reﬁnement to generate a full graph with node and edge labels, including node labels and coordinates.
3D structure prediction.
Our approach is closely related to protein folding (Ingraham et al., 2018;
Yang et al., 2020a; Baek et al., 2021; Jumper et al., 2021).
Inputs to the state-of-the-art models like AlphaFold require a complete protein sequence, its multi-sequence alignment (MSA), and its template features.
These models are not directly applicable because we need to predict the structure of an incomplete sequence and the MSA is not speciﬁed in advance.
Our iterative reﬁnement model is also related to score matching methods for molecular conformation prediction (Shi et al., 2021) and diffusion-based methods for point clouds (Luo & Hu, 2021).
These algorithms also iteratively reﬁne a predicted 3D structure, but only for a complete molecule or point cloud.
In contrast, our approach learns to predict the 3D structure for incomplete graphs and interleaves 3D structure reﬁnement with graph generation.
A NTIBODY S EQUENCE AND S TRUCTURE C O -D ESIGN
Overview.
The role of an antibody is to bind to an antigen (e.g. a virus), present it to the immune system, and stimulate an immune response.
A subset of antibodies known as neutralizing antibodies not only bind to an antigen but can also suppress its activity.
An antibody consists of a heavy chain and a light chain, each composed of one variable domain (VH/VL) and some constant domains.
The variable domain is further divided into a framework region and three complementarity determining regions (CDRs).
The three CDRs on the heavy chain are labeled as CDR-H1, CDR-H2, CDR-H3, each occupying a contiguous subsequence (Figure 1).
As the most variable part of an antibody,
CDRs are the main determinants of binding and neutralization (Abbas et al., 2014).
[Figure]
Following Shin et al.
(2021); Akbar et al.
(2021), we formulate antibody design as a CDR generation task, conditioned on the framework region.
Speciﬁcally, we represent an antibody as a graph, which encodes both its sequence and 3D structure.
We propose a new graph generation approach called
ReﬁneGNN and extend it to handle conditional generation given a ﬁxed framework region.
Lastly, we describe how to apply ReﬁneGNN to property-guided optimization to design new antibodies with better neutralization properties.
For simplicity, we focus on the generation of heavy chain CDRs, though our method can be easily extended to model light chains CDRs.
Notations.
An antibody VH domain is represented as a sequence of amino acids s = s
1 s 2 · · · s n .
Each token s i in the sequence is called a residue , whose value can be either one of the 20 amino acids or a special token (cid:104) MASK (cid:105) , meaning that its amino acid type is unknown and needs to be predicted.
The VH sequence folds into a 3D structure and each residue s i is labeled with three backbone coordinates: x i,α for its alpha carbon atom, x i,c for its carbon atom, and x i,n for its nitrogen atom.
3. G RAPH REPRESENTATION
We represent an antibody (VH) as a graph G ( s ) =
( V , E ) with node features V = { v 1 , · · · , v n } and edge features E = { e ij } i (cid:54) =
j .
Each node feature v i encodes three dihedral angles ( φ i , ψ i , ω i ) related to three backbone coordinates of residue i .
For each residue i , we compute an orientation matrix
O i representing its local coordinate frame (Ingraham et al., 2019) (deﬁned in the appendix).
This allows us to compute edge features describing the spatial relationship between two residues i and j (i.e. distance and orientation) as follows
[Equation]
The edge feature e ij contains four parts.
The positional encoding E pos ( i − j ) encodes the relative distance between two residues in an antibody sequence.
The second term RBF( · ) is a distance encoding lifted into radial basis.
The third term in e ij is a direction encoding that corresponds to the relative direction of x j in the local frame of residue i .
The last term q
( O (cid:62) i O j ) is the orientation encoding of the quaternion representation q ( · ) of the spatial rotation matrix O (cid:62) i
O j .
We only include edges in the K -nearest neighbors graph of G ( s ) with K = 8 .
For notation convenience, we use G as a shorthand for G ( s ) when there is no ambiguity.
3. I TERATIVE R EFINEMENT G RAPH N EURAL N ETWORK (R EFINE GNN)
We propose to generate an antibody graph via an iterative reﬁnement process.
Let G (0) be the initial guess of the true antibody graph.
Each residue is initialized as a special token (cid:104) MASK (cid:105) and each edge ( i, j ) is initialized to be of distance 3 | i − j | since the average distance between consecutive residues is around three.
The direction and orientation features are set to zero.
In each generation step t , the model learns to revise a current antibody graph G ( t ) and predict the label of the next residue t + 1 .
Speciﬁcally, it ﬁrst encodes G ( t ) with a message passing network (MPN) with parameter θ { h ( t )
[Equation]
where h ( t ) i is a learned representation of residue i under the current graph G ( t ) .
Our MPN consists of L message passing layers with the following architecture
[Equation]
where h ( t, 0) i = v i and h ( t ) i
=
h ( t,L ) i .
FFN is a two-layer feed-forward network (FFN) with ReLU activation function.
E ( s j ) is a learned embedding of amino acid type s j .
Based on the learned residue representations, we predict the amino acid type of the next residue t + 1 (Figure 2A).
[Equation]
This prediction gives us a new graph G ( t +0 . 5) with the same edges as G ( t ) but the node label of t + is changed (Figure 2B).
Next, we need to update the structure to accommodate the new residue t +1 .
To this end, we encode graph G ( t +0 . 5) by another MPN with a different parameter ˜ θ and predict the coordinate of all residues.
[Equation]
The new coordinates x ( t +1)
i deﬁne a new antibody graph G ( t +1) for the next iteration (Figure 2C).
We explicitly realize the coordinates of each residue because we need to calculate the spatial edge features for G ( t +1) .
The structure prediction (coordinates x i ) and sequence prediction (amino acid types p t +1 ) are carried out by two different MPNs, namely the structure network ˜ θ and sequence network θ .
This disentanglement allows the two networks to focus on two distinct tasks.
Training.
During training, we only apply teacher forcing to the discrete amino acid type prediction.
Speciﬁcally, in each generation step t , residues 1 to t are set to their ground truth amino acid types s 1 , · · · , s t , while all future residues t + 1 , · · · , n are set to a padding token.
In contrast, the continuous structure prediction is carried out without teacher forcing.
In each iteration, the model reﬁnes the entire structure predicted in the previous step and constructs a new K -nearest neighbors graph
G ( t +1) of all residues based on the predicted coordinates { x ( t +1)
[Equation]
Loss function.
Our model remains rotation and translation invariant because the loss function is computed over pairwise distance and angles rather than coordinates.
The loss function for antibody structure prediction consists of three parts.
• Distance loss : For each residue pair i, j , we compute its pairwise distance between the predicted alpha carbons x ( t ) i,α , x ( t ) j,α .
We deﬁne the distance loss as the Huber loss between the predicted and true pairwise distances
[Equation]
where
distance is squared to avoid the square root operation which causes numerical instability.
• Dihedral angle loss : For each residue, we calculate its dihedral angle (cid:0) φ ( t ) i , ψ ( t ) i , ω ( t ) i (cid:1) based on the predicted atom coordinates x ( t ) i,α , x ( t ) i,c , x ( t ) i,n and x ( t ) i +1 ,α , x ( t ) i +1 ,c , x ( t ) i +1 ,n .
We deﬁne the dihedral angle loss as the mean square error between the predicted and true dihedral angles
[Equation]
• C
α angle loss : We calculate angles γ ( t ) i between two vectors x ( t ) i − 1 ,α − x ( t ) i,α and x ( t ) i,α − x ( t ) i +1 ,α as well as dihedral angles β ( t ) i between two planes deﬁned by x ( t ) i − 2 ,α , x ( t ) i − 1 ,α , x ( t ) i,α
[Equation]
In summary, the overall graph generation loss is deﬁned as L = L seq + L struct , where
[Equation]
The sequence prediction loss L seq is the cross entropy L ce between predicted and true residue types.
3. C ONDITIONAL GENERATION GIVEN THE FRAMEWORK REGION
The model architecture described so far is designed for unconditional generation — it generates an entire antibody graph without any constraints.
In practice, we usually ﬁx the framework region of an antibody and design the CDR sequence only.
Therefore, we need to extend the model architecture to
A
R
G
D
Y Predict amino acid
A
R
G
D
V YY
C … S A … V T
W G Q
G
Message passing
A
R
G
D
Y
Coarsened context sequence
Update structure Update sequence
V YY
C … S A … V T
W G Q
G
V YY
C … S A … V T
W G Q
G
Predict all coordinates 𝒢 ( t ) 𝒢 ( t +0.5)
𝒢 ( t +1) s <
l s > r
A
B
C …TIYPGDGDTGYAQKFQGKATLTADKSSKTVYMHLSSLASEDSAVYYCARGDYYGSNSLDYWGQGTSVT…
D
WGQG TSVT VYYC
Original sequence : s
Context sequence : b l , r ( s )
[Figure]
r
Conditioning via attention.
A simple extension of ReﬁneGNN is to encode the non-CDR sequence using a recurrent neural network and propagate information to the CDR through an attention layer.
To be speciﬁc, we ﬁrst concatenate s <
l
and s >r into a context sequence ˜ s
= s <l ⊕ (cid:104) MASK (cid:105) · · · (cid:104) MASK (cid:105) ⊕ s >r , where ⊕ means string concatenation and (cid:104) MASK (cid:105) is repeated n times.
We then encode this context sequence by a Gated Recurrent Unit (GRU) (Cho et al., 2014) and modify the structure and sequence prediction step (Equation 4 and 6) as
[Equation]
Multi-resolution modeling.
The attention-based approach alone is not sufﬁcient because it does not model the structure of the context sequence, thus ignoring how its residues structurally interact with the CDR’s.
While this information is not available for new antibodies at test time, we can learn to predict this interaction using antibodies in the training set with known structures.
A naive solution is to iteratively reﬁne the entire antibody structure (more than 100 residues) while generating CDR residues.
This approach is computationally expensive because we need to recompute the MPN encoding for all residues in each generation step.
Importantly, we cannot predict the context residue coordinates at the outset and ﬁx them because they need to be adjusted accordingly when the coordinates of CDR residues are updated in each generation step.
For computational efﬁciency, we propose a coarse-grained model that reduces the context sequence length by clustering it into residue blocks .
Speciﬁcally, we construct a coarsened context sequence b l,r ( s ) by clustering every b context residues into a block (Figure 2D).
The new sequence b l,r ( s ) deﬁnes a coarsened graph G ( b l,r ( s )) over the residue blocks, whose edges are deﬁned based on block coordinates.
The coordinate of each block x b i ,e is deﬁned as the mean coordinate of residues within the block.
The embedding of each block E ( b i ) is the mean of its residue embeddings.
[Equation]
Now we can apply ReﬁneGNN to generate the CDR residues while iteratively reﬁning the global graph G ( b l,r ( s )) by predicting the coordinates of all blocks.
The only change is that the structure prediction loss is deﬁned over block coordinates x
b i ,e .
Lastly, we combine both the attention mechanism and coarse-grained modeling to keep both ﬁne-grained and coarse-grained information.
The decoding process of this conditional ReﬁneGNN is illustrated in Algorithm 1.
Require:
Context sequence s <l , s >r 1: Predict the CDR length n 2: Coarsen the context sequence into b
) 3: Construct the initial graph G
) 4: for t = to n do 5:
Encode G ( t ) using the sequence MPN 6:
Predict distribution of the next residue p t +1 7:
Sample s t +1 ∼ categorical( p t +1 ) 8:
Encode G ( t +0 . 5) with the structure MPN 9:
Predict all residue coordinates x ( t +1) i,e 10:
Update G ( t +1) using the new coordinates
Require: A set of antibodies D to be optimized
Require:
A neutralization predictor f .
Require: A set of neutralizing antibodies Q 1: for each iteration do 2:
Sample an antibody s from D , remove its
CDR and get a context sequence b l,r ( s ) 3: for i = 1 to M do 4:
Sample s (cid:48)
i ∼ P Θ ( s (cid:48)
| b l,r ( s )) 5: if f ( s (cid:48) i ) > max( f ( s ) , 0 . 5) then 6:
Q ← Q ∪ { s (cid:48) i } 7:
Sample a batch of new antibodies from Q 8:
Update model parameter minimizing the
3. P ROPERTY - GUIDED SEQUENCE OPTIMIZATION
Our ultimate goal is to generate new antibodies with desired properties such as neutralizing a particular virus.
This task can be formulated as an optimization problem.
Let Y be a binary indicator variable for neutralization.
Our goal is to learn a conditional generative model P Θ ( s (cid:48) | b l,r ( s )) that maximizes the probability of neutralization for a training set of antibodies D , i.e.
[Equation]
[Table]
[Equation]
(15) where f ( s (cid:48) ) is a predictor for P ( Y = 1 | s (cid:48) ) .
Assuming f is given, this problem can be solved by iterative target augmentation (ITA)
(Yang et al., 2020b).
Before ITA optimization starts, we ﬁrst pretrain our model on a set of real antibody structures to learn a prior distribution over CDR sequences and structures.
In each ITA ﬁnetuning step, we ﬁrst randomly sample a sequence s from
D , a set of antibodies whose CDRs need to be redesigned.
Next, we generate M new sequences given its context b l,r ( s ) .
A generated sequence s (cid:48)
i is added to our training set Q if it is predicted as neutralizing.
Initially, the training set Q contains antibodies that are known to be neutralizing ( Y = 1 ).
Lastly, we sample a batch of neutralizing antibodies from Q and update the model parameter by minimizing their sequence prediction loss L seq (Eq.(10)).
The structure prediction loss L struct is excluded in ITA ﬁnetuning phase because the structure of a generated sequence is unknown.
E XPERIMENTS
Setup.
We construct three evaluation setups to quantify the performance of our approach.
Following standard practice in generative model evaluation, we ﬁrst measure the perplexity of different models on new antibodies in a test set created based on sequence similarity split.
We also measure structure prediction error by comparing generated and ground truth CDR structures recorded in the Structural
Antibody Database (Dunbar et al., 2014).
Results for this task are shown in section 4.1.
Second, we evaluate our method on an existing antibody design benchmark of 60 antibody-antigen complexes from Adolf-Bryfogle et al.
The goal is to design the CDR-H3 of an antibody so that it binds to a given antigen.
Results for this task are shown in section 4.2.
Lastly, we propose an antibody optimization task to test whether our method can design new antibodies with desired properties.
Speciﬁcally, we seek to redesign CDR-H3 of existing antibodies in the Coronavirus Antibody Database (CoVAbDab) (Raybould et al., 2021) to improve their neutralization against SARS-CoV-2.
Following works in molecular design (Jin et al., 2020b), we use a predictor to evaluate the neutralization of generated antibodies since we cannot experimentally test them in wet labs.
Results for this task are reported in section 4.3.
Baselines.
We consider three baselines for comparison (details in the appendix).
The ﬁrst baseline is a sequence-based LSTM model used in Saka et al.
(2021); Akbar et al. (2021).
This model does not utilize any 3D structure information.
It consists of an encoder that learns to encode a context sequence ˜ s , a decoder that decodes a CDR sequence, and an attention layer connecting the two.
[Table]
The second baseline is an autoregressive graph generation model (AR-GNN) whose architecture is similar to You et al.
; Jin et al.
(2020b)
but tailored for antibodies.
AR-GNN generates an antibody graph residue by residue.
In each step t , it ﬁrst predicts the amino acid type of residue t and then generates edges between t and previous residues.
Importantly, AR-GNN cannot modify a partially generated 3D structure of residues s 1 · · · s t − 1 because it is trained by teacher forcing.
On the antigen-binding task, we include an additional physics-based baseline called RosettaAntibodyDesign (RAbD) (Adolf-Bryfogle et al., 2018).
We apply their de novo design protocol composed of graft design followed by 250 iterations of sequence design and energy minimization.
We cannot afford to run more iterations because it takes more than 10 hours per antibody.
We also could not apply RAbD to the SARS-CoV-2 task because it requires 3D structures to be given.
This information is unavailable for antibodies in CoVAbDab.
Hyperparameters.
We performed hyperparameter tuning to ﬁnd the best setting for each method.
For ReﬁneGNN, both its structure and sequence MPN have four message passing layers, with a hidden dimension of 256 and block size b = 4 .
All models are trained by the Adam optimizer with a learning rate of 0.001.
More details are provided in the appendix.
4. L ANGUAGE MODELING AND 3D STRUCTURE PREDICTION
Data.
The Structural Antibody Database (SAbDab) consists of 4994 antibody structures renumbered according to the IMGT numbering scheme (Lefranc et al., 2003).
To measure a model’s ability to generalize to novel CDR sequences, we divide the heavy chains into training, validation, and test sets based on CDR cluster split.
We illustrate our cluster split process using CDR-H3 as an example.
First, we use MMseqs2 (Steinegger & S ¨ oding, 2017) to cluster all the CDR-H3 sequences.
The sequence identity is calculated under the BLOSUM62 substitution matrix (Henikoff & Henikoff, 1992).
Two antibodies are put into the same cluster if their CDR-H3 sequence identity is above 40%.
We then randomly split the clusters into training, validation, and test set with 8:1:1 ratio.
We repeat the same procedure for creating CDR-H1 and CDR-H2 splits.
In total, there are 1266, 1564, and 2325 clusters for CDR-H1, H2, and H3.
The size of training, validation, and test sets for each
CDR is shown in the appendix.
Metrics.
For each method, we report the perplexity (PPL) of test sequences and the root mean square deviation (RMSD) between a predicted structure and its ground truth structure reported in
SAbDab.
RMSD is calculated by the Kabsch algorithm (Kabsch, 1976) based on C α coordinate of
CDR residues.
Since the mapping between sequences and structures is deterministic in ReﬁneGNN, we can calculate perplexity in the same way as standard sequence models.
Results.
Since the LSTM baseline does not involve structure prediction, we report RMSD for graphbased methods only.
As shown in Table 1, ReﬁneGNN signiﬁcantly outperforms all baselines in both metrics.
For CDR-H3, our model gives 13% PPL reduction (8.38 v.s. 9.70) over sequence only model and 10% PPL reduction over AR-GNN (8.38 v.s. 9.44).
ReﬁneGNN also predicts the structure more accurately, with 30% relative RMSD reduction over AR-GNN.
In Figure 3, we provide examples of predicted 3D structures of CDR-H3 loops.
Ablation studies.
We further conduct ablation experiments on the CDR-H3 generation task to study the importance of different modeling choices.
First, when we remove the attention mechanism and context coarsening step in section 3.3, the PPL increases from 8.38 to 8.86 (Figure 3C, row 2) and 9.01 (Figure 3C, row 3) respectively.
We also tried to remove both the attention and coarsening
[Figure] [Figure]
modules and trained the model without conditioning on the context sequence.
The PPL of this unconditional variant is much worse than our conditional model (Figure 3C, row 4).
In short, these results validate the advantage of our multi-resolution conditioning strategy.
Moreover, the model performance slightly degrades when we halve the number of reﬁnement steps or increase block size to b = 8 (Figure 3C, row 5-6) .
Lastly, we train a structure-conditioned model by feeding the ground truth structure to ReﬁneGNN at every generation step (Figure 3C, row 7).
While this structure-conditioned model gives a lower PPL as expected (7.39 v.s. 8.38), it is not too far away from the sequence only model (PPL = 9.70).
This suggests that ReﬁneGNN is able to extract a decent amount of information from the partial structure co-evolving with the sequence.
A NTIGEN - BINDING ANTIBODY DESIGN
Data.
Adolf-Bryfogle et al.
(2018) selected 60 antibody-antigen complexes as an antibody design benchmark.
Given the framework of an antibody, the goal is to design its CDR-H3 that binds to its corresponding antigen.
For simplicity, none of the methods is conditioned on the antigen structure during CDR-H3 generation.
We leave antigen-conditioned CDR generation for future work.
Metric.
Following Adolf-Bryfogle et al.
(2018), we use amino acid recovery (AAR) as the evaluation metric.
For any generated sequence, we deﬁne its AAR as the percentage of residues having the same amino acid as the corresponding residue in the original antibody.
Results.
For LSTM, AR-GNN, and ReﬁneGNN, the training set in this setup is the entire SAbDab except antibodies in the same cluster as any of the test antibodies.
At test time, we generate
CDR-H3 sequences for each antibody and select the top 100 candidates with the lowest perplexity.
For simplicity, all methods are conﬁgured to generate CDRs of the same length as the original
CDR.
As shown in Table 1, our model achieves the highest AAR score, with around 7% absolute improvement over the best baseline.
In Figure 4A, we show an example of a generated CDRH3 sequence and highlight residues that are different from the original antibody.
We also found that sequences with lower perplexity tend to have a lower AA recovery error (Pearson R = 0.427,
Figure 4B).
This suggests that we can use perplexity as the ranking criterion for antibody design.
4. SARS-C O V-2 NEUTRALIZATION OPTIMIZATION
Data.
The Coronavirus Antibody Database (CoVAbDab) contains 2411 antibodies, each associated with multiple binary labels indicating whether it neutralizes a coronavirus (SARS-CoV-1 or
SARS-CoV-2) at a certain epitope.
Similar to the previous experiment, we divide the antibodies into training, validation, and test sets based on CDR-H3 cluster split with 8:1:1 ratio.
Neutralization predictor.
The predictor takes as input the VH sequence of an antibody and outputs a neutralization probability for the SARS-CoV-1 and SARS-CoV-2 viruses.
Each residue is embedded into a 64 dimensional vector, which is fed to a SRU encoder (Lei, 2021) followed by
[Table]
average-pooling and two layer feed network
1 and
CDR sequence constraints.
Therapeutic antibodies must be free from developability issues such as glycosylation and high charges (Raybould et al., 2019).
Thus, we include four constraints on a
CDR-H3 sequence s
: 1) Its net charge must be between -2.0 and 2.0 (Raybould et al., 2019).
The deﬁnition of net charge is given in the appendix.
It must not contain the N-X-S/T motif which is prone to glycosylation.
Any amino acid should not repeat more than ﬁve times (e.g. SSSSS).
Perplexity of a generated sequence given by LSTM, AR-GNN, and ReﬁneGNN should be all less than 10.
The last two constraints force generated sequences to be realistic.
We use all three models in the perplexity constraint to ensure a fair comparison for all methods.
Metric.
For each antibody in the test set, we generate 100 new CDR-H3 sequences, concatenate them with its context sequence to form 100 full VH sequences, and feed them into the neutralization predictor f .
We report the average neutralization score of antibodies in the test set.
Neutralization score of a generated sequence s (cid:48) equals f ( s (cid:48) ) if it satisﬁes all the CDR sequence constraints.
Otherwise the score is the same as the original sequence.
In addition, we pretrain each model on the
SAbDab CDR-H3 sequences and evaluate its PPL on the CoVAbDab CDR-H3 sequences.
Results.
All methods are pretrained on SAbDab antibodies and ﬁnetuned on CoVAbDab using the
ITA algorithm to generate neutralizing antibodies.
Our model outperforms the best baseline by a 3% increase in terms of average neutralization score (Table 2).
Our pretrained ReﬁneGNN also achieves a much lower perplexity on CoVAbDab antibodies (7.86 v.s. 8.67).
Examples of generated CDR-H sequences and their predicted neutralization scores are shown in the appendix.
C ONCLUSION
In this paper, we developed a ReﬁneGNN model for antibody sequence and structure co-design.
The advantage of our model over previous graph generation methods is its ability to revise a generated subgraph to accommodate addition of new residues.
Our approach signiﬁcantly outperforms sequence-based and graph-based approaches on three antibody generation tasks.
A M ODEL ARCHITECTURE DETAILS A. R EFINE GNN
Node features.
Each node feature v i encodes three dihedral angles as follows.
[Equation]
Edge features.
The orientation matrix O i =
[ b i , n i , b i × n i ] deﬁnes a local coordinate system for each residue
i (Ingraham et al., 2019), which is calculated as
[Equation]
Attention mechanism.
The attention layer used in Eq.(13) is a standard bilinear attention: attention( c 1: n , h t
[Equation]
= (cid:88)
[Equation]
A. AR-GNN
AR-GNN generates an antibody graph autoregressively.
In each generation step t , AR-GNN learns to encode the current subgraph G 1: t induced from residues { s 1 , · · · , s t } into a list of vectors {
h 1 , · · · , h t } =
MPN θ ( G 1: t ) .
[Equation] [Equation]
For fair comparison, we use the same MPN architecture for both ReﬁneGNN and AR-GNN.
In terms of structure prediction, AR-GNN ﬁrst predicts the node feature ˆ v t +1 of the next residue t + 1 , namely the dihedral angle between its three atoms C α , C, N .
[Equation]
[Equation]
In addition, AR-GNN predicts the pairwise distance between s t +1 and previous residues s 1 , · · · , s t .
[Equation] [Equation]
where FFN is a feed-forward network with one hidden layer and E pos is the positional encoding of t + 1 − i , the gap between residue s t +
and s i in the sequence.
Lastly, AR-GNN predicts the amino acid type of residue s t +1 by
[Equation]
Note that AR-GNN also uses two separate MPNs for structure and sequence prediction.
However, unlike ReﬁneGNN, AR-GNN is trained under teacher forcing — we need to feed it the ground truth structure and sequence in each generation step.
In particular, we ﬁnd data augmentation to be crucial for AR-GNN performance.
Data augmentation is essential because of the discrepancy between training and testing.
The model is trained under teacher forcing, but it needs to decode a graph without teacher forcing at test time.
We ﬁnd mistakes made in previous steps have a great impact on subsequent predictions during decoding.
Speciﬁcally, for every antibody s , we create a corrupted graph (cid:101) G by adding independent random
Gaussian noise to every coordinate:
˜ x i
= x
i + 3 (cid:15), (cid:15) ∼ N (0 , I ) .
In each generation step, we apply
MPN over the corrupted graph instead.
[Equation] [Equation]
The node and edge labels are still deﬁned by the ground truth structure.
Speciﬁcally, let v t and d i,j be the ground truth dihedral angle and pairwise distance calculated from the original, uncorrupted graph G .
AR-GNN loss function is deﬁned as the following.
[Equation]
Similar to ReﬁneGNN, AR-GNN also uses attention mechanism for conditional generation.
Specifically, we concatenate the residue representations (cid:101) h t , (cid:101) g t from MPN with context vectors learned from an attention layer.
[Equation] [Equation]
[Table]
B E XPERIMENTAL DETAILS
Hyperparameters.
For AR-GNN and ReﬁneGNN, we tried hidden dimension d h ∈ { 128 , 256 } and number of message passing layers L ∈ { 1 , 2 , 3 , 4 , 5 } .
We found d h = 256 , L = 4 worked the best for ReﬁneGNN and d h = 256 , L = 3 worked the best for AR-GNN.
For LSTM, we tried d h ∈ { 128 , 256 , 512 , 1024 } .
We found d h = 256 worked the best.
All models are trained by an
Adam optimizer with a dropout of 0.1 and a learning rate of 0.001.
SAbDab data.
The dataset statistics of SAbDab is the following (after deduplication).
For CDRH1, the train/validation/test size is 4050, 359, and 326.
For CDR-H2, the train/validation/test size is 3876, 483, and 376.
For CDR-H3, the train/validation/test size is 3896, 403, and 437.
RAbD conﬁguration.
We provided details of the de novo design setup of RosettaAntibodyDesign (RAbD) here.
For each antibody in the test set, RAbD starts by randomly selecting a CDR from
RAbD’s internal database of known CDR structures.
The chosen CDR-H3 sequence is required to have same length as the original sequence, but it cannot be exactly the same as the original CDR-H sequence.
After the initial CDR structure is chosen, RAbD grafts it onto the antibody and performs energy minimization to stabilize its structure.
Next, RAbD runs 100 iterations of sequence design to modify the grafted CDR-H3 structure by randomly substituting amino acids.
In each sequence design iteration, it performs energy minimization to adjust the structure according to the changed amino acid.
Lastly, the model returns the generated CDR-H3 sequence with the lowest energy.
SARS-CoV-2 neutralization.
Each generative model is pretrained on the SAbDab data to learn a prior distribution over CDR-H3 structures.
Given a ﬁxed predictor f , we use the ITA algorithm to ﬁnetune our pretrained models to generate neutralizing antibodies.
Each model is trained for
ITA steps with M = 100 .
Generated CDR-H3 sequences from our model are visualized in Table 3.
Our neutralization predictor f is trained on the CoVAbDab database.
For simplicity, we only consider two viruses, SARS-CoV-1 and SARS-CoV-2 since other coronavirus have very little training data.
For the same reason, we only consider the spike protein receptor binding domain as our target epitope.
The predictor is trained in a multi-task fashion to predict both SARS-CoV-1 and SARSCoV-2 neutralization labels.
The SRU encoder has a hidden dimension of 256.
The model was trained with a dropout of 0.2, a learning rate of 0.0005, and batch size of 16.
The charge of a residue is deﬁned as C ( s i )
=
I [ s i ∈
{ R, K } ]
+ 0 .
1 · I [ s i
= H ] − I
[ s i ∈
{ D, E } ] (Raybould et al., 2019).
The net charge of a sequence s 1 · · · s n is deﬁned as (cid:80) i C
( s i ) .