Retrieval Augmented Diffusion Model for Structure-informed Antibody Design and Optimization

## Abstract

Abstract Antibodies are essential proteins responsible for immune responses in organisms, capable of specifically recognizing antigen molecules of pathogens. Recent advances in generative models have significantly enhanced rational antibody design.
However, existing methods mainly create antibodies from scratch without template constraints, leading to model optimization challenges and unnatural sequences.
To address these issues, we propose a retrieval-augmented diffusion framework, termed RADAb, for efficient antibody design. Our method leverages a set of structural homologous motifs that align with query structural constraints to guide the generative model in inversely optimizing antibodies according to desired design criteria.
Specifically, we introduce a structure-informed retrieval mechanism that integrates these exemplar motifs with the input backbone through a novel dual-branch denoising module, utilizing both structural and evolutionary information. Additionally, we develop a conditional diffusion model that iteratively refines the optimization process by incorporating both global context and local evolutionary conditions.
Our approach is agnostic to the choice of generative models. Empirical experiments demonstrate that our method achieves state-of-the-art performance in multiple antibody inverse folding and optimization tasks, offering a new perspective on biomolecular generative models.

## 1 Introduction

Antibodies, essential Y-shaped proteins in the immune system, are pivotal for recognizing and neutralizing specific pathogens known as antigens. This specificity primarily arises from the Complementarity Determining Regions (CDRs), which are crucial for binding affinity to antigens (Jones et al., ; Ewert et al., ; Xu & Davis, ; Akbar et al., ). The design of effective CDRs is therefore central to developing potent therapeutic antibodies, a dominant class of protein therapeutics.
However, the development of these antibodies typically relies on labor-intensive experimental methods such as animal immunization or screening extensive antibody libraries, often failing to produce antibodies that target therapeutically relevant epitopes effectively. Thus, the ability to generate new antibodies with pre-defined biochemical properties in silico carries the promise of speeding up the drug design process.

Computational efforts in antibody design have traditionally involved grafting residues onto existing structures (Sormanni et al., ), sampling alternative native CDR loops to enhance affinities(Aguilar Rangel et al., ), and using tools like Rosetta for sequence design improvements in interacting regions (Adolf-Bryfogle et al., ). Many recent studies have focused on applying deep generative models to design antibodies (Luo et al., ; Martinkus et al., ; Zhu et al., ).
They take advantage of geometric learning and generative models to capture the higher-order interactions among residues directly from the data. These innovations provide more efficient methods to search sequence and structure spaces.

Albeit powerful, current generative models struggle to design antibodies that adhere to structural constraints and exhibit desired biological properties. This challenge primarily arises from a lack of diversity in the available training data. Predominantly, research efforts have relied on the SAbDab database (Dunbar et al., ), which comprises fewer than ten thousand antigen-antibody complex structures. The limited scope of this dataset restricts the models’ ability to capture comprehensive high-order interaction information between antigen-antibody residues, thereby increasing the risk of overfitting. Moreover, most existing methodologies attempt to design antibody sequences de novo, without the benefit of template-based guidance. This approach inherently demands a greater volume of data and extensive training or fine-tuning on specific datasets to achieve efficacy in practical applications.

Figure: Figure 1: Illustration of the retrieval-augmented framework.
Refer to caption: x1.png

In this work, we draw inspiration from template-based and fragment-based antibody design to develop a model that fully utilizes protein structural database, effective motif retrieval, and semi-parametric generative neural networks. Our goals are to: (a) leverage template-aware local and global protein geometric information to enhance model generative capability, (b) integrate motif evolutionary signals to prevent overfitting, and (c) necessitate minimal training or fine-tuning for effective use in real-world applications.

To this end, we introduce the Retrieval-Augmented Diffusion Antibody design model (RADAb), a novel semi-parametric antibody design framework. To fully exploit the protein structure space, we first compiled a database of CDR-like fragments from the non-redundant Protein Data Bank (PDB) (Berman et al., ). These CDR-like fragments are linear functional motifs structurally compatible with an antibody CDR loop, found in any protein geometry within the PDB. For a given antibody to be improved, we perform a structural retrieval to obtain motifs with structures similar to the desired CDR framework. As protein sequences capable of folding into similar structures often share homology and consensus, we hypothesize that these retrieved motifs, enriched with evolutionary information, can enhance the model’s generalization.

Unlike traditional rational design methods that optimize by grafting a single CDR fragment, we propose to use a set of structural homologous CDR-like motifs together with the desired backbone for iterative sequence optimization (Figure 1). Our major contributions follow: (1) We propose a first-of-its-kind retrieval-augmented generative framework for rational antibody design. It uses a set of functional CDR-like fragments that satisfy the desired backbone structures and properties to guide generation toward satisfying all the required properties. (2) A novel retrieval mechanism is introduced for integrating these exemplar motifs with the input backbone through a novel dual-branch denoising module, utilizing both structural and evolutionary information. Additionally, we present a coupled conditional diffusion module that iteratively refines the evolution process by incorporating global and local conditions. This allows the model to incorporate more functional information than traditional antibody inverse folding models. (3) Empirical experiments demonstrate that our method improves the state-of-the-art methods in multiple antibody inverse folding tasks, e.g., an 8.08% AAR gains in long CDRH3 inverse folding task and an average of 7 cal/mol absolute $\Delta$$\Delta$G improvements in functionality optimization task, offering a fresh perspective on biomolecular generative models.

## 2 Related work

