AtomSteer improves atomistic model outputs with structural constraints
AtomSteer: Test-Time Constraint Steering for Atomistic Generative Models
Machine LearningArtificial Intelligence
Summary
Diffusion and flow models can create 3D structures like crystals and proteins, but they often miss important details like symmetry and correct chemistry. The authors introduce AtomSteer, a way to fine-tune these models while they generate structures, fixing errors without breaking the rest of the design. By adjusting intermediate steps, AtomSteer helps produce more accurate crystal shapes, molecular details, and protein-ligand fits. This technique works with different existing models without needing to retrain them.
What this means in practice
- •For materials scientists: Improve crystal structure generation accuracy by enforcing symmetry and stereochemistry during model inference without retraining.
- •For medicinal chemists: Enhance protein-ligand docking predictions by reducing stereochemical errors while maintaining valid complex structures.
Authors
Pengfei Cai, Katharina S. C. Jäger, Utkarsh Utkarsh, Nofit Segal, Bowen Deng, Akshay Subramanian, Rafael Gomez-Bombarelli
Abstract
Diffusion and flow matching models offer scalable approaches to generating crystalline materials, molecular conformers, and protein-ligand complexes. However, satisfying crystallographic symmetry, structural geometry, and accurate stereochemistry remains challenging. Satisfying these requirements calls for corrections that preserve the geometry already captured by the pretrained model. In this work, we propose AtomSteer, a test-time inference framework for steering pretrained atomistic generative models with structural constraints. During sampling, AtomSteer corrects intermediate structure predictions with bounded local projections, graph-rigid stereochemical transformations, and projections over periodic symmetries, accepting updates that improve constraint satisfaction while limiting changes to the rest of the structure. We demonstrate AtomSteer across three pretrained backbones, FlowMM, ET-Flow, and Boltz-2, spanning de novo crystal generation and crystal structure prediction, molecular conformer generation, and protein-ligand complex structure prediction. AtomSteer improves space-group distribution coverage in de novo crystal generation and improves target space-group recovery in crystal structure prediction. AtomSteer improves molecular validity and E/Z stereochemical accuracy while retaining conformer ensemble fidelity. For protein-ligand complexes, AtomSteer reduces ligand E/Z errors relative to experimental reference structures while achieving competitive docking success and PoseBusters validity. Our method provides a unified framework for adding chemical and crystallographic control to existing atomistic generative models without retraining or architectural changes.