Stiefel Flow Matching for Moment-Constrained Structure Elucidation

Authors: Austin Henry Cheng, Alston Lo, Kin Long Kelvin Lee, Santiago Miret, Alan Aspuru-Guzik
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: 3d molecular generative models, flow matching, stiefel manifold, structure elucidation

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.

Decision: accept (poster)
Reviewer scores: 8, 6, 6, 8

Reviews and discussion:
<thread>
creation date:
1729536402013

thread_type:
official_review

<thread_object>
creator: reviewer_519h

<content>
summary:
This work proposed a novel flow-based generative framework to address the generation task of molecules with moment constraints. The authors noted the connection between the molecules with fixed moments and the Stiefel manifold and applied flow matching on such a manifold to ensure exact moment constraints. Empirical experiments demonstrated better generation results that satisfy these moment constraints.

strengths:
- The manifold structure of n-body particles with fixed moments was demonstrated to be a Stiefel manifold in the paper, which, to the best of my knowledge, is the very first work.
- Equipped with such a Riemannian structure, the proposed framework can effectively navigate through the Stiefel manifold with the moment constraints exactly satisfied. Empirical experiments also demonstrated zero errors against the moment constraints.
- A thorough mathematical introduction to the Stiefel manifold was provided in the Appendix, together with efficient algorithms to calculate the exponential and logarithm maps on the Stiefel manifold, which makes practical training feasible.

weaknesses:
1. The idea of the moment-constrained generation is not well-motivated in the paper.
    - **Molecules are not rigid bodies**. The dynamics of a molecule (e.g., stretching, twisting of bonds) make it possible for it to have various conformations that vary slightly by their principal moment of inertia. In other words, the moment of inertia measured by rotational spectroscopy is an ensemble property that reflects the average (or effective) moment of inertia of a molecule. In this sense, it seems more reasonable to allow for some violation of the moment constraints.
    - **Practical applications are missing**. It remains unclear what the practical applications of generating molecules with specific moments are. As an example, generating molecules with a specific HOMO-LUMO gap can be significant to the discovery of photoelectric molecules. The moment of inertia, however, does not hold a specific connection to real-world applications.

2. **The number of baselines compared in the paper was small**. Indeed, only KREED was compared. There are other baselines available:
    - The conditional generative model in [1] where molecular properties are fed as additional information during both training and sampling for the conditional generation of molecules with desired properties. As the generative model implicitly learns the connection between the property and the molecule, it does not guarantee the exact satisfaction of constraints, but in practice, it does provide fairly competitive generation results. 
    - As it is always possible to project the generation to the corresponding Stiefel manifold according to Appendix B.6, another natural baseline would be directly projecting the unconstrained generation onto the Stiefel manifold, which also ensures the exact moment constraints.

3. **The experimental results were not convincing enough** to demonstrate the superior performance of the proposed method. Although the moment constraint errors were indeed zero, the validity and stability scores were worse. The RMSD percentage also outperformed the baseline by only a small margin on the GEOM dataset.

4. **The filtering procedure on the GEOM evaluation may give the proposed model with inappropriate advantages**, as it is essentially picking from a larger number of generations with a high likelihood of finding better samples. A similar procedure should be applied to the KREED baseline for a fair comparison.

[1] Hoogeboom, Emiel, et al. "Equivariant diffusion for molecule generation in 3d." International conference on machine learning. PMLR, 2022.

questions:
1. What are the practical applications of the moment-constrained molecule generation? Why is it important to enforce the exact moment constraints? See Weakness 1.
2. What are the performance of additional baselines? See Weakness 2.
3. Can you follow the same postprocessing procedure and compare the proposed model with the KREED-filter baseline performance? See Weakness 3 & 4.
4. What is the empirical training and sampling time of the proposed method compared to a normal unconstrained Euclidean flow matching model? I would expect extra time as the exponential map and logarithm map on the Stiefel manifold are far more complex to compute. Nonetheless, it is always beneficial to know the empirical time.
5. The logarithm map was calculated with 20 iterations. Does it provide a good enough approximation? Can you provide an analysis of the logarithm errors with respect to the number of iterations to better demonstrate your choice of 20 iterations (and 1 iteration for OT)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730146939100

thread_type:
official_review

<thread_object>
creator: reviewer_oa92

<content>
summary:
This paper proposes the use of flow matching generative models to obtain the 3D structure of molecules. This is achieved by explicitly enforcing problem constraints using Stiefel manifold and its properties.

strengths:
- The paper is generally well-written and easy to follow. The math is sound and clearly presented. 

- The use of Stiefel manifold and its connection to the problem and the constraints in (2). 

- The use of geodesics in the Stiefel manifold that provides straight lines for navigating the manifold.

weaknesses:
- How generalizable is the proposed method? Can the authors perform OOD evaluation? 

- Why the need to generate K=10 samples and report only the one with the lowest RMSD? Is this the standard practice? if yes, support is needed. If not, then performing this indicates instability which requires further investigation. 