Antibody Design Computational antibody design primarily follows two paths: conventional energy function optimization methods and machine learning approaches. Early antibody design methods were often limited to sequence similarity and energy function optimization (Lapidoth et al., ; Adolf-Bryfogle et al., ). Recent success of machine learning approaches mainly falls into two directions: antibody sequence design and antigen-specific antibody sequence-structure co-design.
The methods used for antibody sequence design mainly include language-based models (Ruffolo et al., ; Olsen et al., ; Wang et al., ) and inverse folding models (Dreyer et al., ; Høie et al., ).
The other line focuses on antibody sequence-structure co-design mainly taking antibody-antigen complex as a graph, then using graph networks to extract features and predict the coordinates and residue type of antibody CDR (Jin et al., ; Kong et al., ; ; Lin et al., ; Luo et al., ; Zhu et al., ; Martinkus et al., ). Despite powerful, these works typically create antibodies from scratch without template constraints, leading to model optimization challenges and poor generative capability. Instead, our method leverages the power of templates from a structure-informed perspective.

Diffusion generative models Diffusion models (Sohl-Dickstein et al., ; Song et al., ; Ho et al., ) are a class of generative models that have achieved impressive progress on a lot of generation tasks. Denoising diffusion probabilistic models (DDPMs) are a branch of diffusion models, which contain two Markov processes. The forward process perturbs the data into pure noise, and then learns to generate data by reversing the forward Markov process. Because of the diffusion model’s flexibility and controllability, numerous works are focusing on employing retrieval-augmented methods to complement the diffusion framework for text-to-image generation (Sheynin et al., ), image generation (Blattmann et al., ), human motion generation (Zhang et al., ) and small molecule generation (Wang et al., ; Huang et al., ).

Retrieval augmented generative models
Retrieval augmented generation technique was first proposed in the field of natural language processing to enhance the language models by introducing an additional database (Lewis et al., ; Guu et al., ), prompting the language models to generate more realistic and diverse results.
Subsequently, retrieval augmented generation (RAG) has conducted diverse explorations in large fields, including natural language processing (Zhang et al., ; Gao et al., ; Xu et al., ; Yoran et al., ; Caffagni et al., ) and computer vision (Long et al., ; Blattmann et al., ; Hu et al., ; Rao et al., ).
Although there has been some retrieval work in the field of protein design and discovery (Zhou & Grigoryan, ; Aguilar Rangel et al., ), to the best of our knowledge, this is the first-of-its-kind retrieval-based generative framework for antibody design.

## 3 Preliminaries and Notations

### 3.1 Notations

Antibody consists of two heavy chains and two light chains. Each chain’s tip has a complementary site that specifically binds to a unique epitope on the antigen. This site includes six complementary-determining regions (CDRs): CDR-H1, CDR-H2, and CDR-H3 on the heavy chain, and CDR-L1, CDR-L2, and CDR-L3 on the light chain (Presta, ; Al-Lazikani et al., ).

Our work represents each single protein residue in terms of the residue type $s_{i}\in\left\{ACDEFGHIKLMNPQRSTVWY\right\}$, the coordinate ${x}_{i}\in\mathbb{R}^{3}$, and the orientation $\textbf{O}_{i}\in\text{SO}\left(3\right)$,
where $i=1,...,N$ and N is the number of residues in the complex. Concretely, assuming the CDR sequence to be generated includes $m$ amino acids and starts from position $a$, it can be denoted as
$R=\left\{s_{j}\mid j\in\{a+1,...,a+m\}\right\}$. Let $M$ be the length of the antibody, the antibody framework is defined as $C_{ab}=\left\{\left(s_{i},x_{j},\textbf{O}_{j}\right)\mid i\in\left\{1,...,M
\right\}\backslash\{a+1,...,a+m\},j\in\left\{1,...,M\right\}\right\}$. The antibody framework sequence is defined as $S_{ab}=\left\{s_{i}\mid i\in\left\{1,...,M\right\}\backslash\{a+1,...,a+m\}\right\}$. The antigen is defined as $C_{ag}=\left\{\left(s_{i},x_{i},\textbf{O}_{i}\right)\mid i\in\left\{M+1,...,N
\right\}\right\}$.
The retrieved CDR-like fragments are defined as ${\mathbb{A}}=\left\{A_{i}\mid i\in\left\{1,...,k\right\}\right\}$.
The goal of our framework is to extract the antibody CDR structure from the antibody framework complex $C_{ab}$, then input it into the retrieval module to retrieve ${\mathbb{A}}$, and ultimately predict the distribution of $R$ through $C_{ag}$, $C_{ab}$ and ${\mathbb{A}}$.

### 3.2 Diffusion model for antibody design

Due to diffusion models’ excellent performance and controllability, there are now many diffusion-based works that have achieved notable results (Luo et al., ; Villegas-Morcillo et al., ; Kulytė et al., ). To be specific, they are denoising probabilistic diffusion models that transform the amino acid type $s$, the backbone $C\alpha$ atom coordinates $x$, and the amino acid orientation O during the diffusion process. We focus on sequence, of which the forward process perturbs the data in the following ways (Hoogeboom et al., ):

$$ $q\left(\mathrm{~{}s}_{j}^{t}\mid\mathrm{s}_{j}^{t-1}\right)=\operatorname{ Multinomial}\left(\left(1-\beta^{t}\right)\cdot\operatorname{onehot}\left( \mathrm{s}_{j}^{t-1}\right)+\beta^{t}\cdot\frac{1}{20}\cdot\mathbf{1}\right)$ (1) $$

where $\beta^{t}$ is the noise schedule for the diffusion process, as t approaches T, $\beta^{t}$ will approach 1, and the probability distribution will become closer to pure noise. 1 corresponds to a 20-dimensional all-one vector.

