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

• Accurate MOF-adsorbate prediction matters for applications such as carbon capture, where host–guest interactions are critical.Source: §1 · p. 1
• Simulation-based workflows such as GCMC and DFT are expensive for large MOFs and require known host structures, limiting screening of hypothetical frameworks.Source: §1 · p. 1
• Existing generative models use coarse-grained or fixed local geometries and generally generate host structures without adsorbates, missing chemically driven structural variation.Source: §1 · p. 1
• AtomMOF maps building-block identities and adsorbate types directly to equilibrium three-dimensional atomic coordinates and lattice parameters in a fully unconstrained space.Source: §1 · pp. 1–2
• AtomMOF introduces scaling laws for porous-crystal generation and Feynman-Kac steering with MLIP to improve performance, geometric validity, and stability.Source: §1 · p. 2
• On ODAC25, AtomMOF generates diverse adsorption configurations faster than GCMC and finds configurations with adsorption energies below the reference dataset’s lowest energies.Source: §1 · p. 2

2.1. Machine Learning for Crystal Structure Prediction

• 2.1. Machine Learning for Crystal Structure Prediction: Crystal-generation methods include equivariant diffusion, flow matching, Bayesian flow networks, and stochastic interpolants, often addressing periodicity and symmetry.Source: §2.1 · p. 2
• 2.1. Machine Learning for Crystal Structure Prediction: MOF methods manage modularity and scale with rigid-body or fixed-geometry assumptions, whereas AtomMOF enables fully flexible structure prediction.Source: §2.1 · p. 2

2.2. Generative Model for Molecular Complexes

• 2.2. Generative Model for Molecular Complexes: Protein-complex models predict joint structures involving proteins, nucleic acids, and small molecules.Source: §2.2 · p. 2
• 2.2. Generative Model for Molecular Complexes: Generative models for heterogeneous catalysis explicitly capture catalyst–adsorbate interactions, a direction AtomMOF extends to complete MOF–adsorbate systems.Source: §2.2 · p. 2

3.1. Problem Formulation

• AtomMOF represents MOF-adsorbate systems as molecular graphs and predicts atomic coordinates together with lattice parameters for periodic structures.A unit cell is represented by atom identities, Cartesian coordinates, and lattice parameters, while framework and adsorbate components are unions of molecular graphs.Source: §3.1 · pp. 2–3

3.2. Variational Flow Matching for MOFs

• Variational flow matching transforms prior samples into the target structure distribution by predicting clean coordinates and lattice parameters from noisy states conditioned on building blocks.The method jointly models coordinates and lattice parameters along a conditional probability path from the prior at t = 0 to data at t = 1, using an L1 reconstruction objective.Source: §3.2 · p. 3

3.3. Model Architecture

• A Diffusion Transformer encodes atom and pairwise building-block features, models interactions in a deep trunk, and decodes coordinates and lattice parameters through separate heads.The encoder uses atom types, charges, local coordinates, geometry, bonds, and block adjacency; the decoder outputs per-atom coordinates and globally pooled lattice parameters.Source: §3.3 · pp. 3–4

3.4. Feynman-Kac Steering with MLIP

• Feynman-Kac steering uses MLIP energies to bias sampling toward lower-energy structures through an SDE, particle resampling, and an exponentially tilted distribution.The steering strength λ controls the energy-based tilt, while the learned velocity field is converted to an SDE with diffusion and score terms to enable particle resampling.Source: §3.4 · p. 4

4.2. Match Rate & RMSD

• 35.00% higher match rate and 32.68% lower RMSD than MOFFlow-2 make ATOMMOF-L the strongest structure-prediction variant.Both ATOMMOF-M and ATOMMOF-L outperform all baselines across every reported metric, with larger gains under stricter site tolerance.Source: §4.2 · p. 4, Table 1 · p. 6

4.3. Property Evaluation

• ATOMMOF preserves geometric and pore properties more accurately than MOFFlow-2, with the large variant performing best overall.The evaluation covers surface areas, cavity and pore diameters, accessible volume, unit-cell volume, density, and void fraction.Source: §4.3 · p. 5, Table 2 · p. 6

