AtomMOF: All-Atom Flow Matching for MOF-Adsorbate Structure Prediction

TL;DR

AtomMOF presents an all-atom flow-based approach to MOF-adsorbate structure prediction using a Diffusion Transformer and variational flow matching to map building-block graphs directly to equilibrium 3D structures. It introduces Feynman-Kac steering guided by a machine-learned interatomic potential to improve geometric validity and sampling stability, and establishes scaling laws for porous-crystal generation. On the BW dataset, AtomMOF yields $MR$ gains of $35.00\%$ and reduces $RMSD$ by $32.64\%$, while on ODAC25 it achieves substantially faster recovery of adsorption configurations than GCMC and can identify lower-energy adsorption states than the reference. Together, these advances enable efficient, accurate exploration of MOF-adsorbate systems with practical implications for gas separations and carbon capture.

Abstract

Deep generative models have shown promise for modeling metal-organic frameworks (MOFs), but existing approaches (1) rely on coarse-grained representations that assume fixed bond lengths and angles, and (2) neglect the MOF-adsorbate interactions, which are critical for downstream applications. We introduce AtomMOF, a scalable flow-based model built on an all-atom Diffusion Transformer that maps 2D molecular graphs of building blocks and adsorbates directly to equilibrium 3D structures without imposing structural constraints. We further present scaling laws for porous crystal generation, indicating predictable performance gains with increased model capacity, and introduce Feynman-Kac steering guided by machine-learned interatomic potentials to improve geometric validity and sampling stability. On the (MOF-only) BW dataset, AtomMOF increases the match rate by 35.00% and reduces RMSD by 32.64%. On the ODAC25 dataset (MOF-adsorbate), AtomMOF is substantially more sample-efficient than grand canonical Monte Carlo in recovering adsorption configurations and can identify candidates with lower adsorption energies than the reference dataset. Code is available at https://github.com/nayoung10/AtomMOF.

Figures (11)

• Figure 1: MOF-adsorbate structure prediction. Given a set of building blocks $\mathcal{B}$ (comprising linkers, metal nodes, and adsorbate identity), the goal is to predict the equilibrium structure $\mathcal{S} = (\mathbf{A}, \mathbf{X}, \bm{\ell})$, defined by atomic numbers $\mathbf{A}$, Cartesian coordinates $\mathbf{X}$, and lattice parameters $\bm{\ell}$.
• Figure 2: Qualitative comparison of predicted structures. Each structure is shown as a $2\times2\times2$ supercell for visual clarity. Compared to the baselines, AtomMOF-M more closely matches the ground truth and better avoids steric clashes near the metal nodes, illustrating the advantage of all-atom modeling.
• Figure 3: Scaling up AtomMOF improves performance. Both match rate (left) and RMSD (right) exhibit a log-linear relationship with model size, indicating predictable performance gains with increasing number of parameters.
• Figure 4: Effect of model scaling on predicted structures. Structural quality improves with model size: larger models better recover the reference topology and reduce steric clashes.
• Figure 5: Inference efficiency under scaling. Match rate (%, stol=0.3) versus (left) the number of ODE sampling steps and (right) wall-clock sampling time. AtomMOF-L achieves higher match rates than AtomMOF-M at the same step or time budget, and both surpass MOFFlow-2 across budgets. In particular, AtomMOF exceeds MOFFlow-2's peak match rate with only a few steps and largely saturates within $\sim$10 steps.