To reverse the aforementioned forward process and denoise to generate CDR sequence, predictions need to be made by a neural network $F$(·)$[j]$, which takes the antibody-antigen context as condition:

$$ $p\left(\text{s}_{j}^{t-1}\mid\mathcal{R}^{t},\mathcal{C}_{ab},\mathcal{C}_{ag} \right)=\operatorname{Multinomial}\left(F\left(\mathcal{R}^{t},\mathcal{C}_{ab },\mathcal{C}_{ag}\right)\left[j\right]\right)$ (2) $$

As an example, this work uses Diffab (Luo et al., ) as the backbone for the generative model to conduct retrieval augmented generation. Note that the proposed retrieval system is generative model agnostic, and the developed modules can be integrated with any diffusion generative model.

## 4 Methods

We propose RADAb (as demonstrated in Figure 2), a novel structure-informed retrieval-augmented diffusion framework for antibody sequence design and optimization. The model uses a structural retrieval algorithm to search for antibody homologous structures and take their sequences as conditional inputs for the diffusion model to provide homologous patterns and evolutionary information.

Figure: Figure 2: The overall architecture of the proposed RADAb framework. (A) Structural retrieval process, the CDR backbone is input into MASTER and the output is a set of ranked CDR-like fragments. (B) Diffusion process and denoising network which takes antibody-antigen context and retrieved evolutionary information as conditions. (C) Our method restricts the antibody to a small region through fixed structural constraints and retrieval-augmented constraints (functional constraints) to achieve higher fitness.
Refer to caption: x2.png

### 4.1 Structural retrieval of CDR fragments

The structure of a protein is determined by its sequence, and protein sequences that can fold into similar structures exhibit similar properties. These structurally similar protein sequences contain rich evolutionary information. Based on this, we perform retrieval in the PDB database using CDR structures, aiming to obtain fragments that are similar to the real CDR and have homologous sequences, with the expectation that they possess similar functions.

To balance the quality of results and the retrieval speed, we use MASTER (Zhou & Grigoryan, ) for the search. MASTER uses the root-mean-square deviation (RMSD) of backbone atoms as a similarity measure. It queries structural fragments composed of one or more non-contiguous segments and can find all matching fragments from the database within a given RMSD threshold. This allows for fast and accurate searches in the PDB database for protein motifs. Note that MASTER can utilize only the backbone information without any leakage of sequence data during the search process. The retrieval procedure is described in Algorithm 1 and detailed in Appendix A.3.

For the retrieved results, we use the RMSD with the real backbone structure as a score to rank them and filter out the input CDR fragment. For ease of use, we further constructed a CDR-like fragments database (detailed in Appendix A.4). Additionally, to enable the model to learn richer evolutionary information, we filter out identical CDR-like sequences during the training phase. However, to improve the quality of the model’s generation, we do not perform similar filters during generation.

Figure: Algorithm 1 Structural Retrieval Algorithm Overview

### 4.2 Model Architectures

The model takes the antigen-antibody complex’s structure and sequence context, along with the sequences of the CDR-like fragments, as conditional inputs to iteratively denoise.
The first branch of the model learns the global context information of the complex, while another branch takes the local homologous information of CDR-like fragments as input, aiming to learn the functional similarity and evolutionary information of residues with similar structure. The two branches are combined to generate the antibody CDR sequence jointly.

#### 4.2.1 Global geometry context information branch

Context encoder
A protein is formed by the connection of multiple residues. The features of a single residue mainly include the residue type, backbone atom coordinates, and backbone dihedral angles. The features of each pair of residues mainly include the types of both residues, sequential relative position, spatial distance, and pairwise backbone dihedrals. These features are concatenated and then input into two separate MLPs. The output is denoted as $z_{i}$ and $y_{ij}$.

Evolutionary encoder
Recent advances have shown the structure-informed protein language model (PLM) is an excellent tool for creating protein sequence embeddings and providing evolutionary information (Zheng et al., ; Shanker et al., ). Thus, we take ESM2 (Lin et al., ) as an antibody sequence encoder, aiming to capture the evolutionary relations of antibody residues. The state of antibody sequence with CDR at timestep $t$ is fed into it and output is defined as $e^{t}$.

Structure-informed network
The above encoding is used as conditional input to the Structure-informed network. They, along with the CDR sequence and structural state at the current time step, will be input into a stack of Invariant Point Attention (IPA) (Jumper et al., ) layers, and jointly transform into a hidden representation $h_{i}$.
Subsequently, the hidden representation $h_{i}$ is transformed by an MLP to obtain the probability representation $r_{\text{global}}$ of the amino acid type at each CDR site. This probability representation is then input to the local CDR-focused branch.

#### 4.2.2 Local CDR-focused information branch

Post-processing of CDR-like fragments
We first remove the CDR portions from the antibody sequences to obtain the antibody framework sequences. Then, we fill these fragments’ short sequences into the antibody framework sequences, thereby constructing a CDR-like sequence matrix ${\mathbf{E}}$.

CDR-focused Axial Attention
The local CDR-like branch is constituted of a stack of axial attention layers, referred to as CDR-focused Axial Attention. Given that the CDR-like fragments exhibit structures similar to actual CDRs, we employed a tied row attention mechanism used in MSATransformer (Rao et al., ) to leverage these retrieval results. In the standard axial attention (Ho et al., ) mechanism, each row and column are considered independently. However, in MSA (Multiple Sequence Alignment), each sequence exhibits relatively similar structural features. Our matrix format is well-suited for adopting a tied row attention mechanism to fully utilize the structural similarity. When calculating the attention scores for each row, this mechanism simultaneously considers the scores of other rows. This approach not only leverages the structural similarity but also reduces memory usage.

