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

Q1: We added a full analysis of the time complexity of our model in Appendix A.10

Q2: Thanks, we fixed the typos and the format of references.

Q3: In the context of protein modeling, all-atom structures are often interchangeable with "all-heavy atom". Thanks for the note, we made the usage of the term consistent throughout the paper.

Algorithms 1 and 2 are used to encode and decoded the protein into its initial representation which uses degree up to 2. Algorithm 3, however, may be used to produce features of larger degree $l$, through its internal use of tensor square. Because $l_{\textrm{max}}$ is a hyperparameter, the model will build up the geometric representation up to this degree.

Q4: We put the figures together and made a more concise analysis of the two studied interpolations.

Q5: We illustrate the permutation invariance of the side-chain feature parameterization we use. We added a more detailed discussion in Appendix A.1, and incorporated a visual demonstration in Figure 4.

Q6: When choosing our input sequence lengths, we restrict our choices of length to those that will result in perfect alignment of the convolution kernel windows when downsampling down a single representation. For the model settings of our ablation studies, with stride S = 3 and kernel size = 5, the sequence lengths that satisfy this constraint are [1, 5, 17, 53, 160, 485].

We chose ten epochs as we found it to be enough to gather data about training curves and convergence, while still catering to our limited computational resources.

We focused on mini-proteins because our all-atom model performs well at the capacity of our training. For this updated manuscript, we included additional benchmarks on large backbone models that we compare directly to existing architectures.

Q7: To the best of our knowledge, we provide the first demonstration of latent interpolation maintaining consistent chemical validity throughout the trajectory without explicitly enforcing any physical behavior in the latent space. Other methods exist which rely on physics based simulation data for interpolation.

We also added structure recovery comparison in Section 4.1

Thank you for the feedback. Please note the comment above about comparison and baselines.

Q1: We added motivation for the model design at the start of Section 3, as well as at the start description of each component in Sections 3.2-3.4

Q2: We added much more detail on the losses, in Appendix B. In particular, we further elaborate on the chemical losses in Appendix B.2 and include an expression for the final objective function in Appendix B.4

Q3: In Section 4.4, we added comparisons against existing generative models, and added more metrics on diversity and consistency of samples.

Q4: Thank you, we fixed the mistakes.

Thank you for the feedback. Please note the comment above about comparison and baselines.

Q1: In our ablation studies, it is true that a larger downsampling factor carries more parameters and less recoverability, as is expected in producing larger-sized embeddings and decoding bottlenecked representations. Still, the purpose of our ablation was to estimate the cost of training a model to produce those coarsened embeddings at different resolutions, and quantify the quality of their reconstructions for different embedding sizes.

The main motivation for our novel protein autoencoding architecture is to produce single geometric embedding that fully captures the structural features of a protein across multiple resolutions. The inspiration for capturing coarse representations stems from the thermodynamical and evolutionary origin of proteins, and more generally from the need to model building blocks present in structural biology. In our work, we showed that this latent representation of a protein is structured and compact, and can be manipulated while cohesively capturing the structural and chemical features of proteins.

We demonstrate the application of these bottlenecked representations using latent interpolation. To the best of our knowledge, we are the first to explore this form of interpolation for proteins through SO(3)-equivariant features, directly in three-dimensions. The latent interpolation maintains chemical validity throughout its trajectory, as opposed to when the data representation is not downsampled/coarsened and is directly interpolated in the original domain.

The computational gain of our approach starts when the autoencoder model is trained and may be used for downstream tasks. For example, our large backbone model reduces proteins of length 485 into a single geometric tensor representation of 160 scalars and 160 3D vectors, reducing the target size for denoising diffusion. This directly reflects in sampling times orders of magnitude lower than existing models, while performing comparably.

As a final example of the usefulness of our structured latent space, we note that the coarsened embeddings capture sequence labels, all-atom positions, and backbone structure, and our generative model is able to produce all of those simultaneously while diffusing on single representation of features. This contrasts to existing models, which often are required to perform separate diffusion processes on multiple data domains.

Q2: Please refer to the updated manuscript. We added much more detail on the encoding of $\mathbf V^{l=2}$ in Section 3.1, and added diagrams (Figure 4) and further discussion in Appendix A.1.

Q3: In Section 3.1, and Appendices A.1 and A.2, we added details and discussion about the permutation symmetries. In Appendix A.1, Figure 5, we also added a visual demonstration that shows how our proposed feature representations handle the problems that might arise due to the symmetries of residues.

Q4: We added Appendix A.5 where we derive the conditions for coarsening coordinates while satisfying rotation and translation equivariance.

Q5: Thanks, we fixed the minor issues.

Q6: We added a brief description of the HuberLoss in Appendix B.1.