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: This paper proposes a joint generative model (co-designing the sequence and the structure at the same time) for the CDRs of antibodies, with the goal to enhance binding specificity or neutralization capabilities.   The proposed method and two existing baselines are evaluated on 1) perplexity on hold-out set; 2) perplexity and sequence recovery on known antigen-antibody complexes; and 3) redesign CDR-H3 of existing antibodies for coronavirus neutralization as measured by a given neutralization predictor. RefineGNN, the proposed method, shows improved performance in all three tasks.