4.4. Scaling Law

• Match rate and RMSD improve log-linearly as model size increases from 38.4M to 491M parameters.The trend is supported by Pearson correlations of ±0.93 and Spearman correlations of ±1.00 using 10 integration steps and stol=0.3.Source: §4.4 · p. 5, Figure 3 · p. 6
• ATOMMOF-L exceeds ATOMMOF-M at every sampling-step count, while both surpass MOFFlow-2’s peak match rate with only 3 steps.By about 10 steps, match rates largely saturate at approximately 11.7% for ATOMMOF-L and 9.8% for ATOMMOF-M; ATOMMOF-L becomes superior beyond 0.0825 s.Source: §4.4 · p. 6, Figure 5 · p. 7

4.5. Feynman-Kac Steering with MLIP

• Feynman-Kac steering reduces formation-energy error from 0.127 to 0.020 eV/Å and raises MOF validity from 34.16% to 37.87%.Structural validity also increases from 81.70% to 91.70%, indicating improved physical stability and geometric validity.Source: §4.5 · p. 6
• Steering quality improves with more particles and stronger steering, performs best with the immediate potential, and benefits from later application in denoising.The authors attribute the later-timing benefit to MLIP reliability being lower during highly noisy early generation stages.Source: §4.5 · p. 6

5. MOF-Adsorbate Structure Prediction

• The evaluation uses filtered ODAC25 data and compares ATOMMOF with GCMC and random sampling under fixed-framework conditions.The validation split is used as a test set because the official ODAC25 test set is unavailable.Source: §5 · p. 7, §5.1 · p. 7

5.2. Adsorption Energy Success Rate

• Feynman-Kac steering improves MOFChecker validity, structural validity, and formation-energy error relative to an unsteered baseline.The steering study varies particle count, steering strength, resampling interval, start time, and potential mode; later steering generally improves structural quality.Source: Table 4 · p. 7, Figure 6 · p. 8

5.1. Coverage of Adsorption Configurations

• 23.11% coverage in under 1 s versus 21.06% for GCMC in nearly 104 s, while random sampling achieves lower coverage than ATOMMOF.The comparison uses the same time budget framework and evaluates adsorbate placement on a framework fixed to the ATOMMOF-predicted structure.Source: §5.1 · p. 7, Figure 7 · p. 8
• ATOMMOF reaches higher coverage than GCMC and random sampling, indicating better recovery of reference adsorption configurations at comparable or lower runtime.Source: §5.1 · p. 7, Figure 7 · p. 8

5.2. Adsorption Energy Success Rate

• 18.59% success rate with 10 samples versus 12.30% with 1 sample, with lower-than-reference adsorption energies in 14.20% and 6.30% of systems, respectively.For 10 samples, the average minimum energy difference among converged structures is −0.0462 eV; the energy distribution closely matches the reference.Source: §5.2 · p. 8, Table 5 · p. 8, Figure 8 · p. 8

• ATOMMOF models all atoms directly and outperforms prior baselines on BW, supporting high structural accuracy without strong structural constraints.The conclusion contrasts this with rigid-body assumptions and other constrained approaches.Source: §6 · p. 8
• Predictable improvements with model size and compute establish scaling laws for porous-crystal generation.Source: §6 · p. 8
• MLIP-guided Feynman-Kac steering increases the stability and validity of generated structures.Source: §6 · p. 8
• On ODAC25, ATOMMOF recovers diverse adsorption configurations with substantially less wall-clock time than GCMC and random sampling.The generated structures can initialize MLIP relaxation that potentially reaches adsorption energies below the reference.Source: §6 · p. 8

Impact Statement

• Faster structure prediction and sampling may support materials discovery for gas separations and carbon capture.The impact statement also notes potential misuse for harmful chemical design and recommends responsible oversight.Source: Impact Statement · p. 9

A. Model Architecture

