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

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

Code: https://github.com/wengong-jin/RefineGNN

RefineGNN: Iterative sequence and structure generation for CDR design

Abstract

This paper proposes a new generative model for automatically designing antibody complementarity-determining regions (CDRs) to improve binding specificity or neutralization capacity. Conventional generative approaches often treat protein design as sequence generation conditioned on a known three-dimensional structure. That formulation is ill-suited to antibodies because the target 3D structure is usually unknown in advance. To address this, the authors jointly design CDR sequence and structure as a graph. They introduce an iterative refinement graph neural network that autoregressively generates the sequence while iteratively refining its predicted global structure; the predicted structure in turn guides selection of subsequent residues. For efficiency, the model uses a coarse-grained representation when modeling conditional dependencies between CDRs and framework regions. Experiments report higher log-likelihood on held-out data and outperform prior baselines in designing antibodies that neutralize SARS-CoV-2.

Introduction

Antibodies bind pathogens through their CDRs, which determine specificity and neutralization and make antibodies important therapeutics (e.g., for SARS-CoV-2). Computational antibody design aims to generate CDR sequences with desired properties (high binding affinity, stability, etc.), but the combinatorial search space is enormous.

Sequence-only generative methods ignore 3D structure and can yield suboptimal designs, whereas structure-conditioned generators require a predefined target structure—often unrealistic for antibody design.

Core questions:

• How to jointly model sequence and its corresponding 3D structure.
• How to model conditional relationships between CDR sequences and their context (framework regions), including sequence–structure interactions.
• How to optimize antibodies for multiple properties (binding affinity, stability, etc.); classical physics-based methods focus on energy minimization, but the design objective is more complex in practice.

Proposed solution:

• Formulate antibody design as graph generation, with graphs encoding both sequence and 3D structure.
• Introduce RefineGNN, an iterative refinement GNN that dynamically updates predicted 3D structure as new nodes are added during generation.
• Use a coarse-grained graph representation by grouping context sequence into node blocks, improving computational efficiency.

The goal is a generative model that co-designs CDR sequence and 3D structure to optimize target properties (e.g., SARS-CoV-2 neutralization). The iterative refinement framework also relates to score-matching methods for molecular conformation prediction and point-cloud-based diffusion approaches.

Antibody Sequence and Structure Co-Design

Overview

Antibodies bind antigens (e.g., viruses), present them to the immune system, and elicit an immune response. Neutralizing antibodies not only bind antigens but also block their activity.

An antibody comprises heavy and light chains; each chain has a variable region (VH/VL) and constant domains. The variable region splits into framework regions and three CDRs—the most diverse and functional segments, governing binding and neutralization.

Modeling antibody design:

• Antibody design is cast as CDR generation: synthesize CDR sequence conditioned on framework context.
• Graphs encode both sequence and 3D structure.

Core ideas of RefineGNN:

• A graph generative method for joint CDR sequence and 3D structure generation.
• Sequence is generated by iteratively refining a global antibody graph while dynamically adjusting predicted structure.
• During generation, CDR sequence and 3D structure co-evolve for more accurate design.

For simplicity, the experiments focus on heavy-chain CDRs (especially CDR-H3), but the approach extends to light-chain CDR design.

Notations

Amino acid sequence is denoted $s$; <mask> marks unknown residue types; $x$ denotes backbone coordinates.

Graph Representation

Node $v_i$ features are the three backbone dihedral angles of the residue. Edge features have four components; $O_i$ is the local coordinate frame of residue $i$. Edges form a K-nearest-neighbor (KNN) graph with $K=8$ to limit cost.

Iterative Refinement Graph Neural Network (RefineGNN)

Output is a complete antibody graph $G(s)$ with sequence and 3D structure for all nodes.

At each generation step $t$, the pipeline has three stages:

• Graph encoding:

  • A message passing network (MPN) encodes the current graph into per-residue representations.
    
    

  • The MPN updates nodes over multiple message-passing layers.

    

• Sequence prediction:

  • Predict the amino acid type of the next residue from encoder output $h_{t+1}^{(t)}$ via softmax (Fig. 2A).

    

  • Obtain intermediate graph $G^{(t+0.5)}$ with the new residue label filled in (Fig. 2B).

• Structure update:

  • Predict 3D coordinates for all nodes given the current graph (Fig. 2C).

  • A separate structure prediction network (another MPN with parameters $\tilde\theta$) encodes the graph and predicts coordinates.

    
  • Produce $G^{(t+1)}$ with updated node labels and edge features.

Training

• Sequence prediction: Teacher forcing—set the first $t$ residues to ground-truth types and remaining nodes to <MASK>.
• Structure prediction: Fully iterative refinement without teacher forcing; each structure update uses the graph predicted at the previous step.

Loss Function

Overall loss combines sequence and structure terms: $L=L_\mathrm{seq}+L_\mathrm{struct}$.

• $L_\mathrm{seq}$ is cross-entropy for sequence accuracy.

• $L_\mathrm{struct}$ has three parts:

  • Distance loss: Huber loss on predicted vs. true Cα–Cα distances.

    

  • Dihedral loss: MSE on predicted vs. true dihedral angles.

    

  • Cα angle loss: MSE on predicted vs. true angles and dihedrals between Cα atoms.

    

Separate networks: Sequence and structure prediction use distinct MPNs specialized to each task.

Rotation and translation invariance: Distance- and angle-based losses make the model insensitive to rigid transforms of the antibody graph.

Conditional Generation Given the Framework Region

Practical setting: Framework regions are often fixed; only CDR sequence must be designed. The model therefore generates CDRs conditioned on known framework (context) sequence.

Attention propagates framework information into CDR generation.

Conditioning via attention

