Stiefel Flow Matching for Moment-Constrained Structure Elucidation

2024-12-17

arXiv.org

AI Summary

Reconstructing all-atom 3D molecular structures with high accuracy from only the molecular formula and exact rotational inertia tensor remains challenging. Method: This paper introduces the first flow-matching generative framework embedded on the Stiefel manifold St(n,4), where rotational inertia constraints are rigorously modeled as intrinsic geometric constraints on the manifold. The approach integrates Riemannian generative modeling, equivariant optimal transport–inspired flow learning, and rotation-invariance constraints to ensure physically consistent structural sampling. Contribution/Results: Evaluated on the GEOM dataset, our method significantly improves both structural prediction success rate and sampling efficiency. Notably, it maintains high accuracy and computational speed for large molecules—outperforming Euclidean-space diffusion models across all metrics. This work establishes a principled geometric foundation for physics-aware molecular structure generation under rotational inertia constraints.

Technology Category

• Machine Learning: Learning with Manifolds
• Knowledge Representation and Reasoning: Geometric, Spatial, and Temporal Reasoning
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG

Abstract

Molecular structure elucidation is a fundamental step in understanding chemical phenomena, with applications in identifying molecules in natural products, lab syntheses, forensic samples, and the interstellar medium. We consider the task of predicting a molecule's all-atom 3D structure given only its molecular formula and moments of inertia, motivated by the ability of rotational spectroscopy to measure these moments. While existing generative models can conditionally sample 3D structures with approximately correct moments, this soft conditioning fails to leverage the many digits of precision afforded by experimental rotational spectroscopy. To address this, we first show that the space of $n$-atom point clouds with a fixed set of moments of inertia is embedded in the Stiefel manifold $mathrm{St}(n, 4)$. We then propose Stiefel Flow Matching as a generative model for elucidating 3D structure under exact moment constraints. Additionally, we learn simpler and shorter flows by finding approximate solutions for equivariant optimal transport on the Stiefel manifold. Empirically, enforcing exact moment constraints allows Stiefel Flow Matching to achieve higher success rates and faster sampling than Euclidean diffusion models, even on high-dimensional manifolds corresponding to large molecules in the GEOM dataset.

Problem

• Predicting 3D molecular structures from molecular formulas and moments of inertia.
• Leveraging high-precision rotational spectroscopy data for exact moment constraints.
• Improving structure elucidation success rates and sampling speed using Stiefel Flow Matching.

Innovation

• Stiefel Flow Matching for exact moment constraints
• Generative model on Stiefel manifold for 3D structure
• Equivariant optimal transport simplifies flow learning