The input to CDR-focused Axial Attention is a pseudo-MSA matrix ${\mathbf{P}}$ in equation 3. The first row of this matrix is initially filled with the antigen-antibody framework sequence, with the CDR region populated by the noisy sequence $R^{t}_{g}$ (sampled by $r_{\text{global}}$) at the current time step $t$. From the second row to the k-th row (where k is chosen to be 16, meaning the top 15 retrieved CDR-like sequences are used as conditional input), the rows are filled with the CDR-like sequence matrix ${\mathbf{E}}$. The constructed matrix ${\mathbf{P}}$ is then input to CDR-focused Axial Attention to create the homologous embedding and calculate the probability representation $r_{\text{local}}$ (equation 4).

$$ ${\mathbf{P}}=\operatorname{concat}\left(\left(S_{ab}\cup R^{t}_{g}\right),{ \mathbf{E}}\right)$ (3) $$

$$ $\displaystyle r_{\text{local}}[\cdot,j]$ $\displaystyle=G_{\text{col}}\left({\mathbf{P}}_{\cdot,j},t\right)\text{ for all }j\in\text{col},$ (4) $\displaystyle r_{\text{local}}[i,\cdot]$ $\displaystyle=G_{\text{tiedrow}}\left({\mathbf{P}}_{i,\cdot},t\right)\text{ for all }i\in\text{row}$ $$

The row self-attention is computed to capture the internal relationships within the antibody-antigen sequences, while the column self-attention is computed to capture the relationships between the CDR residues and the CDR-like residues.

Skip connection for information fusion
Although the probability distribution of the CDR region created by the antigen-antibody context features has already been fed into the network, to prevent the loss of antigen-antibody context information during forward propagation, the embedding $r_{\text{local}}$ and $r_{\text{global}}$ are added by a skip connection module (He et al., ), then execute softmax to obtain the final probability distribution.

### 4.3 Model training and inference

The overall training objective
The training objective is to minimize the probability distributions predicted by the network under two conditions at each time step and the true posterior distribution at the same time step. Therefore, we choose KL divergence between the two distributions at each residue in the CDR region as the training loss function,

$$ $L_{\text{type\,\,}}^{t}=\mathbb{E}_{\mathcal{R}^{t}\sim p}\left[\frac{1}{m} \sum_{j}{D}_{\text{KL}}\left(q\left(\text{s}_{j}^{t-1}\mid\text{s}_{j}^{t}, \text{s}_{j}^{0}\right)||p\left(\text{s}_{j}^{t-1}\mid\mathcal{R}^{t},\mathcal {C}_{ab},\mathcal{C}_{ag},{\mathbb{A}}\right)\right)\right]$ (5) $$

The training objective of the whole diffusion process is:

$$ $L=E_{t\sim\,\,\text{Uniform\,\,}\left(1...T\right)}L_{\text{type\,\,}}^{t}$ (6) $$

Conditional reverse diffusion process
We employ DDPM to generate sequences. The model starts from time step T, initializing each site of the antibody CDR region as a uniform distribution. Then, through the frozen ESM encoder $E$(·), learned global context network $F$(·)$[j]$ and the local CDR-focused network $G$(·)$[j]$, they predict the noise distribution at each time step jointly and denoise step-by-step:

$$ $e^{t}=E(S_{ab}\cup R^{t})$ (7) $$

$$ $\displaystyle p\left(s_{j}^{t-1}\mid\mathcal{R}^{t},\mathcal{C}_{ab},\mathcal{ C}_{ag},{\mathbb{A}}\right)=\operatorname{Multinomial}\big{[}F\left(\mathcal{R }^{t},\mathcal{C}_{ab},\mathcal{C}_{ag},e^{t}\right)$ (8) $\displaystyle+G\left(F\left(\mathcal{R}^{t},\mathcal{C}_{ab},\mathcal{C}_{ag}, e^{t}\right),{\mathbb{A}}\right)\big{]}\left[j\right]$ $$

During sampling, we remove the CDR region sequences from the antibody structures and fill them with noisy sequence sampled from the uniform distribution. The retrieval process uses the structure of the CDR region as input and outputs a set of CDR-like fragments. Subsequently, this set of CDR-like fragments is fed into the retrieval-augmented diffusion model, serving as a condition along with the antigen-antibody framework context to guide the model in step-by-step denoising and generating the CDR sequences.

## 5 Experiments

To evaluate the performance of our model’s generation, we utilize two tasks: antibody CDR sequence inverse folding (Section 5.1) and antibody optimization based on sequence design (Section 5.2), to compare with the baselines. Additionally, we conducted ablation experiments and further analysis to demonstrate the effectiveness of the retrieval-augmented method (Section 5.3).

The dataset for training the model is obtained from the SAbDab and our established CDR-like fragments dataset. Following the previous work (Luo et al., ), we first eliminated structures with a resolution lower than 4Å and removed antibodies that target non-protein antigens. Chothia (Chothia & Lesk, ) in ANARCI (Dunbar & Deane, ) is used for renumbering antibody residues. We clustered the SADab datasets based on 50% sequence similarity in the CDR-H3 region, and chose 50 PDB files comprising 63 antibody-antigen complex structures as the test set. To ensure distinct training and test sets, we removed structures from the training set that were part of the same clusters as those in the test set.

### 5.1 Antibody CDR Sequence inverse folding

