Ophiuchus: Scalable Modeling of Protein Structures through Hierarchical Coarse-graining SO(3)-Equivariant Autoencoders

A canonical ordering of the atoms of each residue enables the local geometry to be described in a stacked array representation, where each feature channel corresponds to an atom. To directly encode positions, we stack the 3D coordinates of each atom. The coordinates vector behaves as the irreducible-representation of $SO(3)$ of degree $l=1$. The atomic coordinates are taken relative to the $\mathbf{C}_{\alpha}$ of each residue. In practice, for implementing this ordering we follow the atom14 tensor format of SidechainNet [King & Koes (2021)], where a vector $P\in\mathbb{R}^{14\times 3}$ contains the atomic positions per residue. In Ophiuchus, we rearrange this encoding: one of those dimensions, the $\mathbf{C}_{\alpha}$ coordinate, is used as the absolute position; the 13 remaining 3D-vectors are centered at the $\mathbf{C}_{\alpha}$, and used as geometric features. The geometric features of residues with fewer than 14 atoms are zero-padded (Figure 4).

Still, four of the standard residues (Aspartic Acid, Glutamic Acid, Phenylalanine and Tyrosine) have at most two pairs of atoms that are interchangeable, due to the presence $180^{\circ}$-rotation symmetries [Jumper et al. (2021)]. In Figure 5, we show how stacking their relative positions leads to representations that differ even when the same structure occurs across different rotations of a side-chain. To solve this issue, instead of stacking two 3D vectors ($\mathbf{P}^{*}_{u},\mathbf{P}^{*}_{v}$), our method uses a single $l=1$ vector $\mathbf{V}^{l=1}_{\textrm{center}}=\frac{1}{2}(\mathbf{P}^{*}_{v}+\mathbf{P}^{*}_{u})$. The mean makes this feature invariant to the $(v,u)$ atomic permutations, and the resulting vector points to the midpoint between the two atoms. To fully describe the positioning, difference $(\mathbf{P}^{*}_{v}-\mathbf{P}^{*}_{u})$ must be encoded as well. For that, we use a single $l=2$ feature $\mathbf{V}^{l=2}_{\textrm{diff}}=Y_{l=2}\big{(}\frac{1}{2}(\mathbf{P}^{*}_{v}-\mathbf{P}^{*}_{u})\big{)}$. This feature is produced by projecting the difference of positions into a degree $l=2$ spherical harmonics basis. Let $x\in\mathbb{R}^{3}$ denote a 3D vector. Then its projection into a feature of degree $l=2$ is defined as:

$Y_{2}(x)=\left[\sqrt{15}\cdot x_{0}\cdot x_{2},\sqrt{15}\cdot x0\cdot x1,\frac{\sqrt{5}}{2}\left(-x_{0}^{2}+2x_{1}^{2}-x_{2}^{2}\right),\sqrt{15}\cdot x_{1}\cdot x_{2},\frac{\sqrt{15}}{2}\cdot(-x_{0}^{2}+x_{2}^{2})\right]$

Where each dimension of the resulting term is indexed by the order $m\in[-l,l]=[-2,2]$, for degree $l=2$. We note that for $m\in\{-2,-1,1\}$, two components of $x$ directly multiply, while for $m\in\{0,2\}$ only squared terms of $x$ are present. In both cases, the terms are invariant to flipping the sign of $x$, such that $Y_{2}(x)=Y_{2}(-x)$. Equivalently, $\mathbf{V}^{l=2}_{\textrm{diff}}$ is invariant to reordering of the two atoms $(v,u)\rightarrow(u,v)$:

In Figure 5, we compare the geometric latent space of a network that uses this permutation invariant encoding, versus one that uses naive stacking of atomic positions $\mathbf{P}^{*}$. We find that Ophiuchus correctly maps the geometry of the data, while direct stacking leads to representations that do not reflect the underlying symmetries.

Given a latent representation at the residue-level, $(\mathbf{P},\mathbf{V}^{0:{l_{\max}}})$, we take $\mathbf{P}$ directly as the $\mathbf{C}_{\alpha}$ position of the residue. The hidden scalar representations $\mathbf{V}^{l=0}$ are transformed into categorical logits $\bm{\ell}\in\mathbb{R}^{|\mathcal{R}|}$ to predict the probabilities of residue label $\hat{\mathbf{R}}$. To decode the side-chain atoms, Ophiuchus produces relative positions to $\mathbf{C}_{\alpha}$ for all atoms of each residue. During training, we enforce the residue label to be the ground truth. During inference, we output the residue label corresponding to the largest logit value.

