IDFlow

IDFlow is a generative model for 3D molecular structure that reformulates flow matching from an energy-based perspective. It was developed at KTH Royal Institute of Technology with Science for Life Laboratory, Stockholm University, The Alan Turing Institute, and the Max Planck Institute for Polymer Research, and was published at ICML 2025.

Flow matching models learn a velocity field that transports samples from a simple prior to the data distribution, and in structure generation they are typically trained with a data-prediction parameterization: the network sees a noisy interpolant and predicts the clean structure. IDFlow takes the reconstruction error of that prediction to be an energy function, then shapes the energy landscape locally using contrastive samples drawn from the model itself. Minimizing this energy turns out to mean training a network that maps its own output back to itself — an approximately idempotent function.

That idempotency has a direct practical consequence. The same network is both a sampler and a refiner: at inference it predicts a structure, feeds the prediction back through itself, and iterates toward the low-energy region of its own landscape. This is the generative analogue of the recycling step in structure prediction and of predictor-corrector sampling in diffusion models, but it emerges from the training objective rather than being bolted on, and it adds a single hyperparameter to an otherwise standard flow matching setup.

#Key Features

• Energy as reconstruction error: The L2 gap between a prediction and its refinement defines the energy, so no separate energy network, biophysical force field, or confidence model is required.
• Contrastive energy shaping: Negative samples drawn from the flow model itself widen the energy margin around the data manifold, which improves mode coverage instead of collapsing the model onto a narrow set of structures.
• Predictor-refiner sampling: Because the trained map is approximately idempotent, iterating it at inference monotonically improves samples, letting the model reach higher quality at a fixed number of function evaluations.
• Architecture-agnostic objective: The formulation is orthogonal to the backbone. Docking uses an SE(3)-equivariant refinement tensor field network over ligand atom positions; protein backbone generation uses an invariant point attention structure transformer over residue frames.
• Task coverage from one formulation: Single-ligand docking, multi-ligand docking, and unconditional protein backbone generation are all trained with the same objective and sampler.

#Technical Details

IDFlow extends the FrameFlow codebase for protein backbones and follows HarmonicFlow's preprocessing for docking. Docking models are trained on PDBBind v2020 (roughly 19,000 complexes, under both time and 30% sequence-similarity splits) and on Binding MOAD (41,000 complexes) for multi-ligand docking. Backbone generation is trained on a 3,928-structure SCOPe subset of 60-128 residues and on a 19,327-protein PDB set, using 8 A100 GPUs for about 22 hours and 3 days respectively, with a 6-layer network, 256-dimensional node embeddings, 128-dimensional edge embeddings, and 8 IPA heads.

On pocket-level docking at matched inference budgets, IDFlow improves the fraction of predictions under 2 Å RMSD by 5.5% on the sequence-similarity split and 6.4% on the time split, and improves multi-ligand results by 8.1% at the 5 Å threshold. On PDB-trained backbone generation, designability — measured by inverse folding each sample with ProteinMPNN, refolding with ESMFold, and requiring self-consistency RMSD under 2 Å — reaches 0.871 at 100 function evaluations and 0.927 at 200, against 0.824 for FrameFlow at 200 and 0.89 for RFdiffusion, while generating a 100-residue backbone in 3.3 seconds at the lower budget. Diversity and novelty stay at baseline levels, and unlike some higher-designability competitors the samples do not collapse toward all-helical folds.

#Applications

The docking models place ligands in a known binding pocket, the setting that matters for hit triage and pose prediction in structure-based drug discovery, including complexes with multiple bound ligands. The backbone generator produces unconditional protein scaffolds for de novo design pipelines, where designability, structural diversity, and sampling cost jointly determine how many candidates can be pushed to inverse folding and synthesis. More broadly, the objective is a drop-in modification for teams already training flow matching models on molecular structure.

#Impact

IDFlow's contribution is conceptual leverage rather than scale: it connects energy-based modeling, idempotent generative networks, and the refinement heuristics that structure predictors adopted empirically, and shows the connection pays off on two structurally different generation problems under equal compute. The published checkpoint covers unconditional protein backbone generation and is distributed through a cloud drive link in the repository rather than a model hub, alongside a runnable inference script; the repository carries no licence file, so reuse terms are undeclared. Reproducing the docking results requires training from the provided code.