Baselines For traditional methods, we simulated a method of grafting using CDR-like data in the process of rational antibody design. Specifically, we directly graft the retrieved top-1 CDR-like fragment onto the antibody framework, termed Grafting.
For deep learning methods, we selected a series of state-of-the-art protein inverse folding models for comparison with our work, including ProteinMPNN (Dauparas et al., ), a model that utilizes message passing neural network to design sequences with a fixed protein backbone; ESM-IF (Hsu et al., ), a protein inverse folding model that trained on millions of predicted structures; Diffab-fix (Luo et al., ), which can fix the backbone structure and iteratively generate candidate sequences from pure noise in sequence space using diffusion; AbMPNN (Dreyer et al., ), a model fine-tuned ProteinMPNN on antibody sequence and structure. Because it is not open-sourced, we evaluate it on its own test set. For more baseline details, please refer to Appendix A.5.

Table: Table 1: Results of sequence design on SAbDab dataset

Metrics To evaluate the accuracy and rationality of the sequences generated by the model, we selected the following three popular evaluation metrics: (1) Amino Acid Recovery (AAR,$\%$): AAR refers to the ratio of positions where the designed sequence and the true CDR sequence have the same amino acid; (2) Self-consistency RMSD (scRMSD, Å ): To calculate scRMSD, we refold the antibody sequences generated by the model using ABodyBuilder2 (Abanades et al., ). Then, we align the refolded antibody framework with the original antibody and compute the RMSD of the $C_{\alpha}$ atoms in the CDR region. (3) Plausibility: We use pseudo-log-likelihood in an antibody language model, AntiBERTy (Ruffolo et al., ) to calculate plausibility of the generated sequence.

Results As shown in Table 1, RADAb outperforms state-of-the-art methods in each metric and at each CDR region. In particular, in the highly variable and specific CDR-H3 region (Shirai et al., ; Raybould et al., ), our method achieved a great improvement in AAR compared to best-performing methods Diffab-fix and AbMPNN. The evaluation results indicate that the retrieval-augmented method, by introducing structurally similar homologous sequences, has improved the accuracy, consistency, and rationality of the model’s generation.

In addition, CDR-H3 exhibits significant variability in length, sequence, and structure. Typically, deep learning models show decreased performance when generating longer CDR-H3 sequences (Luo et al., ; Høie et al., ). Therefore, we selected a subset of the test set with CDR-H3 lengths longer than 14 to evaluate the generation performance. As shown in Table 3, while the generation performance of all methods declines to some extent, our method demonstrates consistency and significantly outperforms the others, with a larger margin of improvement.

**Table 2: Results of long CDRH3 sequence design performance.**
| Method | AAR(%) | scRMSD | plausibility |
| --- | --- | --- | --- |
| Grafting | 7.79 | 4.05 | -0.785 |
| ProteinMPNN | 46.63 | 2.71 | -0.820 |
| ESM-IF1 | 30.01 | 2.86 | -0.845 |
| Diffab-fix | 42.26 | 3.02 | -0.740 |
| AbMPNN* | 43.27 | 4.39 | -1.012 |
| RADAb | 51.35 | 2.52 | -0.747 |

Figure: Figure 3: Left: Distribution of the samples’ interface energy. Right: Generated CDR-H3 samples and the original structure of PDB: 7d6i antigen-antibody complex. The gray part represents the antibody framework, the red part represents the CDR, and the blue part represents the antigen (with the darker shade indicating the antigen epitope).
Refer to caption: x3.png

### 5.2 Antibody functionality optimization

In this section, we focus on the evolution of antibody sequences and evaluate whether the structure of the evolved sequence has greater functionality compared to the structure of the folded original sequence. To this end, we first fold the designed CDR-H3 sequences with framework sequences and the original real antibody sequences into complete protein structures using ABodyBuilder2. Then, we use FastRelax and InterfaceAnalyzer in PyRosetta (Alford et al., ) to relax the structure and calculate the binding energy $\Delta$G of the antibody-antigen complex.

Metrics
We use various metrics to evaluate the efficacy and functionality of our designed antibodies: (1) $\Delta$$\Delta$G: This metric represents the difference in binding energy between the complex with the designed CDR folded into the structure and the original complex binding energy. (2) $\Delta$$\Delta$G-seq: This metric measures the difference in binding energy between the complex with the designed CDR sequence folded into the structure and the binding energy of the original antibody sequence folded into the structure. It aims to eliminate errors introduced by the folding tool, allowing for a direct comparison of sequence functionality. (3) IMP-seq: This metric indicates the percentage of designed CDR sequences folded into the structure with a lower (better) binding energy than the original antibody sequences folded into the structure.

Results
As shown in Table 3, after folding and relaxing, the antibody sequences we designed show a significant decrease in binding energy compared to other methods, and 37.3% of them have better binding energy than those folded from the original antibody sequences.

To further demonstrate the optimization of antibody sequence functionality by RADAb, we select a specific antigen-antibody complex from the test set (A neutralizing MAb targeting the receptor-binding domain of SARS-CoV-2, PDB: 7d6i). We generate 50 sequences for CDR-H3 and calculate the binding energy $\Delta$G of the folded structures. Among these, 68% of the samples exhibited lower $\Delta$G compared to the original complex. As shown in Figure 3, we select two samples as examples. Although they do not achieve the highest AAR, they demonstrate better binding affinity compared to the native structure.

### 5.3 Analysis

Ablation We conducted a series of ablation experiments for CDR-H3 following the settings described in Section 5.1 to validate the effectiveness and relative contributions of the additional conditions and data we introduced. The specific objectives are: (1) to verify the effectiveness of the retrieval augment module; (2) to assess the validity of the retrieved data; and (3) to evaluate the effectiveness of the evolutionary embedding.

