Energy-Based Flow Matching for Generating 3D Molecular Structure

TL;DR

This paper introduces an energy-based flow matching approach for 3D molecular structure generation, improving training and inference by mapping random configurations to data structures, with strong empirical results in protein applications.

Contribution

It presents a novel energy-based flow matching framework that enhances molecular structure generation by directly learning iterative mappings, with theoretical justification and practical effectiveness.

Findings

Outperforms recent flow matching and diffusion models in protein tasks

Demonstrates effectiveness in protein docking and backbone generation

Provides a simple, theoretically justified approach with empirical success

Abstract

Molecular structure generation is a fundamental problem that involves determining the 3D positions of molecules' constituents. It has crucial biological applications, such as molecular docking, protein folding, and molecular design. Recent advances in generative modeling, such as diffusion models and flow matching, have made great progress on these tasks by modeling molecular conformations as a distribution. In this work, we focus on flow matching and adopt an energy-based perspective to improve training and inference of structure generation models. Our view results in a mapping function, represented by a deep network, that is directly learned to \textit{iteratively} map random configurations, i.e. samples from the source distribution, to target structures, i.e. points in the data manifold. This yields a conceptually simple and empirically effective flow matching setup that is theoretically justified and has interesting connections to fundamental properties such as idempotency and stability, as well as the empirically useful techniques such as structure refinement in AlphaFold. Experiments on protein docking as well as protein backbone generation consistently demonstrate the method's effectiveness, where it outperforms recent baselines of task-associated flow matching and diffusion models, using a similar computational budget.