Relative coordinates are produced directly from geometric features. We linearly project $\mathbf{V}^{l=1}$ to obtain the relative position vectors of orderable atoms $\hat{\mathbf{V}}^{l=1}_{\textrm{ordered}}$. To decode positions $(\hat{\mathbf{P}}^{*}_{v},\hat{\mathbf{P}}^{*}_{u})$ for an unorderable pair of atoms $(v,u)$, we linearly project $\mathbf{V}^{l>0}$ to predict $\hat{\mathbf{V}}^{l=1}_{\textrm{center}}$ and $\hat{\mathbf{V}}^{l=2}_{\textrm{diff}}$. To produce two relative positions out of $\hat{\mathbf{V}}^{l=2}_{\textrm{diff}}$, we determine the rotation axis around which the feature rotates the least by taking the left-eigenvector with smallest eigenvalue of $\mathcal{X}^{l=2}\hat{\mathbf{V}}^{l=2}_{\textrm{diff}}$, where $\mathcal{X}^{l=2}$ is the generator of the irreducible representations of degree $l=2$. We illustrate this procedure in Figure 6 and explain the process in detail in the caption. This method proves effective because the output direction of this axis is inherently ambiguous, aligning perfectly with our requirement for the vectors to be unorderable.

The objective of our Self-Interaction module is to function exclusively based on the geometric features $\mathbf{V}^{0:{l_{\max}}}$, while concurrently mixing irreducible representations across various $l$ values. To accomplish this, we calculate the tensor product of the representation with itself; this operation is termed the "tensor square" and is symbolized by $\left(\mathbf{V}^{0:{l_{\max}}}\right)^{\otimes 2}$. As the channel dimensions expand, the computational complexity tied to the tensor square increases quadratically. To solve this computational load, we instead perform the square operation channel-wise, or by chunks of channels. Figures 7.c and 7.d illustrate these operations. After obtaining the squares from the chunked or individual channels, the segmented results are subsequently concatenated to generate an updated representation $\mathbf{V}^{0:{l_{\max}}}$, which is transformed through a learnable linear layer to the output dimensionality.

To incorporate non-linearities in our geometric representations $\mathbf{V}^{l=0:l_{\textrm{max}}}$, we employ a similar roto-translation equivariant gate mechanism as described in Equiformer [Liao & Smidt (2023)]. This mechanism is present at the last step of Self-Interaction and in the message preparation step of Spatial Convolution (Figure (2)). In both cases, we implement the activation by first isolating the scalar representations $\mathbf{V}^{l=0}$ and transforming them through a standard MultiLayerPerceptron (MLP). We use the SiLu activation function [Elfwing et al. (2017)] after each layer of the MLP. In the output vector, a scalar is produced for and multiplied into each channel of $\mathbf{V}^{l=0:l_{\textrm{max}}}$.

A Sequence Convolution kernel takes in $K$ coordinates $\mathbf{P}_{1:K}$ to produce a single coordinate $\mathbf{P}=\sum_{i=1}^{K}w_{i}\mathbf{P}_{i}$. We show that these weights need to be normalized in order for translation equivariance to be satisfied. Let $T$ denote a 3D translation vector, then translation equivariance requires:

Rotation equivariance is immediately satisfied since the sum of 3D vectors is a rotation equivariant operation. Let $R$ denote a rotation matrix. Then,

Given a single coarse anchor position and features representation $(\mathbf{P},\mathbf{V}^{l=0:l_{\textrm{max}}})$, we first map $\mathbf{V}^{l=0:l_{\textrm{max}}}$ into a new representation, reshaping it by chunking $K$ features. We then project $\mathbf{V}^{l=1}$ and produce $K$ relative position vectors $\Delta\mathbf{P}_{1:K}$, which are summed with the original position $\mathbf{P}$ to produce $K$ new coordinates.

This procedure generates windows of $K$ representations and positions. These windows may intersect in the decoded output. We resolve those superpositions by taking the average position and average representations within intersections.

Training deep models can be challenging due to vanishing or exploding gradients. We employ layer normalization [Ba et al. (2016)] and residual connections [He et al. (2015)] in order to tackle those challenges. We incorporate layer normalization and residual connections at the end of every Self-Interaction and every convolution. To keep roto-translation equivariance, we use the layer normalization described in Equiformer [Liao & Smidt (2023)], which rescales the $l$ signals independetly within a representation by using the root mean square value of the vectors. We found both residuals and layer norms to be critical in training deep models for large proteins.

We handle proteins of different sequence length by setting a maximum size for the model input. During training, proteins that are larger than the maximum size are cropped, and those that are smaller are padded. The boundary residue is given a special token that labels it as the end of the protein. For inference, we crop the tail of the output after the sequence position where this special token is predicted.

