Iterative Refinement Graph Neural Network for Antibody Sequence-Structure Co-design

1 INTRODUCTION . 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 specificity 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 2060 possible CDR sequences and the small solution space which satisfies the desired constraints of binding affinity , stability , and synthesizability ( Raybould et al. , 2019 ) . There are three key modeling questions in CDR generation . The first 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 predefined 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 ) can not 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 refinement 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 efficiently propagate the contextual information to the CDR residues . After pretraining our model on antibodies with known structures , we finetune it using a predefined property predictor to generate antibodies with specific properties . We evaluate our method on three generation tasks , ranging from language modeling to SARS-CoV-2 neutralization 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 . 2 RELATED WORK . Antibody/protein design . Current methods for computational antibody design roughly fall into two categories . The first 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 . ( 2018 ) ; Ingraham et al . ( 2019 ) ; Strokach et al . ( 2020 ) ; 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 specific 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 can not modify a previously generated subgraph when new nodes arrive . While Graphite ( Grover et al. , 2019 ) also uses iterative refinement 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 refinement 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 specified in advance . Our iterative refinement 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 refine 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 refinement with graph generation . 3 ANTIBODY SEQUENCE AND STRUCTURE CO-DESIGN . 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 ) . Following Shin et al . ( 2021 ) ; Akbar et al . ( 2021 ) , we formulate antibody design as a CDR generation task , conditioned on the framework region . Specifically , we represent an antibody as a graph , which encodes both its sequence and 3D structure . We propose a new graph generation approach called RefineGNN and extend it to handle conditional generation given a fixed framework region . Lastly , we describe how to apply RefineGNN 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 = s1s2 · · · sn . Each token si in the sequence is called a residue , whose value can be either one of the 20 amino acids or a special token 〈MASK〉 , meaning that its amino acid type is unknown and needs to be predicted . The VH sequence folds into a 3D structure and each residue si is labeled with three backbone coordinates : xi , α for its alpha carbon atom , xi , c for its carbon atom , and xi , n for its nitrogen atom . 3.1 GRAPH REPRESENTATION . We represent an antibody ( VH ) as a graph G ( s ) = ( V , E ) with node features V = { v1 , · · · , vn } and edge features E = { eij } i 6=j . Each node feature vi encodes three dihedral angles ( φi , ψi , ωi ) related to three backbone coordinates of residue i . For each residue i , we compute an orientation matrix Oi representing its local coordinate frame ( Ingraham et al. , 2019 ) ( defined 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 eij = ( Epos ( i− j ) , RBF ( ‖xi , α − xj , α‖ ) , O > i xj , α − xi , α ‖xi , ,α − xj , ,α‖ , q ( O > i Oj ) ) . ( 1 ) The edge feature eij contains four parts . The positional encoding Epos ( 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 eij is a direction encoding that corresponds to the relative direction of xj in the local frame of residue i . The last term q ( O > i Oj ) is the orientation encoding of the quaternion representation q ( · ) of the spatial rotation matrix O > i Oj . 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.2 ITERATIVE REFINEMENT GRAPH NEURAL NETWORK ( REFINEGNN ) . We propose to generate an antibody graph via an iterative refinement process . Let G ( 0 ) be the initial guess of the true antibody graph . Each residue is initialized as a special token 〈MASK〉 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 . Specifically , it first encodes G ( t ) with a message passing network ( MPN ) with parameter θ { h ( t ) 1 , · · · , h ( t ) n } = MPNθ ( G ( t ) ) , ( 2 ) 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 h ( t , l+1 ) i = LayerNorm ( ∑ j FFN ( h ( t , l ) i , h ( t , l ) j , E ( sj ) , ei , j ) ) , 0 ≤ l ≤ L− 1 , ( 3 ) where h ( t,0 ) i = vi and h ( t ) i = h ( t , L ) i . FFN is a two-layer feed-forward network ( FFN ) with ReLU activation function . E ( sj ) is a learned embedding of amino acid type sj . Based on the learned residue representations , we predict the amino acid type of the next residue t+ 1 ( Figure 2A ) . pt+1 = softmax ( Wah ( t ) t+1 ) ( 4 ) This prediction gives us a new graph G ( t+0.5 ) with the same edges as G ( t ) but the node label of t+ 1 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 . { h ( t+0.5 ) 1 , · · · , h ( t+0.5 ) n } = MPNθ̃ ( G ( t+0.5 ) ) ( 5 ) x ( t+1 ) i , e = W e xh ( t+0.5 ) i , 1 ≤ i ≤ n , e ∈ { α , c , n } . ( 6 ) The new coordinates x ( t+1 ) i define 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 xi ) and sequence prediction ( amino acid types pt+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 . Specifically , in each generation step t , residues 1 to t are set to their ground truth amino acid types s1 , · · · , st , 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 refines 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 ) i , e | 1 ≤ i ≤ n , e ∈ { α , c , n } } . 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 define the distance loss as the Huber loss between the predicted and true pairwise distances L ( t ) d = ∑ i , j ` huber ( ‖x ( t ) i , α − x ( t ) j , α‖ 2 , ‖xi , α − xj , α‖2 ) , ( 7 ) 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 ( φ ( t ) i , ψ ( t ) i , ω ( t ) i ) 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 define the dihedral angle loss as the mean square error between the predicted and true dihedral angles L ( t ) a = ∑ i ∑ a∈ { φ , ψ , ω } ( cos a ( t ) i − cos ai ) 2 + ( sin a ( t ) i − sin ai ) 2 ( 8 ) • 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 defined by x ( t ) i−2 , α , x ( t ) i−1 , α , x ( t ) i , α , x ( t ) i+1 , α . L ( t ) c = ∑ i ( cos γ ( t ) i − cos γi ) 2 + ( cosβ ( t ) i − cosβi ) 2 ( 9 ) In summary , the overall graph generation loss is defined as L = Lseq + Lstruct , where Lstruct = ∑ t L ( t ) d + L ( t ) a + L ( t ) c Lseq = ∑ t Lce ( pt , st ) . ( 10 ) The sequence prediction loss Lseq is the cross entropy Lce between predicted and true residue types .

Summary: The paper proposed a deep generative model named iterative refinement graph neural network to generate antibody CDR for Y-shaped antibodies. Specifically, it sequentially generates the CDR residue sequence and refines the global structure iteratively. Empirical results show superior performance compared with baselines.