SlowGuess/NeurIPS_2023_OCR huggingface dataset:
AbDiffuser: Full-Atom Generation of in vitro Functioning Antibodies
Karolis Martinkus<sup>1</sup>, Jan Ludwiczak<sup>1</sup>, Kyunghyun Cho<sup>1,4</sup>, Wei-Ching Liang<sup>2</sup>, Julien LaFrance-Vanasse<sup>2</sup>, Isidro Hotzel<sup>2</sup>, Arvind Rajpal<sup>2</sup>, Yan Wu<sup>2</sup>, Richard Bonneau<sup>1</sup>, Vladimir Gligorijevic<sup>1</sup>, Andreas Loukas<sup>1</sup>
<sup>1</sup>Prescient Design, Genentech, Roche <sup>2</sup>Antibody Engineering, Genentech, Roche <sup>4</sup>NYU
We introduce AbDiffuser, an equivariant and physics-informed diffusion model for the joint generation of antibody 3D structures and sequences. AbDiffuser is built on top of a new representation of protein structure, relies on a novel architecture for aligned proteins, and utilizes strong diffusion priors to improve the denoising process. Our approach improves protein diffusion by taking advantage of domain knowledge and physics-based constraints; handles sequence-length changes; and reduces memory complexity by an order of magnitude, enabling backbone and side chain generation. We validate AbDiffuser in silico and in vitro. Numerical experiments showcase the ability of AbDiffuser to generate antibodies that closely track the sequence and structural properties of a reference set. Laboratory experiments confirm that all 16 HER2 antibodies discovered were expressed at high levels and that  $57.1\%$  of the selected designs were tight binders.
We focus on the generation of immunoglobulin proteins, also known as antibodies, that help the immune system recognize and neutralize pathogens. Due to their potency and versatility, antibodies constitute one of the most popular drug modalities, with 10 out of 37 newly FDA-approved drugs in 2022 being immunoglobulins [58]. The ability to generate new antibodies with pre-defined biochemical properties in silico carries the promise of speeding up the drug design process.
Several works have attempted to generate antibodies by learning to form new sequences that resemble those found in nature [21; 19; 82]. An issue with sequence-based approaches is that it is hard to determine the properties that render a protein a functional molecule (and an antibody a potent drug) without inspecting a 3D model of the functional state such as an interface or active site. So far, almost all of the first design methods that have enabled novel protein design used carefully curated structure-function information to score the designs [45; 48; 30]. The determinant role of structure on function has motivated numerous works to co-design sequence and structure [4; 24; 89; 34; 43; 3; 51] or to first design the protein backbone and then fill in the amino acid identities [90; 32; 15; 87; 50; 92; 47].
An emerging paradigm for protein generation is that of equivariant diffusion [90; 89; 51]. Protein diffusion models combine ideas from equivariant and physics-informed protein structure representations [37] with advances in denoising diffusion probabilistic models [28] to gradually transform noise to a partial structure. The generated structure is then refined by adding side chains and optimizing atom positions to form a full-atom 3D model of a protein.
A pertinent challenge when generating protein structure is satisfying the appropriate equivariance and physics constraints while also balancing modeling complexity and fidelity. Most current protein models rely on equivariant transformers [37; 73] or graph neural networks [11; 43] to satisfy SE(3)
Correspondence to Andreas Loukas <andreas.loukas@roche.com>.
equivariance and typically represent parts of the protein geometry in angle space. The latter recipe can be set up to build a physics-informed model that respects (many) desired constraints but comes with increased complexity. As a consequence, these models are expensive to train and are often applied to a partial protein, ignoring the side chain placement [43; 51] or even deferring the determination of amino acid identity to a later stage [90; 32]. Alternatively, some works focus on the regeneration of complementarity-determining regions (CDRs) of antibodies that are of particular importance to binding [34; 43; 33; 53; 78; 44], which also helps to reduce the complexity of the problem.
Our work is motivated by the observation that key large protein families, here we focus on the antibody protein family, typically have strong properties, such as an ability to be mapped to a reliable sequence ordinate via sequence alignment. Our main contribution is an equivariant diffusion model called AbDiffuser that is designed to exploit these properties. We show that incorporating family-specific priors into the diffusion process significantly improves generation efficiency and quality.
AbDiffuser relies on a new universal SE(3) equivariant neural network that we call the Aligned Protein Mixer (APMixer). In contrast to existing equivariant architectures used for antibodies and proteins, APMixer models residue-to-residue relations implicitly and is particularly effective in handling sequence length changes. Additionally, its significantly smaller memory footprint makes it possible to generate full antibody structures, including framework, CDRs, and side chains. Our approach to residue representation is made possible by a projection method guaranteeing that bond and angle constraints are respected while operating in coordinate and not angle space, a better match to diffusion with Gaussian noise. We also benefit from the effects of overparameterization by scaling the network size to hundreds of millions of parameters on a single GPU; an order of magnitude improvement over the corresponding (E)GNN architectures. Having a powerful model that generates the full antibody structure is shown to be beneficial for the quality of the designed proteins.
We evaluate AbDiffuser on the generation of antibodies from paired Observable Antibody Space (pOAS) [61], modeling of HER2 binders and antigen-conditioned CDR redesign. Numerical experiments demonstrate that the proposed representation, architecture, and priors enable AbDiffuser to better model the sequence and structural properties of natural antibodies. We also submitted a subset of our proposed HER2 binders for experimental validation in a laboratory. Of the 16 samples submitted,  $100\%$  were expressed as actual antibodies and  $37.5\%$  bound to the target with an average pKD of 8.7. Of these, the subset containing raw/filtered samples achieved a binding rate of  $22.2\% / 57.1\%$  and an average pKD of 8.53/8.78, with our tightest binder slightly improving upon the affinity of the cancer drug Trastuzumab. These results provide the first experimental evidence that a generative model trained on mutagenesis data can reliably (with high probability) create new antibody binders of high affinity, even without post-selection by learned or physics-based binding predictors.
Due to space limitations, we refer the reader to the appendix for a detailed discussion of related work, for method details including proofs, as well as for additional numerical results.
2 Denoising Diffusion for Protein Generation
This section describes how we utilize denoising diffusion to generate antibody sequence and structure. The ideas presented are influenced by previous work on denoising diffusion [26; 5; 28]. The main contributions that distinguish our work are presented in Sections 3 and 4.
We adopt the denoising diffusion framework in which, given a data point  $X_0$ , the forward diffusion process gradually adds noise to form corrupted samples  $X_{t}$ . These samples form a trajectory  $(X_0,X_1,\dots,X_t,\dots,X_T)$  of increasingly noisy data, interpolating from the data distribution  $X_0\sim p(X_0)$  to that of an easy-to-sample prior  $X_{T}\sim p(X_{T})$ , such as a Gaussian. The process is constructed to be Markovian, so that  $q(X_{2},\ldots ,X_{T}|X_{0}) = q(X_{1}|X_{0})\Pi_{t = 2}^{T}q(X_{t}|X_{t - 1})$ .
To generate new data, a neural network learns to approximate the true denoising process  $\hat{X}_0 = \phi(\bar{X}_t, t)$ , which can be achieved by minimizing the variational upper bound of the negative log-likelihood [84; 26; 5]. In our diffusion process, we factorize the posterior probability distribution over the atom positions and residue types:
$$
q (X _ {t} | X _ {t - 1}) = q (X _ {t} ^ {\mathrm {p o s}} | X _ {t - 1} ^ {\mathrm {p o s}}) q (X _ {t} ^ {\mathrm {r e s}} | X _ {t - 1} ^ {\mathrm {r e s}}),
$$
where a Gaussian and a categorical distribution govern the noise applied to atom positions and residue types, respectively. We cover these in detail in Sections B.2 and B.3. The reverse diffusion process is
joint: the model jointly considers atom positions and residue types  $(\hat{X}_0^{\mathrm{pos}},\hat{X}_0^{\mathrm{res}}) = \phi (X_t^{\mathrm{pos}},X_t^{\mathrm{res}},t)$ . Throughout the paper, we use  $X^{\mathrm{pos}}\in \mathbb{R}^{n\times 3}$  as a matrix of antibody atom positions and  $X^{\mathrm{res}}\in \mathbb{R}^{r\times 21}$  as a matrix of one-hot encodings of residue types (20 amino acids and a gap).
3 Aligned Protein Mixer
We present Aligned Protein Mixer (APMixer), a novel neural network for processing proteins from a family of aligned proteins. In particular, we focus on antibodies, although the method can in principle be applied to any sufficiently large protein family by using a family-specific global alignment [81; 57]. We first explain how we ensure SE(3) equivariance and how we account for variable sequence lengths in Sections 3.1 and 3.2. The model is detailed in Section 3.3 and our approach to incorporate physics-based constraints on bond angles and lengths is described in Section 3.4.
3.1 SE(3) Equivariance by Frame Averaging
Any neural network  $\phi$  whose input  $X$  has dimension  $n\times 3$  can be made equivariant or invariant to group transformations by averaging the model outputs over a carefully selected subset of group elements called frames  $\mathcal{F}(X)$  [68]. For example, this has been used successfully to make equivariant self-attention models for the prediction of protein-ligand binding energy [35].
We achieve the desired equivariance to rotations and translations (the SE(3) group) as follows:
$$
X ^ {\text {p o s}} = \frac {1}{| \mathcal {F} (X ^ {\text {p o s}}) |} \sum_ {(R, t) \in \mathcal {F} (X ^ {\text {p o s}})} \phi (X ^ {\text {p o s}} R - 1 t, X ^ {\text {r e s}}) R ^ {T} + 1 t, \tag {1}
$$
where  $t = \frac{1}{n}\mathbf{1}^T X^{\mathrm{pos}}$  is the centroid of our points and the four canonical rotation matrices  $R$  forming the four frames  $\mathcal{F}(X^{\mathrm{pos}}) \subset \operatorname{SE}(3)$  needed to achieve equivariance can be determined based on Principle Component Analysis. More specifically, we obtain three unit-length eigenvectors  $v_{1}, v_{2}, v_{3}$  corresponding to the eigenvalues  $\lambda_{1}, \lambda_{2}, \lambda_{3}$  from the eigendecomposition of the covariance matrix  $C = (X^{\mathrm{pos}} - 1t)^T(X^{\mathrm{pos}} - 1t) \in \mathbb{R}^{3 \times 3}$  and define the four frames as
$$
\mathcal {F} (X ^ {\mathrm {p o s}}) = \left\{\left([ \alpha v _ {1}, \beta v _ {2}, \alpha v _ {1} \times \beta v _ {2} ], t\right) | \alpha , \beta \in \{- 1, 1 \} \right\}.
$$
To respect the fixed chirality of proteins observed in humans, we desire equivariance w.r.t. SE(3) and not E(3) which also includes reflections. As such, when constructing frames the third axis sign is not varied but its direction is determined by the right-hand rule (cross product of the first two axes).
SE(3) invariance can be similarly achieved:
$$
X ^ {\text {r e s}} = \frac {1}{\left| \mathcal {F} \left(X ^ {\text {p o s}}\right) \right|} \sum_ {(R, t) \in \mathcal {F} (X ^ {\text {p o s}})} \phi \left(X ^ {\text {p o s}} R - 1 t, X ^ {\text {r e s}}\right) \tag {2}
$$
We make use of equation 2 when denoising residue types  $X^{\mathrm{res}}$  as they are invariant to the rotation of the antibody, and equation 1 for the prediction of atom positions  $X^{\mathrm{pos}}$ .
3.2 Handling Length Changes by Multiple Sequence Alignment
The version of APMixer investigated in this work is built on top of the AHo antibody residue numbering scheme proposed by Honegger and Pluckthun [27]. This numbering scheme was constructed in a data-driven fashion using information from known antibody structures. For each residue in an antibody chain, it assigns an integer position in [1, 149] based on the structural role of the residue (e.g. being in a particular CDR loop or in the framework region between two particular CDR loops). Essentially all known antibodies fit into this representation [27].
As we are modeling paired antibody sequences (the heavy and light chain), the full representation is a  $2 \times 149$  element sequence, where 149 heavy chain elements are followed by 149 light chain elements. We represent AHo gaps physically as 'ghost' residue placeholders; their position is determined in data pre-processing by linearly interpolating between the corresponding atoms of the nearest existing residues (trivial, due to the use of AHo numbering).
Our experiments confirm that the proposed representation consistently improves generation quality. The reason is two-fold: a) Each antibody chain is now represented as a fixed-length sequence with
149 positions that are either empty gaps or are filled with an appropriate amino acid. This fixed-length representation encompasses antibodies with diverse loop lengths via alignment and thus allows our generative model to internally choose how many residues the generated protein will have. In contrast, many of the non-sequence-based protein, small molecule, and graph generative models [49; 54; 28; 90; 32] require the number of elements in the object to be specified beforehand. b) The sequence position directly implies the structural role that the amino acid needs to perform, which makes it easier for the model to pick up structure-specific rules.
3.3 The APMixer Architecture
Our use of a fixed length representation for the immunoglobin variable domain fold family (Section 3.2) also allows us to forgo traditional architecture choices in favor of a more efficient architecture inspired by the MLP-Mixer [86]. We build the model architecture out of blocks, where each block consists of two MLPs, that are applied consecutively on the columns and the rows
$$
X _ {., j} = X _ {., j} + W _ {2} \rho \left(W _ {1} \text {L a y e r N o r m} \left(X _ {., j}\right)\right) \text {f o r a l l} j \in [ c ]
$$
$$
X _ {i, \cdot} = X _ {i, \cdot} + W _ {4} \rho \left(W _ {3} \text {L a y e r N o r m} \left(X _ {i, \cdot}\right)\right) \text {f o r a l l} i \in [ r ],
$$
of the input matrix  $X \in \mathbb{R}^{r \times c}$ , with  $\rho$  being an activation function, and using the notation  $[k] = (1, \ldots, k)$ . We define the model input as a fixed-size matrix  $X$  combining  $X^{\mathrm{pos}}$  and  $X^{\mathrm{res}}$  with one sequence element per row ( $2 \times 149$  rows in total). In each row, we encode the residue type and all of the atom positions for that residue (e.g., the  $C, C_{\alpha}, N, C_{\beta}$  backbone atoms). These input matrix rows are embedded in higher-dimensional vectors using an MLP. Specifically, using this representation, our input matrix  $X$  has  $r = 2 \times 149$  rows and  $c = 21 + 4 \times 3$  columns.
To achieve equivariance to atom translations and rotations, we can either apply frame averaging (Section 3.1) on the whole model or on each AbDiffuser block individually. We chose the second option as in our preliminary experiments this improved performance. Frame averaging can be applied on high-dimensional embeddings simply by splitting them into three-dimensional sub-vectors and using each to compute the SE(3) frames [68]. To account for that residue types are invariant to Euclidean transformations, while the atom positions are equivariant, we split each block's input and output vectors in half, with one half being treated equivariantly and the other invariantly.
Model complexity and SE(3)-universality. APMixer models pairwise interactions implicitly by operating on rows and columns of the input interchangeably. Thus, its memory complexity grows linearly with the number of residues. This contrasts with the usual quadratic complexity of traditional structure-based models and allows us to do more with a fixed run-time, parameter, and/or memory budget. In Appendix G we prove that the model on top of which APMixer is built is SE(3)-universal, meaning that it can approximate any SE(3)-equivariant function.
We also remark that, in principle, other models, such as a 1D CNN or a transformer, could be used in place of the MLPs in APMixer. With such sequence-length-independent models, we would no longer require a multiple sequence alignment of the given protein family, though this would possibly come at the expense of universality and efficiency.
3.4 Physics-informed Residue Representation by Projection
Atoms within a protein adhere to strong constraints. In principle, a neural network trained with enough data can learn to respect these constraints. However, over a fixed data budget, it can be advantageous to construct a model in a manner that guarantees that its outputs never violate the constraints. Previous work commonly represents proteins in terms of rotation and translation of rigid residue frames and uses idealized residue representations to recover atom positions [37; 73]. Although there is a long tradition of modeling backbone and side chain degrees of freedom in angle space [67; 14], operating in angle space adds modeling complexity and makes diffusion potentially inaccurate [94; 8; 90]. We take a different route and devise a way to respect bond constraints while operating in a global coordinate frame, which works seamlessly with standard Gaussian diffusion for atom positions.
Specifically, inspired by interior-point methods [72] that alternate between optimization and projection steps, we design a new non-parametric projection layer that is applied to both model inputs and outputs. The model and the noising process are allowed to move the atoms freely, and the projection layer then corrects their positions such that the constraints are respected.
Backbone residue projection. We use a reference residue backbone  $(C, C_{\alpha}, N, C_{\beta})$  with idealized bond lengths and angles [17]. As the backbone atoms are rigid, we rely on the Kabsch algorithm [38] to identify the optimal roto-translation between the projection layer's input and the rigid reference residue's atom positions. We apply the transformation to the reference residue and output the resulting atom positions. We also ensure that the distance between the corresponding  $C$  and  $O$  atoms in the output is 1.231 Å while staying as close as possible to the input  $O$  position. The reference residue is also used to represent the AHo ghost residue atoms. Idealizing a backbone in this way usually results in a negligible RMSE to the original structure of  $\sim 5 \cdot 10^{-3} \AA$ .
Side chain projection. We employ a similar idea to constrain the placement of side chains. In contrast to structure prediction [37; 46], we cannot use one idealized side chain per amino acid, since the sequence is unknown during generation. Our solution is to convert all side chains to a generic representation with enough degrees of freedom to account for the exact placement of all atoms. The amino acid side chains have up to 4 bonds that they can rotate around by dihedral angles. These degrees of freedom can be captured by constructing a side-chain template that consists of 4 pseudo-carbon atoms, for which the dihedral angles are rotated in the same way as for the original side chain. If the original side chain has fewer degrees of freedom, we simply set the corresponding dihedral angles to  $180^{\circ}$  such that the atoms forming the dihedral angle lie on a plane. The projection layer then only has to ensure that the bond lengths between the pseudo-atoms are respected. We set the bond length to  $1.54\AA$  because carbon atoms are the most common atoms that form the dihedral angles of real amino acids. This representation can be seen in Figure 1. To recover the full-atom structure, we extract the dihedral angles from the side chain template; the angles are then applied to idealized amino acid-specific templates.
Figure 1: The proposed internal generic side chain representation. The dihedral-defining atoms (orange) from the full-atom representation (top) are used to construct a generic four-atom representation (bottom). If the side chain has fewer than four angles, additional atoms (gray) are placed in the generic side chain to correspond to a  $180^{\circ}$  angle. The full atom representation is recovered by applying matching rotations to an appropriate side chain template.
4 Informative Diffusion Priors
It has been hypothesized that diffusion processes for generation should be tailored to the particular task to increase performance [88; 5]. In Appendix E.1, we present a theoretical argument confirming that choosing a good prior reduces the complexity of the learnable model needed to achieve good-quality denoising. Our bound stated in Theorem E.1 reveals that one cannot learn a simple generative model that fits the data well unless the prior has a small Wasserstein distance to the data distribution.
Armed with this knowledge, we next introduce two types of priors that are incorporated into our diffusion model: a) position-specific residue frequencies that describe sequence conservation of the immunoglobulin fold, and b) conditional dependencies between atom positions.
4.1 AHo-specific Residue Frequencies
A separate benefit of having a fixed-length sequence representation (see Section 3.2) is that we can use the residue frequency at each position as a diffusion prior. To do this, we estimate marginal position-specific categorical distributions  $Q^1, \ldots, Q^{2*149}$  over residue types from the training data and use these to define the noise in the discrete residue-type diffusion (see Appendix B.3). Since AHo aligns residues based on their structural role within the protein, the estimated prior distributions exhibit significant variability across positions and have low entropy in preserved regions of the immunoglobulin fold. In this way, the noise  $q(X_t^{\mathrm{res}}|X_0^{\mathrm{res}}) = \mathrm{Cat}(X_0^{\mathrm{res}}(\beta_tI + (1 - \beta_t)Q^i))$  at every step of the forward process depends on the residue position  $i$  in the fixed-length sequence representation. Gap frequencies are also captured in these noise matrices.
4.2 Encoding Conditional Atom Dependencies
It is a consequence of their chain-like structure that neighboring residue atom positions of proteins are strongly correlated. Theorem E.1 suggests that encoding these correlations within the diffusion process eases the denoising complexity and can free the model to focus on more challenging aspects.
To this end, we capture the conditional independence relations between atom positions by a Gaussian Markov Random Field (GMRF) [42]. The latter corresponds to a Gaussian  $\mathcal{N}(0,\Sigma)$  over atom positions whose precision matrix  $\Sigma^{-1} = L + aI$  is equal to the shifted graph Laplacian  $L = \mathrm{diag}(A1) - A$  associated with the adjacency matrix  $A$ . The GMRF operates under the assumption that node features (e.g., 3D positions) for which there is no edge in the adjacency  $A$  are conditionally independent. Some concurrent works [32; 36] also considered conditional atom dependencies by hand-crafting correlation matrices that capture chain relations between residues. We take a step further by proposing to automatically learn the sparse conditional dependence relations from the training data. Our approach entails estimating a sparse adjacency matrix  $A$  that captures the data variance under the GMRF model. The details are described in Appendix F.