Protein Folding with Neural Ordinary Differential Equations

2025-10-17

AI Summary

The Evoformer architecture in AlphaFold suffers from high computational overhead and inflexible depth due to its discrete 48-layer stack. Method: This work introduces Neural Ordinary Differential Equations (Neural ODEs) into protein structure prediction—first of its kind—proposing a continuous-depth Evoformer. It preserves the original attention mechanisms and geometric modeling while formulating hidden-state evolution over “continuous depth” via an ODE; adaptive step-size solvers balance accuracy and speed, and the adjoint method enables memory-constant backpropagation. Contribution/Results: Trained on a single GPU in just 17.5 hours, the model accurately reconstructs secondary structures such as α-helices. Experiments demonstrate substantial reductions in both computational cost and memory footprint, validating the efficacy and lightweight potential of continuous-depth modeling for protein folding.

Technology Category

• Machine Learning: Deep Neural Architectures and Foundation Models
• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Mixed Discrete/Continuous Search

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Economics, Online Markets and Human Computation: Architectures and workflows that use LLMs for crowd work

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.

Problem

• Reducing computational costs in protein structure prediction models
• Replacing discrete blocks with continuous Neural ODE formulation
• Achieving efficient training with constant memory cost

Innovation

• Continuous-depth Evoformer replaces discrete blocks with Neural ODEs
• Constant memory cost achieved via adjoint method
• Adaptive ODE solvers enable runtime-accuracy trade-off