- The evaluation metric for the proposed method seems to depends on RMSD thresholds defined by the authors. Have these thresholds been employed by previous methods that utilize neural networks for predicting the 3D structure of molecules? If yes, then citations are needed. The authors in Cheng et al. (2024) used correctness and reported the exact RMSD values. 

- Why not reporting the RMSD directly? For example, when predicting a protein 3D structure (e.g., in AlphaFold), the RMSD is reported directly after alignment.

questions:
### Major Comments/Questions: 

- In lines 175 to 177, how practical is the assumption $n\geq 5$? How many examples with $n<5$ are removed from the datasets? 

- To compute the loss in (8), the authors listed 4 requirements. All are well explained other than the second. Further clarification is needed here.

- Many details of the paper is put in the Appendix. This is disrupting the flow of the paper. I suggest that the authors to place an algorithmic procedure for the sampling part. 

### Minor Comments/Questions:

- Based on the definition of $\mathbf{m}$, shouldn't the total number of masses $M$ be equal to $n$?

- Vector $\mathbf{a}$ is not being used in the problem definition other than specifying the number of atoms. 

- The use of $st(n,3)$ instead of $st(n,4)$ in the main body is not well justified. Why Theorem 4 is put in the Appendix? 

- The sampling procedure is not as clearly presented as the training of $v_\theta$. Methmatical procedure of how the pre-trained $v_\theta$ is employed to perform sampling is needed. 

- In practice, the authors use t=0.001 to 1, but the uniform distribution in (8) uses U[0,1]. Shouldn't this be U(0,1]? 

- The last sentence in the Datasets paragraph in Section 4 requires support by a citation. 

- Why the FM model has more parameters than the diffusion model in Cheng et al. (2024)? 

- Ill-sentence in lines 445 to 446.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730687524816

thread_type:
official_review

<thread_object>
creator: reviewer_f4fb

<content>
summary:
The paper proposes a generetive model approach for generating 3D molecular structures conditioned on the molecular formula and moments of inertia. The generative method is an extension of Riemannian flow matching with equivariant optimal transport on the Stiefel manifold.

strengths:
The paper shows the effectiveness of the proposed approach through evaluations over two molecule benchmarks: QM9 and GEOM. The method achieved lower RMSD for the generated molecules with a lower number of function evaluations (NFE) compared to the Euclidean diffusion model.

weaknesses:
Limited model comparisons: The paper only compares its results against one method in the literature, where it might consider other works on Riemannian generative models.

questions:
Does the Stiefel manifold have any constraints? How does it handle the increased complexity of larger molecules and the number of conformers?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730703798140

thread_type:
official_review

<thread_object>
creator: reviewer_locv

<content>
summary:
This paper addresses the molecular structure elucidation problem using a generative approach. This task involves inferring a molecule's 3D structure from its molecular formula and moments of inertia. Traditional generative models conditioned on moments of inertia fail to leverage the precision offered by rotational spectroscopy, which measures moments with high accuracy. The paper introduces Stiefel Flow Matching, a generative model that operates on the Stiefel manifold, where the set of point clouds with fixed moments of inertia is embedded. Empirical results demonstrate higher accuracy and efficiency in generating 3D molecular structures compared to Euclidean diffusion models.

strengths:
- As far as I can tell, the approach used in the paper of embedding the molecular structure elucidation task in the Stiefel manifold is original. This approach leverages manifold geometry to respect exact moment constraints, improving over traditional Euclidean diffusion models. 
- In general, the paper is written and formatted very well. I appreciate the consistent and well-written formalisms and figures, making the otherwise quite abstract paper well-readable. Moreover, the formalisms support the methodology of the work well. 
- The results show some performance improvements, particularly in maintaining moment constraints and reducing sampling cost.
- This work has potential applications in chemistry, pharmacology, and materials science, where precise molecular structures are essential. By leveraging constraints from rotational spectroscopy, this model could significantly improve the accuracy of 3D structure predictions for unknown molecules.

weaknesses:
- Most obviously, for larger datasets like GEOM, the model struggles with validity and stability of generated structures, indicating possible underfitting issues. I think some more concrete suggestions on how to improve these issues would be greatly beneficial.
- Some connections to recent approaches to molecular generation with discrete flow matching are missing in the paper. Even though these problems consider a slightly different task as they operate on discrete domains, some intuition of why the authors' approach does not involve discrete dynamics/generation at all would be useful, as it intuitively would be beneficial for this task. 
- As far as I understand it, some of the alignment and validation steps rely on heuristic checks. Some reflection on how task-specific these are, or some intuition behind the thereby introduced bias would be useful.
- While the authors provide some baseline comparisons, they are somewhat limited and some discussion on how Stiefel Flow Matching specifically outperforms or complements existing diffusion models on continuous or Riemannian spaces could benefit the paper. Moreover, adding more experimental results would aid in this too.
- The paper is longer than 9 pages.

