Aligning Protein Conformation Ensemble Generation with Physical Feedback

Jiarui Lu, Xiaoyin Chen, Stephen Zhewen Lu, Aurelie Lozano, Vijil Chenthamarakshan, Payel Das, Jian Tang

TL;DR: incorporating physical energy feedback into diffusion model for better protein conformation ensemble generation

Abstract: Protein dynamics play a crucial role in protein biological functions and properties, and their traditional study typically relies on time-consuming molecular dynamics (MD) simulations conducted in silico. Recent advances in generative modeling, particularly denoising diffusion models, have enabled efficient accurate protein structure prediction and conformation sampling by learning distributions over crystallographic structures. However, effectively integrating physical supervision into these data-driven approaches remains challenging, as standard energy-based objectives often lead to intractable optimization. In this paper, we introduce Energy-based Alignment (EBA), a method that aligns generative models with feedback from physical models, efficiently calibrating them to appropriately balance conformational states based on their energy differences. Experimental results on the MD ensemble benchmark demonstrate that EBA achieves state-of-the-art performance in generating high-quality protein ensembles. By improving the physical plausibility of generated structures, our approach enhances model predictions and holds promise for applications in structural biology and drug discovery.

Lay Summary:

Proteins are dynamic molecules that constantly change their shape—a process fundamental to their biological functions. Understanding these shape-shifting behaviours is essential for biology and drug discovery. While scientists traditionally rely on physics-based simulations to study them, such simulations are extremely time-consuming, particularly for large proteins. AI models can predict protein structures more quickly, but they often lack a crucial component: the physical principles that govern which shapes are naturally possible. This limitation reduces these models' real-world scientific value.

Our research introduces a new approach called Energy-based Alignment (EBA). This method enables AI models to generate protein structures that are both realistic and physically sound by incorporating feedback from physics-based energy functions during the training process, leading to more accurate predictions.

When tested on benchmark dataset, our approach generates more accurate protein structure ensembles. By bridging the gap between AI and physics-based methods, we provide a faster and more reliable tool for studying protein behavior and supporting drug development.

Application-Driven Machine Learning: This submission is on Application-Driven Machine Learning.

Link To Code: https://github.com/lujiarui/eba

Primary Area: Applications->Health / Medicine

Keywords: Protein, generative models, molecular dynamics, conformation generation, alignments

Submission Number: 14909