Stiefel Flow Matching for Moment-Constrained Structure Elucidation

Thank you for your thoughtful comments. We appreciate your recognition of the novelty of our work in using the Stiefel manifold to exactly satisfy moment constraints and in providing a practical training procedure for an effective Stiefel generative model. We address your comments point-by-point below.

Molecules are not rigid bodies. You are correct that molecules are not perfectly rigid, and that naively converting rotational constants to moments yields effective moments, not equilibrium moments. We have added a discussion of vibrational flexibility and rovibrational correction in Appendix A.4. In summary, exact moment constraints enables the search of the true equilibrium moments up to the precision of calculating the rovibrational correction.

Experiment observes properties that have been vibrationally averaged, including the rotational constants $A(BC)_0$. Vibrational averaging is distinguished from conformational fluctuations such as torsions, which can be frozen out by cooling molecules to their ground vibrational state. However, even in the ground vibrational state, a molecule is still vibrating due to zero-point energy. If one were to a priori guess the equilibrium moments correctly, one could then verify whether they are indeed correct by generating structures, calculating their rovibrational corrections $\alpha$, and then checking their agreement to the experimental rotational constants, via $A(BC)_0 = A(BC)_e - \frac{1}{2}\sum_i \alpha_i^{A(BC)}$. Therefore, given experimental rotational constants $A(BC)_0$, one can use this verification procedure in a fine-grid search for the true equilibrium rotational constants $A(BC)_e$. This is feasible since $A(BC)_e$ consist of only 3 numbers. Experimental precision in $A(BC)_0$ is maintained up to the precision in computing $\alpha$.

To verify that a structure is the true structure, we must know that:

1. its $A(BC)_e$ and $\alpha$ match $A(BC)_0$
2. gradient norm is near 0.

In contrast, without the manifold constraint, samples may have incorrect moment constraints, and we must rely on quantum chemistry geometry optimization to land in a potential energy minimum that just happens to have the correct rovibrational-corrected moments.

Practical application. The practical application is in enabling unknown molecule structure elucidation by rotational spectroscopy, as rotational spectroscopy can measure rotational constants, which are related to the equilibrium moments once you know the rovibrational correction. Our reported success rates are necessarily low due to the need to evaluate on many different test set molecules - in practice, a structure elucidation campaign focuses on a handful of molecules and can therefore devote large computational resources (e.g. 100x more samples, exhaustive energy calculations).

An immediate practical implication of our results is that it actually is possible to identify unknown molecules by just their moments of inertia and molecular formula. The experimental implication of solving this problem is that it could dramatically shift the role of rotational spectroscopy from merely confirming structures to elucidating unknown molecules. This would provide a new method for structure elucidation which does not require mixture separation and may be useful in identifying unknown molecules in the interstellar medium.

Number of baselines is small. See global response. Thank you for the suggestion of KREED-XL-proj.

Experimental results not convincing. We emphasize that validity and stability are reported as averages over all generated samples. Despite poorer averages, Stiefel FM-OT (filter) obtains a greater success rate. The fact that an underfitting model can already provide advantage suggests that the direction of the approach is meaningful, and we discuss potential improvements that overcome the pathologies of Riemannian flow matching in the response to Reviewer Locv.

Thank you for your helpful comments. We appreciate you saying that the paper is well written and easy to follow, as well as your recognition of the novelty of our method in applying the Stiefel manifold and its geodesics. We address your comments below.

Generalizability of the method. Our results are on a molecule-wise train/test split, which means that at test time, the model has never seen any of the unknown molecule's conformers. The increased success rate of Stiefel FM indicates that the model generalizes better than unconstrained diffusion models. We note that the GEOM dataset covers a large area of chemical space which also covers the extent of applicability of rotational spectroscopy: medium-size molecules. Out-of-distribution detection is not perceived to be an issue, since any reported successes will be checked thoroughly using quantum chemistry. Does this answer your question?

Generating $K=10$ samples. We take the minimum RMSD because the important metric is whether the generative model can generate the correct structure at least once. We set $K=10$ based on our computational budget of requiring evaluation on all 29203 test set examples. In actual structure elucidation campaigns, there will only be a handful of molecules of interest, so $K$ should be $>1000$. Correspondingly, significant computational resources will be available so that every generated example can be checked by geometry optimization with quantum chemistry, followed by ensuring that the molecule is stable (gradient norm 0) and has the correct moments.

We emphasize that we are proposing a new problem setting. Other 3D structure prediction problems such as docking usually query $K$ samples and then rank samples based on a confidence head to output a single top-1 predicted structure. We do not rank samples because every generated sample can be evaluated with quantum chemistry.

RMSD thresholds. In our proposed problem setting, the RMSD thresholds are set quite small because the only way to verify that a generated sample is correct is to check that it is stable and has the correct moments. Stability is highly sensitive to tiny changes in conformation. The correctness metric of Cheng et al. is too forgiving, because torsion angles can be rotated, giving the same chemical graph but a different conformer with different moments of inertia.

