EquiJump: Protein Dynamics Simulation via SO(3)-Equivariant Stochastic Interpolants

TL;DR

EquiJump is a novel SO(3)-equivariant deep learning model that efficiently simulates protein conformational dynamics, outperforming existing methods by directly bridging all-atom simulation steps with high accuracy.

Contribution

It introduces EquiJump, a transferable stochastic interpolant framework that unifies and improves protein dynamics simulation across different time steps.

Findings

Achieves state-of-the-art results on fast folding protein trajectories.

Unifies diverse sampling methods for protein dynamics.

Demonstrates transferability across multiple proteins.

Abstract

Mapping the conformational dynamics of proteins is crucial for elucidating their functional mechanisms. While Molecular Dynamics (MD) simulation enables detailed time evolution of protein motion, its computational toll hinders its use in practice. To address this challenge, multiple deep learning models for reproducing and accelerating MD have been proposed drawing on transport-based generative methods. However, existing work focuses on generation through transport of samples from prior distributions, that can often be distant from the data manifold. The recently proposed framework of stochastic interpolants, instead, enables transport between arbitrary distribution endpoints. Building upon this work, we introduce EquiJump, a transferable SO(3)-equivariant model that bridges all-atom protein dynamics simulation time steps directly. Our approach unifies diverse sampling methods and is benchmarked against existing models on trajectory data of fast folding proteins. EquiJump achieves state-of-the-art results on dynamics simulation with a transferable model on all of the fast folding proteins.