Antibody Design and Optimization with Multi-scale Equivariant Graph Diffusion Models for Accurate Complex Antigen Binding

1 Introduction

Antibodies are essential for targeting disease-related antigens in the therapeutic and diagnostic fields [Carter, 2006]. Their binding specificity and affinity are governed by complementarity-determining regions (CDRs), which exhibit high sequence and structural diversity [North et al., 2011]. As a result, designing antibodies to effectively bind complex antigens requires addressing two key challenges: (i) simplifying geometric feature extraction while preserving critical symmetries, and (ii) enhancing model generalization to novel antigen interfaces. This is because current methods often fail to capture precise molecular interactions and preserve structural integrity.

*Corresponding author

First,the complexity of extracting geometric features while accurately modeling symmetries essential for antibody functionality poses a critical challenge in therapeutic antibody design. Thus, accurate CDR modeling is crucial, as geometric attributes such as bond lengths, angles, and torsions regulate antigen-binding capabilities [Carter, 2006]. However, maintaining symmetries such as rotation and translation invariance during feature extraction remains difficult. Traditional methods such as molecular dynamics (MD) and density functional theory provide precise geometric representations but are computationally restrictive in large-scale antibody modeling [Cai et al., 2024; Luo et al., 2022]. Consequently, advances in geometric deep learning (GDL) have introduced symmetrybased inductive biases, enabling accurate molecular modeling with smaller datasets [Shen et al., 2023]. For example, equivariant graph neural networks (EGNNs), central to GDL, efficiently capture molecular properties by embedding symmetries (e.g., rotation and translation invariance) [Garcia Satorras et al., 2021]. In comparison, innovations such as GearBind improve binding affinity predictions through multirelational graph construction and contrastive pretraining [Cai et al., 2024], while IgFold highlights the importance of structural precision in antibody design using advanced deep learning [Ruffolo et al., 2023]. Additionally, models like AbX integrate evolutionary and physical constraints with geometric factors to optimize antibody sequences [Zhu et al., 2024].

Second, generalizing antibody design to out-of-distribution antigen interfaces presents a significant challenge [Li et al., 2024b; Li et al., 2024a], complicating the development of effective antibodies. Antibody-antigen interactions are inherently unique and ad-hoc. This is because antibodies evolve rapidly to counter diverse pathogens without the long coevolutionary histories typical of other protein-protein interactions [Adolf-Bryfogle et al., 2018; Kong et al., 2022]. As a result, predictive models often fail to adapt to novel antigens that differ substantially from the existing training datasets [Joubbi et al., 2024]. As antibody design continues to evolve, integrating advanced computational techniques with a deeper understanding of structural dynamics will be crucial for overcoming these challenges and unlocking the full potential of therapeutic antibodies.

To address these challenges, we propose AbMEGD, an end-to-end solution integrating Multi-scale Equivariant Graph Diffusion for co-designing antibody sequences and structures. Inspired by ViSNet [Wang et al., 2024], which balances computational efficiency with effective geometric information utilization, AbMEGD incorporates geometry calculation and vector-scalar interactive equivariant message passing (ViS-MP). These strategies enable efficient modeling of atomic-level features such as bond angles, dihedral torsions, and improper angles, aligning with classical molecular dynamics force fields while ensuring geometric precision and scalability. AbMEGD combines atomic-level geometric features with residue-level embeddings in a multiscale framework, capturing local atomic details and global sequence-structure interactions to maintain CDR structural fidelity. Leveraging ViS-MP, the model facilitates interactions between vector representations and scalar embeddings, fully exploiting geometric features while preserving E(3)- equivariance. This multi-scale architecture simplifies geometric feature extraction and enhances generalization to complex antigen interfaces. Unlike existing methods such as GEOAB [Lin et al., 2024], which separate structure generation and affinity prediction, AbMEGD employs a unified framework that preserves essential geometric symmetries and addresses the computational challenges of high-throughput antibody design.