Figure: Table 4: Ablation Study. G represents using Ground truth as retrieval results, R represents the Retrieval-augment mechanism, and E represents Evolutionary embedding mechanism.
Refer to caption: x4.png

As shown in Table 4, We demonstrated the retrieval augment module’s effectiveness by inputting the CDR sequence’s ground truth into this module.
We also removed the retrieval augmentation mechanism and the evolutionary embedding mechanism respectively to validate their effectiveness. The experimental results show that both the retrieval augmentation module and the evolutionary embedding module individually improve performance, and using them together maximizes the model’s performance.

Effect of retrieval dataset To further analyze the benefits brought by the retrieval mechanism and retrieved motifs, we conducted a series of comparative experiments on the value $k$ of CDR-like fragments selected as conditions in the diffusion network, as shown in Figure 4.
We found that when the value of $k$ is low, it brings negative benefits to the model, which may be due to overfitting. As the value of $k$ increases, the model performance also gradually improves. When $k$ equals 15, the model achieves the best performance. But when $k$ exceeds 15, the additional information instead introduces noise to the model, leading to performance degradation.

## 6 Conclusion and Future works

In this work, we propose a retrieval-augmented diffusion generative model RADAb for antibody sequence design. This model leverages global geometric information and local template information, incorporating these conditions into the diffusion process to enhance antibody sequence design and optimization. Experimental results demonstrate that RADAb achieves state-of-the-art performance across multiple tasks. The main limitation of this work is that it has not yet been fully validated in wet lab experiments, which will be one of the major tasks in the future. Since we have proposed a comparatively general retrieval method and retrieval-augment framework, another major future task is to extend the model to the design of various protein motifs.

## Appendix A Additional details

### A.1 Model details

For feature dimensions, we set the single residue feature dimension to 128 and the pair feature dimension to 64. We leverage 6 IPA layers to capture geometry information. ESM2 650M is utilized in our model to create the embedding of antibody sequences, and the embedding dimension is 1280.
In the local CDR-focus network, two layers of axial attention were used (two tied row self-attention and two column self-attention). The embedding dimension is 384, the hidden dimension is 1536, and number of attention heads is 6.

### A.2 Implementation details

Our model was developed and executed within the PyTorch framework. For training, We chose the Adam optimizer with a learning rate of 0.0001, weight decay of 0.0, and momentum parameters beta1 and beta2 set to 0.9 and 0.999, respectively. To dynamically adjust the learning rate, we employed plateau as learning rate scheduler. When the validation loss plateaued, the learning rate was reduced by a factor of 0.8, with a minimum learning rate set to 5e-6. The scheduler’s patience was set to 10 epochs. The batch size is 8 during training. We design 8 samples for each CDR in the test set. All experiments are run on a single RTX4090 GPU, with a memory storage of 24GB.

Due to the high variability and specificity of the CDRH3 region, and it is considered the most critical part in determining antigen-antibody binding. We conducted separate training for the sequence design of this region, adding and removing noise only for the CDRH3 region in each training iteration, with a total of 100,000 iterations. The other five regions, being more conserved, were trained together for a total of 250,000 iterations (approximately equivalent to 50,000 iterations per region). The reverse generation process time step t is set to 100.

### A.3 Implementation of structural retrieval

The input consists of the backbone atom coordinates of each amino acid in the CDR region, forming a set of coordinate points $\mathcal{X}$. ${\mathbb{T}}$ represents the structures of all proteins in the PDB database. ${\mathbb{A}}$ represents a set of protein fragments representing CDR-like fragments corresponding to the input CDR structure. $\mathcal{C}$ represents the set of fragments of each structure with a length of m. $J$ represents a linear motif centered on residue $j$ in the structure, with a length equal to the query fragment.

Assume that the coordinates of residue $j$ are aligned with the central residue of $\mathcal{X}$, and then compute the RMSD of $\mathcal{X}$ when aligned onto $\tau$. If the input contains discontinuous multiple structures, cRMSD will be the cumulative RMSD of these structures. MaxA, maxB, and maxC are three different upper-bound thresholds. These thresholds are selected to improve the speed and accuracy of the retrieval algorithm (For detailed proof, please refer to MASTER (Zhou & Grigoryan, )).

### A.4 Implementation of CDR-like database constructing

To eliminate the computational overhead caused by structural retrieval during the model’s training and inferencing, we followed previous work (Aguilar Rangel et al., ) and initially executed the retrieval algorithm on all CDR structures of all antibodies to construct a CDR-like database.

Each of the CDR structures is used as a query to search for structurally similar motifs in the PDB database. The MASTER algorithm is used to match all CDRs against the entire PDB database to find CDR-like structures. This structural search is based on the Kabsch algorithm, using the RMSD of the $C\alpha$ coordinates. For CDR fragments of length 4, the RMSD threshold is 0.4, and the threshold is increased by 0.05 Å for each additional residue (with a maximum threshold set to 1.0 Å).
In this way, we obtain a CDR-like fragments database corresponding to all CDR structures. Except for strictly filtering out results identical to real CDR sequences, no CDR sequence information was leaked in this process.

### A.5 Baseline details

#### A.5.1 Traditional methods

Rosetta-Fixbb (Adolf-Bryfogle et al., ) Rosetta-Fixbb can use energy functions for antibody CDR sequence design. Since DiffAb has already been proven to outperform it (Luo et al., ; Wu & Li, ) on sequence design task, we did not conduct additional comparisons.

Grafting To simulate the rational design commonly used in traditional antibody design methods, which often involves grafting CDR loops. For each CDR region, we directly selected the top-1 fragment(best) from the retrieval database for the structures in the test set and replaced the corresponding original CDR loop sequences with it.