questions:
- The model seems to underfit the GEOM dataset, with low stability and validity rates. Could techniques from recent flow matching techniques or regularization methods improve performance? Alternatively, could a hybrid approach that combines Stiefel Flow Matching and discrete approaches maybe be beneficial here? 
- While optimal transport reduces trajectory length, it appears to slightly reduce success rates. Would a more flexible trade-off strategy that selectively applies optimal transport based on the target molecule’s complexity yield better outcomes? Some comparisons here would be nice. On a similar note, the paper embeds the structure space on the Stiefel manifold, which assumes exact moment constraints. Have the authors considered using approximate embeddings on simpler manifolds for faster computation at the expense of slight constraint violations?
- Do I understand correctly the canonical metric on the Stiefel manifold is used? Given recent developments in Riemannian FM, did the authors explore alternative Riemannian metrics?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732277753042

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We thank all reviewers for their thoughtful and constructive feedback. We are glad that reviewers Reviewers *Locv*, *oa92*, and *519h* recognize our application of the Stiefel manifold to molecular structure elucidation to be “original” and “the very first work”, and that they find our mathematical presentation of the Stiefel manifold and Stiefel geodesics to be “thorough”, “easy to follow”, and “consistent and well-written”. Reviewers *Locv*, *f4fB*, and *519h* also recognize that our contribution empirically improves both the accuracy and efficiency of moment-constrained structure elucidation, with practical “applications in chemistry, pharmacology, and materials science”.

**Empirical utility.** We first make a point that, despite tackling a heavily underconstrained problem, our results show that it is actually possible to take molecular formula and 3 moments of inertia, and then generate the all-atom 3D structure at 0.25 Å resolution, for:
- 27.4% of the test set of QM9 (3580/13033), when combining KREED-XL, Stiefel FM, and Stiefel FM-OT
- 7.9% of the test set of GEOM (2297/29203), when combining KREED-XL, Stiefel FM (filter), and Stiefel FM-OT (filter)

We now respond to common comments.

**Limited baselines.** Reviewers *Locv*, *f4fB*, and *519h* commented that Stiefel FM was compared to limited baselines. We would like to clarify a few items:
1. KREED-XL is a reflection-equivariant diffusion model that is a specialization of approaches like E(3)-equivariant diffusion [1], with differences being that KREED-XL does not change atom types, that it relaxes equivariance from E(3) to reflections, and has a more expressive architecture.
2. KREED-XL is different from KREED (Cheng et al.) because KREED assumes access to unsigned substitution coordinates as extra input, and they have different architectures.
3. KREED-XL-DPS is a significant baseline because Diffusion Posterior Sampling explicitly incorporates the closed-form formula of the moments in an extra guidance term while sampling the diffusion model.
4. We add an additional baseline, KREED-XL-proj, which simply takes all samples of KREED-XL and projects them onto the feasible manifold. The projection satisfies moment constraints exactly and does not change success rate, though it slightly reduces stability, since it has no effect on correct structures while distorting incorrect structures.

[1] Hoogeboom, E., Satorras, V. G., Vignac, C., & Welling, M. (2022, June). Equivariant diffusion for molecule generation in 3d. In International conference on machine learning (pp. 8867-8887). PMLR.

**Underfitting GEOM.** Reviewers *Locv* and *519h* commented that Stiefel FM on GEOM struggles with generating valid and stable structures. We point out that these validity and stability metrics are *averaged over all generated samples*, so validity and stability appear worse because when Stiefel FM fails, it tends to output a very bad structure. But, when Stiefel FM *does* generate the correct structure, it is valid and stable, which makes validity and stability a useful filtering mechanism for Stiefel FM, but not for KREED-XL. We note the greater efficiency (in terms of success rate per NFE) of Stiefel FM-OT for elucidating structure *despite* underfitting the dataset.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The submission studies moment-constrained molecular structure identification, in which we are given the chemical formula of a molecule and the moments of its 3d structure. The goal is to sample from a distribution over structures which are both chemically realistic and agree with the observed moments. In previous work, this problem had been investigated using soft moment constraints. Here, the paper shows how to exactly enforce moment constraints, observing that these constrains can be converted into a manifold constraint in which allowable structures reside on a geodesic submanifold of the stiefel manifold. The paper develops generative models on this Stiefel manifold using flow matching and optimal transport. 

The paper provides a novel approach to moment constrained structure generation. In particular, it provides a novel formulation of this problem as generation of a submanifold of the Stiefel manifold. Experiments demonstrate that the proposed method generates a larger fraction of correct structure compared to existing baselines. While there are some limitations to the method’s performance (this is a highly underdetermined inverse problem; all current methods struggle with the majority of instances). After discussion, reviewers converged to a uniform recommendation to accept, praising the paper’s novel and rigorous formulation and its experimental results.