AtomSteer: Test-Time Constraint Steering for Atomistic Generative Models

Diffusion and flow matching models offer scalable approaches to generating molecular conformers, protein-ligand complexes, and crystalline materials. However, satisfying structural geometry, accurate stereochemistry, and crystallographic symmetry remains challenging. Satisfying these requirements calls for corrections that also 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, ET-Flow, Boltz-2, and FlowMM, spanning molecular conformer generation, protein-ligand complex structure prediction, and de novo crystal generation and 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. AtomSteer improves space-group distribution coverage in de novo crystal generation and improves target space-group recovery in crystal structure prediction. Our method provides a unified framework for adding chemical and crystallographic control to existing atomistic generative models without retraining or architectural changes.