#### A.5.2 Deep learning methods

ProteinMPNN (Dauparas et al., )
ProteinMPNN is a deep learning framework for protein sequence inverse folding. It leverages a message passing neural network to model the complex relationships between amino acids in a protein structure. We use the antibody’s backbone structure as input and keep the sequences outside the CDR regions to be designed fixed. We design sequences for each CDR region separately. The sampling temperature is set to the default value of 0.1.

Esm-IF1 (Hsu et al., )
Esm-IF1 is a protein sequence inverse folding model trained on millions of AlphaFold2 predicted structures. We use the antibody’s backbone structure as input and keep the sequences outside the CDR regions to be designed fixed. We design sequences for each CDR region separately. The sampling temperature is set to the value of 0.2.

Diffab-fix (Luo et al., )
Diffab is a diffusion model that can design sequences of CDR region with a fixed CDR backbone. It takes antigen-antibody framework context as condition to design CDR sequence. For a fair comparison, we retrained it with the default training configuration fixbb.yml.

AbMPNN (Dreyer et al., )
AbMPNN is fine-tuned by antibody structure data and predicted OAS (Observed Antibody Space ) structure data. Its model architecture is consistent with ProteinMPNN but achieves better performance in antibody inverse folding. We use the antibody’s backbone structure as input and keep the sequences outside the CDR regions to be designed fixed.
However, it is not open-sourced yet, so we evaluate it on its own test set. We design sequences for each CDR region separately. The sampling temperature is set to the default value of 0.1.

### A.6 Experiment on CDR-H3’s length

Figure: Figure S1: AAR distribution of different CDRH3 length
Refer to caption: x5.png

We further evaluated each model’s AAR across different CDR-H3 lengths. As shown in Figure S1, although the performance of all models decreases with increasing H3 length, our method still outperforms the others.

## Appendix B Restrict the evolution space by structural retrieval

The search space for protein sequences is vast. If the length of the sequence to be designed is L, the search space is $20^{L}$. General deep learning methods have already significantly reduced the sequence search space, generating sequences of much higher quality than random sequences. However, it is still challenging to find mutation spaces with high fitness. To restrict this space, we first fix the backbone structure of the CDR region, limiting the space to a specific area. Subsequently, we add more constraints to the mutation space by retrieving the PDB database for motifs similar to the CDR backbone structure and taking them as the model’s conditions, as shown in Figure 2 C.

## Appendix C Structure reconstruction

To reconstruct the antibody structure, we use ABodyBuilder2 (Abanades et al., ), a deep learning model capable of predicting antibody light chain-heavy chain complexes. It is significantly faster than AlphaFold2 and offers higher prediction accuracy. We insert the designed CDR sequences into the antibody framework sequence, input it into ABodyBuilder2 to fold, and use OpenMM relax to obtain the structure corresponding to the new CDR sequences. Subsequently, we align the structure to the real antibody framework. Finally, we use the fastrelax function in PyRosetta (Alford et al., ), with the score function set to ref2015 and max iteration set to 1000, to relax the structure.

## Appendix D Training and inference algorithm

In this section, we provide a detailed algorithm for the training (Algorithm 2) and inferencing (Algorithm 3) processes.

### D.1 Training algorithm

Figure: Algorithm 2 Training Procedure of RADAb

### D.2 Sampling algorithm

Figure: Algorithm 3 Sampling Procedure of RADAb

## Appendix E Case Study

We select a portion of the optimized antibodies in Figure S2. They achieved lower binding energy compared to the original antibody structures. Various examples demonstrate that redesigning antibody sequences in the CDR regions by fixing the backbone structure can produce antibodies with lower antigen-binding energy, even if the designed sequences do not have high similarity to natural antibodies. This method is a successful approach to optimizing antibody functionality.

Figure: Figure S2: Optimized antibodies with lower binding energy. The gray parts represent the antibody framework, the red parts indicate the designed CDR regions, and the blue parts represent the antigen.
Refer to caption: x6.png

## Appendix F Limitations

Quality of retrieved data
Due to the inherent limitations of the MASTER algorithms, the retrieved linear motifs may have structural issues. Despite careful screening and filtering, a tiny portion of the data might have lengths that differ from the original CDR loops or may even become discontinuous due to missing residues. These exceptional cases may have a negative impact on our model. We hope that advances in structural retrieval and improvements in alignment will jointly address this issue.

Computational Cost
When constructing the CDR-like fragments database, searching the entire PDB database using all CDR structures from the SabDab database takes approximately 100 hours. Additionally, the ESM2 encoding used to capture antibody sequence evolution information and the axial attention focused on local CDR in the denoising network require more computational resources than typical diffusion methods.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
questions: 1. Is the provided experimental results correct? According to previous works, can $\Delta \Delta G$ be optimized on such a huge scale?  (~100kcal/mol?) Even for the selected examples in 5.2, $\Delta \Delta G$ only changed like ~-40, far away form claimed 109.16.
2. Also, $\Delta \Delta G$ is usually expressed using negative numbers to demonstrate lower energy is achieved. Why $\Delta \Delta G$ is positive in Tab. 3? Suppose you are using the absolute value of $\Delta \Delta G$ , since lower $\Delta G$ is better, why lower $|\Delta \Delta G|$ indicates better performance? Please explain this.
soundness: 2
strengths: 1. I really like the idea of incorporating RAG into antibody design in such an elegant manner, which not only fulfills the relatively blank space of known data but also utilizes evolutionary information. This is intuitive and turned out to be effective.
2. Nice illustrations, which clearly convey the architecture and designs of authors.
summary: This paper combines the popular technique RAG with antibody design by extracting templates from PDB to enrich evolutionary information, leading to better generalization ability. To be specific, the extracted fragments are based on backbone structures and some properties using a novel retrieval mechanism to utilize both structural and evolutionary information. Experiments demonstrate the effectiveness of the proposed RADAb framework.
weaknesses: 1. The experiment part is relatively weak. For example, for antibody functionality optimization task, only one simple experiment is provided.
2. Baselines are inadequate. Large efforts have been devoted to studying antibody design, such as [1][2][3]:. They utilize the same experimental settings and metrics, which should be taken into consideration.
3. If I comprehend correctly, although RADab is capable of capturing structural and sequential information, it is still a sequence design work, rather than sequence-structure co-design or even full-atom generation work, which is a major weakness compared to other methods.

