cs.LGOct 28, 2026

AtomSteer: Test-Time Constraint Steering for Atomistic Generative Models

Authors: Pengfei Cai, Katharina S. C. Jäger, Utkarsh Utkarsh, Nofit Segal, Bowen Deng, Akshay Subramanian, Rafael Gomez-Bombarelli

Organizations: Massachusetts Institute of Technology

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.

Figures & tables

Appendix figures & tables19 assets

Supplementary material from the paper’s appendix.

Appendix

Explore similar work

CardsList
  1. uFlowCSP: Crystal Structure Prediction using Mean flow generative models

    Sep 9, 2026Sourin Dey, Dipannoy Das Gupta, Lai Wei +2Crystal Structure PredictionOne-Step Generative Modeling

  2. Multimodal Crystal Flow: Any-to-Any Modality Generation for Unified Crystal Modeling

    Feb 23, 2026Kiyoung Seong, Sungsoo Ahn, Sehui Han +1Crystal Structure GenerationUnified Multimodal Models

  3. Crystalite: A Lightweight Transformer for Efficient Crystal Modeling

    Apr 2, 2026Tin Hadži Veljković, Joshua Rosenthal, Ivor Lončarić +1Diffusion TransformerTransformer