Stiefel Flow Matching for Moment-Constrained Structure Elucidation

TL;DR

This paper introduces Stiefel Flow Matching, a novel generative model that accurately predicts 3D molecular structures from moments of inertia by leveraging the geometry of the Stiefel manifold, outperforming existing methods.

Contribution

The paper formulates the structure elucidation problem on the Stiefel manifold and develops a flow matching approach that enforces exact moment constraints for improved molecular structure prediction.

Findings

Achieves higher success rates in structure prediction.

Enables faster sampling compared to Euclidean diffusion models.

Effectively handles high-dimensional molecular data.

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.