Stiefel Flow Matching for Moment-Constrained Structure Elucidation
Authors: Austin Henry Cheng
Alston Lo
Kin Long Kelvin Lee
Santiago Miret
Alan Aspuru-Guzik
Venue: ICLR 2025 Poster (2025)
Keywords: 3D molecular generative models
flow matching
Stiefel manifold
structure elucidation
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We propose Stiefel flow matching as a generative model on the Stiefel manifold and apply it for molecular structure elucidation by rotational spectroscopy.

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.

Published as a conference paper at ICLR 2025
STIEFEL FLOW MATCHING FOR MOMENT-CONSTRAINED STRUCTURE ELUCIDATION
Austin H. Cheng 1.2 Alston Lo Kin Long Kelvin Lee
Santiago Miret
Alán Aspuru-Guzik 1,2,4 1,2,4
Uniivesiity of Toronto Vector Institute Intel Labs *Acceleration Consortium
https: /ttp::///gthbb..///////p com/ camm/apuuu--a gww.kiggouapssaeeeeeee group/stiefflFM
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 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 St,n, 4). We then propose Stiefel Flow Matching 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 Fuclidean diffusion models, even on high-dimensional manifolds corresponding to large molecules in the GEOM dataset.
INTRODUCTION
Elucidating the structure of unknown molecules a central task in chemistry, important for analyzing environmental samples (Moneta al., 2023), identifying novel drugs (Sonstrom al., 2023), and determining potential building blocks of life the interstellar medium (McGuire et al., 2016). The challenge is to aggregate information from multiple sources of analytical data to unambiguously determine a molecule's s structure. Rotational spectroscopy holds a unique capacity to provide precise measurements of molecule's rotational constants, which are closely related to its moments of inertia. In turn, the connection between these moments and 3D structure has routinely provided the highest quality gas-phase 3D structures attainable from experiment (Domingos et al., 2020). Typically, structure clucidation with rotational spectroscopy proceeds by confirming whether a known structure's 's moments match with experiment (Lee & McCarthy, 2019; McCarthy al., 2020). However, this approach is inherently restricted to molecules whose structures have already been catalogued, and leaves no prescription for undiscovered molecules such as novel natural products and key reactive intermediate species that cannot be easily isolated (Womack et al., 2015).
To overcome this limitation, we apply generative modeling to infer candidate structures based on the moments and molecular formula alone. By themselves, moments of inertia are a 3-number summary of how a molecule's 's mass distributed in space. Going from moments to 3D structure (3 + 3n values) is therefore severely underconstrained inverse problem. Nevertheless, deep generative models such as diffusion (Ho et al., 2020; Song et al., 2020) and flow matching (Lipman et al., 2023; Liu et al., 2022) have shown promise in solving various inverse problems (Song et al., 2022; Chung et al., 2023; Song et al., 2023a). Indeed, previous work has applied Euclidean diffusion models sample 3D structures conditioned given of moments of inertia (Cheng al., 2024).
"Equal contribution. <austi nêcs. toronto.edu, alaton. 1 .utoronto.ca>

---

Published as a conference paper at ICLR 2025
Px 378.264
Px 378.264
Px 378.264
Py 163.007
163.007
Py 163.007
P2 9.747
0.747
P2 0.757
X
X
X1
Te,St(n,4)
Output: True structure
Input:
U
Moment- satisfying point clouds (Ours)
molocular formula CgHgO2N
U
moments inertia Ps) (378.254, 163.007, 9.747)
Uo.
M St 4)
Chemically valid point
Figure Stiefel Flow Matching learns to elucidate 3D molecular structure from moments and molecular formula alone by transforming uniform Stiefel noise Xo into valid molecular structures X1. Generative modelling on the Stiefel manifold St(n, 4) guarantees that samples always have the correct moments of inertia, which allows the network focus only on generating chemically stable structures. Within the intersection of these spaces lies the true 3D structure.
However, analytic formulas for the moments of inertia are sufficiently simple that we can define the feasible space where these constraints are always exactly satisfied. Such precise adherence to moments can potentially leverage the many digits of precision provided by rotational spectroscopy (Shipman et al., 201 1) in order to constrain the space of plausible structures. We first show that M, the set of n-atom point clouds with fixed moments of inertia, is embedded in the Stiefel manifold St(n, 4). We then propose Stiefel Flow Matching as a generative model on the Stiefel manifold for solving the moment-constrainee structure elucidation problem (Figure 1). Our approach augments Riemannian flow matching (Chen & Lipman, 2024) with equivariant optimal transport (Klein et al., 2023; Song et al., 2023c), which simplifies and shortens generation paths. Concretely, our contributions are:
1. We propose the task of moment-come consetriind structure elucidation as challenging generative modelling problem the Stiefel manifold.
2. To solve this problem, we present Stiefel Flow Matching, Riemannian flow matching approach. Furthermore, we formulate an objective for equivariant optimal transport on the Stiefel manifold, which obtains shorter and simpler flows.
Stiefel Flow Matching predicts 3D structure with greater success rate and lower cost than Fuclidean diffusion models on both the QM9 (Ramakrishnan et al., 2014) and GEOM (Axelrod Gomez-Bombarelli, 2022) datasets.
2 BACKGROUND AND APPROACH
We consider 3D molecule as point cloud of n atoms with atomic numbers E No and 3D coordinates X = (I;, Vi, > E Rnx3 We also refer to X as the molecule's 3D structure. Molecules have translational, rotational, and permutation symmetry. However, and X are stored on a computer, they necessarily have a node ordering and orientation. Therefore, when we refer to structure, really mean the equivalence class containing X under these symmetries.
2

---

Published as a conference paper at ICLR 2025
2.1 PROBLEM STATEMENT
We are given a molecule's molecular formula and moments of inertia, and wish to infer the 3D structure of the molecule. The molecular formula provides with the number of atoms n, atomic numbers a E N", and atomic masses m E I". The moments of inertia' (Px; Py: Pz) are three nonnegative numbers that summarize the mass distribution the point cloud. They are defined the cigenvalues of the planar dyadic P E R2x3 (Kraitchman, 1953), calculated from masses m and the 3D atomic coordinates X
(Iz M/s 2:) 2i) E Rxx as
maa me 2
m;3;2; mayi2 m,3
P x (diag m)
Σ
mayia m.yi mizsi
This definition assumes coordinate system whose origin is the weighted center of mass of the molecule, which gives the constraint m X = O. We assume that Px > Py > Pz 0, which eliminates rare edge cases (Appendix A.2). Note that P is symmetric and positive semidefinite and can be thought of as the weighhed covariance matrix the point cloud, as principal component analysis (PCA).
Diagonalizing P yields three cigenvectors, referred to as the principal axes of rotation, which orient molecule in canonical representation up sign-flips. The principal axes correspond to directions which "explain" the most molecular mass, in analogy to PCA. The principal axis system then the coordinate system whose origin is the center-of-mass and whose axes are the principal axes of rotation. We fix our coordinate system to be the principal axis system by construction, which gives the following constraints X:
Px
m,a =
Σ
0 =
m,J;;
py m.y?
m
m,lii
(2)
Pz
2 m
0=
0
m;zi-
These constraints also canonicalize a 3D structure to sign-flips of the axes (e.g. x + -x). The center-of-mand constraint removes translational degrees of freedom, while the off-diagonal constraints remove rotational degrees of freedom.
The goal of molecular identification from moments and molecular formula is find all molecular structures X which are consistent with these constraints and are thermodynamically stable, i.e., local minima of the potential energy surface. Structures which satisfy these criteria can then be compared to experimental measurements from rotational spectroscopy (Appendix A.3).
2.2 THE FEASIBLE SPACE OF MOMENT-CONSTRAING STRUCTURES
The Stiefel manifold is the set of orthonormal n x matrices, defined as
St(n,p) = (U E R2 U = Ip)
where I, the p p identity matrix, and n 2 p. An element of St(n,, p) can thought as the first p columns of some n-ddimensional (improper) rotation matrix, can be thought of as collection of p orthonormal n-dimensional vectors. We provide more background Appendix
Momenn-constrained structures X can be mapped into (1,,, 4) by scaling rows by masses and columns by moments. Letting M = Di= my be the total mass, consider the following construction:
mi Px
F1
My Py y1
mi Pz
T M
U=
E Rxx4
ma Fe
Px
m Py
Vn
ma Pz
2n
mm
M
This overview slight simplification of rotational spectroscopy. Full details outlined in Appendix A.

---