We analyse the time complexity of a forward pass through the Ophiuchus autoencoder. Let the list $(\mathbf{P},\mathbf{V}^{0:l_{\textrm{max}}})_{N}$ be an arbitrary latent state with $N$ positions and $N$ geometric representations of dimensionality $D$, where for each $d\in[0,D]$ there are geometric features of degree $l\in[0,l_{\textrm{max}}]$. Note that size of a geometric representation grows as $O(D\cdot l^{2}_{\textrm{max}})$ in memory.

• •

  The cost of an SO(3)-equivariant linear layer (Fig. 7.b) is $O(D^{2}\cdot l_{\textrm{max}}^{3})$.

• •

  The cost of a channel-wise tensor square operation (Fig. 7.c) is $O(D\cdot l^{4}_{\textrm{max}})$.

• •

  In Self-Interaction (Alg. 3), we use a tensor square and project it using a linear layer. The time complexity is given by $O\left(N\cdot(D^{2}\cdot l_{\textrm{max}}^{3}+D\cdot l_{\textrm{max}}^{4})\right)$ for the whole protein.

• •

  In Spatial Convolution (Alg. 5), a node aggregates geometric messages from $k$ of its neighbors. Resolving the $k$ nearest neighbors can be efficiently done in $O\left((N+k)\log N\right)$ through the k-d tree data structure [Bentley (1975)]. For each residue, its $k$ neighbors prepare messages through linear layers, at total cost $O(N\cdot k\cdot D^{2}\cdot l^{3}_{\textrm{max}})$.

• •

  In a Sequence Convolution (Alg. 4), a kernel stacks $K$ geometric representations of dimensionality $D$ and linearly maps them to a new feature of dimensionality $\rho\cdot D$, where $\rho$ is a rescaling factor, yielding $O((K\cdot D)\cdot(\rho\cdot D)\cdot l^{3}_{\textrm{max}})=O(K\cdot\rho\cdot D^{2}\cdot l^{3}_{\textrm{max}})$. With length $N$ and stride $S$, the total cost is $O\left(\frac{N}{S}\cdot K\cdot\rho\cdot D^{2}\cdot l^{3}_{\textrm{max}}\right)$.

• •

  The cost of an Ophiuchus Block is the sum of the terms above,

  $$O\left(D\cdot l_{\textrm{max}}^{3}\cdot N\cdot(\frac{K}{S}\cdot\rho D+kD)\right)$$

• •

  An Autoencoder (Alg. 7,8) that uses stride $S$ convolutions for coarsening uses $L=O\left(\log_{S}(N)\right)$ layers to reduce a protein of size $N$ into a single representation. At depth $i$, the dimensionality is given by $D_{i}=\rho^{i}D$ and the sequence length is given by $N_{i}=\frac{N}{S^{i}}$. The time complexity of our Autoencoder is given by geometric sum:

  $$O\left(\sum_{i}^{\log_{S}(N)}\left(l_{\textrm{max}}^{3}\rho^{i}D\frac{N}{S^{i}}(K\rho^{i+1}D+k\rho^{i}D)\right)\right)=O\left(Nl_{\textrm{max}}^{3}(K\rho+k)D^{2}\sum_{i}^{\log_{S}(N)}\left(\frac{\rho^{2}}{S}\right)^{i}\right)$$

  We are interested in the dependence on the length $N$ of a protein, therefore, we keep only relevant parameters. Summing the geometric series and using the identity $x^{\log_{b}(a)}=a^{\log_{b}(x)}$ we get:

  $$=O\left(N\frac{(\frac{\rho^{2}}{S})^{\log_{S}(N)+1}-1}{\frac{\rho^{2}}{S}-1}\right)=\begin{cases}O\left(N^{1+\log_{S}\rho^{2}}\right)&\text{for }\rho^{2}/S>1\\ O\left(N\log_{S}N\right)&\text{for }\rho^{2}/S=1\\ O\left(N\right)&\text{for }\rho^{2}/S<1\end{cases}$$

  In most of our experiments we operate in the $\rho^{2}/S<1$ regime.

The Huber Loss [Huber (1992)] behaves linearly for large inputs, and quadratically for small ones. It is defined as:

We found it to significantly improve training stability for large models compared to mean squared error. We use $\delta=0.5$ for all our experiments.