• The denoiser uses an encoder, trunk, and decoder to transform noisy atomic and lattice inputs into predicted clean coordinates and lattice parameters.The encoder builds single-atom and pairwise representations, the trunk models atomic interactions, and the decoder produces per-atom outputs.Source: A. Model Architecture · pp. 12–13

B. Feynman-Kac Steering Algorithm

• Feynman-Kac steering biases sampling toward low-energy structures by weighting and resampling an ensemble of K particles using an MLIP potential.The algorithm exposes steering strength λ, particle count K, and resampling interval as hyperparameters.Source: B. Feynman-Kac Steering Algorithm · p. 14

C. Training Details

• Training uses 8 NVIDIA B200 GPUs, with architecture-specific hyperparameters for ATOMMOF-S, ATOMMOF-M, and ATOMMOF-L.Source: C. Training Details · p. 15, Table 6 · p. 15
• Dynamic batching sets batch size using a squared-atom limit to accommodate high variability in atomic counts.Source: C. Training Details · p. 15
• AdamW with a linear warmup scheduler is used for optimization, with loss weights and optimizer settings specified separately.Source: C. Training Details · p. 15, Table 7 · p. 15
• Multi-stage curriculum learning progressively increases the maximum number of atoms used in each training stage to improve convergence and stability.Source: C. Training Details · p. 15, Table 8 · p. 15

D. Additional Results for Scaling Law (Section 4.4)

• Scaling from 38.4M to 491M parameters preserves strong log-linear improvements in match rate and RMSD under stol=0.5.The reported correlations are Pearson 0.92 and −0.94, respectively, with Spearman ±1.00.Source: D. Additional Results for Scaling Law… · p. 16, Figure 9 · p. 16
• >20% match rate is reached by both ATOMMOF variants in 2 integration steps, versus <10% for MOFFlow-2, although ATOMMOF-M saturates slightly below it beyond 10 steps.ATOMMOF-L remains the top-performing model overall, indicating a better sample-quality versus inference-cost trade-off at low budgets.Source: D. Additional Results for Scaling Law… · p. 16, Figure 10 · p. 16

E.1. Metric Definitions

• The steering evaluation measures MOFChecker validity, structure validity, and formation energy error.Structure validity requires pairwise distances of at least 0.5 Å and unit-cell volume of at least 0.1 Å3; formation energy error compares predicted and ground-truth formation energies.Source: §E.1 · p. 17

E.2. Comparison of Energy Distributions

• MLIP-guided steering shifts ATOMMOF energy distributions toward the reference range.The comparison uses energy distributions computed with eSEN-ODAC25-Full.Source: §E.2 · p. 17, Figure 11 · p. 17

E.3. Sensitivity Analysis

• Each sensitivity experiment varies one Feynman-Kac steering parameter while keeping all others at their default values.The analyzed settings include particle count, steering strength, resampling interval, steering start time, and potential.Source: §E.3 · p. 17, Table 9 · p. 18, Table 10 · p. 18
• Table 11 reports the numerical results underlying Figure 6 for the steering sensitivity analysis.The table summarizes the outcomes associated with the tested hyperparameter variations.Source: §E.3 · p. 17, Table 11 · p. 18

F. Experimental Details for MOF-Adsorbate Structure Prediction (Section 5)

• GCMC uses MOF force fields combining Dreiding and UFF, adsorbate-specific force fields, DDEC6 charges, Ewald electrostatics, and simulations at 298 K.The simulations use a 12.8 Å interaction cutoff and a force tolerance of 10−6 kcal/mol/Å for long-range electrostatics.Source: F. Experimental Details for… · p. 19
• For each fixed adsorbate type and loading, GCMC performs 10,000 initialization cycles followed by 10,000 production cycles and uses the final production configuration.Translation, rotation, and reinsertion moves are attempted with equal probability.Source: F. Experimental Details for… · p. 19
• Random sampling rejects configurations when any periodic atom pair satisfies d_ij ≤ 0.8(r_i + r_j).The criterion is applied to MOF-adsorbate and inter-adsorbate distances using the minimum image convention.Source: F. Experimental Details for… · p. 19