Published as a conference paper at ICLR 2025
It can be verified by inspection and comparison the Equation 2 constraints that the columns of U are orthonormal. That is, U E St(n, 4). However, we cannot yet freely convert between X and U:
the last column of U is not free as it must equal the unit mass vector in = vm/M to satisfy the zero center-of-mass constraint. To convert an arbitrary U St(n, 4) to an X satisfying all constraints, we first apply a rigid n-dimensional rotation R to U so that its last column is aligned to vi (Appendix B.5), before finally unscaling the rows and columns of U.
The feasible space M of moment-constraiin structures is therefore the subset of St(n, 4) whose last column is fixed m, i...
M = {U E St(n,,4) U,, }
(5)
In fact, we show in Appendix B.7 that M is totally geodesic submanifold of (t 4). which means that shortest paths between points in M stay in M. The first three columns of elements in M form the intersection between St(n, 3) and the orthogonal complement to span(m)) (m)) which is in turn equivalent to St(n 1,3). Hence, the dimension of M 3n 9, which corresponds to removing 3 translational, 3 rotational, and 3 moment degrees of freedom, consistent with the 9 constraints in Equation 2. Going forward, we assume that the molecule of interest contains n 2 5 atoms, so that we always deal with Stiefel manifolds of strictly rectangular matrices.
2.3 NAVIGATING THE STIEFEL MANIFOLD
The Stiefel manifold provides rich structure for navigating the feasible space of molecular structures. As manifold, it is locally Euclidean but globally curved. This means that every point U E St(n, p) is attached a vector space called its langent space TeySt (n, p). For the Stiefel manifold, these tangent spaces are given
SuSt(n. p) (A R
UU A A 00-
Then, equipping every tangent space with an inner product (., ")u turns St(n,,) p) into Riemannian manifold, giving rise to notions of angles and distances. The collection of inner products for each tangent space is called the Riemannian metric. One such metric for the Stiefel manifold is the canonical metric (Edelman et al., 1998),
(A, Auu trace A (L,, - JUU A,
(7)
which we exclusively use for this work. The canonical metric induces a norm ||||| (A, A) on each tangent space, which gives the length of a curve [0, 1] + St(n. p) as L(7) - So (Il( Curves that are locally length-minimziing are called geodesics, providing notion of "straight lines" for efficiently navigating around the manifold. Geodesics are defined by their starting point and initial velocity. Indeed, the exponential map expo (A) takes in a starting point U E St(n, p) and initial velocity A E SttSt p), and outputs the final manifold point after following this geodesic for unit time. The exponential map is locally invertible, which gives the existence of the logarithmic map (U). The logarithmic map takes in starting point U and target point Ü, and outputs the tangent vector needed to travel from U to U. Algorithmss for computing exponential and logarithmic maps under the canonical metric for the Stiefel manifold are given in Appendix B.3.
3 STIEFEL FLOW MATCHING
Having shown that the feasible space of conttranne structures is a Stiefel manifold, we can now formulate the problem of moment-consttrined structure elucidation as an unconstrained generative modeling problem on the Stiefel manifold. An attractive approach to this is flow matching (FM), which trains network as a time-dependent velocity field 24 that transforms samples from a prior noise distribution xo Po (z())) into samples which approximately match the data distribution In ~ ~pp (x1) R Pama (x1) (Lipman et al., 2023). The integration of the velocity field over time is then a continuous normalizing flow 4/, which generates marginal probability densities Pt over time by the pushforward operation Pt = [et] (Po). In practice, this is realized by sampling initial conditions Io ~ Pu())) and evolving them from time to f according to the ODE de de uz (Appendix C.4). The goal of training is to approximate this u, using neural network vo (f, x) parameterized by 0.
Flow matching is readily generalized to distributions on Riemannian manifolds (Chen & Lipman, 2024). When closed-form geodesics available, Riemannian flow matching provides a simulation-
4

---

Published as a conference paper at ICLR 2025
free training objective for learning ue, called Riemannian conditional flow matching,
CRCFM ()) = ((t)
[I -i) -.lle.]
(8)
Intuitively, this loss trains the model to interpolate between many pairs of (noise, data). To compute this loss, we require (1) sampling Ug ~ p(U)) from a prior noise distribution, (2) geodesic interpola- tion between Uo and U, to get U,, (3) computing the time derivative of the interpolant U,, and (4) evaluating the norm -l1,.
To sample uniformly from the feasible space M, we can sample U uniformly from St(n, 4) (Ap- pendix B.4) and then rigidly rotate U so that its last column aligns with the unit mass vector rn (Appendix B.5). Then, geodesics can be computed from the exponential (Edelman et al., 1998) and logarithmic (Zimmermann & Hüper, 2022) maps, which are efficient to compute for St(n,
U, expr, logg (U1)),
making our training objective simulation-free. Once we have the interpolant U, which will be the in- put to the neural network, we can then calculate the network's 's target U.. Instead of autodifferentiation, which introduces unnecessary overhead, we compute U, using the logarithmic map,
U, = logu, (U1);
(10)
observing that U, is a unit length tangent vector along the geodesic from U. to UI. Finally, compute the norm of (t, U,) U, following Appendix B.2. Additionally, Appendix Theorem shows how we can compute the logarithm in St 3), rather than St(n2, 4), which slightly saves time.
3.1 REFLECTION AND PERMUTATION EQUIVARIANCE
the principal axes of rotation by construction. This canonicalizes the 3D structure, removing translational symmetries and reducing the rotational symmetries to sign-flip symmetries of the eigenvectors of P. which are the coordinate axes (Puny et al., 2022; Duval et al., 2023). Hence, the flow needs to be equivariant with respect to sign-flips of the coordinate axes, which we call cenffeetioo-equivariiinet et al., 2024). In addition, because 3D structure is invariant under node order permutations, the learned velocity should be equivariant to permutations. Together, the explicit equivariance constraints on the network are given as vt(1, IU, = Noo(1, U,) R. for all node permutations II and reflections R = diag(=1, #1,1), We satisfy these constraints with a reflection -qquvarian graph neural network architecture described in Appendix C.I.
As mentioned earlier, we set the coordinate
for brevity (Lim et al., 2023; Cheng
3.2 EQUIVARIANT OPTIMAL TRANSPORT
Flow matching learns a velocity field which transports samples from the noise distribution to the data distribution, but there is guarantee that samples will follow paths y that are optimal with respect to transport cost L(r). Optimal paths are desired because they afford more efficient training and faster generation (Pooladian et al., 2023; Tong et al., 2024). Since molecules have permutation and rotational symmetries, optimal paths should connect structures whose equivalence classes are close to each other. This corresponds to finding optimal node permutations and rotations to align noise samples Xo to data samples X1. In contrast to the Euclidean case (Klein et al., 2023; Song et al., 2023c), only reflections are needed, not rotations, because the coordinate axes are already fixed in place as the principal axes of rotation by Equation 2.
In addition, transport cost on St(n, 4) must be measured using the Riemannian distance. Thus, the optimal transport map from Un to Un minimizes the following cost over coom-type-poee sype-preserving node permutations II (i.e., Ha a. and reflections R = diag( +1, +1,1
d(U.. U,) = || logmu, ||TUR While searching for optimal alignments, we do not compute the full Stiefel logarithm and instead approximately calculate distance using only one iteration of the inner loop described in Appendix Algorithm 2, which we justify in Appendix E. This approximate distance is then heuristically optimized over atom permutations and reflections with a greedy random local search (Appendix D). Appendix Algorithm 4 outlines this procedure, which samples several permutations for each reflection to approximately identify the best reflection, and then refines the permutations by a local search of random atom sype-preserving index swaps (Appendix Figure 9).
(11)
5

---

Published as a conference paper at ICLR 2025
Table 1: Experimental results on QM9. Stiefel FM shows no violation of moment constraints as shown in the Error metrics, and has the highest success rate for structure elucidation, with the lowest computational cost.
< RMSD f
Method
Error +
Valid ?
Stable
Diverse
NFE +
0.25 A
0.10A
Stiefel Random
0.00
0.00
0.00
0.061
nan
2.640
KREED
11.22:10.28
9.55:10.26
5.18
0.878
-1.335
1.429
1000
KREED-XL
13.65=00.3
10.94=0.27
3.64
0,933
-1.048
0,870
1000
12.36++0.29
9.40=10.26
1.33
0.744
0.826
1.060
1000
KREED-XL XL-proj
13.67=00.00
10.93=0..27
0.00
0.924
0.905
0.871
1000
Stiefel FM
15.17:1..11
13.82:10.30
0.00
0.882
1.125
1.040
200
Stiefel FM-OT
13.99+0.30
02.68=0.29
0.00
0.835
1.039
1.045
200
Model
KREED-XL
600
Mode:
KREED-XL-DPS
KSTED-X
Stiafol FM
KREED XL DPS
Steff FH
Startel FN-OT
Count
300
200
200
100
100
8.00
0.10
0.20
0.30
Figure 2: Histograms minimum RMSD for predicted QM9 examples show two distinct clusters for RMSD. The 0.25A threshold captures molecular structures that useful for structure elucidation.
EXPERIMENTS
We evaluate Euclidean diffusion models and Stiefel Flow Matching on the QM9 and GEOM datasets. For each example, the model takes in moments and molecular formula and produces K = 10 samples.
Datasets. For QM9 (Ramakrishnan et al., 2014), dataset. We abbreviate -rOMM (Axeirod & Homeez-oombarelle 2022) as GEOM. We use the same training, validation, and test splits Cheng et al. (2024), except we remove examples that unstable, have less than 5 atoms, or have exactly zero-valued moments, which drops 3122 examples from QM9 and drops no examples from GEOM. QM9 has train/val/test splits of (04265/13056/1300) molecules, while GEOM has splits of 233625/29203/2903 molecules, or 5537598/9203/29223 conformers. We only predict 3D structures for the lowest-energy conformers of GEOM. This reflects experimental reality, as the lowest-energy conformer typically has the highest proportion in the population after cooling by supersonic jet expansion (Ruoff et al., 1990).
the conformers provided by the GEOM
Structure elucidation. The only measure of model's success is its ability to generate the correct 3D structure X with high accuracy at least once, where accuracy is measured by root-mean-squee -quaredd- deviation (RMSD) of the predicted coordinates to the ground truth. High accuracy is needed because the only confirming evidence available to us is (I) agreement with moments and (2) thermodynamic stability by quantum chemistry, For this reason, we use stringent thresholds of RMSD 0.25 Á and RMSD < 0.10 to ensure that the predicted structure is the same potential energy basin as the true structure. Success is reported if the minimum RMSD over K 10 generated samples satisfies these thresholds. The success rate is the percentage the test set whose generated samples has minimum RMSD which satisfies each threshold. Error bars are standard emrrr of the mean.

---

Published as a conference paper at ICLR 2025
Stiefel FM
1000
Moddd
Statel EM
Steffi OT
1000
1000
Stiefel FM-OT
1000
0.75
1.50
2.00
Corve Longth
Figure 3: (Left) Learned sampling trajectories for Stiefel FM and Stiefel FM-OT on QM9. Each column begins generation from the same noise. (Right) Histogram of curve lengths of all QM9 sampling trajectories for Stiefel FM and Stiefel FM-OT. Permutation and reflection alignment lead to simpler and shorter paths.
Evaluating RMSD is nontrivial because the generated and ground truth structure must first be aligned under same-atom-type permutations and reflections. We use the same RMSD procedure as Cheng al. (2024): We first align node permutations by solving a linear assignment problem whose cost matrix is squared Euclidean distance, and repeat this for all reflections, taking the minimum. Then we compute RMSD between the ground truth coordinates and the aligned coordinates. This is similar to the alignment procedure used in (Klein et al., 2023; Song et al., 2023c).
Auxiliary metries. To characterize the samples generated by each model, report additional metrics. In contrast to success rate, which checks the minimum RMSD of generated samples for each example, these auxiliary metries are averaged over all generated samples (except diversity). Error measures how much the generated structure violates the moment constraints. If P is the computed planar dyadic the generated structure, this is computed as / ||riuu P P) where triu takes the upper triangular part of the matrix (contains 6 elements). Validity is a heuristic check based Deteeeeeennn rmi neConnect vi vvity (Landrum, 2013; Kim & Kim, 2015). Stability is the log norm of the gradient of energy with respect to coordinates, as reported by the xtb quantum chemistry program (Bannwarth et al., 2019). Stable structures should have gradient norm close to zero (log norm very negative), assuming the structure is at local minimum and not a saddle point. Diversity is calculated as the average pairwise RMSD of all generated samples of single example. NFE is the number of function evaluations used during generation, and measures computational cost. We do not set boldface for these metrics because they do not correspond directly to success criteria.
on bond detection with rdDete wwwsstermiiBBBodce
Baselines. Given the novelty the problem, the number of available baselines is limited. We compare the performance of Stiefel FM to KREED (Cheng et al., 2024). KREED is reflection- cquivariant diffusion model trained to generate 3D structure conditioned on molecular formula and moments of inertia, and is a specialization of Euu-eiauiaar approaches like EDM (Hoogehoom et al., 2022). Since our architecture for Stiefel FM is much larger (parameter-wist than the model architecture used in KREED, and because KREED is tailored for slightly different task, also train another reflection-eqqiiaariat diffusion model with an identical neural network architecture to Stiefel Flow Matching, which we label KREED-XL. The planar dyadic is computable at every step of the generation process, which means that this task can be treated as nonlinear inverse problem: On top of KREED-XX XL, we apply Diffusion Posterior Sampling (DPS) (Chung al., 2023), which guides generation with an additional drift term for minimizing the planar dyadic error. As a simple baseline which exactly satisfies moment constraints, we report performance for uniform random sampling the Stiefel manifold. We also report the results of KREED-XL -XL after projecting samples onto the feasible manifold (Appendix B.6). Relevant hyperparameters for methods are provided in Appendix C.3.
Results. Table shows our experimental results on QM9, reporting Stiefel FM with and without optimal transport (OT). We note that the average number of atoms in QM9 is 18, meaning that on average the model must infer 3(18) 9 = 39 values from 3 moments. We find that Stiefel FM can generate the correct structure with greater success rate than all Euclidean diffusion models. We
7

---

Published as a conference paper at ICLR 2025
Table 2: Experimental results on GEOM. Stiefel FM generates few valid structures on its own due the increased difficulty of manifold-constranne generative modelling. When adjusted to generate the same number of valid molecules, Stiefel FM-OT (filter) obtains the highest success rate, without surpassing baselines in computational cost.
% < RMSD ?
Method
Error +
Valid t
Stable +
Diverse f
NFE +
0.25 A
0.10 Á
Stiefel Random
0.00++000
0.00++0.00
0.00
0.000
nan
4.104
0
KREED
0,0410.01
0.02:10.01
58.36
0.353
-0.583
2.286
1000
KREED-XL
3.54+0.11
2.02++008
30.71
0.907
0.900
2.190
1000
BREED-XL-pru
3.54
2.04
0.00
0.904
00.52
2.188
1000
Stiefel FM
2.17:0.09
0.24
0.00
0.388
-0.66
2.212
200
Stiefel FM-OT
2.44
1.49=0.07
0.00
0.376
-0.002
2.195
200
Stiefel FM (filter)
3.57=0.11
2.06=#0.08
0.00
0.889
-0.437
2.183
600
Stiefel FM-OT (filter)
3.94
2.42+
0.00
0.869
-0.352
2.165
600
175
KREED-XL
500
Model
Stiefol FM (fiitrr
Stretel FM (fiited
Streful PS6OT (Riter)
400
125
360
COurt
100
50-
100
25-
8.00
0.05
0.10
0.15
0.20
0.25
0.30
Min RMSD (A]
Figure 4: Histograms of minimum RMSD for predicted GEOM examples.
also see that incorporating the analytic formula of the moments via DPS does improve agreement with the moments, but the cost of accuracy. In contrast, Stiefel FM does suffer from this tradeoff. Projecting samples onto the manifold does not change success rate, because the projection leaves correct structures untouched, while only distorting incorrect structures. In addition, Stiefel FM uses only 20% of the computation used by diffusion models. Euclidean diffusion models such as KREED-XL can produce more valid and stable structures, though they may not necessarily generate the correct structure. Figure reveals that when Stiefel FM's predictions extremely accurate, achieving RMSD even below 0.05 Á. Training with equivariant optimal transport helps learn simpler generation paths and reduces the average curve length of generation trajectories from 1.696 to 1.547 (Figure 3). Interestingly, training with optimal transport slightly reduces the success rate of Stiefel FM on QM9, though this trend is reversed for GEOM. Appendix Table reports additional trials experimenting with optimal transport, noritll timestep sampling (Esser al., 2024), and stochasticity (Bose al., 2024).
correct, is likely be
The task of structure clucidation poses an even harder challenge on GEOM, with an average of 46 atoms, giving a task of inferring 3( (6) problem is hopelessly underconstrained. Nevertheless, Table 2 shows that both diffusion models and flow matching can obtain a nontrivial success rate for structure elucidation. However, the low accuracy, validity, and stability of Stiefel FM(-OT) suggests that the model is underfitting the dataset, even as KREED-XL is able to generate numerous valid samples. We observe that less than half of the struetures generated by Stiefel FM are valid, which aligns with the fact that there are fewer valid structures on the feasible manifold M than in regular Euclidean space. This suggests that a fairer evaluation utilizing the strength of the manifold constraint should generate similar number of
129 values. Naively, one would expect that this
8

---

Published as a conference paper at ICLR 2025
valid structures for each model. This is justified by the fact that validity does not use ground truth labels and can be computed at only nominal cost. Therefore, we generate 30 samples for Stiefel FM and Stiefel FM-OT before filtering to retain up to K 10 valid samples. Note that the combined computational cost of generating 3x as many samples (3 x 200 NFE) is still lower than KREED-XL (1000 NEE). After filtering. Stiefel FM-OT obtains similar validity rate to KREED-XI, but obtains the highest success rate for structure elucidation. Even after filtering, the mediocre stability of generated samples suggests that success rate can be improved further, though at the slightly higher cost of quantum chemistry calculations. Now, we see that optimal transport does help Stiefel FM to predict accurate structures, while also reducing average curve length from 1.421 to 1.344 (Appendix Figure 8). However, biasing generation by selecting for valid samples seems to also select for longer generation trajectories. We find that generation paths that land on the correct structure are usually longer than generation paths that land on incorrect structures (Appendix Figure 7). This may be explained by the fact that initial points are sampled anywhere uniformly on the manifold, but for success they must end up on the single true structure. In contrast, there many incorrect structures all over the manifold, which may end up on average closer to random initial points.
One should expect success rates to be of this magnitude for solving heavily underconstrained problem. This is simply due to the fact that there exist many stable structures with the same molecular formula and very similar moments of inertia. Indeed, trained models usually generate realistic-looking molecules - see generated examples in Appendix Figure 5 and Figure 6. Furthermore, only 10 to 30 samples were queried for each molecule, but an actual structure elucidation campaign would have a much larger compute budget for generating thousands of samples. As highlight, our results show that it actually possible at 0.25 Å resolution to elucidate 27, 27.4% of the test set of QM9 (580/13033) when combining KREED-XL -L,, Stiefel FM, and Stiefel FM-OT; and 7.9% of the test set of GEOM (2297/29203), when combining KREED-XL XL, Stiefel FM (filter), and Stiefel FM-OT (filter).
RELATED WORK
Generative models for 3D molecules. While generative models have been actively explored for 3D structure prediction, few works have applied these methods to the task of structure elucidation (Cheng et al., 2024). Adjacent work in applying diffusion and flow matching models to 3D molecules include molecular generation (Hoogeboom et al., 2022; Song et al., 2023c). conformer search (Jing et al., 2022; Xu al., 2022), docking (Corso et al., 2022), biomolecular assembly (Abramson et al., 2024) and Boltzmann generators (Klein et al., 2023). Generative models on Riemannian manifolds have also been applied to protein design (Bose et al., 2024; Yim al., 2023) and crystal structure prediction (Jiao et al., 2024). Recent work has applied generative modelling to predict 3D molecular structure from powder X-ray diffraction patterns (Lai et al., 2024; Riesel et al., 2024) and 3D protein structure from cryo-EM density maps (Levy et al., 2024).
Deep learning for molecular identification using rotational spectroscopy. A limited number of works provide other parts of the complete workflow needed for structure elucidation using rotational spectroscopy. Zaleski & Prozument (2018) propose RAINet as a forward modeling approach for assigning rotational spectra, where a set of peaks is fed into a classifier, that categorizes the spectra to an appropriate multilayer perceptron that outputs spectroscopic parameters, including moments. McCarthy & Lee (2020) adopt a "mixtureeof-expprts approach that maps spectroscopic parameters and approximate molecular formula into a set of complementary experimental observables and SMILES strings, though with limited success. Cheng al. (2024) present a Euclidean diffusion model KREED for determining 3D structure from moments, molecular formula, and unsigned substitution coordinates, the latter of which is also measurable from rotational spectroscopy, but can be difficult and expensive to obtain. Most recently, Schwarting et al. (2024) provide a thorough analysis into the inverse problem, examining the frequency with which different molecules have moments of inertia that very close in value. Stiefel Flow Matching could disambiguate these structures by providing a diversity of structures that satisfy moment constraints exactly.
Statistics on the Stiefel manifold. While the Stiefel manifold is often studied in the context of optimization (Absil et al., 2008; Chen et al., 2021; Kong et al., 2022), a number of works study the Stiefel manifold (Chakraborty & Vemuri, 2019; Chikuse, 1990). One
probability distributions distribution on the Stiefel manifold is called the matrix von Mises-Fisher distribution or matrix Langevin distribution (Pal et al., 2020; Chikuse, 2003; Jupp& & Mardia, 1979). Wang & Solo (2020)

---

Published as a conference paper at ICLR 2025
propose particle filtering algorithm the Stiefel manifold, with the first application of optimal transport on the Stiefel manifold. Yataka et al. (2023) propose a continuous normalizing flow on the Grassmann manifold, which is closely related to the Stiefel manifold.
CONCLUSION
We propose Stiefel Flow Matching, Riemannian generative model for generating samples subject to exact orthogonality constraints, and apply it to the challenging inverse problem of structure elucidation from moments of inertia and molecular formula. Empirically, Stiefel Flow Matching achieves a higher success rate than Euclidean diffusion approaches. Satisfying the constraints exactly will enable future advances in Riemannian generative modelling to directly transfer generating more stable molecules, without needing to consider agreement to the moments.
6.1 FUTURE RESEARCH DIRECTIONS
Improving Stiefel generative models. Riemannian flow matching empirically shows degradation compared to Riemannian diffusion (Lou et al., 2024; Zhu et al., 2024). This has been attributed to two pathologies of Riemannian flow matching for compact manifolds: (1) the geodesic- geodesic-based velocity field is discontinuous at the cut locus (Lou et al., 2024; Zhu et al.. 2024), and (2) the probability density has a shrinking support (Stark et al., 2024; Holderrieth et al., 2024). These pathologies may explain the difficulty of Stiefel FM in fitting GEOM, and motivate the development of alternative probability paths for Stiefel flow matching. such as Stiefel diffusion (De Bortoli et al., 2022), diffusion mixtures (Jo & Hwang, 2023), or flows which asymptotically land on the Stiefel manifold (Ablin & Peyré, 2022; Gao et al., 2022). Training Stiefel diffusion with the denoising score matching objective requires the heat kernel as training targets for the neural network, but these targets are expensive to compute (Azangulov et al., 2022: Lou et al., 2024). Alternatively, the matrix Langevin distribution may be amenable to modelling with Star-Shaped Denoising Diffusion Probabilistic Models (Okhotin et al., 2024). Future work can also explore varying the metric used on the Stiefel manifold, since the Stiefel manifold admits a a 1-parameter family of metrics generalizing the canonical metric (Hüper et al., 2021). These metrics have efficient numerical algorithms for the exponential and logarithm (Mataigne et al., 2024). Sample-time advances in flow matching, such as corrector sampling (Gat al., 2024) enhancing the flow with jump process (Holderrieth al., 2024), are also orthogonal avenues for improvement.
key limitation of Stiefel Flow Matching is the requirement of molecular formula as input. Recent approaches in discrete flow matching (Campbell et al., 2024; Gat et al., 2024) could enable a multimodal flow simultaneously vary both continuous atom positions and discrete atom types. The jump processes of generator matching (Holderrieth et al., 2024) are particularly natural for this problem, as they could allow the flow to jump between Stiefel manifolds of different sizes.
Modeling on real-world chemical data. Since we require only minimal information in the moments and molecular formula, another direction is to incorporate other conditioning information, such as energy and force information, fragments of 2D graphs, dipole moments, or other sources of analytical chemistry data. We can do so through MCMC sampling (Du et al., 2023), guidance (Song et al., 2023a;b; Mardani et al., 2023), or diversity sampling approaches (Corso et al., 2024). This presents an opportunity in mass spectrometry, as the 3D structural information provided by the moments can distinguish molecules which have the same mass. Controllable diversity can also be leveraged here since, in a molecule identification campaign, we can eliminate candidates once experiment confirms that they are not correct. We then want sample structures which are different from these eliminated candidates.
Additional applications of Stiefel Flow Matching, Stiefel generative models can also be applied to other domains of orthogonality constrained data, such as molecular orbitals (Mrovec & Berger, 2021; Aoto & da Silva, 2021), orthogonal neural network weights (Kong et al., 2022), and covariance matrices in neural data (Nejatbakhsh et al., 2024).
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
We thank Luca Thiede, Kevin Xie, Adamo Young, and Karsten Kreis for helpful discussions. Molecular visualizations were created with py3Dmol (Rego & Koes, 2015). This research was undertaken thanks part to funding provided to the University of Toronto's 's Acceleration Consortium from the Canada First Research Excellence Fund in preparing this research were provided by the Acceleration Consortium. A.A.-G. thanks Anders G. Froseth for his generous support. A.A.GG. also acknowledges the generous support of the Canada 150 Research Chairs program.
Computational resources used
We acknowledge the Python community (Van Rossum et al., 1995; Oliphant, 2007) for developing the core set of tools that enabled this work, including Py Py'oorch (Paszke et al., 2019), PyTorch Lightning (Falcon & The PyTorch Lightniing team, 2019), PyTorch Geometric (Fey & Lenssen, 2019b), Geomstats (Miolane et al., 2020; 2024), einops (Rogozhnikov, 2022), RDKit (Landrum et al., 2024), Py3Dmol (Rego & Koes. 2015), Jupyter (Kluyver al., 2016). Matplotlib (Hunter, 2007), seaborn (Waskom, 2021). NumPy (Harris et al., 2020), SciPy (Virtanen et al., 2020), pandas (The pandas development team), pybind (W. Jakob, Rhinelander, D. Moldovan and others, 2017), and Eigen (Guennebaud, Jacob, et al., 2010).
REFERENCES
Pierre Ablin and Gabriel Peyré. Fast and accurate optimization on the orthogonal manifold without retraction. In International Conference Artificial Intelligence and Statistics, pp. 5636-5657.
PMLR, 2022.
Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew Ballard, Joshua Bambrick, et al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, pp. 1-3, 2024.
P-A Absil, Robert Mahony, and Rodolphe Sepulchre. Optimization algorithms on matrix manifolds. Princeton University Press, 2008.
Yuri Alexandre Aoto and Márcio Fabiano da Silva. An optimisation on the grassmannian with applications to quantum chemistry. Proceeding Series of the Bruzilian Society of Computational and Applied Mathematics, 8(1), 2021.
Simon Axelrod and Rafael moannz--mmaarlli.. GEOM, energy-annotated molecular conformations for property prediction and molecular generation. Scientific Data, 9(1):1-14, 2022.
Iskander Azangulov, Andrei Smolensky, Alexander Terenin, and Viacheslav Borovitskiy. Stationary kernels and gaussian processes on lie groups and their homogencous spaces i: the compact case. arXiv preprint arXiv:2208. /4960, 2022.
Christoph Bannwarth, Sebastian Ehlert, and Stefan Grimme. Gfn2-xtb accurate and broadly parametrized self-consistent tight-binding quantum chemical method with multipole electrostatics and dependent dispersion contributions. Journal chemical theory and computation, 15 (3):1652-1671) 2019.
Thomas Bendokat, Ralf Zimmermann, and P-A Absil. A grassmann manifold handbook: Basic geometry and computational aspects. Advances in Computational Mathematics, 50(1):1-51, 2024.
Joey Bose, Tara Akhound Akhound-Sadegh, Guillaume Huguet, Kilian FATRAS, Jarrid Rector- Brocor-Brooks, Cheng- Hao Liu, Andrei Cristian Nica, Maksym Korablyov, Michael M. Bronstein, and Alexander Tong. SE(3)-stochasti flow matching for protein backbone generation. In The Twelfth International Conference on Learning Representations, 2024. URL https: id=kJFIH23hhx
httpps///ppeeeeeieee hwtp//ppeeeeee.eeeeoee net forum?
Gordon G Brown, Brian C Dian, Kevin O Douglass, Scott M Geyer, and Brooks H Pate. The rotational spectrum of epifluorohydrin measured by chirped-pulse fourier transform microwave spectroscopy. Journal of Molecular Spectroscopy, 238(2)2000-12 2006.
11

---

Published as a conference paper at ICLR 2025
Gordon G Brown, Brian C Dian, Kevin O Douglass, Scott M Geyer, Steven T Shipman, and Brooks H Pate. A broadband fourier transform microwave spectrometer based on chirped pulse excitation. Review of Scientific Instruments, 79(5), 2008
Andrew Campbell, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola. Generative flows on discrete state-spaces: Enabling multimodal flows with applications to protein co-design. arXiv preprint arXiv:2402.049997 2024.
Rudrasis Chakraborty and Baba C Vemuri. Statistics on the stiefel manifold: theory and applications. 2019.
Ricky T. Q. Chen and Yaron Lipman. Flow matching general geometries. In The Twelfth International Conference on Learning Representations, 2024. URL net/ forum? htte://oommmmmm77moim ddg
Shixiang Chen, Alfredo Garcia, Mingyi Hong, and Shahin Shahrampour. Decentralized riemannian gradient descent on the stiefel manifold. In International Conference on Machine Learning, pp- 1605. PMLR, 2021.
Austin H Cheng. Alston Lo, Santiago Miret, Brooks H Pate, and Alán Aspuru-Guzik Determining 3d structure from molecular formula and isotopologue rotational spectra in natural abundance with reflection-equivarian diffusion. The Journal of Chemical Physics, 160(12), 2024.
Yasuko Chikuse. Distributions of orientations on stiefel manifolds. Journal of multivariate analysis, 33(2):247-264, 1990.
Yasuko Chikuse. Concentrated matrix langevin distributions. Journal of Multivariate Analysis, 85(2): 375-394, 2003.
Hyungjin Chung, Jeongsol Kim, Michael Thompson Mccann, Mare Louis Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. In The Eleventh International Conference on Learning Representations, 2023. URL /tttt::::///ppeeeeee id=OnD 92G.GTO..
net forum?
Gabriele Corso, Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi Jaakkola. Diffdock: Diffusion steps, twists, and turns for molecular docking. arXiv preprint arXiv:2210.01776, 2022.
Gabriele Corso, Yilun Xu, Valentin De Bortoli, Regina Barzilay, and Tommi S. Jaakkola. Particle guidance: non-i.i.d. diverse sampling with diffusion models. In The Twelfth International Confer- ence on Learning Representations, 2024. URL https: openreviev net / forum?iree ?id= KqbCvIFBY7.
Valentin De Bortoli, Emile Mathieu, Michael Hutchinson, James Thornton, Yee whye Teh, and Arnaud Doucet. Riemannian bcore-based generative modelling. Advances in Neural Information Processing Systems, 35:2406-2422,, 2022.
J. Demaison, James E. Boggs, and Attila G. Császár (ods.). Equilibrium Molecular Structures: From Spectroscopy Quantum Chemistry. CRC Press, Boca Raton, 2011. ISBN 978-1-4398-113244
Prafulla Dhariwal and Alexander Quinn Nichol. Diffusion models beat GANs on image synthesis. In A. Beygelzimer, Y. Dauphin, Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021.
Sander Dieleman, Laurent Sartran, Arman Roshannai, Nikolay Savinov, Yaroslav Ganin, Pierre H Richemond, Arnaud Doucet, Robin Strudel, Chris Dyer, Conor Durkan, et al. Continuous diffusion for categorical data. arXiv preprint 1. 1./129., 2022.
Sérgio R. Domingos, Cristóbal Pérez, Mark D. Marshall, Helen O. Leung, and Melanie Schnell. Assessing the performance of rotational spectroscopy in chiral analysis. Chemical Science, 11(4)): 10863-10870, October 2020. ISSN 2041-6539, doi: 10. 10..039/D0SC03752..
Yilun Du, Conor Durkan, Robin Strudel, Joshua B Tenenbaum. Sander Dieleman, Rob Fergus, Jascha Sohl-Dickstein, Arnaud Doucet, and Will Sussman Grathwohl. Reduce, reuse, recycle: Composi- tional generation with energy-based based diffusion models and meme. In International conference on machine learning, pp. 8519-8510. PMLR, 2023.
12

---

Published as a conference paper at ICLR 2025
Alexandre Agm Duval, Victor Schmidt, Alex Hernández-Gariit Santiago Miret, Fragkiskos D Malliaros, Yoshua Bengio, and David Rolnick. Faenet: Frame averaging equivariant gnn for materials modeling. In International Conference Machine Learning, pp. 9013-9033. PMLR, 2023.
Alan Edelman, Tomás A Arias, and Steven T Smith. The geometry of algorithms with orthogonality constraints. SIAM journal on Matrix Analysis and Applications, 20(2):03-3353 1998.
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. arXiv preprint arXiv:2403.0322 2024.
William Falcon and The PyTorch Lightning team. PyTorch Lightning, March 2019. URL https:
//github. com /In wttp.///w...iiiinnnnngtaaeeeeeee
Matthias Fey and Jan Eric Lenssen. Fast Graph Representation Learning with Py Geometric, May 2a. URL https::tt//.... / / thub. com/ httpe//teeetttttteheeeeeeeee
Matthias Fey and Jan Eric Lenssen. Fast graph representation learning with pytorch geometric. arXiv preprint 19XXv:1903.02428 2019b.
Bin Gao, Simon Vary, Pierre Ablin, and P-A Absil. Optimization flows landing on the stiefel manifold. IFAC- PuC--PppesOOnLiee, 55(30):25-300 2022.
Itai Gat, Tal Remez, Neta Shaul, Felix Kreuk, Ricky TQ Chen, Gabriel Synnaeve, Yossi Adi, and Yaron Lipman. Discrete flow matching. arXiv preprint UrXhv:2407.15555, 2024.
Walter Gordy, Rohert I. Cook, and Arnold Weissberger. Microwave molecuaar spectra, ley New York, 1984. URL httpst // app. knovel wctm//////eeoouuceeeeeeeeeee / toc.
spectra, volume 18. Wi-
Gaël Guennebaud, Benoît Jacob, et al. Eigen v3. http://cigen.tuxffaily.rrr 2010.
Charles R. Harris, K. Jarrod Millman, Stéfan J van der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, Robert Kern, Matti Picus, Stephan Hoyer, Marten H. van Kerkwijk, Matthew Brett, Allan Haldane, Jaime Fernández del Río, Mark Wiebe, Pearu Peterson, Pierre
Kevin Sheppard, Tyler Reddy, Warren Weckesser, Hameer Abbasi, Christoph Gohlke, and Travis E. Oliphant. Array programming with NumPy. Nature, 585:357-362, 2020. doi: 10. 0068//41566-00022222
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6855 1851, 2020.
Peter Holderrieth, Marton Havasi, Jason Yim, Neta Shaul, Itai Gat, Tommi Jaakkola, Brian Karrer, Ricky TQ Chen, and Yaron Lipman. Generator matching: Generative modeling with arbitrary markov processes. arXiv preprint arXiv:24 2052887 2024.
Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3d. In International conference machine learning, pp. 8867 8887. PMLR, 2022.
John Hunter. Matplotlib: A 2d graphics environment. Computing in science & engineering, 9(33): 90-95, 2007.
Knut Hüper, Irina Markina, and Fátima Silva Leite. lagrangian approach extremal curves on stiefel manifolds. AIMS, 2021.
Rui Jiao, Wenbing Huang, Peijia Lin, Jiaqi Han, Pin Chen, Yutong Lu, and Yang Liu. Crystal structure prediction by joint equivariant diffusion. Advances Neural Information Processing Systems, 36, 2024.
Bowen Jing, Gabriele Corso. Jeffrey Chang, Regina Barzilay, and Tommi Jaakkola. Torsional diffusion for molecular conformer generation. Advances in Neural Information Processing Systems, 2022.
13

---

Published as a conference paper at ICLR 2025
Jachyeong Jo and Sung Ju Hwang. Generative modeling manifolds through mixture of riemannian diffusion processes. arXiv preprint arxiv:2310.02216, 2023.
Peter E Jupp and Kanti Mardia. Maximum likelihood estimators for the matrix von mises-fishe and bingham distributions. The Annals of Statistics, 7(3):599-66 1979.
Yeonjoon Kim and Woo Youn Kim. Universal structure conversion method for organic molecules: from atomic connectivity to three-dimensional geometry. Bulletin of the Korean Chemical Society, 36(7):1769-177 2015.
Leon Klein, Andreas Krämer, and Frank Noe. Equivariant flow matching. In Thirty-seventh Conference on Neural Information Processing Systems. 2023, URL stttps://opeeeeee net/forum? forum? dd=eLH2NF0011.
Thomas Kluyver, Benjamin Eagan-Kelley Fernando Pérez, Brian Granger, Matthias Bussonnier, Jonathan Frederic, Kyle Kelley, Jessica B Hamrick, Jason Grout, Sylvain Corlay, et al. Jupyter notehooks-a publishing format for reproducible computational workflows. Elpub, 2016:87-90, 2016.
Lingkai Kong, Yuqing Wang, and Molei Tao. Momentum stiefel optimizer, with applications to suitably-orthogonm attention, and optimal transport. arXiv preprint arXiv:2205.14173. 2022.
J Kraitchman. Determination molecular structure from microwave spectroscopic data. American Journal of Physies, 21(1):17-24, 1953.
Qingsi Lai, Lin Yao, Zhifeng Gao, Siyuan Liu, Hongshuai Wang, Shuqi Lu, Di He, Liwei Wang, Cheng Wang, and Guolin Ke. End-to-end crystal structure prediction from powder x-ray diffraction. arXiv preprint arXiv:2401.03822 2024.
Greg Landrum. Rdkit documentation. Release, 1(1-79):4, 2013.
Greg Landrum, Paolo Tosco, Brian Kelley, Ric, David Cosgrove, sriniker, Riccardo Vianello, gedeck, NadineSchneider, Gareth Jones, Eisuke Kawashima, Dan N, Andrew Dalke, Brian Cole, Matt Swain. Samo Turk, Aleksandr Savelev, Alain Vaucher, Maciej Wójeikowski, Ichiru Take, Vincent F. Scalfani, Daniel Probst, Kazuya Ujihara, guillaume godin, Axel Pahl, Rachel Walker, Juuso Lehtivarjo, Francois Berenger, strets and jasondhiggs. rdkit/rdkit: 2024. URL httpt::///// http:://ooooooooooo
February
/ http://eeoooooooo.33... 10633624.
John M. Lee. Smooth Manifolds, 1-29. Springer New York, New York, NY, 2003. ISBN 978- 0-387-21752-9, doi: 10.10077/978.0333--2172222 978-0- 3000-38877222229-3
URL https: http:////do...0/0.01010007/1
Kin Long Kelvin Lee and Michael McCarthy. Study of Benzene Fragmentation, Isomerization, and Growth Using Microwave Spectroscopy. The Journal of Physical Chemistry Letters, 10(10);
April 2019. ISSN 1948-7185. doi: (0.102//acs.jpceet.hhe 9000566....
Kin Long Kelvin Lee and Michael McCarthy. Bayesian Analysis of Theoretical Rotational Constants from Low-Cost Electronic Structure Methods. The Journal of Physical Chemistry A, 124(5): 898-98 910, February 2020. ISSN 1089-5639. doi: 10. (ttllacs.jpca.9b0990
Axel Levy, Eric Chan, Sara Kridovich-Keil, Frédéric Poitevin, Ellen D Zhong, and Gordon Wetzstein. Solving inverse problems in protein space using diffusion- based priors. arXiv preprint 04X2:24006042399 2024.
Derek Lim, Joshua Robinson, Stefanie Jegelka, and Haggai Maron. Expressive sign equivariant networks for spectral geometric learning. In Thirty-seventh Conference Neural Information Processing Systems, 2023. URL https: https::/ opw.revieeeeeeeeeeeeeeeeeeeeeeee net forum? www.rrummiddwwddyyddo UWdy
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference Learning Repre- sentations, 2023. URL https: /ttpp//ppeeeeeee/e//eeeeeeeeeeee .tt.we..weeeeee.e..............e forum? d= dedd==/gvvcDDC....
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.0300 2022.
14

---

Published as a conference paper at ICLR 2025
Aaron Lou, Minkai Xu, Adam Farris, and Stefano Ermon. Scaling riemannian diffusion models. Advances in Neural Information Processing Systems, 36, 2024.
Morteza Mardani, Jiaming Song, Jan Kautz, and Arash Vahdat. A variational perspective solving inverse problems with diffusion models. arXiv preprint arXiv:2305.0439), 2023.
Alan G Marshall and Christopher L Hendrickson. High-resolution mass spectrometers. Anmu. Rev. Anal. Chem., 1 1:579-599, 2008
Simon Mataigne, Ralf Zimmermann, and Nina Miolane. An efficient algorithm for the riemannian log- arithm on the stiefel manifold for family of riemannian metrics. arXiv preprint arXiv:2403.11730, 2024.
Michael McCarthy and Kin Long Kelvin Lee. Molecule Identification with Rotational Spectroscopy and Probabilistic Deep Learning. The Journal of Physical Chemistry A, 124(15):3002-3017, April 2020. ISSN 1089-5639, 1520-5215. doi: 10.102 10.1022/acss /tt.aaapaaa..e01336..
Michael C. McCarthy, Kin Long Kelvin Lee, P. Brandon Carroll, Jessica P. Porterfield, P. Bryan Changala, James H. Thorpe, and John F. Stanton. Exhaustive Product Analysis of Three Benzene Discharges by Microwave Spectroscopy. The Journal of Physical Chemistry A, 124(25):110--5181, doi:
June 2020. ISSN 1089-5639, 1520-5215.
Brett A McGuire. 2018 census of interstellar, circumstellar, extragalactic, protoplanetary disk, and exoplanetary molecules. The Astrophysical Journal Supplement Series, 239(2):17, 2018.
Brett A. McGuire, P. Brandon Carroll, Ryan A. Loomis, Ian A. Finneran, Philip R. Jewell, Anthony J. Remijan, and Geoffrey A. Blake. Discovery of the interstellar chiral molecule propylene oxide (CH3CHCH20), Science, 352(6292); 144--914: June 2016. doi: 10. 126/science.aae0:3. a320328.
Nina Miolane, Nicolas Guigui, Alice Le Brigant, Johan Mathe, Benjamin Hou, Yann Thanwerdas, Stefan Heyder, Olivier Peltre, Niklas Koep. Hadi Zaatiti, Hatem Hajri, Yann Cabanes, Thomas Gerald, Paul Chauchat, Christian Shewmake, Daniel Brooks, Bernhard Kainz, Claire Donnat, Susan Holmes, and Xavier Pennec. Geomstats: python package for riemannian geometry in machine learning. Journal of Machine Learning Research, 21(223):1-9, 2020. URL http: //jmlr. org. www.ooggggppg///eee capprssvv1/// 19 027 html.
Nina Miolane, Luis F. Pereira, Saiteja Utpala, Nicolas Guigui, Alice Le Brigant, Hzaatiti, Yann Cabanes, Johan Mathe, Niklas Koep, elodiemaignant, ythanwerdas, xpennec, tgera168, Christian, Tra My Nguyen, Olivier Peltre, John Harvey, pehauchat, julesdeschamps, Quentin Barthélemy, mortenapedersen. Maya95assal, Abdellaoui-Souhai,, Adele Myers, Felix Ambellan, Florent-Michel, Stefan Heyder, Shubham Talbar, Yann de Mont-Marin, and Marius. geom- stats/geomstats: Geomstats v2.8.0, September 2024. URL https: zenodo. aeroodo.173377 13737807.
org/10.../..ooo 5281/
Benedetta Giannelli Moneta, Sara Elsa Aita, Elena Barbaro, Anna Laura Capriotti, Andrea Cerrato, Aldo Laganà, Carmela Maria Montone, Susy Piovesana, Federico Scoto, Carlo Barbante, al. Untargeted analysis of environmental contaminants in surface snow samples of svalbard islands by liquid chromatography-hige resolution mass spectrometry. Science of the Total Environment, 858: 159709, 2023.
Martin Mrovec and JA Berger. diagonalization free optimization algorithm for solving kohn-shan equations of closed-shell molecules. Journal of Computational Chemistry, 42(7):492-504, 2021.
Amin Nejatbakhsh, Isabel Garon, and Alex Williams. Estimating noise correlations across continuous conditions with wishart processes. Advances in Neural Information Processing Systems, 36, 2024.
Andrey Okhotin, Dmitry Molchanov, Arkhipkin Vladimir, Grigory Bartosh, Viktor Ohanesian, Aibek Alanov, and Dmitry P Vetrov. Star-shaped denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 36, 2024.
Travis E Oliphant. Python for scientific computing. Computing in science engineering, 9(3):10-20, 2007.
15

---

Published as a conference paper at ICLR 2025
Subhadip Pal, Subhajit Sengupta, Riten Mitra, and Arunava Banerjee. Conjugate priors and posterior inference for the matrix langevin distribution on the stiefel manifold. 2020.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. PyTorch: An imperative style, sesh-performacce deep learning library. Advances in Neural Information Processing Systems, 32, 2019.
Aram-Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo-Enrich, Brandon Amos, Yaron Lipman, and Ricky I. Q. Chen. Multisample flow matching: Straightening flows with minibatch couplings. In Andreas Krause, Emma Brunskill, Kyunghyun Cho. Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference Machine Learning. volume 202 of Proceedings of Machine Learning Research, pp. 28100 28100-28127, PMLR, 23-29 Jul 2023. URL https: / /tteee//prroceeeeeee pttpppooooodioeooccooooo eeg.ediddi......e.ee. ngs mlr www...p.es.......... www.pprsssvv000 2012220/pooolaloo anbooladiaan2aa an23a html.
Omri Puny, Matan Atzmon, Edward J. Smith, Ishan Misra, Aditya Grover, Heli Ben-Hame and Yaron Lipman. Frame averaging for invariant and equivariant network design. In International Conference on Learning Representations, 2022. URL https: http:///pppeeeeeeeeeeeeeee net. ntt//fruuttttt forumm?
/
1d=z2Uyy55nXRR.
Cristina Puzzarini and John F Stanton. Connections between the accuracy of rotational constants and cquilibrium molecular structures. Physical Chemistry Chemical Physics, 25(3):1421-1429, 2023.
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):1-7,2014
Nicholas Rego and David Koes. 3Dmolljs: Molecular visualization with WebGi... Bioinformatics, 31 12224 2015.
Eric A Ricsel, Tsach Mackey, Hamed Nilforoshan, Minkai Xu, Catherine K Badding, Alison B Altman, Jure Leskovec, and Danna E Freedman. Crystal structure determination from powder diffraction patterns with generative machine learning. Journal of the American Chemical Society, 2024.
Alex Rogozhnikov. Einops:
Clear and reliable tensor manipulations with einstein-like notation. In International Conference on Learning Representations, 2022. URL forum? htt:://oommmdddoddadoaee id==oppsss22bbb...
RS Ruoff, TD Klots, T Emilsson, and HS Gutowsky. Relaxation conformers and isomers in seeded supersonic jets of inert gases. The Journal of chemical physics, 93(5):3142-311 1990.
Marcus Schwarting, Nathan A. Seifert, Michael J. Davis, Ben Blaiszik, Ian Foster, and Kirill Prozument. Twins in rotational spectroscopy: Does rotational spectrum uniquely identify molecule?, April 2024.
Steven Shipman, Justin Neill, Richard D Suenram, Matt T Muckle, and Brooks Pate. Structure determination of strawberry aldehyde by broadband microwave spectroscopy: Conformational stabilization by dispersive interactions. The Journal of Physical Chemistry Letters, 2(5):443-448, 2011.
Jiaming Song, Arash Vahdat, Morteza Mardani, and Jan Kautz. Pseudoinverse-guidded diffusion models for inverse problems. In International Conference on Learning Representations, 2023a. URL
net wtw.eee...neffttttteeeeeeoe id= 9_MMAAMKKMMM.M
Jiaming Song, Qinsheng Zhang, Hongxu Yin, Morteza Mardani, Ming- Yu Liu, Jan Kautz, Yongxin Chen, and Arash Vahdat. Loss-guided diffusion models for plug-and-play controllable generation. In International Conference on Machine Learning, pp. 32483-32498. PMLR, 2023b.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:201 134466, 2020.
16

---

Published as a conference paper at ICLR 2025
Yang Song, Liyue Shen, Lei Xing, and Stefano Ermon. Solving inverse problems in medical imaging with score-based generative models. In International Conference on Learning Representations. 2022. URL https: wwwp//peeevvieeee wwwp:/wwwee.....eeeeeeeeeeeeee wtt.:etetooomm.ddaaaaatayaa
Yuxuan Song, Jingjing Gong, Minkai Xu, Ziyao Cao. Yanyan Lan, Stefano Ermon, Hao Zhou, and Wei- Ying Ma. Equivariant flow matching with hybrid probability transport for 3d molecule generation. Thirty-seventh Conference on Neural Information Processing Systems, 2023c.
Reilly Sonstrom, Zoua Pa Vang, Haley N. Scolati, Justin L. Neill, Brooks H. Pate, and Joseph R. Clark. Rapid Enantiomeric Excess Measurements of Enantioisotopomers by Molecular Rotational Resonance Spectroscopy. Organic Process Research & Development, ((07):1185-113) July 2023. ISSN 1083-6160. doi: 10. (0.1021/accs.oorr330000022
Jean-Thibaut Spaniol, Kin Long Kelvin Lee, Olivier Pirali, Cristina Puzzarini, and Marie-Aline Martin-Drumed A rotational investigation of the three isomeric forms of seanoethhnnylbenzeen (HCC-C H -N): Benchmarking experiments and calculations using the "Lego brick" approach. Physical Chemistry Chemical Physics, 25(8):6397-6405 2023. doi: 10.1039/D2CP448255
Hannes Stark, Bowen Jing, Chenyu Wang, Gabriele Corso, Bonnie Berger, Regina Barzilay, and Tommi Jaakkola. Dirichlet flow matching with applications to dna sequence design. arXiv preprint arXiv:2402.05841, 2024.
The pandas development team. pandas-dew/pandas:: Pandas. URL httph:://gtttubb.... pandas- setrdassdev/par /pee/saands
Alexander Tong, Kilian FATRAS, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector- Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research, 2024. ISSN 2835- 8856. URL https: : / /ppprevveeewweeee net/ forum? (ccrummmd=CDDScTT 3AW. Expert Certification.
Guido Van Rossum, Fred L Drake, et al. Python reference manual, volume 111. Centrum voor Wiskunde en Informatica Amsterdam, 1995.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, F. ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus. S. Vishwanathan, and Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Pauli Virtanen, Ralf Gommers, Travis F. Oliphant, Matt Haberland, Tyler Reddy, David Cournapeau, Evgeni Burovski, Pearu Peterson, Warren Weekesser, Jonathan Bright, Stéfan J. van der Walt, Matthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. J. Nelson, Eric Jones, Robert Kern, Eric Larson, C J Carey, ilhan Polat, Yu Feng, Eric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, E. A. Quintero, Charles R. Harris, Anne M. Archibald, Antônio H. Ribeiro, Fabian Pedregosa, Paul van Mulbregt, and SciPy 1.0 Contributors. SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python. Nature Methods, 17:261-274 2020. doi: 10. (.ddi..org/1i 10.00..03//4440 011-220000000-2----111111---1-10
URL https:
Natalaja Vogt, Jürgen Vogt, and Jean Demaison. Accuracy of the rotational constants. Journal of molecular structure,
W. Jakob, J. Rhinelander, D. Moldovan and others. pybind!1 Seamless operability between C++11 and Python, 2017.
Zhichao Wang and Victor Solo. Particle filtering on the stiefel manifold with optimal transport. In 2020 59th IEEE Conference on Decision and Control (CDC), pp. 4111-4116. IEEE, 2020.
Michael L. Waskom. Seaborn: statistical data visualization. Journal of Open Source Software 6(60): 3021, 2021.
Caroline C Womack, Marie-Aline Martin- Drummn, Gordon G Brown, Robert W Field, and Michael C McCarthy. Observation of the simplest criegee intermediate ch2oo in the gas-phas, ozonolysis of ethylene. Science advances, 12)xe1400105, 2015.
17

---

Published as a conference paper at ICLR 2025
Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. arXiv preprint arXiv:2203.02923, 2022.
Ryoma Yataka, Kazuki Hirashima, and Masashi Shiraishi. Grassmann manifold flows for stable shape generation. Advances in Neural Information Processing Systems, 36:72377-72411, 2023.
Lia Yeh, Lincoln Satterthwaite, and David Patterson. Automated, context-free assignment of asym- metric rotor microwave spectra. The Journal of chemical physics, 150(20), 2019.
Jason Yim, Brian L Trippe, Valentin De Bortoli, Emile Mathicu, Arnaud Doucet, Regina Barzilay, and Tommi Jaakkola. Se (3) diffusion model with application to protein backbone generation. arXiv preprint arXiv:2302.02277, 2023.
Daniel P. Zaleski and Kirill Prozument. Automated assignment of rotational spectra using artificial neural networks. The Journal of Chemical Physics, 149(10):104166, September 2018. ISSN 0021-9606. doi: N0.1063/1.5033711
Yuchen Zhu, Tianrong Chen, Lingkai Kong, Evangelos A Theodorou, and Molei Tao. Trivialized mo- mentum facilitates diffusion generative modeling on lie groups. arXiv preprint arXiv:2405. 16381, 2024.
Ralf Zimmermann and Knut Hüper. Computing the riemannian logarithm on the stiefel manifold: Metries, methods, and performance. SIAM Journal on Matrix Analysis and Applications, 43(2): 953-980, 2022.
18

---

Published as a conference paper at ICLR 2025
A ROTATIONAL SPECTROSCOPY & MOLECULAR STRUCTURE
A.1 NOTATION
(Px, ,y, P2) are really called the planar moments of inertia. We refer to them as the moments of inertia as shorthand. The actual moments of inertia (IA, IB. Ic) are uniquely related to the planar moments by a simple linear transformation (Kraitchman, 1953).
A.2 INERTIA EDGE CASES
We assume that the molecule of interest has moments of inertia Px > Py Py> 22>00 0. or in other words, is nonplanar asymmetric rotor. This assumption holds for the vast majority of molecules. We now discuss edge cases, such as perfectly symmetric, planar, or linear molecules. It is worth noting that, owing to their rarity and symmetries, these edge cases have significant overlap with the set of molecules that have already been studied (McGuire, 2018).
When structures have two equal eigenvalues (Py- P2z) their inertia matrix, there is no longer unique choice of these two principal axes. These axes now sweep out a plane of possibilities, and numerical diagonalization will arbitrarily pick two orthogonal axes from this plane. But, making an arbitrary choice does not break the mapping between X and U in Equation (4). The only issue is that to respect this additional symmetry, the flow should invariant to in-plane rotoreflections the molecule. However, these examples are so rare that they can be ignored: 76 examples in QM9 and 8 examples in GEOM. If needed, this symmetry can be handled using data augmentation.
Stiefel Flow Matching cannot handle exactly planar (Pz 0) and exactly linear (Py Pz 0) molecules due to divide by zero in Equation (4). For this reason, 3109 examples are removed from QM9. The vast majority of "planar" molecules are actually slightly nonplanar and therefore pose issue. In rare cases where molecules are truly planar, Stiefel Flow Matching can be reformulated for 1 and 2 dimensions.
A.3 EXPERIMENTAL WORKFLOW
For an unknown molecule, we assume that its molecular formula be measured by high-resolution mass spectrometry (Marshall & Hendrickson, 2008) and that its moments can be measured from rotational spectroscopy (Gordy et al., 1984). When not available, molecular formula may be guessed by brute force, informed by the size of the moments.
At a high level, rotational spectroscopy observes how molecules freely rotate in the gas phase. The rotation of molecules is quantized, giving rise to discrete set of rotational states. Molecules can absorb or emit radiation characteristic wavelengths transition between these energy levels. Rota- tional spectroscopy measures the energies of these transitions. A broadband microwave spectrometer can simultaneously measure thousands of these transitions, producing a spectrum of many sharp peaks (Brown et al., 2006; 2008). Rotational transition energies are in the microwave to far infrared region, which is why rotational spectroscopy is also known as microwave spectroscopy.
For asymmetric rigid molecule, and neglecting effects like centrifugal distortion and hyperfine structure, the molecule's rotational energy levels are essentially determined by three unique rotational constants, A(BC), A > B > C. Rotational constants inversely proportional the principal A²
moments of inertia, i.c., A = 27Ą Each energy level is also indexed by quantum numbers. Transitions are the differences in these energy levels. The molecule must have an appreciable dipole moment for these transitions to be measured.
Once a spectrum measured, the rotational spectroscopist tasked with assigning each transition to its quantum numbers and ultimately assigning rotational constants A A(BC). Spectral assignment is challenging problem tackled in other works (Zaleski & Prozument, 2018; Ych et al., 2019). Rotational constants then simply be inverted to obtain effective moments of inertia.
Using these effective moments, a spectroscopist can now search for the true structure using any of the models developed in this work. In an actual structure elucidation campaign, only few targets are considered, which permits querying K > 1000 samples for each target, and also leaves enough computational resources to evaluate the stability of every generated sample by quantum chemistry.
19

---

Published as a conference paper at ICLR 2025
A.4 EXPERIMENTAL PRECISION
While the proposed method generates structures which satisfy the moment constraints exactly, and while rotational spectroscopy can measure experimental rotational constants many digits of precision (Vogt et al., 2011), 1), it is unfortunate that the experimental rotational constants do not directly translate to moments inertia. This because molecules are not perfectly rigid: Experiment observes properties that have been vibrationally averaged, including the rotational constants. Conformational fluctuations such as torsions 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. As a result, the experimental rotational constants A(BC)a are proportional to (p))) whereas equilibrium rotational constants A(BC), are proportional to (r) where (-) denotes vibrational average.
Structures in QM9 and GEOM have been geometry optimized to reach equilibrium structures (r)- This error due to zero-point vibration effects is the major source of uncertainty between equilibrium and experimental rotational constants, the order of 1% relative error (Vogt al., 201 1; Puzzarini & Stanton, 2023).
Typically, experimental rotational constants ABC)) be corrected into equilibrium rotational constants A A(BC)e by a rovibrational calculation. Experimental and equilibrium rotational constants are related by perturbative expansion (Demaison et al., 2011)
A(BC), A(BC),
1 Σ A(BC)
A(BC) 16s
(12)
2
where j index normal modes of vibration (i.j {1,2, ................................ 3n 6} for n atoms), and and y are first and second order interaction constants that couple rotational and vibrational motions together. As Taylor series, is generally much larger than %, y is usually neglected. The rovibrational correction a can be calculated using an electronic structure method with a good compromise between computational time and accuracy (Puzzarini & Stanton, 2023; Spaniol et al., 2023). However, this calculation requires knowing the structure in the first place.
Alternatively, experimental rotational constants be approximately corrected equilibrium rotational constants by simple empirical scaling factors (Lee & McCarthy, 2020).
But, exact moment constraints do allow application in the following sense: if one were a priori guess the equilibrium moments correctly, one could then verify whether they are indeed correct by generating structures, calculating their rovibrational corrections, and then checking their agreement to the experimental rotational constants. Therefore, given experimental rotational constants A(BC"). one can use this verification procedure constants A ABCC).. This is feasible since A(BC), consist of only 3 numbers. Experimental precision in A(BC)o is maintained up to the precision in computing a. To verify that structure the true structure, we must know that (1) its A(BC), and match A(BC), and (2) its gradient norm is O.
fine-grid search for the true equilibrium rotational
B THE STIEFEL MANIFOLD
In this section, we discuss various facts about the Stiefel manifold and computations thereon. We refer readers to (Lee, 2003; Edelman et al., 1998; Bendokat et al., 2024) for further details.
The Stiefel manifold (t. is the set of rectangular orthonormal matrices of shape n x pt
St(n, p) (U E Raxp UT== Ip).
(13)
St(n, p) is manifold of dimension np =p(p + 1). An element of St(n,p) can be thought of as collection of p orthonormal r--dimensional vectors (i.e. p-frame living in r--dimensional space), or as the first columns of an n X orthogonal matrix. In fact, St(n. p) generalizes some well-known spaces. For example, St(n, n)) is the orthogonal group O(n), and St(n,, n 1) is diffeomorphic to the special orthogonal group SO(n), and St(n, 1) recovers the unit nt-phere 5n-1
B. TANGENT SPACE
The tangent space of St(n, p) a point U is identified with the subspace
StSt((n p) (A E RMxP| UTA A + A A - 0}. In other words, this is the set of p matrices A for which U A is skew-symmetric.
(14)
20

---

Published as a conference paper at ICLR 2025
B.2 CANONICAL METRIC
Smoothly equipping every tangent space with an inner product turns St(n,p) into a Riemannian manifold. We exclusively consider the canonical metric,
(A, = trace T (In (In-JUU
(15)
Any inner product induces a norm by ||al|o VIA, A)u. For the canonical metric, the squared norm can be rearranged as:
| - A) trace A A trace(U A) A -
VIIUT|13S (16)
where || || is the Frobenius norm. Computing the canonical norm in this way is much faster for the node-wise graph batching approach used by PyTorch Geometric (Fey & Lenssen, 2019a).
B.3 EXPONENTIAL AND LOGARITIMM
The Stiefel exponential be computed by the algorithm presented by Edelman et al. (1998), who give the closed-form expression for geodesic with initial conditions 7()) = U and 7(0) = A. Here, reproduce the algorithm for computing expu (A) 7(1).
Algorithm 1 Computing a Stiefel geodesic 7(t). (Edelman al., 1998)
Require: Base point U Stnn, p). initial direction E TySt(n, p). time 1: QR f- (In UU ), where Q E Roxp and R E Rpxp
QR decomposition
2: A - U A
R
E R2px2p
R
0
(()
3:
+ expm (tA), where M(t),N( E R
> take submatrices
(()
4: return (M + Qt(() E St(n.
An efficient algorithm for computing the Stiefel logarithm is reproduced and simplified here from Zimmermann Hüper (2022). We implement the logarithm in C++ for speed on the CPU, must be called on every fetch of a data example. This does not present a bottleneck: computing logarithmic maps takes an average of 0.1 ms for molecules both QM9 and GEOM. The logarithm has cost of O(nz²). The thin QR decomposition and initial matrix multiplications have cost of O(np²), while the inner loop's 's Schur decomposition, Sylvester solve, matrix multiplications, and matrix exponential have cost of (pp)). We only execute the inner loop maximum of 20 times (see Appendix E for convergence).
Algorithm 2 Computing the Stiefel logarithm. (Zimmermann & Hüper,, 2022)
Require: Base point U E St(n,,) p), target point U E St((n,))
1: M + U E Rexp
2: QN + UM E Raxp
thin QR
M
3: OR N
2pxp
O E
R2
compute orthogonal completion via full QR
M
4: V
N
.p:pp
ESO(2p)
flip sign if det
5: for k
1....
20 do
6:
A
-
B
log (V)
use Schur matrix logarithm
B
C
7:
St
S
BBB
1
11
12
8:
solve C ST TS for I E RP*P
symmetric Sylvester equation > matrix exponential, or use Cayley transform rotate last p columns of V
9:
$ # exp, (T) 12px
10:
V. V.p.pp + V ppp2 $
11:
return U QB E St,
21

---

Published as a conference paper at ICLR 2025
B.4 UNIFORM SAMPLING
To sample uniformly on St(n,,) p) with respect to the Haar measure, we can compute U Z Z)-1/2 for a random matrix Z E Raxp whose elements are drawn i.i.d. from standard Gaussian (i.e. Z = randn (n, p)).
B.5 HOUSEHOLDER REFLECTIONS
To align the final column of a matrix U E St(n,,)) unit vector y. we can rotate its columns under a transformation E SO(n) such that RU_, = y. Let U. æ. For any unit vector E S the Householder matrix (v) I, 2ve has determinant Then,
y R H a + yll
E (O(m)
(17)
is our desired rotation. In particular, verify that
+ y x + yl|
(2 2x y)
H
y) -0-1x+ ()+
(18)
so that Rx= (z)((()) y. desired.
B.6 ORTHOGONAL. PROJECTIONS
The orthogonal projection of point Un E Rux onto St(n, p) is special case of the well-known orthogonal Procrustes problem:
arg UI, Uolle, subject U U Ip-
(19)
Letting Uo = AEB T under a singular value decomposition, the solution to this problem is AB The orthogonal projection of point Z E Tx onto the tangent space St, p) is computed by Z 2-- U syymnU U Z), where sym( AA)
2 (A + A is symmetrization operation.
B.7 THE ZERO CENTER MENTER-OF-AEE MASS SUBMANIFOLD
For 2 5, we are interested in the subset St,n, 4) obtained by fixing the final column fixed unit vector E 5m-1;
= = (U E St(n, p) U.,1 = a).
(20)
Note that S n - ({a)) a level set the projection map
:: Sn 4) + 5n U+
(21)
In fact, æ is a smooth surjective map of constant rank, so a submersion by the global rank theorem. Hence, M is embedded submanifold by the submersion level set theorem. The tangent vectors to M are exactly those whose final column zero, since
TrMM ker dæ (A St 4) 10 =0.
(22)
If M further inherits the canonical metric from St(n, 4), then it becomes a Riemannian manifold. Note that M is homeomorphic to St(n 1, 3), so is connected and compact. Hence, is geodesi- cally complete and any two points in M can be connected with Length-minimizinn geodesic on M. Theorem shows that M is totally geodesic, i.e., any geodesic on MM a geodesic on St(ni, 4). This allows perform simpler computations in the ambient space. Theorem 5 computes projections of arbitrary matrices onto TuM
Theorem 1. As defined above, M is totally geodesic.
Proof. A sufficient condition is that for any U E M and A TuM, the geodesic y with initial conditions 70) = and (0)) = A stays within M. Fortunately, is given in closed-form by Edelman et al. (1998) in Algorithm 1. Since A = the fourth columns of R and U A are zero. Since U A is skew-symmetric by Equation 14, its fourth row is also zero. Then, the fourth and
22

---

Published as a conference paper at ICLR 2025
column of A and expm (t.) zero, but expm((t)4., 1. It follows that (t)., (0,0,0,1) and (t)., = O, so that
(M
(23)
Hence, 7(1) E M as desired.
We now provide some lemmas which are useful for proving that the Stiefel logarithm on M can be computed using St 3) rather than St,n, 4) (Theorem 4).
Lemma 2. Let X be square matrix whose tth row and column are O. Positive powers of X retain zeros in the 8th row and column. Specifically, suppose
E
F
E
F H
== 0 0 0
(24)
G
H
G
G
GOH
Then we have
0 F
X k
0 0
(25)
0 H
Proof. We
show this by induction over k. It is clear this is true for k 1. Then, assume
E
F
X
0
0
and aim to show for k + 1. We know that
0
H
E
F
F
E
F H
ÉE FG ÉF FH
(26)
G
H
G
H
G
GE+HG HG GF HH
At the time, have that
E
F
E
F
FG 0
X k++
0 0
0
0
0
0
=
0 0 0
(27)
G
0
H
G
H
GE+HG
0
GF+HN HH
which completes the proof.
Lemma 3. Let X be a square matrix whose ith row and column are O. The matrix exponential "ignores" these zeros. Specifically, suppose
E
0
E
F
X
0
expm
G
H
G
H
(28)
G
0
H
Then we have
E
0
F
exp, (X)
0 1 0
(29)
0
H
Proof. The matrix exponential is given by
exp m (X)
=
x
(30)
k!
km
By Lemma 2, all non-identity terms of this series are identical to terms in the series for E F
H but with zeros inserted in the tth row and column. The identity term then contributes
expm. G
the extra in the main diagonal.
23

---

Published as a conference paper at ICLR 2025
Theorem 4. Let M, and let U be its first three columns so that U (Ü a).
The exponential on M can be computed as the exponential on St(n, 3) after discarding the last column of A A 0).
expu ( (cxPo a)
(31)
Similarly, the logarithm on M can be computed as the logarithm on St(n, followed by concatenat- ing a zero column:
logu, (U1) (logrr, (Ü1) 0)
(32)
Proof. We go line-by-ling bine through Algorithm to show that the Stiefel exponential on M is equivalent to the exponential for St(n, 3).
The last column of (In UU ) is because the last column of A is 0. Then,
QR -L.-UU (Q q)
R 0 0
(33)
where q is some orthogonal vector the 3 columns of Q and It is the first 3 columns and of R. This corresponds to the QR decomposition QR (I, ÜÜ ÜÜT) )
In addition, we have that UTA Ü 0
0
Thus, the block matrix A in Algorithm is given as
0
R
0
0
0
0
0
(34)
0
0
0
0
By Lemma 3, the matrix exponential of "ignores" the extra rows and columns of zeros. Therefore,
M()
0
R
0
(() 1"". N(t) ....
0
0
expm.(tA) expm
0
0
(
0
0
0
0
(35)
where tilde terms are equal to their counterparts in the exponential of St(n, 3).
Finally, the output of Algorithm is
+ = (Ú a) (M()) - 9) ((())
0
(36)
= (ŪM(( (QN(() 0)
(37)
=(UM() QN(() a) (AXPo( a).
(38)
(39)
Since the Stiefel exponential is locally invertible, this also shows that the Stiefel logarithm can be computed using Uo and Ü,
Theorem 5. Let U E M, and let U be its first three columns so that U Ü a). Given E RAx3
the minimum-norm projection of 0) onto TyM is given by [ Z) 0), where
T (Z) = skew TT
( + (I, - T Z E p
(40)
Ú skew Ü 7 (I, UU Z gax3
(41)
Note there is no tilde in the second term.
24

---

Published as a conference paper at ICLR 2025
Proof. We can rewrite
TuM = [[A0] Ü A 0, A,,
(42)
= {[āo] Á E St.St 3) n) }
(43)
Edelman et al. (1998) give an orthogonal projection of Z onto StSt n, 3) as:
(Z) Ú skew(U Z (I, ÚÚ )2
(44)
where skow (A) = (4 A A) The orthogonal projection of Z onto span(a) is T2(Z)= (I, aa )z
(45)
We can check that иT and m2 commute with π = π o 122 = 22 mi Hence, π is an orthogonal projection onto St,,),,) n span(a)+ as desired.
C EXPERIMENTAL DETAILS
C.I ARCHITECTURE
Explicitly, the neural network takes in moments Pr Pz), time atom types a, and coordinates X, and outputs a Stiefel tangent vector. Moments are embedded using sinusoidal features with wavelength geometrically spaced from 0.0001 10,000. Time is similarly embedded but with wavelength range from 0.001 to 1. We use reflection reflection-equiiariaatt network. Note that given infection-invariann function f. the mapping X + sign( (X) f(X) is reflection -quuiaionn-quivarriant,, where sign(X X) gives the element-win wisee signs of X, with (0) = O. Thus, the problem is reduced to constructing a network that reflection-invariant with respect to the input coordinates X.
As input to our network, we begin by featurizing the molecule in obtain invariant node features h E Robe by using the molecule's ' unsigned coordinates and atom types, and edge features Eij E Rdaty are computed from the unsigned differences between pairs of atomic coordinates. These features are passed through a Transformer backbone (Vaswani et al., 2017). We use the PreLN layout with an adaptive version of LayerNorm (Dieleman et al., 2022: Dhariwal & Nichol, 2021) that conditions on the timestep and molecule's ' moments. In addition, the attention module is replaced with message-passing block that jointly updates the node and edge features:
manner. We
(Vij, aij. e + MLP(h;, hjs eij); for all i, Bg t Bij + tely for all i,
(46)
(47)
y: + Σ exp(aa)) j=1 k=1 exp(axx) Vij. for all
(48)
h; + h, + Linear(y:) for
(49)
The last two equations are reminiscent of self-attention in Transformers, and along the same lines, we multi-heade extension of them. Across all experiments, We use 16 Transformer-lik blocks, 12 update heads, SiLU/Swish activations, and a 4x expansion in each block's feed-forward module. In total, model has 1 154M trainable parameters.
dnode 768 and dedge 192.
This architecture is distinct from the pretrained network of KREED (Cheng et al., 2024), which was tailored for predicting 3D structure given molecular formula, moments of inertia, and substitution coordinates. KREED was trained using random dropout of input substitution coordinates. For QM9, the model sometimes observed examples with substitution coordinates during training. However, for GEOM, the model always received at least some substitution coordinates during training. Therefore, KREED is operating -ut-of-distribution when provided with no substitution coordinates all on GEOM.
C.2 STOCHASTICITY
Diversity important for identifying unknown molecules outside the training set. Therefore, we experiment with adding stochasticity to the dynamics of the flow, leading individual paths to be stochastic. To do so during training, we apply an exponential map to a Gaussian variable with noise scale y. cos(x()+1 in the tangent space of the interpolant before calculating Ú again. During sampling, 2 add Gaussian noise of the same noise schedule to the tangent vector at every time step.
25

---

Published as a conference paper at ICLR 2025
Table 3: General training and sampling hyperparameters.
Hyperparameter
QM9
GEOM
Epochs
1000
60
Batch size per GPU
256
24
Optimizer
Adam
Adam
Learning
10-1
10-1
Training
Learning rate warmup steps
2000
2000
Weight decay
0.01
0.01
Gradient clipping
yes
yes
EMA decay
0.9995
0.9995
Timesteps
1000
1000
KREED
Schedule
polynomial
polynomial
Stiefel FM
Timesteps
200
200
Table Training and sampling hyperparameters for Stiefel Flow Matching.
Dataset
Model
Timestep sampling
OT
stochasticity 7
Stiefel FM
uniform
no
0.00
Stiefel FM-OT
uniform
yes
0.00
QM9
Stiefel FM-OT-stoch
uniform
yes
0.10
Stiefel FM-In
logit-normal
no
0.00
Stiefel FM-In-OT
logit-normal
yes
0.00
Stiefel FM
uniform
0.00
GEOM
Stiefel FM-OT
uniform
yes
0.00
Table 5: Extended ablation study for QM9, Stochasticity and logit-normal timestep sampling do not provide clear advantage.
%
RMSD
Method
Error
Valid
Stable |
Diverse t
NFE J
0.25 A
0.10 Á
Stiefel FM
15.17 0.31
13.82 0.30
0.00
0.882
-1.125
1.040
200
Stiefel FM-OT
13.99 + 0.30
12.68 + 0.29
0.00
0.835
-1.039
1.045
200
Stiefel FM-stoch
15.13 0.31
13.83 0.30
0.00
0.877
-1.116
1.045
500
Stiefel FM-In
15.74 0.32
11.45 + 0.28
0.00
0.880
-0.600
0.982
200
Stiefel FM-In-OT
14.90 0.31
12.45 0.29
0.00
0.875
-0.687
1.026
200
Training (min epoch)
Training (it s)
Sampling (seconds K=10 samples)
Method
Dataset
KREED-XL
QM9
1.02
6.7
13.9
Stiefel FM
QM9
1.32
5.1
2.9
Stiefel FM-OT
QM9
3.48
1.9
2.9
KREED-XLL -XL
GEOM
225.6
17.0
71.3
Stiefel FM
GEOM
224.4
17.1
15.0
Stiefel FM-OT
GEOM
229.8
16.7
15.0
Table 6: Training and sampling timing QM9 and GEOM. A CPU bottleneck of computing logarithms and optimal transport exists for QM9 due to a large batch size of 256, but this CPU bottleneck disappears for GEOM due small batch size of 24.
C.3 TRAINING
Models were trained on 4 NVIDIA A100 40GB GPUs. Tables and 4 gives our hyperparameters. We the adaptive gradient clipping strategy from Hoogeboom al. (2022).
26

---

Published as a conference paper at ICLR 2025
C.4 SAMPLING
During sampling, we sample uniformly from M and iteratively query the trained model for tangent vector A at every step. If stochasticity is turned on, Gaussian noise with scale y cos(mt)+1 is added to this tangent vector. At every step, U, is projected onto the manifold (Appendix B.6) before projecting the tangent vector to the tangent space on the manifold St,n,,(n4 4) (Theorem 5). Integration proceeds by applying expu, to this tangent vector, scaled by dt.
2
Algorithm 3 Sampling under Stiefel Flow Matching.
Require: Size n. atom types a, moments Pxyz (Px. Pr, Pz), stochasticity y, timesteps T
1: U ~ Uniform(St(n, 4)) using Appendix B.4
2: Project U to M (defined in Equation 5) using Appendix B.5
3: t 8+0
4: ta + 1/T
5: for step in 1.....
. do do x XU, m, Pxyz) by inverting Equation Á + (t X, Pxyz, a) E Rax3
6:
7:
8:
if step then
9:
A < lA + 7 (cos(TT)) 1)e, where Et NO0 1)
10:
else
11:
A + t
12:
Project A to T using Theorem 5
13:
U + cxpv(()
1+ 1
return X XUU m. Pxyz) by inverting Equation
27

---

Published as a conference paper at ICLR 2025
KREED-NL
Sfee PM
Figure 5: Selected QM9 examples. Best viewed zoomed in. Examples are sorted by RMSD ground truth, which is shown in cyan. Green panels indicate meeting the threshold of 0.25 A.
KREED XL
Stiefel FM -T (fitte)
Figure 6: Selected GEOM examples. Best viewed zoomed in. Examples are sorted by RMSD to ground truth, which is shown in cyan. Green panels indicate meeting the threshold of 0.25 A.
28

---

Published as a conference paper at ICLR 2025
Success (0.25A)
3.0
Success (0.25A)
2.5
False
False
True
2.5
True
2.0
Density
Density
2.0
1.5
1.0
0.5
0.0
1.00
1.25
1.50
1.75
2.00
2.25
2.50
0.5
1.0
1.5
2.0
Curve Length
Curve Length
Figure 7: Normalized histogram of curve lengths of success and failure examples of Stiefel FM the QM9 (left) and GEOM (right) datasets. Success cases are more often the result of longer generation paths.
2500
Stiefel FM
Stiefel FM-OT
Stiefel FM (filter)
2000
Stiefel FM-OT (filter)
1500
Count
1000
0.25
0.50
0.75
1.00
1.25
1.50
1.75
2.00
Curve Length
Figure 8: Equivariant optimal transport shortens generation trajectories for GEOM. However, gener- ating more samples and filtering by validity has slight bias towards longer generation trajectories.
D GREEDY RANDOM OPTIMAL ASSIGNMENT ALGORITHM
The greedy random local search looks for the best reflection and permutation to minimize the Stiefel distance between a given Un and U, d(U,, U1) = || log (Ui)l | This distance is approximated using only one iteration of the logarithm. first decides on the best reflection R by trying many random permutations for each reflection. Then, the best reflection is kept, and the permutation is optimized using random swaps of indices. The computational cost of computing the optimal transport map is O(np?), ), it relies fixed number of computations of the logarithm.
29

---

Published as a conference paper at ICLR 2025
Algorithm 4 Heuristic alignment algorithm.
Require: Uo. 11 St(n,,)) p), atom types a, number of restarts R, local search budget L
1: c" < 20
D lowest cost seen so far find good reflection
2: for reflections R (diag(() & E (-1,+1}33 do
3: for k: = 0, ,..... R do
4: sample random alom-type-preservingg node permutation II
5: Ucand f- R
6:
c d(Uuand: U1). if c³ then
approximate distance
7:
8:
c² t
9:
Unes + Ueand
10:
for k = 0.1...., do
local search over node swaps
11:
sample (i, such that a; = aj. without repeating pairs Ueand swap (ubep i,
12:
d(Uennd UI), if e" then
approximate distance
14:
t c
16:
Ubest f- Ucand
17:
return Uest
interpolate
to X1
Xo
Align to X;
reflect z-axis, nodes
X
Interpolate
X
X;
Figure 9: Equivariant optimal transport aligns noise samples X, to training examples X1 over reflections and permutations, leading to smoother and shorter paths shown to the model during training.
E STIEFEL LOGARITHM EMPIRICAL ANALYSIS
We empirically analyze the convergence of the Stiefel logarithm (Algorithm 2). We sample 100k training examples from QM9 or GEOM to used as Un and for each example sample one random point Ug. We compute the true logarithm A loogr, (U, using large number of iterations and compare it to the 20-iteration truncated logarithm. We record error as the infinity norm of the difference between the true and approximate logarithms. We set the convergence threshold to be le-6.
30

---

Published as a conference paper at ICLR 2025
For QM9, the 20-iteration logarithm converges 97.8% of the time (median iterations converge), and the median error in case of nonconvergence is 2.5e-4. The Spearman correlation between the 1 1-iteration approximate distance and the true distance is p 0.87. For GEOM, the 20-iteration logarithm converges 99.7% of the time (median 9 iterations to converge), and the median error in case of nonconvergence is 8.8e-5. The Spearman correlation between the 1-iteration approximate distance and the true distance is p = 0.83.
These results empirically validate that the I-iteration approximate distance used in Algorithm 4 an upper bound on the true distance, and that it is valid heuristic which generally maintains the same relative ordering as the true Stiefel distance.
16000
20000
14000
12000
8000
10000
4000
28 3M 20--0
20 27 20>30
Number Harations converge
Number Rarations to converge
Figure 10: Histograms of the number of inner iterations of the Stiefel logarithm required to converge to an error of le-6 for 100k random training examples of QM9 (left) and GEOM (right).
Spearman
Spearnsn
3.50
3.25
Nngrosimate Distance
3.0
3.00
lgproeimate
2.75
2.5
2.50
2.25
2.0
1.5
1.5
2.0
2.5
3.0
3.5
1.75
2.00
2.25
2.50
2.75
3.00
3.25
Irue Distance
Figure 11: Parity plots comparing the 1-iteration approximate Stiefel distance the true Stiefel distance of QM9 (left) and GEOM (right).
31


Reviews:
[{"review_id": "StH3dZO99y", "reviewer": "Reviewer_Locv", "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.", "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? \n- 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?\n- 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?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "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. \n- 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. \n- The results show some performance improvements, particularly in maintaining moment constraints and reducing sampling cost.\n- 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.\n- 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. \n- 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.\n- 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.\n- The paper is longer than 9 pages.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qlee9fwxNK", "reviewer": "Reviewer_f4fB", "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.", "questions": "Does the Stiefel manifold have any constraints? How does it handle the increased complexity of larger molecules and the number of conformers?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "b5Cfpdj6CW", "reviewer": "Reviewer_oa92", "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.", "questions": "### Major Comments/Questions: \n\n- In lines 175 to 177, how practical is the assumption $n\\geq 5$? How many examples with $n<5$ are removed from the datasets? \n\n- To compute the loss in (8), the authors listed 4 requirements. All are well explained other than the second. Further clarification is needed here.\n\n- 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. \n\n### Minor Comments/Questions:\n\n- Based on the definition of $\\mathbf{m}$, shouldn't the total number of masses $M$ be equal to $n$?\n\n- Vector $\\mathbf{a}$ is not being used in the problem definition other than specifying the number of atoms. \n\n- 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? \n\n- 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. \n\n- 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]? \n\n- The last sentence in the Datasets paragraph in Section 4 requires support by a citation. \n\n- Why the FM model has more parameters than the diffusion model in Cheng et al. (2024)? \n\n- Ill-sentence in lines 445 to 446.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is generally well-written and easy to follow. The math is sound and clearly presented. \n\n- The use of Stiefel manifold and its connection to the problem and the constraints in (2). \n\n- 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? \n\n- 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. \n\n- 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. \n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XB1JejBQr0", "reviewer": "Reviewer_519h", "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.", "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.\n2. What are the performance of additional baselines? See Weakness 2.\n3. Can you follow the same postprocessing procedure and compare the proposed model with the KREED-filter baseline performance? See Weakness 3 & 4.\n4. 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.\n5. 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)?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "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.\n- 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.\n- 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.\n    - **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.\n    - **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.\n\n2. **The number of baselines compared in the paper was small**. Indeed, only KREED was compared. There are other baselines available:\n    - 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. \n    - 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.\n\n3. **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.\n\n4. **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.\n\n[1] Hoogeboom, Emiel, et al. \"Equivariant diffusion for molecule generation in 3d.\" International conference on machine learning. PMLR, 2022.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)