The vector map loss measures differences of internal vector maps $V(\mathbf{P})-V(\hat{\mathbf{P}})$ between predicted and ground positions, where $V(\mathbf{P})^{i,j}=(\mathbf{P}^{i}-\mathbf{P}^{j})\in\mathbb{R}^{\times 3}$. Our output algorithm for decoding atoms produces arbitrary symmetry breaks (Appendix A.2) for positions $\mathbf{P}_{v}^{*}$ and $\mathbf{P}_{u}^{*}$ of atoms that are not orderable. Because of that, a loss on the vector map is not directly applicable to the output of our model, since the order of the model output might differ from the order of the ground truth data. To solve that, we consider both possible orderings of permutable atoms, and choose the one that minimizes the loss. Solving for the optimal ordering is not feasible for the system as a whole, since the number of permutations to be considered scales exponentially with $N$. Instead, we first compute a vector map loss internal to each residue. We consider the alternative order of permutable atoms, and choose the candidate that minimizes this local loss. This ordering is used for the rest of our losses.

We consider bonds, interbond angles, dihedral angles and steric clashes when computing a loss for chemical validity. Let $\mathbf{P}\in\mathbb{R}^{N_{\textrm{Atom}}\times 3}$ be the list of atom positions in ground truth data. We denote $\hat{\mathbf{P}}\in\mathbb{R}^{N_{\textrm{Atom}}\times 3}$ as the list of atom positions predicted by the model. For each chemical interaction, we precompute indices of atoms that perform the interaction. For example, for bonds we precompute a list of pairs of atoms that are bonded according to the chemical profile of each residue. Our chemical losses then take form:

where $\mathcal{B}$ is a list of pair of indices of atoms that form a bond. We compare the distance between bonded atoms in prediction and ground truth data.

where $\mathcal{A}$ is a list of 3-tuples of indices of atoms that are connected through bonds. The tuple takes the form $(v,u,p)$ where $u$ is connected to $v$ and to $p$. Here, the function $\alpha(\cdot,\cdot,\cdot)$ measures the angle in radians between positions of atoms that are connected through bonds.

where $\mathcal{D}$ is a list of 4-tuples of indices of atoms that are connected by bonds, that is, $(v,u,p,q)$ where $(v,u)$, $(u,p)$ and $(p,q)$ are connected by bonds. Here, the function $\tau(\cdot,\cdot,\cdot,\cdot)$ measures the dihedral angle in radians.

where $H$ is a smooth and differentiable Heaviside-like step function, $\mathcal{C}$ is a list of pair of indices of atoms that are not bonded, and ($r_{v},r_{u}$) are the van der Waals radii of atoms $v$, $u$.

When training autoencoder models for latent diffusion, we regularize the learned latent space so that representations are amenable to the relevant range scales of the source distribution $\mathcal{N}(0,1)$. Let $\mathbf{V}^{l}_{i}$ denote the $i$-th channel of a vector representation $\mathbf{V}^{l}$. We regularize the autoencoder latent space by optimizing radial and angular components of our vectors:

The first term penalizes vector magnitudes larger than one, and the second term induces vectors to spread angularly. We find these regularizations to significantly help training of the denoising diffusion model.

We weight the different losses in our pipeline. For the standard training of the autoencoder, we use weights:

For fine-tuning the model, we increase the weight of chemical losses significantly:

We find that high weight values for chemical losses at early training may hurt the model convergence, in particular for models that operate on large lengths.

To visualize the learned latent space of Ophiuchus, we forward 50k samples from the training set through a large 485-length model (Figure 10). We collect the scalar component of bottleneck representations, and use t-SNE to produce 2D points for coordinates. We similarly produce 3D points and use those for coloring. The result is visualized in Figure 10. Visual inspection of neighboring points reveals unsupervised clustering of similar folds and sequences.

To compute the scTM scores, we recover 8 sequences using ProteinMPNN for 500 sampled backbones from the tested diffusion models. We used a sampling temperature of 0.1 for ProteinMPNN. Unlike the original work, where the sequences where folded using AlphaFold2, we use OmegaFold [Wu et al. (2022b)] similar to [Lin & AlQuraishi (2023)]. The predicted structures are aligned to the original sampled backbones and TM-Score and RMSD is calculated for each alignment. To calculate the diversity measurement, we hierarchically clustered samples using MaxCluster. Diversity is defined as the number of clusters divided by the total number of samples, as described in [Yim et al. (2023)].

In Figure (13) we show generated samples from our miniprotein model, and compare the marginal distribution of our predictions and ground truth data. In Figure (14.b) we show the distribution of TM-scores for joint sampling of sequence and all-atom structure by the diffusion model. We produce marginal distributions of generated samples Fig.(14.e) and find them to successfully approximate the densities of ground truth data. To test the robustness of joint sampling of structure and sequence, we compute self-consistency TM-Scores[Trippe et al. (2023)]. 54% of our sampled backbones have scTM scores > 0.5 compared to 77.6% of samples from RFDiffusion. We also sample 100 proteins between 50-65 amino acids in 0.21s compared to 11s taken by RFDiffusion on a RTX A6000.

We include metrics on designability in Figure 12.