SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We evaluate AbDiffuser's ability to design antibodies. After describing the experimental setup and metrics, Section 5.1 presents in silico tests illustrating the effect of our proposed changes on the generation quality. Section 5.2 details in vitro results showcasing the ability of our method to design new expressing antibodies that bind to a known antigen. Further analyses can be found in the Appendices K and L, whereas experiments on CDR redesign in SAbDab are in Appendix N.
Baselines. We compare APMixer with: a) a sequence transformer based on BERT [12; 74] whose output is folded to recover the structure; b) an E(n) Equivariant Graph Neural Network (EGNN) [75] which is a popular choice for tasks such as 3D molecule generation [28], antibody CDR loop inpainting [43], and antibody generation [89]; and c) a FA-GNN [68], corresponding to a standard GNN with SE(3) equivariance attained by frame averaging. We also evaluate the proposed informative antibody-specific priors using all of these architectures. To ensure the comparison is performed fairly and under similar settings, we always use the projection layer and our diffusion implementation only varying the denoising model architecture. In the OAS task, we also compare our diffusion-based approach with d) the IgLM [80] antibody language model. To ensure that it generates paired OAS-like sequences, we condition the generation on the subsequences of the first and last few amino acids taken from pOAS sequences (something that our approach does not need). It is also important to note that the comparison is not fair, since the publicly available IgLM model was trained on 558M sequences that also include the whole paired OAS (105k sequences) and the test set we use. So, in many ways, IgLM's performance represents the best results we could ever hope to achieve with a sequence-only approach. We also compare with e) dyMEAN [44], which is the only other antibody generative model previously shown to be able to jointly generate the full strucuture and sequence. In the binder generation task, we compare with f) RefineGNN [34], g) MEAN [43] and h) DiffAb [53], three state-of-the-art geometric deep learning methods for CDR redesign.
Metrics. The quality of generated sequences is measured in terms of their naturalness (inverse perplexity of the AntiBERTy [74] model), closeness to the closest antibody in the training set in terms of edit distance, and stability estimated by IgFold. We also verify that the generated antibody sequences satisfy the appropriate biophysical properties using four additional structure-based metrics [69]: CDR region hydrophobicity (CDR PSH), patches of positive (CDR PPC), and negative charge (CDR PNC), and symmetry of electrostatic charges of heavy and light chains (SFV CSP). The metrics applied to generated structures focus primarily on the estimated free energy  $\Delta G$  using Rosetta [2] and RMSD for backbone heavy atom positions as compared to IgFold [73] predictions. More details can be found in Appendix J.
As we want to estimate how well the entire data distribution is captured, in all cases except RMSD, we report the Wasserstein distance between the scores of the sequences in the test split and the scores of the generated sequences. As a reference, we also report the baseline metrics achieved by the sequences and structures in the validation set. A generative model that approaches or matches these values is effectively as good at modeling the distribution of the specific metric in consideration as i.i.d. sampling. The test set and the generated set always have 1000 examples.
Further details on training and implementation can be found in Appendix M.
Model	W1(Nat.)↓	W1(Clo.)↓	W1(Sta.)↓	W1(PSH)↓	W1(PPC)↓	W1(PNC)↓	W1(CSP)↓	W1(ΔG)↓	RMSD↓
Validation Set Baseline	0.0150	0.0043	0.0102	0.8301	0.0441	0.0176	0.4889	1.0814	—
Transformer	0.5308	0.4410	1.2284	25.8265	0.2324	0.2278	2.7925	—	—
Transformer (AHo)	0.4456	0.3474	0.5351	6.4490	0.1641	0.0593	2.3472	—	—
IgLM* [80]	0.1103	0.0484	0.0577	11.0675	0.0413	0.0671	1.9274	—	—
dyMEAN [44]	0.1319	0.1600	0.0423	3.9145	0.1566	0.2929	2.3711	601.1153	3.8157
EGNN	0.3988	0.2655	0.3547	2.1115	0.1486	0.1085	1.9881	1586.0160	9.8231
EGNN (AHo)	0.3329	0.2229	0.2904	8.1620	0.1263	0.1075	0.7978	1714.2734	10.0628
EGNN (AHo & Cov.)	0.3482	0.2374	0.2443	2.5632	0.1190	0.0462	1.2184	1015.8926	9.4814
FA-GNN	0.4141	0.2822	0.4302	2.5330	0.1696	0.1164	1.7886	22.7988	0.8617
FA-GNN (AHo)	0.3407	0.2263	0.2344	2.3272	0.1411	0.1306	1.6046	8.7506	0.8321
FA-GNN (AHo & Cov.)	0.2785	0.1669	0.0815	5.4440	0.0493	0.0212	0.7768	15.3670	0.8814
AbDiffuser (uniform prior)	0.2837	0.1419	0.2188	3.1364	0.0727	0.1691	1.3874	38.8417	0.8398
AbDiffuser (no projection)	0.2378	0.1529	0.0694	2.3530	0.0637	0.0793	0.7376	6313.2495	11.1431
AbDiffuser (no Cov.)	0.2309	0.1107	0.1235	1.2392	0.0664	0.0511	0.6453	17.7322	0.6302
AbDiffuser	0.1979	0.0921	0.0662	2.3219	0.0314	0.0285	0.6662	13.3051	0.5230
AbDiffuser (side chains)	0.0916	0.0520	0.0186	6.3166	0.0209	0.0754	0.8676	16.6117	0.4962
Table 1: Antibody generation based on Paired OAS [61]. AHo denotes models that use AHo numbering and position-specific residue frequencies. Cov denotes models that use the learned covariance. IgLM is denoted by * since it was trained on significantly more data (including the test set) and was given part of the sequence to bootstrap generation. The top three results in each column are highlighted as First, Second, Third.
Model	Parameters ↑	Memory (training) ↓	Memory (generation) ↓	Generation time ↓
Transformer	84M	14GB	15GB	3.2 min
EGNN	39.3M	78GB	16GB	22.6 min
FA-GNN	9.4M	75GB	38GB	9.5 min
AbDiffuser	169M	12GB	3GB	2.3 min
Table 2: Number of parameters, model memory consumption during training with a batch size of 4 and memory consumption with the time taken to generate a batch of 10 examples for paired OAS.
5.1 Paired OAS Generation
We focus on matching the distribution of 105k paired sequences from the Observed Antibody Space database [61] folded with IgFold and optimized with Rosetta [2].
Table 1 summarizes the results. Generated samples by AbDiffuser improve upon baselines on nearly all fronts, even compared to the IgLM language model which was trained on magnitudes more data (especially when concerning structure-related metrics). dyMEAN GNN-based model struggles in this distribution modeling task and is the only model tested that does not achieve the perfect uniqueness of the generated samples (58.  $2\%$  unique). The experiment also corroborates our analysis (Theorem E.1) on the benefit of informative priors to diffusion: using a position-specific residue type frequency (AHo) and encoding conditional atom dependencies through a learned covariance matrix (Cov.) helps to improve the ability of most models to capture the pOAS distribution. Interestingly, including the learned covariance can sometimes noticeably improve the quality of the generated sequences (FA-GNN), but its strong benefit to structure quality is only felt when the model is powerful enough to model the structure well (APMixer). Inclusion of AHo numbering and position-specific frequencies improves all models. We perform a similar ablation for APMixer by setting the prior distribution to uniform (uniform prior) and observe a similar performance drop.
To interpret the fidelity of the generated structures, we recall that IgFold uses an ensemble of 4 models followed by Rosetta optimization and that, on average, individual IgFold models (before ensembling) achieve an RMSD of 0.4239 on the test set. Therefore, in this regard, the structures created by AbDiffuser are nearly indistinguishable from the test set structures (RMSD of 0.4962). A more detailed analysis of per-region RMSD can be found in Appendix K. We further observe that when no projection layer is used and instead one uses standard diffusion to predict the noise added to the atom positions [28], the training becomes less stable and the model can fail to learn to generate good structures. Encouragingly, forcing AbDiffuser to model side chain positions alongside the backbone tends to improve the similarity of the generated sequences (Naturalness, Closeness, Stability). This is likely due to the strong coupling between feasible side chain conformations and residue types. The generated side chains get an average Rosetta packing score of 0.624, whereas folded and Rosetta-optimized structures have a mean packing score of 0.666. Recalling that a packing score above 0.6 is widely considered good [77; 66; 3], we deduce that AbDiffuser is able to generate physically plausible side chain conformations. When available, we also use the side chain positions
Model	W1(Nat.) ↓	W1(Clc.) ↓	W1(Sta.) ↓	W1(PSH) ↓	W1(PPC) ↓	W1(PNC) ↓	W1(CSP) ↓	W1(ΔG) ↓	RMSD ↓	pbind ↑	Uniq. ↑
Validation Set Baseline	0.0011	0.0003	0.0061	1.3183	0.0196	0.0114	0.3280	2.0350	—	0.8676	100%
MEAN [43]	0.0072	0.0009	0.0267	8.4184	0.0184	0.0231	0.4108	8.5981	0.7792	0.7767	38.9%
DiffAb [53]	0.0074	0.0014	0.0498	0.5481	0.0097	0.0067	3.4647	6.7419	0.4151	0.8876	99.7%
RefineGNN [34]	0.0011	0.0004	0.0026	0.5482	0.0053	0.0046	0.1260	—	—	0.7132	100%
Transformer (AHo)	0.0014	0.0031	0.0097	1.3681	0.0171	0.0140	0.2657	—	—	0.3627	100%
EGNN (AHo & Cov.)	0.0013	0.0030	0.0102	1.1241	0.0123	0.0236	0.2441	1967.5280	9.2180	0.3626	100%
FA-GNN (AHo & Cov.)	0.0018	0.0030	0.0063	0.4158	0.0107	0.0056	0.2644	76.7852	3.1800	0.4576	100%
AbDiffuser	0.0013	0.0018	0.0028	0.4968	0.0205	0.0113	0.1588	6.4301	0.3822	0.5761	100%
AbDiffuser (side chains)	0.0010	0.0005	0.0062	1.2909	0.0115	0.0029	0.0948	32.0464	0.4046	0.6848	100%
AbDiffuser (τ = 0.75)	0.0005	0.0011	0.0054	0.3934	0.0148	0.0129	0.1785	6.2468	0.3707	0.6382	100%
AbDiffuser (s.c., τ = 0.75)	0.0005	0.0004	0.0126	1.8510	0.0126	0.0017	0.0917	12.8923	0.3982	0.7796	100%
AbDiffuser (τ = 0.01)	0.0008	0.0014	0.0265	2.5944	0.0206	0.0053	0.2378	15.2200	0.3345	0.9115	99.7%
AbDiffuser (s.c. τ = 0.01)	0.0015	0.0024	0.0159	1.5043	0.0210	0.0126	0.5173	114.4841	0.6795	0.9436	91.4%
Table 3: Generating Trastuzumab mutants based on the dataset by Mason et al. [55]. The top three results in each column are highlighted as First, Second, and Third. Multiple approaches can generate sequences similar to the test set, but generating predicted binders (large  $p_{\mathrm{bind}}$ ) is considerably harder.
predicted by the model for the  $\Delta G$  energy estimation. Even though this procedure is expected to generate results slightly different when compared to backbone-only models (in the latter case missing side chain atoms are first repacked based on the rotamer libraries before the energy minimization step), we still observe a high overlap between the resultant energies. This further highlights the quality of side-chain prediction.
It should be noted that  $\Delta G$  energy computation and some additional structure-based metrics [69] (i.e., CDR PSH, CDR PPC, CDR PNC, SFV CSP) are inherently susceptible to even minor changes in the geometry of the modeled structures. Thus, in line with the discussions by Raybould et al. [69], the overall trends of these metrics can be used to assess the generated samples as similar or dissimilar to the reference distributions, but one should not fall into the trap of overly focusing on the specific values attained. From this perspective, most structure-based models do sufficiently well on these metrics, perhaps with the exception of EGNN  $\Delta G$ .
In Table 2 we show that APMixer is able to use an order of magnitude more parameters with a smaller memory footprint during training and offers more efficient sample generation, compared to the baseline architectures, on Nvidia A100 80GB GPU.
5.2 Generating HER2 Binders
Antibody generative models can be used to explore a subset of the general antibody space, such as the binders of the target antigen. The purpose of modeling and expanding a set of binders is twofold: a) it allows us to rigorously validate our generative models in a setup more tractable than denovo design; b) from a practical standpoint, it sets the ground for the optimization of properties that render binding antibodies drugs, such as developability, immunogenicity, and expression, allowing efficient exploration of the binder space [21; 63]. Note that sufficiently large libraries consisting of antibodies of variable binding affinity are usually discovered during the very early stages of the drug design process by means of high-throughput experiments [64; 76; 55]. Thus, the data used here can be sufficiently easier to obtain than the crystal structures usually assumed in CDR redesign experiments.
We use the Trastuzumab CDR H3 mutant dataset by Mason et al. [55] which was constructed by mutating 10 positions in the CDR H3 of the cancer drug Trastuzumab. The mutants were then evaluated using a high-throughput noisy assay to determine if they bind the HER2 antigen. After discarding all duplicate sequences and sequences that appear in both the binder set and the non-binder set, we are left with 9k binders and 25k non-binders. The generative models were trained only on the binder set. Separately, we train a classifier based on the APMixer architecture to distinguish binders from non-binders. The classifier achieves  $87.8\%$  accuracy on a randomly selected test set of 3000 antibodies, implying that the predicted binding probability is a somewhat informative metric for how tightly the binder distribution is approximated by the models.
The computational results summarized in Table 3 evidence that AbDiffuser closely approximates the binder distribution. Although most models achieve results close to those achieved by the validation set<sup>1</sup>, the sequence-based baseline struggles to capture the features that convey binding; evidenced by the low predicted binding probability. GNN-based baselines struggle to predict the correct structure
<sup>1</sup>The Wasserstein distance of the validation set is likely overestimated here due to the validation set being smaller than the test set (100 vs 1000 samples).
(RMSD above 3) producing structures with unreasonably high Amber energy. As a result, the binding probability of their designs is lower than that of the more powerful structural models.
Next, we look at the two baselines that only redesign CDR H3 instead of generating the whole antibody. RefineGNN [34] manages to closely capture the sequence and biophysical characteristics and generates sufficiently good binders. DiffAb [53], which is another diffusion-based model, achieves the best binding probability out of the CDR redesign baselines. MEAN [43] generated few unique sequences (38.9% uniqueness) and failed some of the distribution-based metrics (Naturalness, CDR PSH). dyMEAN [44] collapsed to always generating a single sample, thus we did not report its results. The overfitting behavior of MEAN and dyMEAN can be attributed to the reconstruction-based training objective and the fact that they were designed for a slightly different task of CDR redesign with different antigens, instead of the single antigen we have here.
Due to the use of iterative refinement and reconstruction-based training objectives, RefineGNN and MEAN focus on the generation of the most likely examples. Focusing on high-likelihood modes is especially beneficial here, as the experiment that was used to create the dataset is noisy. Diffusion models can also be adjusted to focus on prevalent modes by reducing the temperature  $\tau$  of the denoising process. In the Gaussian case, we specified the temperature as the scaling factor of the noise added during the reverse process (Appendix B.2), whereas in the discrete case we specified it as the temperature of the model output softmax function (Appendix B.3). We observe that a slight reduction of the temperature helps to improve the general distribution fit across almost all metrics. Reducing the temperature further boosts the binding probability, but, as expected, can result in a slight loss of diversity. Using a higher temperature slightly increases the Stability Wasserstein distance while improving Stability. The phenomenon occurs because the model is no longer concerned with fitting low-likelihood modes of the real distribution that contain structures of poor stability.
In contrast to MEAN, RefineGNN and DiffAb, which only redesign the CDR H3 of the test-set structures, AbDiffuser generates full antibodies and still achieves better sequence similarity. MEAN [43] also produced structures of noticeably worse RMSD, which can be explained as follows: a) as we see in Appendix K, MEAN does not predict CDR H3 positions as well; b) changing CDR H3 can impact the most-likely conformation of the overall antibody, something that the CDR inpainting models cannot account for. We do not report the RMSD and  $\Delta G$  for RefineGNN as it does not place the generated CDR H3 loop in the reference frame of the original antibody.
The AbDiffuser model that also generates side chains generally achieved better sequence similarity (Naturalness, Closeness) and better binding probability than the backbone-only model, but a worse similarity of sequence stability. Furthermore, while the side chain model achieved a worse overall structure quality ( $\Delta G$ , RMSD), as we see in Appendix K it predicted CDR H3 positions more precisely, which is the main desideratum in this experiment.
5.3 In Vitro Validation
We further validate our designs through an in vitro experiment. As shown in Figure 2, all submitted designs were expressed and purified successfully (average concentration of  $1.25\mathrm{mg / ml}$ ) and an average of  $37.5\%$  of the designs were confirmed binders with  $\mathrm{pKD}\in [8.32,9.50]$  (higher is better) whose average was 8.70. The binding rate was improved (from  $22.2\%$  to  $57.1\%$  ) when considering designs that were additionally filtered so that they were better than the bottom 25th quantile in every metric (naturalness, RMSD, etc.) and a classifier trained to distinguish binders from non-binders predicted that they bound with high confidence. This increase in binding for filtered samples suggests that our selected metrics indeed correlate with desirable in vitro characteristics. Our best binder belonged to the latter set and its pKD of 9.50 was slightly above Trastuzumab (not in the training set) while differing in 4 positions of the CDR H3 loop. Further examination of the top binders is performed in Figure 3 and Appendix L.
In contrast to the only other comparable study for ML-based Trastuzumab CDR H3 mutant design by Shanehsazzadeh et al. [76], our best binder had  $\mathrm{pKD} = -\log_{10}(3.17^{-10}) = 9.50$ , while the best binder found by Shanehsazzadeh et al. had  $\mathrm{pKD} = 9.03$ . Two important differences between the two studies are that: we trained on binders and tested 16 samples in the wet lab, while Shanehsazzadeh et al. used a simpler sequence infilling model trained on a large set of generic antibodies to generate 440k candidate sequences that were filtered in the wet-lab using high-throughput screening to identify 4k binders, of which 421 were selected to be tested using precise SPR measurements. The fact that our
Figure 2: In vitro validation of AbDiffuser designs in terms of their ability to express (left), binding affinity (center), and binding rate (right). The 'raw' column corresponds to randomly selected generated antibodies, whereas 'filtered' designs were additionally filtered by in silico screening.
Figure 3: HER2 structure and the 3D structures of generated binders highlighted in different colors. AbDiffuser has learned to redesign part of the binding interface while maintaining affinity. Our tightest binder achieved a pKD of 9.5 which, accounting for experimental variability, is akin to Trastuzumab whose average measured pKD was 9.21.
approach required 26x fewer precise SPR wet-lab experiments to find a better binder hints toward a major improvement in efficiency. This highlights the importance of new and more efficient generative models for antibodies and more efficient, and thus powerful architectures such as APMixer.
The work of Shanehsazzadeh et al. [76] also shows that high-throughput wet-lab experiments can be used to build sufficiently large datasets of target-specific binders without starting from a known binder. This provides a straightforward way to train AbDiffuser on new targets.
After good binders are determined through high-throughput and SPR experiments, at the end of the drug discovery process, an experimental 3D structure of the bound complex is often produced. In recent works, redesigning CDRs in such co-crystal structures has been a popular task [53; 43; 34]. To that end, Appendix N investigates how AbDiffuser can be adopted for it and shows that it offers much better sequence recovery rates than the current state-of-the-art diffusion model for the CDR redesign [53].
We propose a denoising model that deviates from conventional practice in deep learning for protein design. APMixer enjoys the benefits of the best protein models (SE(3) equivariance, universality), while being significantly more memory efficient. We also show how to tailor the diffusion process to antibody generation by incorporating priors that enable the handling of variable-length sequences and allow the model to navigate strong bond constraints while operating in an extrinsic coordinate space that supports Gaussian noise. In our future work, we want to apply our method to protein families beyond antibodies, as organized in CATH [81] and Pfam [57]. Promising examples include the TIM-barrel fold and enzyme families, such as SAM-dependent methyltransferases, long-chain alcohol oxidases, and amine dehydrogenases [83].
[89] R. Vinod, K. K. Yang, and L. Crawford. Joint protein sequence-structure co-design via equivariant diffusion. In NeurIPS 2022 Workshop on Learning Meaningful Representations of Life, 2022.
[90] J. L. Watson, D. Juergens, N. R. Bennett, B. L. Trippe, J. Yim, H. E. Eisenach, W. Ahern, A. J. Borst, R. J. Ragotte, L. F. Milles, et al. Broadly applicable and accurate protein design by integrating structure prediction networks and diffusion generative models. bioRxiv, 2022.
[91] D. N. Woolfson. A brief history of de novo protein design: minimal, rational, and computational. Journal of Molecular Biology, 433(20):167160, 2021.
[92] K. E. Wu, K. K. Yang, R. v. d. Berg, J. Y. Zou, A. X. Lu, and A. P. Amini. Protein structure generation via folding diffusion. arXiv preprint arXiv:2209.15611, 2022.
[93] K. Xu, W. Hu, J. Leskovec, and S. Jegelka. How powerful are graph neural networks? In International Conference on Learning Representations, 2019.
[94] J. Yim, B. L. Trippe, V. De Bortoli, E. Mathieu, A. Doucet, R. Barzilay, and T. Jaakkola. Se (3) diffusion model with application to protein backbone generation. arXiv preprint arXiv:2302.02277, 2023.
A Related Work
The literature on computational methods for protein design is vast. We focus here on machine learning methods for protein sequence and structure generation and refer the interested reader to the relevant reviews by Woolfson [91]; Ferruz et al. [18] for broader context.
A.1 Protein Language Models
The first works on protein generation focused on language modeling [16; 21; 71; 85]. Motivated by the observation that the natural forces driving antibody sequence formation, namely V(D)J recombination and the selective pressure in the presence of an antigen referred to as somatic hypermutation, are distinct from other proteins, recent works have also built antibody-specific language models. Therein, large-scale models trained on millions of sequences, such as AntiBERTy [74] and IgLM [80], have been used to generate full-length antibody sequences across different species or restore missing regions in the antibody sequence [62; 56].
A.2 Protein Structure Generation
These approaches use equivariant architectures to generate proteins and antibodies in 3D. We broadly distinguish two strategies: backbone-first and sequence-structure co-design.