We validate AbMEGD using the SAbDab database [Dunbar et al., 2014], which includes diverse antigen-antibody complexes. Compared to DiffAb, a leading model in antibody design, AbMEGD achieves a 10.13% increase in amino acid recovery (AAR), a 3.32% rise in improvement percentage (IMP), and a 0.062 ˚A reduction in root mean square deviation (RMSD) within the critical CDR-H3 region. These results highlight its capacity to generate accurate structures and highaffinity antibodies, particularly for complex antigens. AbMEGD also excels in CDR design and affinity optimization tasks. Furthermore, its robust generalizability supports highaffinity antibody design for novel and out-of-distribution antigens, establishing a new benchmark in balancing structural conservation with improved functionality.

In summary, our contributions are as follows:

• Unified Generative Framework with Multi-Scale Equivariant Graph Diffusion: We propose AbMEGD, an end-to-end framework that unifies sequence-structure co-design and affinity optimization. By leveraging a multi-scale architecture, AbMEGD captures global sequence-structure interactions and local atomic details, preserving geometric symmetry while ensuring accuracy and efficiency.
• Superior Performance and Broad Applicability: Evaluations on the SAbDab database confirm AbMEGD's superior performance in sequence-structure co-design and affinity optimization. The model achieves structural fidelity and enhances functionality, showcasing adaptability across antibody design tasks and setting a new benchmark in therapeutic antibody development.

2 RelatedWork

2.1 Computational Antibody Design

Traditional antibody design methods rely on statistical energy functions and Monte Carlo sampling for sequence-structure optimization [Adolf-Bryfogle et al., 2018; Ruffolo et al., 2021], but suffer from inefficiency and local optima due to the rugged energy landscape. They also struggle with the vast sequence-structure space, especially for complex antigens, limiting design diversity. Deep learning has enabled sequence-based generative models that treat proteins as language [Alley et al., 2019; Shin et al., 2021], offering efficient sequence generation but often neglecting structural cues critical for antigen binding. The first CDR co-design model [Jin et al., 2021], targeting SARS-CoV-2, incorporated structure but depended on antigen-specific predictors, limiting generalizability. Moreover, it focused on backbone atoms while ignoring side-chain orientations essential for protein-protein interactions. In contrast, our method captures both backbone and side-chain features via multi-level representation, enabling a more complete modeling of the antigen-antibody interface for high-affinity binding.

2.2 Diffusion-based Generative Model

Diffusion probabilistic models (DDPMs) have gained traction for generating protein sequences and structures by denoising samples from a prior distribution [Sohl-Dickstein et al., 2015; Song and Ermon, 2019; Ho et al., 2020; Martinkus et al., 2024]. Models like DiffAb [Luo et al., 2022] and DIFFFORCE [Kulyt˙e et al., 2024] have demonstrated superior efficiency in generating antibody sequences and structures. However, these models primarily focus on residue-level features and often fail to capture atomic-level interactions essential for high-specificity antibody-antigen binding. In comparison, our approach enhances existing DDPMs using VisNetinspired graph-based modules [Wang et al., 2024] to model atomic-level geometric features. This facilitates the effective representation of complex three-dimensional (3D) spatial relationships among the atoms in antibodies and antigens, optimizing antibody specificity and affinity through detailed atomic interaction modeling.

3 Methodology

This section is organized as follows: Section 3.1 introduces the notations used throughout this paper and formally defines the problem. Section 3.2 presents the multi-level feature generation and unified representation mechanism, integrating residue- and atom-level features. Finally, Section 3.3 details the diffusion processes tailored for modeling antibodies.

3.1 ProblemDefinitions and Notations

In this study, we represent each amino acid in a protein complex using three components: type, Cα atom coordinates, and orientation. Specifically, the amino acid type is denoted as sri ∈ {ACDEFGHIKLMNPQRSTVWY }, where i = 1, . . . , N and N is the total number of amino acids in the complex. The spatial position of each residue is described by the Cα coordinates xri ∈ R3 , and the orientation is denoted as Ori ∈ SO(3), which captures the geometric directionality in 3D space.