[1]Kong X, Huang W, Liu Y. End-to-end full-atom antibody design[J]. arXiv preprint arXiv:2302.00203, 2023.

[2]Xiangzhe Kong, Wenbing Huang, and Yang Liu. Conditional antibody design as 3d equivariant
graph translation. arXiv preprint arXiv:2208.06073, 2022.

[3]Shitong Luo, Yufeng Su, Xingang Peng, Sheng Wang, Jian Peng, and Jianzhu Ma. Antigen-specific
antibody design and optimization with diffusion-based generative models for protein structures.
Advances in Neural Information Processing Systems, 35:9754–9767, 2022.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: N/A
soundness: 3
strengths: Originality: See the weaknesses below

Quality: The paper overall has good quality. The experiment is sound.

Clarity: The paper is written clearly.

Significance: The paper significantly improves upon the baseline models in terms of the metrics. It has good significance.
summary: This paper introduces a retrieval-augmented diffusion framework for anti-body design and optimization. Its key contribution is that it retrieves structurally-similar protein segments to the CDR, and use them as conditions in the diffusion generation. Experiments demonstrates impressive results, showing improved performances in antibody CDR sequence inverse folding and antibody functionality optimization. Ablation experiments show that both the retrieval augmentation and evolutionary embedding (ESM2) contribute to the improved performance.
weaknesses: Originality: There has been retrieval-augmented diffusion model proposed in other applications, such as in [1] and [2] (protein-specific 3D molecule generation). Although the paper is the first to introduce it in the antibody design, the method novelty is fair due to the prior works. The paper needs to also cite previous retrieval-augmented generative models, and clearly state the contribution in terms of the method.


[1] Blattmann, Andreas, et al. "Retrieval-augmented diffusion models." Advances in Neural Information Processing Systems 35 (2022): 15309-15324.

[2] Huang, Zhilin, et al. "Interaction-based Retrieval-augmented Diffusion Models for Protein-specific 3D Molecule Generation." Forty-first International Conference on Machine Learning.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: Did you investigate using a similar approach for structure and sequence re-design (as done by DiffAb)? This would expand the scope of the paper.

When you ablate the numberk of reference loop sequences fed to the model (Figure 4), you show that the optimal number of sequences to use is 15 "But when k exceeds 15, the additional information instead introduces noise to the model, leading to performance degradation.". Do you have any motivation why this would be the case? I'd imagine that as long as the mined loops satisfy some RMSD/quality threshold, they should always be beneficial to the model.
soundness: 2
strengths: Inclusion of the data-mined motifs from PDB with similar loop structures is a quite nice and intuitive improvement. That helps the performance of the model. The paper is clearly written and easy to follow. Ablation is performed on the key changes introduced, showing that they are all beneficial.
summary: In this paper authors propose a method to improve upon inverse folding methods by mining PDB for similar structural motifs and use an MSATransformer to take them into account when predicting the loop sequences, along side the usual structure-conditioned diffusion model as used in DiffAb. Inclusion of the mined structures improved the performance of the model in the chosen experiments inverse folding and energy optimization experiments.
weaknesses: Authors state that other models "typically create antibodies from scratch without template constraints, leading to model optimization challenges and poor generative capability". This is a rather harsh critique and I would like such sentences to be backed either by citations or experiments directly showing that. It is true that the proposed approach did generally better than baselines in the chosen experiments, but the difference to DiffAb was marginal and there can be a multitude of reasons for that difference.

In the energy optimization experiment authors say that authors use structures "folded from the original antibody sequences" to evaluate the energy. Folding can be arbitrarilly bad in determining the loop strucutres (esp. CDR H3). So this might very much end up measuring just the agreement between the folding method and the method being tested (how common the folding method finds the sequence to be) instead of how good the energy trully would be. If sequence recovery is reasonably high, one could use grafting and MD to get higher quality/more plausable complexes. Actually, what is the RMSD to original crystal structure, of the reference structure your approach produces (RMSD of original co-crystal vs folding original Ab and runing Rosetta fast relax)?

While in general I like the idea of data mining the PDB for similar loops and using them as reference. I find that the paper is a marginal improvement on existing literature and is of somewhat limited scope for a ICLR paper (only doing inverse folding, while say DiffAb can also do complete loop sequence + structure redesign, and having only two experiments, one of which is not fully convincing to me).

Meta-review: The paper introduces a retrieval-augmented diffusion model, RADAb, for antibody design using structural homologous motifs to guide the generative process. It claims state-of-the-art performance in antibody optimization tasks.

Strengths include innovative integration of retrieval mechanisms with diffusion models and potential for enhancing biomolecular design.

Weaknesses could be limited experimental validation.

The decision to accpet is based on the paper's innovative approach in antibody generation.

Decision: Accept (Poster)