Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Protein Folding with Neural Ordinary Differential Equations

Arielle Sanford, Shuo Sun, Christian B. Mendl

TL;DR

This work tackles the computational burden of AlphaFold 2's Evoformer by replacing the 48 discrete blocks with a Neural ODE that evolves the MSA and pair representations over a continuous depth $t\in[0,1]$, enabling constant memory via the adjoint method. The proposed Evoformer ODE defines a vector field $f_\theta(m,z,t)$ that reconstitutes core Evoformer operations with gating functions $\sigma_m(t)$ and $\sigma_z(t)$, and uses an RK4 solver to integrate the dynamics. Training relies on ground-truth trajectories derived from OpenFold, employing a two-phase scheme to learn the full transformation from input to the final Evoformer block output, achieving substantial speedups (≈$4.85$ s per protein on a single GPU) while preserving broad structural features such as $\alpha$-helices. The results demonstrate a promising, lightweight, and interpretable alternative for biomolecular modeling, with future work pointing toward larger MSAs, higher-capacity representations, adaptive solvers, and longer sequences to close the gap with the traditional block-based model.

Abstract

Recent advances in protein structure prediction, such as AlphaFold, have demonstrated the power of deep neural architectures like the Evoformer for capturing complex spatial and evolutionary constraints on protein conformation. However, the depth of the Evoformer, comprising 48 stacked blocks, introduces high computational costs and rigid layerwise discretization. Inspired by Neural Ordinary Differential Equations (Neural ODEs), we propose a continuous-depth formulation of the Evoformer, replacing its 48 discrete blocks with a Neural ODE parameterization that preserves its core attention-based operations. This continuous-time Evoformer achieves constant memory cost (in depth) via the adjoint method, while allowing a principled trade-off between runtime and accuracy through adaptive ODE solvers. Benchmarking on protein structure prediction tasks, we find that the Neural ODE-based Evoformer produces structurally plausible predictions and reliably captures certain secondary structure elements, such as alpha-helices, though it does not fully replicate the accuracy of the original architecture. However, our model achieves this performance using dramatically fewer resources, just 17.5 hours of training on a single GPU, highlighting the promise of continuous-depth models as a lightweight and interpretable alternative for biomolecular modeling. This work opens new directions for efficient and adaptive protein structure prediction frameworks.

Protein Folding with Neural Ordinary Differential Equations

TL;DR

This work tackles the computational burden of AlphaFold 2's Evoformer by replacing the 48 discrete blocks with a Neural ODE that evolves the MSA and pair representations over a continuous depth , enabling constant memory via the adjoint method. The proposed Evoformer ODE defines a vector field that reconstitutes core Evoformer operations with gating functions and , and uses an RK4 solver to integrate the dynamics. Training relies on ground-truth trajectories derived from OpenFold, employing a two-phase scheme to learn the full transformation from input to the final Evoformer block output, achieving substantial speedups (≈ s per protein on a single GPU) while preserving broad structural features such as -helices. The results demonstrate a promising, lightweight, and interpretable alternative for biomolecular modeling, with future work pointing toward larger MSAs, higher-capacity representations, adaptive solvers, and longer sequences to close the gap with the traditional block-based model.

Abstract

Recent advances in protein structure prediction, such as AlphaFold, have demonstrated the power of deep neural architectures like the Evoformer for capturing complex spatial and evolutionary constraints on protein conformation. However, the depth of the Evoformer, comprising 48 stacked blocks, introduces high computational costs and rigid layerwise discretization. Inspired by Neural Ordinary Differential Equations (Neural ODEs), we propose a continuous-depth formulation of the Evoformer, replacing its 48 discrete blocks with a Neural ODE parameterization that preserves its core attention-based operations. This continuous-time Evoformer achieves constant memory cost (in depth) via the adjoint method, while allowing a principled trade-off between runtime and accuracy through adaptive ODE solvers. Benchmarking on protein structure prediction tasks, we find that the Neural ODE-based Evoformer produces structurally plausible predictions and reliably captures certain secondary structure elements, such as alpha-helices, though it does not fully replicate the accuracy of the original architecture. However, our model achieves this performance using dramatically fewer resources, just 17.5 hours of training on a single GPU, highlighting the promise of continuous-depth models as a lightweight and interpretable alternative for biomolecular modeling. This work opens new directions for efficient and adaptive protein structure prediction frameworks.
Paper Structure (11 sections, 5 equations, 3 figures, 8 algorithms)

This paper contains 11 sections, 5 equations, 3 figures, 8 algorithms.

Table of Contents

  1. Introduction
  2. Background
  3. AlphaFold 2 Architecture
  4. NeuralODE
  5. Evoformer ODE Architecture
  6. Training
  7. Results
  8. Conclusions and Next Steps
  9. Details of the Evoformer ODE Architecture
  10. Notation and Shapes
  11. Learned Parameters

Figures (3)

  • Figure 1: ODE Evoformer block. Arrows show the information flow. The shape of each array is shown in parentheses, where $s$ is the number of MSA sequences, $r$ is the number of residues, and $c$ is the channel dimension.
  • Figure 2: Inference runtime as a function of protein length. Linear and quadratic fits are shown for the NeuralODE and OpenFold Evoformer, respectively.
  • Figure 3: Predicted protein structures visualized using PyMOL pymol. Columns correspond to different proteins (1fv5, 4epq, 1ujs). Rows show predictions from (top) the NeuralODE Evoformer, (middle) OpenFold ground truth, and (bottom) a model using only 24 Evoformer blocks. Color denotes per-residue confidence (pLDDT; red = high, blue = low).