We consider the antigen structure and the antibody framework as given, focusing our design efforts on the CDRs of the antibody. Specifically, let the generated CDR consist of m amino acids, indexed from c + 1 to c + m, the set of residues in the CDR is denoted as R = {(srj , xrj ,Orj ) | j = c + 1, . . . , c +m}. Our objective is to model the conditional distribution P(R | C), where C = {(sri , xri ,Ori ) |

ri ∈ {1, . . . ,N}\{c+1, . . . , c+m}} represents the antigenantibody framework excluding the CDR regions.

3.2 Multi-level Feature Generation and Unified Representation Mechanism

The overall workflow of AbMEGD is illustrated in Figure 2. This study presents an innovative multi-level feature generation and unified representation mechanism that seamlessly integrates residue- and atom-level features. This approach leverages the complementary nature of these features to simultaneously capture global structural information and finegrained geometric details, enabling comprehensive characterization of antibody-antigen complexes.

Residue-Level Feature Generation. At the residue level, we employ a multi-layer perceptron (MLP) module to process features of individual residues and residue pairs. We define the following variables to represent the feature states of the CDRs:{strj , xtrj ,Otrj }c+m j=1 ∪ {stri , xtri ,Otri }i∈{1,...,N}\{c+1,...,c+m}, where strj denotes the amino acid type of residue j at time step t, xtrj represents the Cα backbone coordinates, and Otrj ∈ SO(3) indicates the orientation.

The single-residue features encapsulate each residue's chemical properties, geometric characteristics, and local structural environment. Thus, the features for residue ri are defined as [Luo et al., 2022]:

where tri represents the amino acid type features, obtained by embedding the amino acid type sri into a low-dimensional vector. gri denotes the local geometric features, derived by constructing a local reference frame based on the backbone atomic coordinates pri . dri encodes the backbone dihedral angles through angular encoding of ϕri and ψri , which describe the torsional conformation of the protein backbone. Finally, fri represents the fragment type features, capturing local fragment-level structural information.

Pairwise residue features describe the interactions and geometric relationships between residues. The features for a residue pair (ri, rj) are represented by [Luo et al., 2022]:

where tri,rj represents the residue-pair type features, encoded by embedding the combination of the amino acid types sri and srj . rri,rj indicates the relative positional features of the residue pair, derived from the sequence index difference ∆ri,rj = |ri -rj|. dri,rj denotes the geometric distance features, calculated as the Euclidean distance between the backbone atoms of the two residues ∥pri -prj ∥ and further represented using Gaussian distribution encoding. Lastly, ori,rj refers to the orientation features of the residue pair, based on the relative orientation angles (ϕri,rj , ψri,rj ). Specifically, ϕri,rj and ψri,rj capture spatial orientation and torsional relationships, ensuring rotationally invariant geometric properties.

Atom-Level Feature Generation. To enhance complex representation, we focus on heavy atoms (N, Cα, C,O, and side-chain Cβ, excluding glycine which lacks Cβ). The atomlevel feature extraction process utilizes the positions of selected atoms X ∈ RA×3 and their atomic indices Z ∈ N A, where A is the number of heavy atoms.

Initially, an embedding block encodes the atomic indices and pairwise edge distances into a high-dimensional embedding space. Then, a series of ViS-MP blocks iteratively update the heavy atoms' scalar and vector representations based on their interactions. Subsequently, residual connections are incorporated between consecutive ViS-MP blocks [Wang et al., 2024]. Thus, the final output includes node and edge features, representing the molecular structure at the atomic level.

Furthermore, the embedding process starts with the heavy atomic indices zai of atom ai, embedded along with the pairwise distances eaiaj using radial basis functions within a cutoff distance. The initial node and edge embeddings are defined as [Wang et al., 2024]:

where N(ai) denotes the set of first-order neighbors of atom ai, and aj represents one of its neighbors. The initial vector embedding v0ai is initialized to 0 and projected into the embedding space as v0ai ∈ R3×K , where K is the hidden dimension size.

Moreover, the ViS-MP mechanism facilitates the interaction between scalar and vector features, enhancing the geometric representation of the molecular structure. Therefore, the mechanism's key operations expressed as follows [Wang et al., 2024]:

Here, hlai denotes the scalar embedding of atom ai at layer l, f lai,aj stands for the edge feature between atoms ai and aj at layer l, and vlai represents the vector embedding of atom ai at layer l. Meanwhile, mlai is the aggregated scalar message for atom ai at layer l, computed by the scalar message function ϕsm. Similarly, mlai is the aggregated vector message for atom ai at layer l, computed by the vector message function ϕvm. The update functions ϕsun and ϕvun respectively update the scalar and vector embeddings based on the aggregated messages and current state. Likewise, the edge features f l+1 ai,aj are updated using ϕsue , which incorporates the vector rejection of embeddings vlai and vlaj along the relative position vector raiaj , ensuring rotationally invariant geometric representation.

Notably, this approach updates node and edge features during message passing, providing a more comprehensive geometric representation compared to previous studies [Sch¨utt et al., 2021; Th¨olke and De Fabritiis, 2022].

Residue and Atom Feature Integration. The residue features (eri , zri,rj ) and atom features (hai , vai ) are integrated and input into an invariant point attention (IPA) [Jumper et al., 2021] module to capture interactions between residue and atom features:

Furthermore, the IPA module facilitates learning local-toglobal relationships within molecular structures by leveraging the interplay between residue and atom features.

3.3 Diffusion Processes

Unlike traditional models focused on denoising molecular states, our approach incorporates multi-level molecular features, enhancing accuracy and structural fidelity by leveraging residue- and atom-level information. Let Rt = {stj , xtj , Otj }c+m j=c+1 denote the intermediate states of residue j at time step t. At t = 0, the system represents real data, including observed CDR sequences and structures, while at t = T, it corresponds to samples from the prior distribution. Thus, the forward diffusion process adds noise incrementally from t = 0 to t = T , while the generative process denoises in reverse. Hence, the diffusion processes for amino acid types stj , coordinates xtj , and orientations Otj are defined as follows:

Multinomial Diffusion for Amino Acid Types. The forward diffusion process for amino acid types follows a multinomial distribution, formulated as [Hoogeboom et al., 2021]:

where the onehot function maps an amino acid type into a 20-dimensional one-hot vector, and 1 represents an all-one vector. The parameter βt type denotes the probability of uniformly resampling an amino acid across 20 types. As t → T , βt type approaches 1, driving the distribution closer to uniformity. Thus, the following probability density offers an efficient mechanism to perturb stj at each timestep t during training [Hoogeboom et al., 2021]:

where ˜αt type = ∏tτ=1 (1 - βτ type ).

The generative diffusion process is defined as:

where F(·)[j] represents a neural network model that takes as input both the structural context (comprising the antigen and antibody framework) and the CDR state from the previous step. This model predicts the probability distribution over the 20 amino acid types for the j-th residue in the CDR.

Unlike forward diffusion, the generative process relies on the structural context C and the CDR state from the previous step, which includes positions and orientations. Generative diffusion aims to approximate the posterior q(st-1 j |stj , s0j ), derived from Eq.10 and Eq.11, for denoising purposes. Thus, the objective of training the generative diffusion process to predict amino acid types is to minimize the expected Kullback-Leibler (KL) divergence between Eq.12 and the posterior distribution. Let DjKL denote the KL divergence for the j-th term. Then, we define the loss function as:

where DjKL = DKL (q(st-1 j | stj , s0j )∥p(st-1 j | Rt , C)).

Cα Coordinate Diffusion. To accurately capture the hierarchical spatial configurations between residues and atoms, our framework models the forward diffusion of normalized Cα coordinates xtj as a Gaussian distribution:

where the diffusion rate βtpos determines the noise level at each timestep t, increasing from 0 to 1 as the process progresses. The cumulative scaling factor, ¯αtpos = ∏tτ=1 (1 - βτpos ). By applying the reparameterization trick [Ho et al., 2020], the generative diffusion process is formalized as:

where the mean µp(Rt , C) is computed as:

where G(·)[j] is a neural network model that predicts the standard Gaussian noise ϵj ∼ N(0, I), which is added to the scaled coordinate of amino acid j, √¯α0 posx0j , following the reparameterization process described in Eq.15: xtj = √¯α0 posx0j + √1 - ¯α0 pos ϵj .

The objective function for training the generative process is the expected mean squared error (MSE) between G and ϵj. This is derived by aligning the distribution p with the posterior q(xt-1 j | xtj , x0j ) [Ho et al., 2020]:

SO(3)-Based Orientation Diffusion. We propose a diffusion framework based on the SO(3), extending the method [Leach et al., 2022] by incorporating multi-level features. The diffusion process on SO(3) is modeled as an isotropic Gaussian distribution, with dynamics governed by the diffusion parameter βtori . This process transitions molecular orientations from well-defined real states at t = 0 to near-uniform distributions at t = T, simulating the gradual addition of noise [Leach et al., 2022]:

IGSO(3) represents the isotropic Gaussian distribution on SO(3), parameterized by a mean rotation and a scalar variance [Leach et al., 2022; Matthies et al., 1988; Nikolayev and Savyolov, 1997]. Then, the ScaleRot operation modifies the rotation matrix by scaling its rotation angle while keeping the rotation axis fixed [Gallier and Xu, 2003]. ¯αtori = ∏tτ=1 (1 - βτ ori ), where βτ ori represents the variance that increases with the diffusion step t. The conditional distribution used for the orientation generation process is defined as:

Here, H(·)[j] represents a neural network designed to denoise orientations, producing the denoised orientation matrix for amino acid j. Training the conditional distribution involves aligning the predicted orientation from H(·) with the real orientation. Consequently, the training objective is formulated to minimize the expected discrepancy, measured using the inner product between the real and predicted orientation matrices:

where ˆOt-1 j = H(·)[j] represents the predicted orientation matrix for amino acid j.

The Overall Training Objective. By summing Eq.13, 18, and 22, and taking the expectation over t, the final training objective function is formulated as:

Training begins by selecting a random time step t and generating noisy states {stj , xtj ,Otj }c+m j=c+1 ∼ p through the diffusion process, which incorporates noise into the data as described in Eq.11, 15, and 20. The model parameters are updated by computing the loss based on the perturbed data and applying backpropagation.

4 Experiments

We propose AbMEGD, an end-to-end framework that integrates multi-scale equivariant graph diffusion for codesigning antibody sequences and structures. It addresses two key tasks: sequence-structure co-design (Section 4.1) and antibody optimization (Section 4.2).

4.1 Sequence-Structure Co-design

We evaluated our model using data from the SAbDab [Dunbar et al., 2014], collected up to September 2024. Antibodyantigen complexes with resolutions worse than 4 ˚A and those binding non-protein antigens were excluded. Antibodies were clustered by 50% CDR-H3 sequence identity, as in [Luo et al., 2022], and five clusters (19 antibody-antigen complexes) formed the test set, including antigens from SARSCoV-2, MERS, and influenza. The remaining clusters formed the training set.

For evaluation, we removed the CDRs from the test set, generating sequences and structures for the removed regions. We fixed the CDR length to the original value (though other lengths could be explored). Each model generated 100 samples per CDR, with all generated and native structures refined using AMBER99 force field [Lindorff-Larsen et al., 2010] in OpenMM [Eastman et al., 2017] and Rosetta [Alford et al., 2017].

Baseline Models. To evaluate performance, as described in Section 2.1, we selected three baseline models for comparison. DiffAb [Luo et al., 2022] represents a diffusionbased generative model specifically designed for antibody structure and sequence modeling. DIFFFORCE [Kulyt˙e et al., 2024] leverages a diffusion framework with force-guided sampling to enhance the quality of generated molecular structures. Additionally, we included the traditional energy-based method RABD [Adolf-Bryfogle et al., 2018], which utilizes energy minimization to design antibodies and optimize binding affinities.