Context construction by concatenation: $\tilde{s}=s_{<l}\oplus\mathrm{}\cdots\mathrm{}\oplus s_{>r}$.

A GRU encodes context; sequence is predicted and structure updated step by step.

Multi-Resolution Modeling

Attention captures framework sequence but not spatial interactions between framework and CDR. Multi-resolution modeling addresses this:

Coarse-grained context:

• Group context residues into blocks of $b$ residues to reduce cost.
• Block embeddings and coordinates are averages over constituent residues.

Edges on the coarse context graph use block-averaged coordinates. During generation, RefineGNN refines both the coarse graph and the fine local CDR graph.

Conditional generation is summarized in Algorithm 1.

Property-Guided Sequence Optimization

Generate antibody sequences with target properties (e.g., antigen neutralization), formulated as an optimization problem.

1. Pretrain the generative model on real antibody structures to learn a prior over CDR sequence and structure.
2. Sample an antibody sequence $s$ from the training set, remove its CDR, keep framework context, and generate $M$ candidate sequences with the generative model.
3. Score each candidate $s’$ with a pretrained property predictor $f(s’)$; if $f(s_i^{\prime})>\max(f(s),\mathrm{threshold})$, add the candidate to optimization set $Q$.
4. Sample a batch from $Q$ as training targets and minimize generative loss $L_\mathrm{seq}$.

Experiments

Setup

• Language modeling and 3D structure prediction:

  • Perplexity (PPL) of generated antibody sequences.
  • Structural error vs. ground truth via RMSD on generated CDR structures.

• Antigen-binding antibody design:

  • Evaluate RefineGNN on a standard antibody design benchmark.
  • Goal: generate CDR-H3 sequences that bind a given antigen.

• Antibody optimization:

  • Redesign CDR-H3 on SARS-CoV-2 data to improve neutralization.
  • Fixed framework with CDR-H3 optimization—a common antibody engineering setup.

Baselines

• LSTM (sequence-only):
  • No 3D structure; encoder for context, decoder for CDR, with attention between them.

• AR-GNN (autoregressive graph generation):

  • GNN generates the antibody graph residue by residue.
  • Cannot revise partially generated 3D structure during generation.

• RAbD (Rosetta Antibody Design):

  • Physics-based design: CDR fragment selection, energy minimization, iterative sequence design.
  • Used only for antigen-binding design; not for SARS-CoV-2 optimization, which would require predefined 3D structures.

Hidden dimension 256; block size $b=4$.

Language Modeling and 3D Structure Prediction

Data

• SAbDab (Structural Antibody Database): 4994 antibody structures.
• Train/validation/test splits by CDR sequence clusters (e.g., CDR-H3).

Metrics

1. Perplexity (PPL):

   • Measures plausibility of generated sequences.
   • Lower PPL indicates better match to the data distribution.

2. RMSD (Root Mean Square Deviation):

   • Deviation between generated and native structures.
   • Lower RMSD indicates more accurate structures.

Visualizations show RefineGNN predictions aligning more closely with native antibody structure.

Antigen-Binding Antibody Design

Data

• Benchmark from Adolf-Bryfogle et al. (2018).
• 60 antibody–antigen complexes for design evaluation.
• Given framework structure, design CDR-H3 to bind the antigen.
• Simplified setting: no antigen structure used during CDR-H3 generation; antigen-conditioned generation is left for future work.

RefineGNN achieves 34.14% AAR, substantially above all baselines (~7% gain); Table 1.

• Fig. 4A shows example generated CDR-H3 sequences in antigen-bound complexes.
• Predicted structures are aligned and grafted via Kabsch; residues differing from the parent sequence are highlighted in red.
• Lower perplexity correlates with lower amino-acid recovery error (Pearson $R=0.427$).
• Perplexity is a useful ranking criterion for candidate sequences in antibody design.

SARS-CoV-2 Neutralization Optimization

Data — CoVAbDab (Coronavirus Antibody Database):

• 2411 antibodies with binary labels for neutralization at specific epitopes (SARS-CoV-1 and SARS-CoV-2).
• Split by CDR-H3 clusters into train, validation, and test (8:1:1).

Neutralization predictor

• Goal: Predict SARS-CoV-2 neutralization.
• Input: Heavy-chain variable sequence (VH).

• Architecture:

  • Per-residue embeddings (64-D).
  • SRU encoder (Simple Recurrent Unit).
  • Average pooling followed by a two-layer feed-forward network.
  • Outputs probabilities $p_1$ and $p_2$ for neutralizing SARS-CoV-1 and SARS-CoV-2.
  • Score: $f(s)=p_2$ for SARS-CoV-2 neutralization.

CDR sequence constraints

Therapeutic design constraints:

1. Net charge:

   • Net charge in $[-2.0, 2.0]$.
   • Defined from residue types (e.g., R, K positive; D, E negative).

2. Avoid glycosylation:

   • No N-X-S/T motifs (X = any amino acid).

3. Amino acid repeats:

   • No single amino acid repeated more than five times in a row (e.g., “SSSSS” is invalid).

4. Perplexity constraint:

   • Generated sequence PPL must be below 10.
   • PPL computed with LSTM, AR-GNN, and RefineGNN for fair comparison.

Evaluation metrics

• Neutralization score:

  • For each test antibody, generate 100 new CDR-H3 sequences.
  • Concatenate with framework context to form full VH sequences.

  • Score with predictor $f$:

    • If all constraints are satisfied, use $f(s’)$.
    • Otherwise, use the parent antibody score for comparison.

• Perplexity (PPL):

  • Pretrained SAbDab models evaluated on CoVAbDab test set.

After pretraining on SAbDab, the model is fine-tuned on CoVAbDab with ITA (Iterative Target Augmentation).