Report RMSD directly. We do not report mean/median RMSD because these aggregate metrics do not distinguish the performance of each method. Most examples have high RMSD, owing to the high difficulty of the problem, which washes out the mean or median. We therefore show RMSD histograms for transparency on the choice of threshold.

$n \geq 5$ assumption. There are 6 examples of molecules with $n < 5$ in QM9, and none in GEOM. There are relatively very few small molecules compared to larger molecules. The vast majority of already-known molecules are also small, because they are easy to study. In contrast, we focus on the setting of elucidating unknown molecules.

Explaining loss computation. Step (2) is geodesic interpolation, which is given by Eq. (9). We refer to the Appendix for numerical algorithms of the exponential and logarithm.

Appendix interrupts flow, add sampling. We have done our best to make the main text of the paper self-contained in conveying important details related to method and experimental results. We intentionally leverage the appendix for showing greater details for interested readers. We are happy to discuss which specific details you think would be a better fit for the main text. We have also added an algorithm box describing sampling as you suggested.

Masses. You are correct that there are $n$ masses, one for each atom. $M=\sum_i \boldsymbol{m}_i$ is the total mass of the molecule.

Vector $\boldsymbol{a}$. This is defined this way just so it is clear that atom types are discrete. There is an implicit mapping between atom types and the mass of the most abundant isotope. $\boldsymbol{a}$ is used in the appendix to specify that input features to the neural network include the embedded discrete atom types.

Thank you for your thoughtful review of the paper. We appreciate your support of the work for its originality, readable presentation, and practical utility. We address your comments below:

Underfitting GEOM, concrete directions for improvement. The difficulty of Stiefel FM may be attributed to pathologies in Riemannian flow matching with the geodesic distance on compact manifolds:

1. The velocity field under the geodesic distance suffers from discontinuity at the cut locus, as argued in [1].
2. The time-dependent probability density suffers from a shrinking support, which has been discussed for cases like the simplex [2] or $SO(3)$ (Appendix G.2 of [3]).

Other works empirically report degraded results of Riemannian flow matching versus Riemannian diffusion on the torus [1,4]. These pathologies motivate the development of other probability paths for Stiefel flow matching, such as Stiefel diffusion, or flows which asymptotically land on the Stiefel manifold [6], both of which are discussed for future work.

At sample time, some potential improvements for flow matching are to use corrector sampling [5], or to enhance the flow with a jump process [3]. At training time, we have tried other flow matching techniques, such as a logitnormal timestep schedule [7] and stochastic flow matching [8], but with limited improvement. Minibatch optimal transport [9] is technically challenging to apply here because each example has a variable number of atoms.

Connections with discrete flow matching. In our setting, the only degrees of freedom are the continuous atom positions, which does not include the discrete atomic identities because we assume a fixed molecular formula. In future work where molecular formula is unknown, the jump processes of generator matching [3] could enable a multimodal flow which simultaneously varies continuous atom positions and discrete atom types. This process would jump between Stiefel manifolds corresponding to different molecular formulae. Importantly, jump processes would not require initial knowledge of the number of atoms.

Heuristic alignment and validation. We discuss each method of alignment and validation.

• The 3D alignment used in evaluating RMSD does not use a heuristic: it finds the node-permutation and reflection which yields the smallest least-squares Euclidean deviation. This is done by solving 8 linear assignment problems, one for each reflection, and taking the minimum.
• The alignment of noise and data samples for optimal transport minimizes Stiefel distance using greedy randomized search. We cannot rely on a linear assignment solver since we must minimize the nonlinear Stiefel distance. We suspect that solving this assignment problem is more difficult than the quadratic assignment problem, which is NP-hard already, and leave full validation of this hypothesis to future work. We have observed that decreasing the number of iterations of the heuristic can yield suboptimal paths, but can still provide useful learning signal. This is in line with other works where optimal transport approximated using mini-batches can still provide valuable learning signal [9].
• The validity metric uses rdkit's rdDetermineConnectivity, which assigns bonds based on a lookup of bond lengths and valency considerations.

Limited baselines. See global response.

Longer than 9 pages. The maximum page length is 10 according to the call for papers:

the main text must be between 6 and 10 pages (inclusive). ... We encourage authors to be crisp in their writing by submitting papers with 9 pages of main text. We recommend that authors only use the longer page limit in order to include larger and more detailed figures. However, authors are free to use the pages as they wish, as long as they obey the page limits.

Thank you for your comments on the paper. We appreciate your recognition of the effectiveness of our evaluation for our proposed method and the fact that we achieve better RMSD with a lower number of function evaluations. We address your comments below.

Limited baselines. See global response.

Stiefel manifold constraints. The Stiefel manifold imposes the constraints of moments of inertia. The only other constraint is that the molecule must be energetically stable, which should be learned by the model.

Large molecules. The performance of Stiefel flow matching on large molecules is shown by its performance on GEOM. All conformers are used during training, but evaluation considers only the lowest energy conformer. This is because in a rotational spectroscopy measurement, the sample is cooled, and so the lowest-energy conformer is likely to have the highest population, which therefore shows the strongest signal in the spectrometer.

[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.