Evaluation Metrics. The models are evaluated using four key metrics: (1) AAR measures how accurately the generated sequences match the native amino acid sequences; (2) RMSD quantifies the structural deviation between the generated and native CDRs by comparing the Cα atom coordinates;

and (3) IMP indicates the proportion of designed antibodies that achieve improved binding energy (∆G) compared to the baseline, calculated using the InterfaceAnalyzer module in Rosetta.

Table 1 compares four antibody design models across three metrics. AbMEGD excels overall, leading in AAR for H2 and H3, with competitive results in H1. It achieves the lowest RMSD in H1, indicating high structural accuracy, and performs well in H2 and H3 RMSD. While DIFFFORCE slightly outperforms AbMEGD in IMP for H2 and H3, AbMEGD leads in H1, balancing binding energy improvement with sequence recovery. DIFFFORCE demonstrates its strength in affinity optimization, particularly in H3, while DiffAb performs well in AAR and achieves the lowest RMSD in H2 but lags in IMP. RAbD shows moderate IMP in H2 but low AAR, indicating poor sequence recovery, with favorable RMSD only in H3 due to reliance on the Rosetta energy function. In contrast, AbMEGD achieves robust energy improvements without Rosetta supervision. Building on DiffAb, which is widely recognized as a leading model in antibody design, AbMEGD significantly improves performance in the critical CDR-H3 region, achieving a 10.13% increase in AAR (from 28.70% to 38.93%), a 3.32% boost in IMP (from 20.42% to 23.74%), and a 0.062 ˚A reduction in RMSD (from 3.481 ˚A to 3.419 ˚A). These results highlight AbMEGD's capability to refine sequence-structure co-design and affinity optimization, surpassing current models. Figure 3 illustrates three designed examples of CDR-H3 interacting with the SARS-CoV2 RBD. Among these, Sample 1 exhibits the strongest binding affinity, as reflected by its lowest binding energy, and demonstrates superior structural complementarity with the antigen. In contrast, Sample 3 shows a higher binding energy compared to the reference, and its CDR shape appears less compatible with the antigen, indicating weaker interactions.

4.2 Antibody Optimization

Our model is designed to optimize existing antibodies, providing an important application in the pharmaceutical domain. The optimization process begins by perturbing the CDR sequence and structure for t steps using the forward diffusion process. Afterward, starting from the (T - t)-th step (with t steps remaining) of the generative diffusion process, the sequences are denoised to produce a set of optimized antibodies. This procedure is illustrated in Figure 4a.

We focus specifically on optimizing the CDR-H3 antibody region in the test set with varying t values. For each antibody and t, the CDR is independently perturbed 100 times, yielding 100 optimized CDRs distinct from the original CDR. We evaluate the performance of our model and DiffAb by reporting the percentage of antibodies with IMP, RMSD, and AAR of the optimized CDRs relative to the original antibody. Additionally, we compare the optimized antibodies with de novo antibodies (t = T = 100) introduced in Section 4.1.

As depicted in Table 2 and Figure 4b , our optimization method produces antibodies with superior binding energy (IMP) as measured by the Rosetta energy function. At t = 4, the optimized CDRs reach an IMP score comparable to that of de novo antibodies, while maintaining significantly higher structural similarity to the original CDRs. The RMSD and AAR results further demonstrate that our model achieves a better trade-off between binding energy improvement and structural conservation compared to DiffAb, making it more suitable for practical applications.

5 Conclusion

In this study, we propose the AbMEGD, an advanced diffusion-based generative model addressing key challenges in antibody design. By integrating residue- and atomic-level features into a unified framework, AbMEGD achieves high performance in sequence-structure co-design and antibody optimization. The experimental results confirm its ability to balance structural conservation and enhance functionality during antibody refinement.

Ethical Statement

There are no ethical issues.

Acknowledgments

This work was supported in part by the Natural Science Foundation of China (No.62476203), and the Guangdong Provincial Natural Science Foundation General Project (No.2025A1515012155). Key Program of Hubei Natural Science Foundation Traditional Chinese Medicine Innovation and Development Joint Fund (2025AFD470).