AssembleFlow: Rigid Flow Matching with Inertial Frames for Molecular Assembly
Authors: Hongyu Guo
Yoshua Bengio
Shengchao Liu
Venue: ICLR 2025 Poster (2025)
Keywords: rigid flow matching
inertial frame
quaternion representation
material
molecular assembly
molecular crystallization
Primary area: generative models
TL;DR: We introduce the first generative framework that enforces molecular rigidity in the assembly of a set of molecules for innovative material design.

Abstract: Molecular assembly, where a cluster of rigid molecules aggregated into strongly correlated forms, is fundamental to determining the properties of materials. However, traditional numerical methods for simulating this process are computationally expensive, and existing generative models on material generation overlook the rigidity inherent in molecular structures, leading to unwanted distortions and invalid internal structures in molecules. To address this, we introduce AssembleFlow. AssembleFlow leverages inertial frames to establish reference coordinate systems at the molecular level for tracking the orientation and motion of molecules within the cluster. It further decomposes molecular $\text{SE}(3)$ transformations into translations in $\mathbb{R}^3$ and rotations in $\text{SO}(3)$, enabling explicit enforcement of both translational and rotational rigidity during each generation step within the flow matching framework. This decomposition also empowers distinct probability paths for each transformation group, effectively allowing for the separate learning of their velocity functions: the former, moving in Euclidean space, uses linear interpolation (LERP), while the latter, evolving in spherical space, employs spherical linear interpolation (SLERP) with a closed-form solution. Empirical validation on the benchmarking data COD-Cluster17 shows that AssembleFlow significantly outperforms six competitive deep learning baselines by at least 45\% in assembly matching scores while maintaining 100\% molecular integrity. Also, it matches the assembly performance of a widely used domain-specific simulation tool while reducing computational cost by 25-fold.

Published as a conference paper at ICLR 2025
ASSEMBLEFLOW: RIGID FLOW MATCHING WITH INER- TIAL FRAMES FOR MOLECULAR ASSEMBLY
Hongyu Guo
Yoshua Bengio
National Research Council Canada University of Ottawa hongyu. tawa ca
Mila Québec AI Institute
Université de Montréal
CIFAR AI Chair
yoshua www.bbennniiiiii.m waw.ioimmmaa....mm
Shengehao Liu
Université de Montréal shengchao. umint real ca
ABSTRACT
Molecular assembly, where cluster of rigid molecules aggregated into strongly correlated forms, is fundamental to determining the properties of materials. How- ever, traditional numerical methods for simulating this process are computationally expensive, and existing generative models material generation overlook the rigidity inherent in molecular structures, leading to unwanted distortions and in- valid internal structures molecules. To address this, introduce AssembleFlow. AssembleFlow leverages inertial frames to establish reference coordinate systems at the molecular level for tracking the orientation and motion of molecules within the cluster. It further decomposes molecular SE(3) transformations into translations in IR and rotations in SO(3), enabling explicit enforcement of both translational and rotational rigidity during each generation step within the flow matching framework. This decomposition also empowers distinct probability paths for each transforma- tion group, effectively allowing for the separate learning of their velocity functions: the former, moving in Euclidean space, uses linear interpolation (LERP). while the latter, evolving in spherical space, employs spherical linear interpolation (SLERP) with a closed-form solution. Empirical validation on the benchmarking data COD- Clusterl7 shows that AssembleFlow significantly outperforms six competitive deep learning baselines by at least 45% in assembly matching scores while maintaining 100% molecular integrity. Also, it matches the assembly performance of widely used domain-spectt simulation tool while reducing computational cost by 25-fold.
INTRODUCTION
Deep learning methods have been revolution- izing scientific research across various do- mains, enabling breakthroughs in fields such drug discovery (Yu al., 2024), material science (Merchant et al., 2023), and molecular design (Loeffler et al., 2024). For instance, protein folding systems have demonstrated un- precedented accuracy and creativity in design- ing protein structures (Jumper et al., 2021; Back et al., 2021), driving innovation in drug discovery. These advancements underscore the transformative potential of machine learn- ing in tackling complex scientific problems.
Initial Structure Weakly-comellttd
Firel Structure Strongly-corrlattd
Figure 1: Illustration of assembly a cluster of three molecules transitioning from a weakly correlated structure (left) to strongly correlated crystal structure (right). key challenge for existing generative models in material generation preserving the rigidity of each molecule throughout this transformation 3D space.
Molecular assemble or crystallization one such complex process where rigid molecules transition from weakly correlated arrangement highly ordered, strongly correlated structure. During this

---

Published as a conference paper at ICLR 2025
(a) R3 Euclidean space
(c) 803 rotation sphere
(d) rotation sphere
LERP
(b) R3 Euclidean space
SLERP
LERP
Molecule
Mass center of molecule
Principal axes of inertial frame
Figure 2: (a. b): The molecule's center of mass (CoM) is used the reference point for translation and interpolation in Fuclidean space. (c, d): An inertial frame serves as the reference coordinate system, tracking rigid molecular rotation- rotation operations are implemented using quaternion representation- and enabling interpolation in spherical space.
process, each molecule approximated maintain its shape and structure unchanged as it moves in the 3D space, as illustrated in Figure 1. Crystallization plays pivotal role in determining the physical properties of materials, including their mechanical strength, electrical conductivity, and thermal stability (Porter et al., 2009; Carter & Norton, 2013), making it key process in material science (Ashby & Jones, 2012), pharmaceuticals (Hilfiker, 2006), and nanotechnology (Gonsalves et al., 2000). For example, the crystalline form of drug can affect its solubility and bioavailabil- ity (Byrn et al., 1999; Healy et al., 2017); Similarly, precise control over molecular arrangements is key to optimizing the electronic and catalytic performance of organic semiconductors, polymers, and molecular catalysts (Saparov & Mitzi, 2016).
Traditional numerical methods have long been employed to simulate the crystallization pro- cess (Martínez et al., 2009; Van Der Spoel et al., 2005), but they are often computationally expensive and inefficient, limiting their scalability and practical use in large-scale applications. On the other hand, despite the importance of crystallization, existing machine learning methods struggle to capture the physical constraints critical this process. A major limitation is the failure account for the inherent rigidity of molecular structures during crystallization (Liu et al., 2024c), often leading to un- wanted distortions and invalid internal structures, i.e., non-rigid molecules. an assembly, molecules must retain their rigid atomic structures, as this rigidity is essential for producing meaningful packing arrangements. However, current generative models for molecular crystallization treat molecules as flexible entities (Liu et al., 2024c), resulting in physically unrealistic packing structures and atomic arrangements, failing to retain individual molecule's 's structure intact.
To address these limitations, introduce AssembleFlow, novel framework specifically designed to incorporate the rigid body constraints inherent in molecular assembly or crystallization. As illustrated in Figure 2, AssembleFlow leverages inertial frames establish reference coordinate systems for assembling molecules. Because SE(3) is the semi-direct product the rotation group SO(3) and the translation group IR3, we can further decompose the group SE(3) transformations into translations in IR (Figure 2 (a, b)) and rotations in SO(3) (Figure 2 (c, d)). Such decomposition allows the explicit enforcement of both translational and rotational rigidity the molecular level effectively, ensuring that each molecule in the cluster moves as a unified, rigid body throughout the crystallization process. During such enforcement, AssembleFlow employs a distinct approach for learning the SE(3)-equivariant velocity functions associated with translations and rotations. For translations, uses linear interpolation (LERP) Fuclidean space, while for rotations, leverages spherical linear interpolation (SLERP) in spherical space, with closed-form solution. This distinction in handling the translation and rotation groups allows AssembleFlow to accurately model the rigid transformations of molecules during each prediction and generation step within the flow matching framework (Lipman et al., 2022; Liu et al., 2022; Albergo & Vanden- Vanden-Eijndenn 2022).
We empirically evaluate AssembleFlow using the benchmarking crystallization dataset COD- Cluster17. The quantitative results reveal that AssembleFlow significantly outperforms six competi- tive deep learning baselines by at least 45% in terms of assembly matching score. Also, AssembleFlow
2

---

Published as a conference paper at ICLR 2025
exhibits strong assembly performance compared to widely used domain-speciffi simulation tool for molecular assembly, achieving this with a 25-fold reduction in computational cost. Furthermore, we present qualitative results, including atomic collision properties of predicted crystals, which further demonstrate AssembleFlow's effectiveness in preserving and modeling the rigidity of the molecular crystallization and assembly process. Our work is the first to implement rigid generation in SE(3) space for molecular assembly. We also want to mention that in what follows, we use molecular assembly, crystallization, and molecular packing interchangeably.
PRELIMINARIES
Molecular crystallization. Molecular crystallization is transition of molecules from weakly correlated structures strongly correlated structures, e-g., from liquid gas phase to solid phase, as illustrated in Figure 1. One example is liquid water freezing into ice, transitioning from a liquid phase to solid phase. The crystallization from a gas phase directly solid phase is called deposition.
-quivvarrree For geometric modeling for crystallization, one critical property of the target function is rotation-equivariant and maanlation-equivariinn (i.e., -quiivaiuia))iant We here provide a brief introduction on the 3E(3))-quivariane and for more detailed discussions of -qquva-equi refer the reader (Smidt et al., 2018; Brandstetter et al., 2021; Liu et al., 2023; Zhang et al., 2023). SE(3) equivaaraae the property for the geometric modeling function f X Y as:
((x((a)x) = pr (a)) Va E G
(1)
where (a) and (a are the SE(3) group representations on the imput and output space, re- spectively. Squiv-iquiaaiiant modeling Equation (1) is essentially saying that the designed deep learning model is modeling the whole SE(3) group transformation trajectory on the molecule conformations, and the output is the transformed ý accordingly. One concrete example is that when we rotate the input molecular system by certain angle, the predicted forces by SE(3)-equivariant models will also rotate accordingly.
Conditional flow matching. Conditional flow matching (CFM) (Lipman et al., 2022) and two parallel works (Rectified Flow (Liu et al., 2022) and Stochastic Interpolants (Albergo & Vanden Vanden-Eijnden, 2022)) formulate the distribution modeling problem learning a vector field that can generate probability path mapping from simple distribution at = 0 to the target distribution at f = 1. Please refer the original papers for more detailed discussion (Lipman et al., 2022).
In the crystallization processes, our geometric data are atomic coordinates in the 3D Euclidean points E RR3, and the atomic type is fixed during the whole crystallization process, so we may well
ignore that. Then we define time- time-dependent vector field v : [o, 1] x R + IR3. A dime-dependent vector field defines a time-dependent diffcomorphic map. called flow, Φ [0, 1] R R3, The
vector field defines flow via an ordinary differential equation as
(r) = (r))) dt
(2)
probability density path is denoted as p [0. 1] R + Rso. Existing flow model (Chen et al.. 2018) maps a prior distribution Po to another distribution Pt with push-forward equation or change
dat
of variable rule: (r(r) = (o))- = Po Po ((0=(((r)) (r))det
do.(r) Thus, modeling the likelihood of ds
data distribution at can be transformed into modeling the velocity field matching problem with parameterized velocity field vog i.e., flow matching:
CPM = ((ve) - vr(r)))
(3)
With the continuity equation (Villani et al., 2009), we can further derive an equivalent objective by considering the conditional vector field conditioned on the empirical data rı, i.e., ve(r(ri), and the resulting objective is the conditional flow matching:
Lam = E. ( - vo(r, rr,))=
(4)
3 METHOD: ASSEMBLEFLOW
Problem Formulation. AssembleFlow is designed to model rigid transformations during crystalliza- tion, ensuring that each molecule in the cluster remains rigid throughout the transformation process.
3

---

Published as a conference paper at ICLR 2025
Rigorously, modeling P(fry) (ra)), where {ry} and {r:} are the atom conformations in final and initial positions respectively. During this process, we assume the rigidity of molecules. Noticeably, will use preprocessed dataset where the prior conformations are geometrically optimized and fixed (Liu et al., 2024c).
This section outlines the five key steps in the algorithm's development. Specifically, in Section 3.1, explain how inertial frames can be leveraged to provide a stable reference for tacking the orientation of multiple assembling molecules in the Euclidean space, Such a reference perspective can guarantee rigid structures during molecular rotations throughout the assembly process. Building on these frames, there are multiple ways for rotation representation, so Section 3.2, we illustrate how to use quaternion representation for capturing the rotation transformation induced from inertial frames. This followed by a detailed discussion of AssembleFlow, a rigid flow matching method, in Section 3.3. AssembleFlow decomposes the assembly probabilistic paths into SO(3) group path and R group path to guarantee the rigidity, and learns the time-dependent vector fields through a How-matching framework the two path spaces respectively. In Section 3.4, we employ the reparameterization trick to make AssembleFlow more numerically stable. Finally, in Section 3.5, we present the two types of SEii)-equivariant flow matching velocity functions specifically designed for use in AssembleFlow. Note: the pseudo algorithm of our AssembleFlow is provided in Appendix E.3.
3.1 SO(3) GROUP AND INERTIAL FRAME FOR RIGID PACKING
The core of AssembleFlow lies in the utilization of the inertial frame as the reference frame. Within this frame, the rotation matrix in the SO(3) group defines how the molecular system rotates rigidly. This serves as the key step in Assemble for modeling rigid transformations in SO(3) group.
SO(3) group. The special orthogonal group, denoted (0)),, is group of rotation matrices that represent rotations in n-dimensional Euclidean space. In this paper, we are interested in n = 3 dimensional space, and every rotation matrix used to perform a rotation in 3D space can be represented as an element of S3(3) group. The SO(3) group consists of all orthogonal matrices with determinant 1 R E 3x3 such that R2 R = I, where R is the transpose of R and is the identity matrix.
Inertial frame as the reference frame. An incrtial frame is a reference frame such that it can provide consistent basis for describing molecule's motion, including rotation. Here, utilize the inertial frame to build a basis for explaining how each molecule rotates in the Euclidean space. One example is illustrated in Figure 2. Im- portantly, an inertial frame provides co- ordinate system such that a molecule stays rigid and does not deform over the crys- tallization or modeling process; we here assume the system is not influenced by ex- ternal forces. Next, we will detail how iner- tial frames are used to represent the rotation matrix for rigid molecules.
(A) Three principal axes
(b) Three principal axss
Rotatable
Initial Axes
Final Aves
Inttal Aves
Final Axes
(c) Fourth non-coplaner point
Fourth non-coplaner
Rotatable
Non-rotatable
Initial Axes
Final Axes
Aves
Axes
Figure 3: (a, b) show two potential rotational alignments hetween two coordinate systems (axes). (c. d) show that only
one unique rotation is possible for four non-coplanar points.
First, we employ the following four sequential steps to derive the reference frames that construct the rotation matrix from N atomic positions
Calculate the mass center: c = + Σ ri.
Adjust position relative to the center r; = = -r-
Compute the inertia tensor i = Σ, ||: rir?' where T is the unit diagonal matrix.
Obtain the principal axes of inertia by applying eigen-decomposition on i. We have
i = QAQT, where Q is the orthogonal matrix whose columns are the eigenvectors of and A is the diagonal matrix whose elements are the eigenvalues X; of i, representing the principal moments of inertial along the principal axes.
The above steps yield the three principal axes in Q. To adopt this for modeling the crystallization process (Figure 1), we build inertial frames for each molecule in the cluster. In case, for each
4

---

Published as a conference paper at ICLR 2025
molecule need to build two inertial frames, for the weakly correlated and strongly correlated structures, respectively. We call these two frames initial (inertial) frame F; and final (inertial) frame Fy.
Second, we apply the axes Q,, respectively. As illustrated in Figure 2 (d), can only perform the rotation based on the aligned principal axes or the aligned coordinate systems. However, as we conduct the eigen- decomposition of frames F; and Fy, it is not guaranteed that the corresponding principal axes Q; and Q2 are aligned. To align the two coordinate systems, we aim to match both the directions and the orders of the corresponding principal axes in each set. We will detail this alignment process next.
to obtain the initial principal axes Q, and final principal
Align the directions between initial and final coordinate systems. First, for a given inertial frame F. we have three axes in Q composing a coordinate system. This can lead eight possible directions. To align the directions between the initial system Q: and final system 2,, we add a first constraint the three axes must form a right-handed coordinate system. Such a filtering step can be achieved by using cross-product: if the local frame system is right-handed, then the cross-product between any two axes should match the third axis or share the same direction as the third axis. Otherwise, the coordinate system is left-handed, then we randomly revert one basis out of three. This reduces to four potential combinations of directions. To further align the directions, we introduce Lemma and Theorem to help provide us with theoretical guidance.
Lemma 1. For an initial inertial frame F; and a final inertial frame Ff, we build up the correspond- ing right-handed principal axes as coordinate systems, Qi and Q2, respectively. Suppose we have to change the directions of Q, to match Qu, then we should change the directions of two buses in Q2-
Proof. There are three bases in Q5- If we change one or three basis directions in Qy. then Q, will change from hagdedaande to deft-handedness, which violates the assumption. Thus, if we need to change the directions of Q2 to match Q,, we should change the directions of two bases in Q5. O
With Lemma 1, we find that using three base vectors is insufficient to determine the direction aligmment of three axes. To define the directions that can match between the initial and final coordinate systems, we need to incorporate extra node an auxiliary, in Theorem 1.
Theorem 1. For an initial inertial frame F, and a final inertial frame Fj, we build up the corre- sponding right-handed principal axes coordinate systems, Qi and Q respectively. Then we need to incorporate a fourth point that is not coplanar with the three basis vectors, to align the directions of two coordinate systems with one unique rotation transformation matrix.
In Figure (a, b). we can see least two possible rotation matrices to transform from initial axes to final axes. However, when we add a fourth non-coplanar point in Figure 3 (c), the rotation transformation becomes unique, and the corresponding rotation in Figure 3 (d) is invalid. Then more rigorously, the proof includes the two key steps: (1) Using Lemma 1, we can find multiple rotation matrices for alignment between coordinate systems. (2) After introducing the fourth non-coplanar point, the contradiction proves that there exists only one unique rotation for alignment. For more rigorous proof, please refer to Appendix D. O
Proof Sketch. We first provide intuitive examples in Figure 3.
Align the ordering between initial and final coordinate systems. We can typically sort the eigenvectors (as for principal axes) through the corresponding eigenvalues. The main challenge comes when there a tie in eigenvalues. Because this is rare case, we propose doing a depth- search to enumerate all the possible combinations of basis orderings of Q2, to match Qi
Outputs and engineering issue: numerical stability. Without loss of generality, we can assume that we do not change the axis direction or ordering in the initial coordinate system Qi, and we only change Q, to Q,. so as to align with Q. The ultimate rotation matrix is thus R = Qf Q,- Meanwhile, we would like to point out that multiple numerical stability issues exist. This can arise in the following scenarios: (1) when the sampled points near the origin, (2) when checking if the eigenvalues are tied or not, (3) when extracting a fourth non- non-planar point for alignment, (4) when verifying whether rotating the initial atoms (points) matches the final atoms. To mitigate these issues, we carefully select a threshold value and clamp the (reconstructed) coordinates to this minimum threshold.
Summary. In this section, first introduce the basic concepts of SO(3) group and inertial frame. Then we present how we construct the initial and final inertial frames for each molecule. i.e., F, and F,, in molecular crystallization. Next, by applying the eigen-decomposition the constructed
5

---

Published as a conference paper at ICLR 2025
inertial frames, obtain the initial and final principal axes (right-handen Qi and Q, respectively. Finally, we align the Q to Q; by checking the directions and ordering of three axes in Q,. This results in two aligned bases (Q; and Q2) and a rotation matrix R such that Q = R RQ; As a result, this enables rigid molecule-level rotations during their transformations in the assembling processes.
3.2 QUATERNION REPRESENTATION FOR ROTATION
For the SO(3) group in Section 3.1, there are multiple ways to represent rotation transformation in addition to the rotation matrix R. If we want to model the rotation matrix directly, we must guarantee that the generated matrix variable satisfies the two properties discussed in Section 3.1, namely Orthogonality and Determinant. Such a constrained modeling is challenging. Thus, an alternative way of rotation representation with a more flexible formulation preferred. To attain this goal, utilize quaternion representation defined through the inertial frame, as described below.
Definition. A quaternion is defined as:
q + yj + zk (w,v),
(5)
where w, x, y, 2 real numbers, and i, lo are the fundamental quatemnoo units. Equivalently, the w is called the real part, while v = (z, y, 2) is a 3D vector representing the imaginary part.
Rotation quaternion. A rotation quaternion is a unitary quaternion, i.e., w² + +2² + y² 2² 1, and this can be easily achieved by taking the normalization of the quaternion variables. In what follows, will assume the quaternion is rotation quaternion unless otherwise specified. Notice that for each rotation matrix, there are two equivalent quaternions, and -q. Here we manually enforce the real number part of the generated quaternion to be non-negative.
Tranformation from rotation matrix to rotation quaternion. There are multiple ways to extract the quaternion from the rotation matrix, and we provide more detailed discussion in Appendix B. In this work, we adopt eigendecomposition (Horn, 1987; Bar-Itzhack, 2000) to extract the initial and final quaternion (i.e., q: and qÿ) from the initial and final coordinate systems (i.e., Q; and Qy).
Spherical interpolation (SLERP) for quaternion interpolation. One of the main advantages of using quaternion is that it is friendly to interpolation on the SO(3) space, i.e., the spherical interpolation (SLERP) between two quaternions Qo and q1:
SLERP(go. q1, sin((1 - f) + sin(tw)qı
(6)
sin(w)
where w the angle between qu and qı, and t E [0, I] is interpolation parameter. Thus, we can that SLERP provides smooth and uniform rotation between two quaternions. An example is provided in Figure 2.
We provide comprehensive discussion of various rotation representations in Appendix B, including quaternion multiplication and vector rotation using quaternion. Please consult that section for details.
3.3 PATH INTERPOLATION IN ASSEMBLEFLOW
To model the crystallization process, our AssembleFlow method integrates the inertial frames and quaternion representation for rotation, as discussed in Sections 3.1 and 3.2, into a conditional flow matching framework. We note that, unlike most existing flow matching methods that focus solely on atom-level diffusion paths in the Euclidean space, which suffices for non-rigid transformation, AssembleFlow operates within the full SE(3) group space in the molecule level due to the rigidity requirement.
Recall that the crystallization process involves the movement molccule in the cluster, AssembleFlow jointly models the rotational transformations in SO(3) space and translational transformations in R space. Such decomposition ensures the preservation of rigid molecular structures throughout the crystallization process. We next detail these two transformations.
cluster of molecules, and for each
Modeling translations in R The goal is model the molecule-level translations in each cluster, and AssembleFlow achieves this by modeling the translations on each molecule's mass center, as depicted in Figure 2. For notation simplicity, we will R represent the translation vector.

---

Published as a conference paper at ICLR 2025
We adopt the flw-matchhing framework, and the goal here is to learn the probability of final mass center æ from the initial mass center 34 i.e., p(æy (zt)). ) To this end, assume that we use linear interpolation (LERP) for path interpolation, by treating æo = Dú and Φ - ay then for interpolation parameter t E [0, 1/. we have the interpolated translation
LERP (2o, æ1, = two + (1 t)x1.
Next, we introduce an SEuii-eqivvaiian function ug t) the core module to learn the velocity at time t. Thus the objective function is defined as:
Egs
æ - 2o u (as, q t)[l"
(8)
Modeling rotations SO(3). For modeling the rotations the SO(3) group, recall that we can find the initial bases Q; and final principal bases , from inertial frames in Section 3. and then transform them into the rotation quaternions qi and q, as introduced in Section 3.2. The task here is model ((gllg))
Thus, is natural adopt the spherical interpolation (SLERP) as the smooth translation between two quaternions. We treat qo = q. and qı = q, and plug them into Equation (6). This gives us the interpolated rotations at time l. The first-order derivative of SLERP has an analytical formula:
w cos(tw)qq cos((1 - t)w)qo)
SLERP(QU: qz,t)
(9)
sin(w)
Similarly here, we then introduce an SE(3)-equivariant function v0,so(3) (gr: f) to model the molecule- level rotation velocity time t. The objective function becomes:
d SLERP((oo q1.t) (0,5o13) (xt, q., qr,t)] t)l
C5o(3)
(10)
dt
Inference. For inference, AssembleFlow conducts the sampling step in SO(3) and g alternatively:
Tt++ x. öt 19,23 (Z:, qr, g++1 q: 5t (0,50(3) (æt, qe,
(11)
Thus, both the molecule- translation 2z+1 and molecule-leved rotation q++1 applied on each molecule, and we repeat Equation (11) for steps to obtain the predicted strongly correlated molecule position. However, in Equation (11), it remains open question how to obtain the w for SO(3) generation, since w is the angle between qo and qı, and qı is unknown during the inference process. We address this by proposing the reparameterization trick will be discussed next in Section 3.4.
3.4 REPARAMETERIZZTTIN FOR STRONGLY CORRELATED STRUCTURES
We here leverage the reparameterization trick to directly model the SE(3) action at time T instead of the velocity each time t. Equivalently, the velocity of SE(3) action can be written as:
uogs (ee qe, t) ( (t q:, t) x.,) / t)
w cos(tw) a1 (2., t).t) - cos f)w)
(12)
(0) (wi, q.
sin(w)
(x,, In other words, we directly estimate the translation & f (t, qr, () and rotation Ộ (x qt, () in the final step or the strongly correlated structure. The objectives over the two spaces thus become:
E2s reparameter E| L (mx, ge, f)
Lso(3) appammmttr = Rillq - ς1.0 q:, ().)
(13)
The final objective function the summation of two terms. Besides, such a reparameterizatton enables us to conduct inference using the Euler algorithm:
ah+1 24 & (#1.0 t) xz)((1 t),
@ cos(tw) &T. o ow q: - cos((1 - t))))
(14)
+ d5
sin
where @ is the angle between go and a,, q: t). The velocity functions of & $1,0 qe. f) and q (2;, qi, are Squuieaquivariant and will be discussed next in Section 3.5.
7

---

Published as a conference paper at ICLR 2025
Table AssembleFlow against six generative models on COD-Cluster17 with 5K, 10K, all samples. The best results are marked in bold.
Packing Matching
Validity
PM (atom)
PM (center)
Collision
Separation
Compactness
CODD-Custer\7-5K
GNN-MD
13.67 0.06
13.80
0.07
27.53 0.49
0.22 0.11
100.00 0.00
CrystalSDE-VE
15.52 = 1.88
16.46 0.99
1.20 1: 0.08
27.17 + 0.86
57.47 2: 7.76
CrystaISDE: VP
18.15
4.46
34.00
CrystaFFlowVV VE
14.87 + 7.07
13.08
4.51
1.37 + 0.04
35.70 0.73
8.40 4.17
CrystalFlow-Vy
15.71 2.69
17.10
1.89
1.38 1: 0.04
35.43
0.88
4.87 1.09
CrystaFFlow-LFRR
13.59 0.09
13,26
0.09
0.34 0.01
97.38
0.10
100.00 0.00
AssembleFlow (ours)
7.27 + 0.04
6.13
0.10
0.33 0.00
97.64
0.36
100.00 0.00
Dataset: COD-Cluster17-00K
GNN-MD
0.06
0.05
27.88 0.49
0.23
100,00 0.00
CrystalSDE-V VE
17.25 + 2.46
17.86 1.11
0.99 0.27
32,99
10.72
34,93 + 14.99
CrystaIDDE: VP
22.20
3.29
21.39
1.50
0.53 0.35
52.48
15.44
16.83 18.09
CrystalFlow-VE
16.41 + 2.64
16.71
+ 2.35
1.42 + 0.03
33.79
+ 0.51
5.47 0.47
CrystalFlow- VP
19.39 1 4.37
16.01
3.13
1.44 1: 0.03
33.35
0.55
4.23 0.48
CrystalFlow-I
97.32
AssembleFlow (ours)
7.38 + 0.03
6,21 0.05
0.31 + 0.00
97.73
0.16
99.93 0.05
Dataset: COD-ClusterI7-Ate
GNN
22.30 + 12.04
14.51 0.82
24.29 4.58
4.13 5.60
98.77
1.73
CrystalSDE-VE
17.28 0.73
18.92 + 0.03
0.19 +
15.47
12.42
2.51
2.37
CrystalDE: VP
18.03 11 4.56
20.02 1: 3.70
0.55 1: 0.19
48.78
1.70
6.88
2.82
CrystalFlow-yy
12.80
141
35.34 0.28
2,90
CrystalFlow-VP
13,50
13.28 + 0.48
1.51 0.02
33.06
1.31
6.61
3.17
CrystalFlow-LFF
13.61 + 0.00
13.28 0.01
0.34 + 0.00
97.34 0.02
99.99
0.01
AssembleFlow (ours)
7.37 0.01
6,21 0.01
0.31 + 0.00
98.15
0.22
0.00
3.5
MULTI-GRAINED VELOCITY FUNCTION
Recall that the data structure considered here is the cluster of molecules, thus is natural split the modeling into intra- intra-molecule and inter-molecule modeling, as introduced below. For intra- molecule modeling, we adopt PaiNN (Schütt et al., 2021), which one of the most widely used SE(3)-equivariant models. It can encode the inherent geometric structural information of individual molecules. Then for inter-molecule modeling, we consider two options equiivaaaar models: (1) Atomic-level modeling that utilizes all the atoms' positions for learning the molecular- rotation and translation for the next step. (2) Molecular-level modeling that directly utilizes the molecular-level rotation and translation for next-step prediction. This concludes discussion AssembleFlow, and more details are provided in Appendix E. A high-level overview and pscudo algorithm are provided in Algorithms and 2 in Appendix E.3.
4 EXPERIMENTS
4.1 EXPERIMENT SETUP
Implementation. The codes and checkpoints are available at this GitHub repository.
Datasets. We evaluate our method using the crystallization dataset COD-Cluster17 (Liu et al., 2024c). This COD-Cluster17 contains 133K crystals and is a curated subset derived from the Crystallography Open Database (COD) database (Grazulis et al., 2009). We consider three versions of C with 5k, 10k, and all data, respectively. Detailed discussion on this dataset is provided in Appendix G.
Evaluation metrics. We evaluate the performance of the compared approaches using a comprehensive set of metrics tailored to assess the quality of crystallization packing. These metrics include: (I) Packing Matching (PM) (Chisholm & Motherwell, 2005): This metric measures how well the generated molecular assemblies match the reference crystal structures in terms of spatial arrangement and packing density. Following (Liu et al., 2024c), we employ packing matching on both the atomic level (PM-atom) and the mass-center-level (PM-center) (Chisholm & Motherwell, 2005). (2) Atomic Collision: This follows (Cordero al., 2008). It measures the percentage collided atom pairs the predicted assemblies. Atoms must maintain minimum covalent distance governed by the balance of attractive and repulsive forces. (3) Separation: We extend the metric from (Xie et al., 2022: Yang al., 2024) setting. cluster of molecules is valid if the minimum distance between molecules
8

---

Published as a conference paper at ICLR 2025
is above 0.5À (Court et al., 2020). This metric is referred to separation to measure the validity to avoid unphysical interactions at the molecular level. (4) Compactness: We propose this measure by calculating the percentage of simulated clusters where the maximum atomic pairwise distances are below 100Å. A higher compactness value suggests more efficient arrangement, where the intermolecular spaces are minimized, leading to a denser crystalline structure. Detailed discussions on these metrics are provided in Appendix G.
Baselines. We compare our method with two categories
Table 2: Ablation studies of PackMol and AssembleFlow variants.
PackMol Assembleelow-Aton Assembli Molecule
of baselines: state-of-the. art deep generative models and an established domain- specific simulation tool.
Dataset: COD-Chuster17-55
PM (atom) J
7.10 = 0.05
7.27 = 0,04
7.67 0.10
PM (center) +
6.05 + 0.04
6.13 = 0.10
6.77 = 0.10
Collision +
0.32 + 0.00
0.33 0.00
0.37 0.01
(I) Deep generative base- lines. For generative models, we evaluate our approach against GNN- MD (Liu et al., 2024c), CrystalSDE (Liu et al., 2024c), CrystalFlow (Liu et al., 2024c), and differ- ent variations of them, including CrystalSDE- VE, CrystalSDE-VP CrystalFlow-Vir VV, CrystalFlow- VP, and CrystalFlow-LERP. These models employ various mecha- nisms handle the challenges of molecular crystallization. CrystalSDE-V2 VE and CrystalSDE-Va use stochastic diffcrential equations to model diffusion processes under different parameterizations. CrystalFlow- VE and CrystalFlow- VP apply flow matching principles for diffusion- based interpolation path, with the latter focusing on variance- preeeeriiign methods. CrystalFlow-LERG utilizes linear interpolation to handle molecular transformations, striking a balance between computational complexity and performance.
Separation
99.56 1 0.08
97.64 0.36
92.95
Dataset: COD-Cluster17-10m
PM (atom) J
7.16
7.38 = 0.03
7.65 = 0.17
PM (center) 4
6.11 + 0.01
6.21 = 0.05
6.69 + 0.20
Collision &
0.30 0.00
0.31 0.00
0.35 0.01
Separation
99.45 +: 0.10
97.73 0.16
92.67 1: 0.32
Dataset: COD-ClusterI7-AII
PM (atom) J
7.15 0.01
7.37 0.01
7.47
PM (center) 4
6.09 0.01
6.21 0.01
6.36 0.06
Collision +
0.30 + 0.00
0.31 0.00
0.33 + 0.00
Separation
99.42 -: 0.03
98.15 0.22
95.66 1: 0.08
(2) Domeii-sppecffic simulation baseline. We also compare our method with PackMol (Martínez et al., 2009), well-established simulation tool widely used in the field for molecular packing. PuckMol has long been go-to solution for chemistry and material experts due to its ability to generate initial molecular configurations for follow-up simulations, making it an important and relevant baseline for evaluating molecular assembly tasks. More detail on this baseline is in Appendix F.
4.2
MAIN RESULTS
The comparison results with the generative modeling baselines and the simulation model are presented in Tables and 2, respectively.
As shown in Table 1, AssembleFlow significantly outperformed all six deep generative models across almost all metrics. For example, AssembleFlow improved Packing Matching by at least 45% compared to other models. Notably, most baselines struggled with rigid packing, leading very low Separation scores, except for CrystalFlow-LERP For example, AssembleFlow achieved a Separation rate of 97.64%, while GNN-MD only reached 0.22%.
As shown in Table 2, when compared to the domain domain-specific tool PackMol, our data-driven approach demonstrates strong assembly performance relative this widely used simulation method. Remarkably, our method achieves 100% validity. matching that of the tool. While this outcome highlights the promise of AssembleFlow. it is expected for PackMol, as it leverages well-established domain knowledge and heuristic physical rules to determine molecular orientations. Promisingly, both methods achieved a very low Collision rate. The results indicate that the data-driven AssembleFlow performs comparably to the domain simulation tool PackMol.
simulation
4.3 ABLATION STUDIES

---

Published as a conference paper at ICLR 2025
Velocity function. As dis- cussed in Section 3.5, As- sembleFlow can utilize two types of -quiiaaqiaa velocity functions. We here evaluate their performance and, as shown in the last two columns of Table 2, AssembleFlow-Atom generally outperforms AssembieFloo-Moleeuue as atomic-level information
stomE
collasion
atomic colision
Figure 4: Atomic collisions (red circles) in predicted assemblies.
provides more detailed geometric insights. Despite this, both variations of AssemblcFlow exhibit strong results, significantly surpassing other deep learning baselines.
Atomic collision. To show atomic collisions of the assemblies, visualize, in Figure 4, two atomic collisions in assemblies predicted by AssembleFlow, where two molecules collide into each other (indicated by the red circle).
Computational cost. In Figure we present the computation time needed for AssembleFlow and PackMol. It reveals that our data-driven method achieves 25-fold reduction in computational costs. This suggests that our method can be scaled effectively for larger molecular systems and datasets.
5 CONCLUSION AND OUTLOOK
We introduced AssembleFlow, a generative model that maintains the inherent rigidity of molecular structures during assembly. By us- ing inertial frames for positional references at the molecular level, AssembleFlow accurately tracks molecular orientation and motion. It de- composes transformations into translations and rotations, enforcing rigidity throughout the gen- eration process. This innovative approach en- ables the model to separately learn velocity func- tions using linear and spherical interpolation for accurate rigid molecular assembly.
864
35
PackMol
AssembleFlow
Figure 5: Comparison of computational cost in hours for 10,543 molccule clusters from COD-Cluster17. Pack- Mol requires around 864 hours, while AssembleFllw requires 35 hours.
Empirical results COD-Clusterl7 show that AssembleFlow outperforms six state-of-the-art deep learning models while maintaining molecular integrity and achieves comparable assembly performance established simulation tool, significantly redueing computational costs. This suggests that AssembleFlow offers an efficient solution for large-scale molecular simulations and opens avenues for faster material generation molecular engineering.
To the best of our knowledge, AssembleFlow is the first to implement rigid generation in SE(3) space for molecular assembly. It has the potential to be generalized to more complex and challenging scenarios, such as simulating the crystallization process of polymorphs with diverse configurations and structures. These scenarios are critical for accurately modeling real-world materials, where different polymorphic forms can exhibit vastly different properties. However, the current dataset, COD-Cluster17, lacks the necessary information to fully explore this aspect. As such, extending AssembleFlow to handle these intricate processes remains an exciting direction for future research, where more specialized datasets could enable deeper insights into material formation and polymorph behavior. We aim to address these challenges in future work.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
We thank the anonymous reviewers for their valuable suggestions and feedback. YB acknowledges support from NRC AI4D, CIFAR, and the CIFAR AI Chair program. This project's computational resources are provided by NRC and the Digital Research Alliance of Canada.
REFERENCES
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. 15
Michael S Albergo and Eric Handen-Fijnden, Building normalizing flows with stochastic interpolants. arXiv preprint (XXiv:2209.15777 2022. 2, 15
Marloes Arts, Victor Garcia Satorras, Chin-Wer Wei Huang, Daniel Zugner, Marco Federici, Cecilia Clementi, Frank Noe, Robert Pinsler, and Rianne van den Berg. Two for one: Diffusion models and force fields for coarse-grained molecular dynamies. Journal of Chemical Theory and Computation, 2023. 15
Michael F. Ashby and David RH Jones. Engineering Materials 2. Butterworth-Heiineaam. 4 edition, 2012. 2
Minkyung Back, Frank DiMaio, Ivan Anishchenko, Justas Dauparas, Sergey Ovchinnikov, Gyu Ric Lee, Jue Wang, Qian Cong, Lisa N. Kinch, R. Dustin Schaeffer, Claudia Millán, Hahnbeom Park, Carson Adams, Caleb R. Glassman, Andy DeGiovanni, Jose H. Pereira, Andria V. Rodrigues, Alberdina A. van Dijk, Ana C. Ebrecht, Diederik J. Opperman, Theo Sagmeister. Christoph Buhlheller, Tea Pavkov-Keller, Manoj K. Rathinaswamy, Udit Dalwadi, Calvin K. Yip, John E. Burke, K. Christopher Garcia, Nick V. Grishin, Paul D. Adams, Randy J. Read, and David Baker. Accurate prediction of protein structures and interactions using three-track neural network. Science, 373(6557):811-87, 2021. doi: 10. (0.1166/sieencc.ab8877555 //www. wwwence.ecc.............. org/ doi /tw////s////// 1 sciencee abj j8744.
URL https:
Itzhack Y Bar-Itzhack. New method for extracting the quaternion from rotation matrix. Journal of guidance, control, and dynamics, 23()):1085-1087 2000. 6, 21
Avishck Joey Bose, Tara Achound-Sadeeae Kilian Fatras, Guillaume Huguet, Jarrid Rector-Brovd Brooks, Cheng- Caang-HHool Liu, Andrei Cristian Nica, Maksym Korablyov, Michael Bronstein, and Alexander Tong. Se (3)-shocc flow matching for protein backbone generation. arXiv preprint 10.00:2310.02331,, 2023. 15
Johannes Brandstetter, Rob Hesselink, Elise van der Pol, Erik J Bekkers, and Max Welling. Ge- ometric and physical quantities improve e (3) equivariant message passing. arXiv preprint vrXiv:2110.0290: 2021. 3
WH Brooks, WC Guida, and KG Daniel. The significance of chirality in drug design and development. Current topics in medicinal chemistry, 11(7):760, 2011. 15
Stephen R. Byrn, Robert R. Pfeiffer, and John G. Stowell. Solid-State Chemistry of Drugs. SSCI, Inc., West Lafayette, IN, 1999. 2
C. Barry Carter and M. Grant Norton. Ceramic Materials: Science and Engineering. Springer, edition, 2013. 2
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
James Alexander Chisholm and Sam Motherwell. Compack: a program for identifying crystal structure similarity using distances. Journal of applied crystallography, 38(1):228-22 2,122, 2005. 8, 31
Beatriz Cordero, Verónica Gómez, Ana E Platero-Prats, Mare Revés, Jorge Echeverría, Eduard Cremades, Flavia Barragán, and Santiago Alvarez. Covalent radii revisited. Dalton Transactions, (21):2832-2838, 2008. 31
11

---

Published as a conference paper at ICLR 2025
Callum J. Court, Batuhan Yildirim, Apoorv Jain, and Jacqueline M. Cole. 3-d inorganic crystal structure generation and property prediction via representation learning. Journal of Chemical Information and Modeling, 60(10):4518 4535, 2020. 9, 31
Stefan Doerr, Maciej Majewski, Adrià Pérez, Andreas Kramer, Cecilia Clementi, Frank Noe, Toni Giorgino, and Gianni De Fabritis. Torchmd: A deep learning framework for molecular simulations. Journal of chemical theory and computation, 17(4):2355-2363, 2021. 15
Daniel Flam-Shepherd and Alán Aspuru-Guzik Language models generate molecules, materials, and protein binding sites directly in three dimensions as xyz, cif, and pdb files. arXiv preprint arXiv:2305.055 05500,, 2023. 15
Octavian Ganea, Lagnajit Pattanaik, Connor Coley, Regina Barzilay, Klavs Jensen, William Green, and Tommi Jaakkola. Geomol: Torsional geometric generation of molecular 3d conformer ensembles. Advances Neural Information Processing Systems, 34:13757-13769, 2021. 15
Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International conference on machine learning, pp.
1263-1272. PMLR, 2017. 28
K. Elsabett Gonsalves, C.-P. Wong, Lisa T. Kuhn, and R. Shukla. Nanotechnology: Molecularly Designed Materials. American Chemical Society, 2000. 2
Saulius Grazulis, Daniel Chateigner, Robert T. Downs, Alex F. T. Yokochi, Manuel Quirós, Luca Lutterotti, Elena Manakova, Justinas Butkus, Peter Moeck, and Armel Le Bail. Crystallography open database an open-access collection of crystal structures. Journal of Applied Crystallography, 42(4):726-72) 2009. doi: 10. 100//000 0117700022888800110000 1898090166900. 8, 31
Jiaqi Guan, Wesley Wei Qian, Xingang Peng, Yufeng Su, Jian Peng, and Jianzhu Ma. 3d equivariant diffusion for target-aware molecule generation and affinity prediction. In The Eleventh International Conference on Learning Representations, 2022. 15
Jiaqi Guan, Xiangxin Zhou, Yuwei Yang, Yu Bao, Jian Peng, Jianzhu Ma, Qiang Liu, Liang Wang. and Quanquan Gu. DecompDifft Diffusion models with decomposed priors for sructure-based drug design. Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 Proceedings of Machine Learning Research, pp. 1827-11884.. PMLR, 23-29 Jul 2023. URL https: /ttpp/eecccegaaaee/eepepp/
v 22 gww.uaa.nnnnannnm 15
Anne Marie Healy. Zelalem Ayenew Worku, Dinesh Kumar, and Atif M. Madi. Pharmaceutical solvates, hydrates and amorphous forms: special emphasis on cocrystals. Advanced Drug Deliv- ery Reviews, 117:25-46, 2017. URL /api. www.nticschoooo.ooooo.....
38541447. 2
Rolf Hilfiker. Polymorphism in the Pharmaceutical Industry. Wiley- 2006.
Jonathan Ho, Xi Chen, Aravind Srinivas, Yan Duan, and Pieter Abbeel. Flow++: Improving flow- based generative models with variational dequantization and architecture design. In International conference on machine learning, pp. 2722-2730. PMLR, 2019. 17
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6855 2020. 15
Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion
for molecule generation 3d. In International conference on machine learning, pp. 8867-8887.
PMLR, 2022. 15
Berthold KP Horn Closed-form solution of absolute orientation using unit quaternions. Josa a, 4(4): 629-642, 1987. 6,
Guillaume Huguet, James Vuckovic, Kilian Fatras, Eric Thibodeau-Lanee Laufer, Pablo Lemos, Riashat Islam, Cheng Hao Liu, Jarrid Rector-Brooks Tara Akhound- -aaoeh,, Michael Bronstein, et al. Scquence-augmentted se (3)-flow matching for conditional protein backbone generation. arXiv preprint 205.1:2405.20333.. 2024. 15
12

---

Published as a conference paper at ICLR 2025
Rui Jiao, Wenbing Huang. Peijia Lin, Jiaqi Han, Pin Chen, Yutong Lu, and Yang Liu. Crystal structure prediction by joint equivariant diffusion. Advances in Neural Information Processing Systems, 36, 2024. 15
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool,, Russ Bates, Augustin Zidek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873))583--888 2021. 15
Leon Klein, Andreas Krämer, and Frank Noé. Equivariant flow matching. Advances in Neural Information Processing Systems, 36, 2024. 15
Jonas Köhler, Michele Invernizzi Pim De Haan, and Frank Noé. Rigid body flows for sampling molecular crystal structures, In International Conference on Machine Learning, pp. 17301-17326.
PMLR, 2023. 15, 16, 17
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint 22Xi::2200.02447, 2022. 2, 15
Shengchao Liu, Weitao Du, Yanjing Li, Zhuoxinran Li, Zhiling Zheng, Chenru Duan, Zhi-Ming Ma. Omar M. Yaghi, Anima Anandkumar, Christian Borgs. Jennifer T Chayes, Hongyu Guo, and Jian Tang. Symmetry-informed geometric representation for molecules, proteins, and crystalline materials. In Thirty-seventh Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2023. URL httpstss: openreview n forum? 1d=.ggsNe 3, 15, 27
Shengchao Liu, Weitao Du, Yanjing Li, Zhuoxinran Li, Vignesh Bhethanabotla, Nakul Rampal, Omar Yaghi, Christian Borgs, Anima Anandkumar, Hongyu Guo, et al. multi-grained sym- metric differential equation model for learning protein-ligand binding dynamics. arXiv preprint 2xXi::2401.11122 2024a. 15
Shengchao Liu, Divin Yan, Weitao Du, Weiyang Liu, Zhuoxinran Li, Hongyu Guo, Christian Borgs, Jennifer Chayes, and Anima Anandkumar, Manifold constraed nucleus-level denoising diffusion model for gructure-based drug design. arXiv preprini arXiv:2409. 1054409.0588. 10584, 2024b. 15, 31
Shengchao Liu, Divin Yan. Hongyu Guo, and Anima Anandkumar. Equivariant flow match- ing framework for learning molecular cluster crystallization. In ICML 2024 Workshop on Geomelry-groundle Representation Learning and Generative Modeling, 2024c. URL https: //openreview. net. wawprrummmddcccepmee
1. 2, 8, 9, 15,
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03003 2022. 2, 15
Hannes H Loeffler, Jiazhen He, Alessandro Tibo, Jon Paul Janet, Alexey Voronov, Lewis H Mervin, and Ola Engkvist. Reinvent 4: Modern di-driven generative molecule design. Journal of Chemin- formatics, 16(1):20, 2024.
Youzhi Luo, Chengkai Liu, and Shuiwang Ji. Towards symmetry-awand generation of periodic materials. Advances in Neural Information Processing Systems, 36, 2024. 15
Leandro Martínez, Ricardo Andrade, Ernesto G Birgin, and José Mario Martínez. Packmol: A package for building initial configurations for molecular dynamics simulations. Journal of computational chemistry, 30 30(13):1157,,164, 2009. 2, 9
Rebecca U Ve and Niamh M O' Boyle. Chirality of new drug approvals (2013-2022): trends and perspectives. Journal of Medicinal Chemistry, 67(4):2305 2024. 15
Amil Merchant, Simon Batzner, Samuel S. Schoenholz, Muratahan Aykol, Gowoon Cheon, and Ekin Dogus Cubuk, Scaling deep learning for materials discovery. Nature, 2023. doi: 10. 5415866-02--0735-9
David A. Porter, Kenneth E. Easterling, and Mohamed Y. Sherif. Phase Transformations in Metals and Alloys. CRC Press, third edition, 2009. 2
13

---

Published as a conference paper at ICLR 2025
Bayrammurad Saparov and David B. Mitzi Organic-inorganin hybrid perovskites: Structural versatility for functional materials design. Chemical Reviews, 116(7):4558 4596, 2016.
Kristof I Schüitt, Oliver T Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. arXiv preprint arXiv:2102.03150, 2021.
8. 27, 28, 30
Tess Smidt, Nathaniel Thomas, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and translation-ego -equivviinniuiviiiant neural networks for 3d point clouds. arXiv preprint arXiv: w0Xiv018020082113 2018. 3. 15
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019. 15
David Van Der Spoel, Erik Lindahl, Berk Hess, Gerrit Groenhof, Alan E Mark, and Herman JC Berendsen. Gromacs: fast, flexible, and free. Journal of computational chemistry, 26(16):
1701-1718, 2005. 2
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017. 27, 28
Cédric Villani al. Optimal transport: old and new, volume 338. Springer, 2009.
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computa- tion, 23(7):1661-1674, 2011. I. 15
Tian Xie, Xiang Fu, Octavian- Eugen Ganea, Regina Barzilay, and Tommi S. Jaakkola. Crystal diffu- sion variational autoencoder for periodic material generation. In International Conference on Learn- ing Representations, 2022. URL httpps:///penneveee..eee www.eeeee/tooedeeeeeeeeeee -t 8,
Sherry Yang, Simon Batzner, Ruiqi Gao, Muratahan Aykol, Alexander L. Gaunt, Brendan McMorrow, Danilo J. Rezende, Dale Schuurmans, Igor Mordatch, and Ekin D. Cubuk. Generative hierarchical materials search. 2024. URL https:/ wwwaaaai.oogg............a..a abs / 24009 06762. 8, 31
Jason Yim, Andrew Campbell, Andrew YK Foong, Michael Gastegger, José Jiménez-Luna Sarah Lewis, Vietor Garcia Satorras, Bastiaan Veeling, Regina Barzilay, Tommi Jaakkola, al. Fast protein backbone generation with se (3) flow matching. arXiv preprint 2rXiv:2310.05297 2023a. 15
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 20Xx:2330..0227, 2023b. 15
Min Yu. Weiming Li, Yunru Yu, Yu Zhao, Lizhi Xiao, Volker Lauschke, Yiyu Cheng, Xingcai Zhang, and Yi Wang. Deep learning large-scale drug discovery and repurposing. Nature Computational Science, 4:1-15, 08 2024. doi: 10. 10/8//44358-00444000999994
Xuan Zhang, Limei Wang, Jacob Helwig, Youzhi Luo, Cong Fu, Yaochen Xie, Meng Liu, Yuchao Lin, Zhao Xu, Keqiang Yan, et al. Artificial intelligence for science in quantum, atomistic, and continuum systems. arXiv preprint arXiv:2307.0842. 2023, 3, 15
14

---

Published as a conference paper at ICLR 2025
A RELATED WORKS
A.1 GEOMETRIC MODELING
Geometric modeling of molecules has predominantly been applied 3D Euclidean space (Smidt et al., 2018). It requires that the representation and generation function over the molecular system remain cquivariant to the rotations and translations, i.e., the -quivvaqui, ensuring that molccular properties are preserved regardless of orientation or position in space (Smidt al., 2018; Zhang et al., 2023; Liu et al., 2023).
We note that reflection (or chirality) is also important factor in geometric molecule modeling. For multi-component molecular systems like protein-ligand binding complexes, each individual molecule lead different energies and corresponding properties with distinct chiralities (Brooks et al., 2011; Vickerer & O' D'Boyle, 2024). Also, as shown in AlphaFold2, the natural molecules should be sensitive to the chirality (Jumper et al., 2021). Thus, more physically accurate geometric modeling should be SE(3) -qquivaqui and dellectio--ntisyymmen AssembleFlow satisfies this in Section 3.
and we have shown how our proposed
We also want to highlight another line of research that downplays the importance of symmetry in molecular modeling, as demonstrated by models like AlphaFold3 (Abramson et al., 2024) and XYZTransformer (Flam-Shepherd & Aspuru-Guzik 2023). These models avoid incorporating geometric symmetries because enforcing group symmetry constraints, such as SE())-equivariane,, can limit model's expressiveness. In some domain-sseccfff tasks, relaxing these constraints has resulted in strong performance. However, in the case of crystallization, maintaining molccular rigidity-a key symmetric property-is crucial. As demonstrated in previous work (Liu et al., 2024c), neglecting these equivariance constraints during molecular crystallization leads to unrealistic molecular structures.
A.2 GENERATIVE MODELS ON GEOMETRIC DATA
The geometric modeling on the continuous 3D Euclidean space can be naturally merged with deep generative frameworks, where the goal is to learn the geometric data distribution and generate new molecules. The deep generative models include but are not limited to denoising score matching (Vin- cent, 2011; Song & Ermon, 2019), denoising diffusion probabilistic model (Ho et al., 2020), and flow matching (Lipman et al., 2022; Liu et al., 2022; Albergo & Landen-Eijnden 2022). Such geometric generative models have been widely adopted for molecule and material generation (Hoogehoom et al., 2022; Jiao et al., 2024; Luo et al., 2024), structure-base drug design (Guan et al., 2023; 2022; Liu al., 2024b), and molecular dynamics simulation (Doerr et al., 2021; Arts et al., 2023; Liu et al., 2024a).
A.3 GENERATIVE MODELS ON RIGID GEOMETRIC DATA
Though our work is the first to apply rigid generation in SE(3) space for molecular pack- ing/crystallization, similar ideas have been carried out for protein backbone generation (Jumper al., 2021; Yim et al., 2023a;b; Bose al., 2023; Huguet et al., 2024) and relevant works (Köhler et al., 2023; Ganea et al., 2021; Klein et al., 2024)) In the protein backbone generation and folding setting, there exists a well-defined local frame structure for each residue: the backbone atom pairs (C (C-N) N) and (C. C") form two bases, and their cross product leads to the third basis, which is normal vector perpendicular to the two bases. Thus, the rotation matrix can be easily constructed based on such local frame structure. Such modeling paradigm is appealing for macromolecules like proteins reduce the computational cost.
In this field, AlphaFold2 adopts this local-frame idea to model the folding process (Jumper et al., 2021), while Frff applies this idea and denoising diffusion model for protein structure gen- eration (Yim et al., 2023b). Similarly, FrameFlow and FoldFlow integrate local frames with flow matching to learn the protein dynamics (Yim et al., 2023a; Bose et al., 2023; Huguet et al., 2024). However, this approach cannot be easily extended to crystallization tasks, constructing reliable lo- cal frames to establish positional references for assembling molecules is not straightforward. Instead, innovatively introduce inertial frames to achieve this goal.
15

---

Published as a conference paper at ICLR 2025
A.4 COMPARISON BETWEEN SPHERICAL INTERPOLATION AND EXPONENTIAL MAP INTERPOLATION ON S0(3)
In addition spherical interpolation, could also consider the exponential map interpolation, as used in (Riemannian FM, FrameFlow, and FoldFlow). In this section, we would like to compare the theoretical differences. We further conduct experiments for empirical comparison, and please check Appendix H.2 for more details.
We mark exponential map interpolation (EMLERP), and is defined as:
EMLERP(ro: r1, exPro (1 (r1)),,(r))).
(15)
We then discuss how FoldFlow and FoldFlow2 utilized this equation since they the latest work along this line of using EMLERP. By utilizing the tangent space of the manifold and axis-angle representation, existing works (FoldFlow, FoldFlow2) have been using approximated closed-form solution for the derivative:
EMIERP(PO rı, () = logr. To
(16)
Thus, their objective function on SO(3) is:
d Cso(3) EMLERP(ro, r1, - (0,,0033 Te, t)l?
dt To t
(17)
logr.
(0,50(3) rt, t)ll?.
To sum up, these two methods do not have a clear methodological advantage over each other; however, the EMLERP considers more approximation tricks in implementation. We summarize the main differences Table 3.
Table 3: Comparison between the interpolation paths in AssembleFlow and FrameFlow.
AssembleFlow (ours)
FoldFlow
Reference Frame
Eigenvectors of inertial frames
Sram--chhiidt on N-C-C,
SO(3) Interpolation
Spherical Interpolation
Exponential Map Interpolation
Equation
sini (o(((--jee 1/1/0 tain
exp,, (t log,, (r1))
amto] Equation (9)
Derivative
Equation (16)
Rotation Representation (for Velocity / Objective Function)
Quaternion Yes
Rotation Matrix or Axis-angle
Reparameterization
No
A.5
COMPARISON BETWEEN OUR WORK AND RIGID STRUCTURE SAMPLING (KÖHLER ET AL., 2023)
We would like to emphasize that our work differs from (Köhler et al., 2023) in the following fundamental aspects. Noticeably, would like to emphasize that our work can be seen as an extension of (Köhler et al., 2023), yet addressing more practical and challenging problems, including rigid modeling on arbitrary molecules, SEu3)-equivaaraan modeling, and interpolation modeling.
1. Experiment and Data Difference: (Köhler et al., 2023) targets at modeling the transition, e.g., water molecules at different temperatures (no code or comments related to methane rigid modeling were found in the GitHub repository). The MD simulation can provide samples at the stable or cquilibrium status. Our work models the transition from unstable to stable conformations.
2. Dynamic Transition Modeling vs. Stable Structure Sampling: Unlike (Köhler et al., 2023) which focuses on stable structure sampling, our work models the dynamic transition process from weakly correlated (unstable) structures to strongly correlated (stable) structures. Notably, dynamic transition modeling toward stability is identified as a nontrivial next step in the ICML work (Köhler et al., 2023), which states "... flow model for the positions that can handle phase transitions', We also want to emphasize that the transition of molecular cluster from weakly correlated (unstable) structures to strongly correlated (stable) structures is a special case of the general dynamics. The limitation here is the data insufficiency (lack of intermediate snapshots), not modeling.
16

---

Published as a conference paper at ICLR 2025
3. Objective Difference:
(Köhler al., 2023) first introduces molecular equilibrium sampling with Boltzmann distribution in Eq 1,2. But this is not the goal, (Köhler et al., 2023) changes to estimating log-ratio VF = log( Zo, /Zog) between two configurations in Eq 3. This measure estimates the energy difference and tells which state is more stable. (Köhler et al., 2023) says using trackable priors like Gaussian can be biased, thus it takes the insight from previous work, and targets solving the problems Eq 3 and 7 (learned free energy perturbation). (Notice that other equations Eq 2,4.5 are preliminaries, not directly related to the core method in this work.) This reveals another theoretical difference between this work and our work: (Köhler et al., 2023) is estimating the upper bound of log-ratic between two stable states (with MD simulations), and ours is directly modeling p(X (X)) from unstable to stable transition.
4. Use of Inertial Frames for Rigid Modeling: frame for H-O-H (codes here) (no code or comments related to methane rigid modeling were found in the GitHub repository). In contrast, our approach is more generalizable, as the inertial frame can serve
(Köhler et al., 2023) specifically models the
reference frame for any molecule. Thus, (Köhler et al., 2023) cannot directly applied to our dataset, as is limited few types of constructed frames, and our work can be viewed as an extension of (Köhler et al., 2023) to more general setting.
5. -qqiieeqivarrint Symmetry Modeling: (Köhler et al., 2023) states that it has limitation on not exploiting the SE(3) symmetry of jointly moving all rigid bodies'. We solve this issue by introducing the SEuu)-equivarian modeling from two granularities.
6. Limitations of Coupling Layers in Normalizing Flow: (Köhler al., 2023) does modeling with an extra bijectivity constraint in coupling layers, limiting the model capacity (Ho et al., 2019). Flow matching enables flexible velocity functions under the interpolation framework.
We list the comparison in the table below:
Table 4: Comparison between (Köhler et al., 2023) and AssembleFlow.
Paper (Kither al., 2023)
AssembleFlow (ours)
Experiments
Transition between two stable conformations
Transition from unstable to stable
Data
Water and methane (experiment missing GitHub molocules
COD organic molecules
Frame construction
H-O-H frame specific for water molecale code
Incrtial frame for any organic molccule
comments related methane rigid modeling
were found in the GitHub repository)
Objective function
Upper bound log-ratio between stable
Direct estimation of conditional density from stable to stable
formations for energy difference
Modeling SFC3) symmetry moving
No
Yes
rigid bodics
Avoid bijectivity constraint in coupling layers
No
Yos
17

---

Published as a conference paper at ICLR 2025
B ROTATION REPRESENTATION
In this section, we will be mainly discussing three types rotation representations. It is important to note that "representation" here refers to the data structure commonly used in the machine learning community, rather than the concept of a representation space.
In Appendix C.1, we will explain how to use inertials to represent the rigid structures of molecules in a cluster. With such a rigid representation, we can then decompose SE(3) transformation into tuple of SO(3) and R$ transformations.
natural way to represent the SO(3) transformation is by using the rotation matrix, as will be introduced in Appendix B.1. However the rotation matrix is not flexible and must satisfy specific mathematical properties to make sure it is valid rotation in space, so we need a more flexible and efficient representation. To this end, we would like introduce axis-angle representation in Appendix B.2 and quaternion representation in Appendix B.3. The axis-angle representation and quaternion representation are closely related, and their transformation is discussed in Appendix B.4. Last but not least, the transformation between the rotation matrix and quaternion is in Appendix B.5, and the transformation between axis-angle representation and rotation matrix representation is in Appendix B.6. An overview of the rotation representation and the corresponding transformations are listed Figure 6.
Tubb
Ri + Rip Rip Ry Rig Ru Bi
Rig Riy Ri Rip Rys Mat
Riy Rg +
Roy Ra + My
d oor
Rotation Motrik R
Rg R23
Rig Ri
Ma Ria
Precipal Axees of Eigendwsompostton
Ny
2hind
2ay Zec
R
2xy Das
L-227 22?
20z 2ary
2ox Zere
C-xps
K=
u0
-M
(ron 4,80 y.a ut
Axis-angla
Rotation Quaternion
Rotation angle (1 and rotetion socis
(w: K.W.K,N k.
(M) H., H,),
Figure 6: Transformation between rotation matrix, quaternion, and axis anile representation.
B. ROTATION REPRESENTATION WITH ROTATION MATRIX
Definition Rotation matrix R E R3333 satisfying two conditions:
1. Orthogonality: The rows and columns of R are orthonormal vectors. R RRT = II
2. Determinant: The determinant R must be 1. det(F) =.
The set of all orthogonal matrices of size 3 with determinant is representation of a group known as the special orthogonal group SO(3). To generate a rotation matrix SO(3), certain properties for the rotation matrix must be satisfied, and such constrained generation is challenging. Thus, an alternative way of representing the rotation matrix is preferred. To this end, we consider axis-angle representation, described below.
Rotation with Rotation Matrix rotation operation is:
rotation point, p = (x) y/ 2), the 3D Euclidean space, the
x y
20
R11
R12
R13
rr
R
Rai
Rga
R23
y
(18)
Rai
Rgz
Rag
Properties
If we have multiple rotation matrices, and we want to yield a single matrix that combines these rotations into one, then have options:
18

---

Published as a conference paper at ICLR 2025
Extrinsic rotations. All rotations refer to fixed and global coordinate system, and the rotation matrices are ordered from the right to the left. we apply rotations R1. R2, and Rg, then it is written R R1 Intrinsic rotations. A rotation refers to the last rotated coordinate system, and the rotation matrices ordered from the left the right. If apply rotations R1, R2, and Rz. then it is written as R = Rí RgRR3.
B.2 ROTATION REPRESENTATION WITH AXIS-ANGLE REPRESENTATION
The axis-angle represents the rotation by its angle 0 and the rotation axis u. For example, a rotation
of 180 degrees around the Y-Axis would be represented as 0 = 180, (0, 1,0). The representation very intuitive, but for actually applying the rotation, another representation required, such as quaternion rotation matrix.
Definition
The axis-angle representation rotation is defined by two components:
1. Rotation axis: unit vector (uzy Uy Ug) that specifies the direction of rotation, ||]]]]
1.
2. A scalar / specifies the angle of rotation around the rotation axis.
For annotation, an axis-angle rotation can thus presented by four numbers as (0, L, 0. 2).
B.3 ROTATION REPRESENTATION WITH QUATERNION REPRESENTATION
Quaternion represents rotation by 4D vector and is more concise representation than a rotation matrix. It requires more math and is less intuitive, but is much more powerful representation. Quaternion representation has been widely used in rigid motion modeling, robotics modeling, and quantum mechanics (e.g., the spin of an electron and the polarization of photon). In this work, we are focusing on the first case, rigid motion modeling.
Definition
A quanternion is: w ri yj + 2k (w, v),
(19)
where x, y, 2 real numbers, and j, l are the fundamental quaternion units. The w is the real part, and = (x, y, 2) is a 3D vector representing the imaginary part.
Multipliation of basis elements
The multiplication for the basis elements i, lo is defined below:
f = j² k ij -ji = k: jk kjj
(20)
ki ik:
ijk -1
Quanternion multiplication: Hamilton product a.k.a., Hamilton product. For two quaternions r (ro. ri, T2, and (so, 81, 82, 83):
This can give us the quaternion multiplication,
t=r8, r8,
(21)
where
tg = To80 F181 P282 T383
4 To + T180 T283 T382
(22)
t2 = To82 + r183 T280 T381
t3 To83 T182 r2$1 + T380-
Point rotation with quaternion We rotate point (Uz, By U2) by the quaternion q (w, xr y, z) using the following three steps:
1. Transform quaternion p (0, Va, Ug U2).
19

---

Published as a conference paper at ICLR 2025
2. Construct the conjugate quaternion (w,- -2, y. -2).
3. There are two types rotation operations:
(a) Active rotation: p' q pq. when the point is rotated w.rt. the coordinate system. (b) Passive rotation: p' = qpQ", when the coordinate system is rotated w.r.t. the point. Notice that the two rotations opposite from each other. In case, the passive rotation.
4. The resulting vector v' is extracted from the imaginary part of p'.
Spherical Interpolation (SLPER) for Quaternion Interpolation Quaternions are often pre- ferred for interpolating between rotations because they offer smoother interpolation than axis-angle representation. The spherical interpolation defines the geodesic over the rotation group.
sin((1 )) + sin(tw)q:
SLERP(qu- q1 t)
(23)
sin(w)
where w is the angle between qu and qı-
Properties
A quaternion is unit quaternion if ||و|| w² + x² +
All rotation quaternions must be unit quaternions.
A rotation of Qa followed by rotation of qu be combined into single rotation:
aoqa. Note that order matters.
The conjugate
quaternion q" (w, -I, - 22).
The inverse of rotation quaternion is q' Then we can see that qq qq"
(1, 0, 0, 0).
QQattemminn multiplication is associative: abc a(bc).
Quaternion multiplication not commutative: ab + ba.
B.4 TRANSFORMATION BETWERN AXIS-ANGLE AND QUATERNION
Axis-angle representation to quaternion representation Axis-angle representation is u (ugs Uy Uz) equipped with a rotation angle 0, and the rotation quaternion is unitary, i.e., w² x² + y² + +2 The quaternion is thus:
0 + sin ( (u, + Ug + u,k),
q cos
(24)
2
or equivalently in the vector form:
q (cos
sin
We sin
My sin 2 Uz).
(25)
2
2
Quaternion representation axis-angle representation convert it to axis-angle representation:
Quaternion is q
(w, x, y. 2). and to
1. Compute the angle 0 2 cos (w).
2. Compute the axis u.
(0. 2) V 22 22 22
(x, 16
u
(26)
sin 8 2
B.5 TRANSFORMATION BETWEEN QUATERNION AND ROTATION
Quaternion to rotation matrix
Given a quaternion
(wi, v. y, z), the corresponding rotation
matrix is
1 2y²
22 2ry + 2wz
2ry 2wz
2r 2wy
R =
1 2x
2x
2yz 2wr
(27)
2rz 2wy
2yz + 2wx
1 2x 2ny
20

---

Published as a conference paper at ICLR 2025
R11
R12
R13
Rotation matrix to rotation quaternion Given a rotation matrix R
R21
Rz2
R23
then
R3ı
R32
R33.
we first calculate the magnitude of four quaternion components as below:
R1 + Roz + R33
lgol
R11 Ro2 R33
1ga1
(28)
R11 + Ry2 R333
lual
Rn
R22 + R33
lgsl
To find the signs of the four elements, we can find the largest magnitude:
If Igol is the largest, then
Ra2 R23
R13 Rsi
R21 R12
(29)
4w
4w
4w
If 191| is the largest, then
Raz R2a
R12 + R21
R13 + Rai
x 41
(30)
4x
4x
4x
If [92| the largest, then
R13 Rai
R12 + R2i
R23 + R32
y 92
(31)
4y
4y
4y
If 19s] is the largest, then
R21 R12
R13 + R31
R23 + R32
2 43
x=
y
(32)
42
42
4%
The sign is ambiguous because any rotation has two possible quaternion representations, If one is known, the other one can be found by taking the negative of all four terms.
Besides, there exist other solutions, for instance, extracting quaternion from rotation matrix using cigendecomposition (Horn, 1987; Bar- Itzhakk, 2000). We first construct a matrix K with:
R + R2z + Raa
R32 R23
Ria Rai
R21 R
Rg R23
Ru R22 Ra3
R12 Rai
R13 + Rai
K
(33)
3
R1 Rs1
R12 + R21
Ryg R11 Rgs
R23 + R32
Re1 R12
R13 + R31
R23 Ra2
Ra3 R11 Ro2
Then we perform eigendecomposition K VAVT, where a diagonal matrix with eigenvalues and V is the matrix with eigenvectors as columns. Finally, pick up the eigenvector w.r.t. the largest eigenvalue, and this eigenvector is the unit quaternion.
B.6 TRANSFORMATION BETWEEN AXIS-ANGLE REPRESENTATION AND ROTATION MATRIX
Axis-angle representation to rotation matrix
Construct the skew-symmetric matrix:
0
U2
Uyy
K
U
0
(34)
U,y
Ug
0
According to the Rodrigue's rotation formula, we have:
R
sin(O) ()
(35)
21

---

Published as a conference paper at ICLR 2025
Rotation matrix to axis-angle representation of R, i.e.,
The rotation angle can be computed using the trace
Tr R
0 = cos
(36)
2
where Tr( is the trace, Tr(R)= Rui + R22 + R33.
Then we calculate the rotation direction or rotation axis u
(tez My: Uz), as
R32
R23
R13 R31
R21 R12
Ug
U3 2sin 0
(37)
ssin
sin 0
Notice that when = 0, there is no definition of rotation axis and the whole rotation matrix R is unitary. When 0 n, because sin 0 O, then we need other methods (e.g., eigendeccompostionn) to determine the rotation axis.
22

---

Published as a conference paper at ICLR 2025
C INERTIAL FRAME FOR RIGID BODY
C.I RIGID REPRESENTATION OF MOLECULE
We employ the following four sequential steps to derive the reference frames that construct the rotation matrix from N atomic positions
Calculate the mass center: с - + Eir:. Pi.
Adjust position relative to the center T;
Compute the inertia tensor 1 = E. |1:11 - r.r... where / is the unit diagonal matrix.
Obtain the principal axes of inertia by applying sheendecomposition on Î. We have î QAQ", where is the orthogonal matrix whose columns are the eigenvectors of f, and diagonal matrix whose elements are the eigenvalues A; of i, representing the principal moments of inertial along the principal axes.
C.2 ORTHOGONAL MATRIX
In linear algebra, orthogonal matrix or orthonormal matrix is square matrix whose columns and rows are orthonormal vectors. This can be written as
Q'Q QQ" =1.
(38)
This leads the equivalent characterization: matrix Q is orthogonal if its transpose is equal its inverse:
Q" Q-1
(39)
Notice that when discussing matrices, the two terms (orthogonal and orthonormal) can be used interchangeably.
If Q is square matrix, then the conditions RR = / and R R are equivalent. Proof sketch: Rd R = and RD RR R. so R R-1, This can give us RR I.
C.3 ROTATION MATRIX FROM RIGID BODY
From Appendix C.I, we can construct the inertial tensors. Then we employ eigenvalue decomposition on the inertial tensor. The normalized eigenvectors v1, U2 vg form orthonormal basis, which can be used to construct the rotation matrix, i.e.,
R
U2
vz].
(40)
Figendecomposition Inertial Tensors For inertial tensor J, the decomposition is: with Tu; , where A; are eigenvalues and E R3x1 are eigenvectors. Thus, we can have
lu X
IR RA
(41)
I RAR-1
Coo
0
0
where
0
A1
0
Because the inertial tensors are symmetric matrices, we have that
0
A2
matrices orthonormal matrices.
C.4 EXPLORATIONS ON OTHER REFERENCE FRAME OPTION
One critical question is in addition the Inertial frame, do have other options for modeling the rigidity? One simple solution is to directly apply the eigendecomposition as principal component analysis (PCA) the point clouds of centered molecules.
First, we would like to clarify that there are two roles and one important property the inertial frame and its eigenvectors:
23

---

Published as a conference paper at ICLR 2025
Role 1:
The three bases in inertial frames act as reference canonical pose.
Role 2: The three bases enable modeling the velocity function in SO3 space.
What's 's more important, expect the three bases to be numerically stable.
Then getting back to the question, though both can help build up the reference frame or canonical pose, there are certain aspects would like to emphasize when comparing PCA and Inertial frames.
Canonical pose: The key question is defining canonical pose. If we just do PCA on the point clouds, this cannot guarantee the group symmetry in SE(3). However, if we remove the center point, then the group symmetry can be guaranteed; then we can use SVD to get find this is somewhat similar the
the three principal components. Till this step, inertial frame construction (Sec 3.1).
Difference between PCA and Inertial Frame as reference frames: Though both can be used for building up the reference frame or canonical poses, SVD (for PCA) on the set of point clouds N 3 (N the number of atoms) can be less numerically unstable, while eigendecomposition Inertial Frame 3 x 3 can be more numerically stable. We conduct an experiment to verify this.
Experiment setup: Suppose we have weakly-correlan correlated structures X1. and we find the corresponding bases using either eigendecom- position centered X or inertial frame I. The two bases marked Bo 3x3 and B1 E 3x3
structures Xo and strongly-
Objective: We obtain the rotation matrix with Rd BJJ B1. and then the whole molecular system as Xi Xo Rt and are measuring the reconstruction errors as MSE(X; X1). ). We mark the MSE using inertial frame and PCA as d and ÔPCA. respectively. If óinettill < ÓPCA, then we can conclude that using the inertial frame is more stable than using PCA, and vice versa. Notice that since the MSE reconstruction is meaningless when it is too small, so we only compare these two frames when least one of them has reconstruction greater than or equal a threshold 0.
rotate
Results: The comparison results are in Table 5, and we can observe that in general, using the inertial frame is more stable than PCA. We are listing multiple reconstruction threshold 9 in Table 5, and we are using 1e in the main article.
Table 5: Comparison of using inertial frame and PCA for reconstruction (%).
le-3
le-4
le-5
le-6
le-7
1e-8
P ((Sneniill < Ôpcx)
0,434
0.884
1.338
1.794
2.254
2.727
P P(Oteniaal SPCA)
0.371
0.756
1.147
1.539
1.934
2.345
C.5 KABSCH ALGORITHM
Kabsch algorithm is way to compute the optimal rotation matrix that minimizes the root-mean- square deviation (RMSD) between two sets of points (atoms in our case). However, it guaranteed in the 10D-Cluster17 dataset that the molecules in weakly correlated structures rotate molecules in strongly correlated structures; in other words, the RMSD can be approximately 0 if we use the Kabsch algorithm, which is equivalent to calculating the rotation matrix directly after we fix the poses. We have shown how to calculate the rotation matrix in the experiment above.
24

---

Published as a conference paper at ICLR 2025
D PROOF OF THEOREM
Proof. For three vectors, can easily find counter-example as illustrated in Figure 3 (a, b). Figure (a, b) describes two cases where we have the same initial frame, and we can rotate to two different final frames with rotation matrices, yet the righthandness still matches. We can easily see that there are four options of rotation matrices in this case, and we cannot uniquely determine the final inertial frame in this case.
More rigorously, let us first assume that there exists a rotation transformation R that can transform the initial coordinate system Q; the final coordinate system Qy. as:
TQ T Q
Q0,07
R
Q
(42)
Q
Q1,
First, Lemma 1, we should change either zero or two directions for direction alignment.
Then without loss of generality, we can assume the two directions to be the last two axes. Thus, we obtain rotation matrix Rº such that H' is rotating R. along vector Q,,0 with 180 degrees, We can represent It' using Rodrigue's rotation formula, Rt = (2Q I)R. Thus, we can have:
Qu.07 T Q Q2.2
Q
R'.
=2Q
1)
Q
Q1111
(43)
122
This is essentially saying starting from one initial frame, we can have multiple matched final frames. Thus, using only three vectors cannot uniquely determine the direction matching. We provide two examples Figure (a, b).
For the four vectors, we introduce an extra atom into the inertial frame system, and such an extra atom point is nonplanar to the three base axes. Then the problem becomes: starting from an initial frame and extra point, can we find multiple rotation matrices such that the final frames have reflected directions? To be more rigorous, let us have the following formulation.
First, let assume we have this rotation matrix:
T
T
Q
T
Q
R
(44)
Q v
Qi,2
v
As discussed above, we need to guarantee the right- heht-handdedees property, thus, without loss generality. here we also assume the last two axes are reflected. The question turns to: does exit another rotation matrix R', such that:
Q
T
QQ,07
T
Qr11
Q
= R'.
(45)
Q1.2
Q
We now use contradiction. Since we still have the two axes rotated 180 degrees around the first axes. Qro. so R² = (2Qr.0Q7 R)R Then given the two conditions v" Ro" and v" RR'"" we have (2Qr,0Q7. T)oT" T
25

---

Published as a conference paper at ICLR 2025
If we let y [K1, N2, kis and
[Ui, U2, 03), then have
I)v"
v²
Thy k1
kg
[11
2/1
ki
kik2
k2ka
1/2
A2233
kikk
kaka_
Loa
033
[k (1101 u1 + k2 + k3tz)
2/1
(46)
k2 (2 + k2 + k3t3)
2/2
(kivi + k2 + kgdz))
033
Tki
(k1 + k2 + k3v3)
22
2/2
K3
L33
After calculation, obtain that Qr,0 cu, where is coefficient. However, claimed in the condition, v does not lie in the same line as Q thus, there does not exist such another rotation matrix R' * R satisfying Equation (45). We also provide examples in Figure 3 (c, d).
By contradiction, we can tell that there only one unique rotation mapping from the initial inertial frame to the final inertial frame.
To sum up, three points cannot form rigid structure in Fuclidean space, thus there can exist multiple reflection transformations, leading to opposite inertial frames. Four points can form rigid structure, thus there exists only one reflection transformation.
26

---

Published as a conference paper at ICLR 2025
E PROBLEM FORMULATION AND MORE DETAILS OF ASSEMBLEFLOW
E. PROBLEM FORMULATION
We would like to emphasize that previous works aim at atomic level modeling, while our proposed AssembleFlow focuses on molecular level modeling. Meanwhile, both models need satisfy the SEvuv-equiv,riance, as detailed below.
Atomic Level Modeling Existing deep learning frameworks have been using atomic level modeling. For each atom r, the inference step is:
T++1 = Pt + 30,t,
(47)
s.t. o. is -quiieaquivaiinn.
Thus, we can observe that such problem formulation cannot guarantee the rigidity of each molecule during the crystallization process.
Molecular Level Modeling our proposed AssembleFlow, it learns the translation and rotation at the molecule level. For each atom r, the inference step is:
r++1 (r, + æ0,t), s.t. Dg, and Rø.: are SE(3)-equivariant
(48)
The Ee.s and RØ., are molecular level modeling. Notice that this also holds after we take the reparameterization, discussed in Section 3.4. In Appendix E.2, we will explore how define the SE(()--quivariant models on top of that.
E.2 EQUIVARRAT VELOCITY FUNCTION
We consider types of SEu()aaquivariaat models the velocity function. As shown in Figure the inputs are the same for learning: the positions at the initial step and the final step, respectively, We take the position of the mass center for each molecule in the cluster to obtain the translation in E (æ; and æ f), and we take the first principal axes of inertial frames to obtain the reference coordinate system for rotation in SO(3) (q, and a, with alignment). Then we adopt Equations (6) and (7) for the interpolation on SO(3) and R group respectively, which gives translation æ, and rotation q at interpolation time / E [0, 1].
Recall that the data structure considered here is the cluster of molecules, thus it is natural to split the modeling into intra-molecule and inter-molecule modeling, as introduced below.
Intra-molecule Modeling. For each molecule the cluster, we adopt the SE(3)-equivariant PaiNN (Schütt et al., 2021) to obtain the representation for each atom. Such an atomic representation can encode the inherent geometric structural information of individual molecules, which can be passed to inter-molecule modeling in the next step.
Inter-molecule Modeling. This step aims to model the inter-molecule interactions during the molecular crystallization process based on the intra-molecule representation. We can have two options for SE(3)-equivariant inter-molecule modeling: (1) to project æ, and 9. back obtain the atom-wise position and do modeling, as in Figure 7(a), (2) to directly perform molecular level modeling on molecular-level translation æ; and rotation Qt. as in Figure 7(b).
Atoiic level modeling. This means we build up the SE(3)--quivariant models on top of atom positions r; at time f, and the outputs are the rotation and translation for step t + 1 or step T' (if we use reparameterizationn..
Obtain intra-molecule representation h. using PaiNN (Schütt et al., 2021).
Obtain time embedding h2 with positional encoding (Vaswani et al., 2017).
Build up the vector frame basis (Liu et al., 2023) for each atom F., based its neighborhoods.
Then update the atomic representation h the summation of ha and h
27

---

Published as a conference paper at ICLR 2025
Inlitlatom pooaion
Final stom posiion
Intal stom posttion
Firdl stom sootion
Transition Rosszon Og
Trcrkition 4
Taandation 22
Transation X Roteton Quaaeroo
Reaaten Q
Nottton Cb
Quarmio 5j
Cuntersion g)
Quanteriin
rrrmiltton 4 and quatemson
trur trarslatoon quatemicn
coo pootion
netecuie rppreseraation
Prpliatant Geomenii
Erplasaart Geometio
Repertinleiad fira
Reperemeteetzed tranolaion
Popaoneeriaod rotaton
(a) Atomic level molecle modeling.
(b) Molecular level molecle modeling.
Figure 7: Illustration of two types of SE(3)-equivarian inter-molecule velocity functions.
Apply message passing (Gilmer al., 2017) to exchange the information between each atom and its neighborhood on top of the atomic representation h., and vector frame basis Fa-
The outputs include the atomic scalar representation ha, E R and atomic vector representation ha." E IR x3xd where N is the total number of atoms in the cluster. Then conduct the aggregation to obtain the molecular level predicted rotation velocity q E R x3 and predicted translation velocity i E Rxx3, where M is the number molecules in the cluster.
This holds similarly if we are going to predict the final rotation âı.0 and final translation vector 110 when using the reparameterization
Molecular level modeling. This means we build up the SE(3) models top of Ø; and q. at time I, and the outputs are the rotation and translation for step + I or step T (if we are using reparameterization).
Obtain intra- molecule representation hg using PaiNN (Schütt et al., 2021).
Obtain time embedding h, with positional encoding (Vaswani et al., 2017).
Then we update the atomic representation h, as the summation of ha and h4-
Aggregate the atomic representation get the molecular representation h... AAply the message passing (Gilmer et al., 2017) to exchange the information between each molecule and other molecules the cluster, where the interactions are treated as the forces in the moter-molecule level in the cluster.
-
-The outputs include molecular level predicted rotation velocity Åo Mx3 and predicted translation velocity ba IR M x33 where M is the number of molecules in the cluster.
This holds similarly if we are going to predict the final rotation Qı.0 and final translation vector L when using the reparamcterization.
Notice that for the predicted quaternion qo. we only predict the imaginary part. We then concat the real part as 1, followed by normalization to make a rotation quaternion.
E.3 ALGORITHM
In this section, we provide the pseudocodes for AssembleFlow. Following Section 3, we illustrate the reparameterized version here. The training and inference algorithms are in Algorithms I and 2, respectively. We would also like to emphasize the following prior steps: (1) We first construct the inertial frames discussed in the main article. (2) Then construct the coordinate system (in rotation matrix) and conduct the alignment between the initial and final frames. (3) Last but not least, transform the coordinate system from rotation matrix quaternion as discussed in Appendix B.5.
28

---

Published as a conference paper at ICLR 2025
Algorithm Learning of Assembleflow.
1: Inputs: For N atoms in M molecules, we have atomic level initial position ro E R molecular level initial rotation quaternion To xx and translation 2o E IR atomic level final position R molecular level final rotation quaternion qi E IR and translation
E RMxx3 timestep T IR, epoch E R, coefficients oo, 01 E IR.
2: for epoch e E] do
3: Sample t E [1.
4: Conduct I.ERP to obtain translation æ at time I, following Equation (7).
5: Perform SLERP to obtain quaternion q at time i, following Equation (6).
6: Predict the final quaternion Ii = Q1,0 (Zt, qz, t) and translation # = 11 (we, q: t) using
SEuiv-aqui modeling discussed in Appendix E.2. Minimize loss = Ckp ,eppaameee
7:
reparameter + 0t2 Lso(3) reparameters as defined in Equation (13).
8: end for
Algorithm 2 Inference of AssembleFlow.
1: Inputs: For N atoms in M molecules, have atomic level initial position ro R-Nx3
molecular level initial rotation quaternion god and translation 20 RRxx3, timestep T
R, learned SE(3)-equivariant models final rotation quaternion fn,o((;, qe: and translation
L1,0( Qt: Qt-t). t).
2: for timestep é E [] do
3: Predict the final quaternion 9i - (3t, gt, () and translation á 10 (2 qt, () using SE(3)--quivariant modeling as discussed in Appendix E.2.
4: Calculate the next-step quaternion qi41 and translation क141 as Equation (14).
5: Move the cluster of molecules w.r.t. at++ and ø;+1-
6: Obtain the corresponding atomic positions T++1 time + following Equation (48).
7: end for
8: The final predicted crystal structure fT..
E,4 HYPER-PARE PARERRRRR
We provide the key dyper-parametenn of AssembleFlow in Table 6.
Table 6: Hyperparameter specifications for AssembleFlow.
Model
Hyperparameter
Value
Intra-modeling
PaiNN
embedding dim
[128]
num of layers cutoff
(3]
[5)
read out
[mean] J
Intra-modeling Atomic Level
of layers
[2.5]
num of convolution
(2)
num of head num of timesteps
[4 8)
500, 200)
Og (a)
[)
[1. 10)
Intra-modeling Molecular Level
of layers
[4,5)
num of head
[. [4
num of timesteps
(50, 200)
(1)
[1, 10)
Optimization
seed epochs cutoff e learning rate
42, 123)
1000, 2000)
20, 20,50)
le-4, Sc-4)
optimizer
Adam
29

---

Published as a conference paper at ICLR 2025
F DETAILS OF BASELINES
F.1 DEEP LEARNING BASELINES AND PARAMETERS
Notice that for all the baselines listed below, also adopt the PaiNN for atomic level representa- tion (Schütt et al., 2021), and the hyperparameters are the same Appendix E. We list the remaining hyperparameters of baselines in Table 7.
Table 7: Hyperparameter specifications for deep learning baselines.
Model
Hyperparameter
Value
GNN-MD
seed epochs num timesteps
[0, 42, 123)
3000
(50, 200}
cutoff
(00, 50)
learning rate
lc-4, 5c-4)
optimizer
[Adam }
CrystalSDE
SDE type seed epochs num of timesteps
(VE, VP [0. 42, 123)
3000
[50, 200]
cutoff learning rate optimizer
200 50}
le-4, Se-4]
Adam
CrystalFlow
interpolation type
(VE, VP. LERP) [0, 42, 123)
seed epochs num of timesteps
3000
[50, 200)
cutoff lcarning rate
(20, 50)
1c-4, 5c-4] (Adam }
optimizer
F.2 PACKMOL
We want mention that using PackMol for evaluation with the packing matching metrics is non- trivial. This is because the atom ordering in each molecule and the molecule ordering in each cluster are different between PackMol simulated results and the ground-trutt truth results. We that deep learning methods do not have this issue because the orderings of initial atoms/molecules match with the final atoms/mmolecules.
Thus, to address this issue, we first use the Hungarian algorithm to match the mass centers of simulated results obtain least matching distance with the ground truth mass centers, i.e., PM (center). This gives us the molecule ordering mapping from simulated clusters to the ground-truth clusters. Then for each molecule simulated-and-gree anndgrouund-truhh pair, we apply the Hungarian algorithm again to obtain the minimum distance for alignment.
30

---

Published as a conference paper at ICLR 2025
G DATASET AND EVALUATION METRICS
G.1 DATASET
We evaluate our method using the crystallization dataset COD-Cluster17 from (Liu et al., 2024c). This COD-Cluster17 is curated subset derived from the COD database (Grazulis et al., 2009). We note that some single molecules/substanees can crystallize in different forms, known as polymorphs. This arises due to changes in the configurations during the process, such as the environment temperature, pressure, and solvent. COD-Cluster17 simplifies this setup by ignoring the configuration information and treats the crystallization problem density estimation problem.
G.2 DETAILS OF EVALUATION METRICS
We illustrate five types of evaluation metrics below. Notice that in the original dataset, the dynamics trajectories of molecules missing. Thus, evaluation is based on the ground truth cluster geometry at the last step.
Packing Matching (PM) This metric quantifies how well the generated molecular assemblies match the reference crystal structures in terms of spatial arrangement and packing density (Chisholm Motherwell, 2005). is a key indicator of how accurately model can replicate real-world crystallization patterns. We provide both the atomic MP, denoted as "PM (atom)" and mass-center- level PM, denoted as "PM (center)"
Collision This follows (Cordero al., 2008; Liu et al., 2024b). measures if there is any atomic collision in the predicted assemblies. Atoms must maintain minimum pairwise distance governed by the balance of attractive and repulsive forces. More concretely, we are using covalent radii as the most strict metric for atomic collisions in molecular generation. This is because it provides a precise lower bound for the distances between atoms when they are bonded. other words, covalent radii represent the distance at which two atoms form stable covalent bond, which is very close and well-defined interaction compared to non-covalent interactions. However, other types of atomic radii, such van der Waals radii or ionic radii, can be used for different purposes, depending the nature of the interaction you're modeling.
Separation We extend the metric from (Xie al., 2022; Yang al., 2024) setting. A cluster of molecules is valid if the minimum distance between molecules is above 0.5A (Court et al., 2020). This metric is referred to as separation to measure the validity to avoid unphysical interactions at the molecular level.
Compactness We propose this measure by calculating the percentage of simulated clusters where the maximum atomic pairwise distances are below 100A. This assesses the spatial efficiency of the molecular assemblies, indicating how closely the constituent molecules are packed together. higher compactness value suggests a more efficient arrangement, where the intermolecular spaces are minimized, leading to a denser crystalline structure.
31

---

Published as a conference paper at ICLR 2025
H
ABLATION STUDIES
H.1 ABLATION STUDIES ON RANDOM SAMPLING
Here, we add another baseline by randomly sampling translation and rotation.
For (3), we can do random sampling.
For IR3 we first obtain the range of atom positions in the training data, and then we just do uniform sampling within this range.
The results are in Table 8. As observed in Table 8, the Random baseline performs exceptionally well across all three validity metrics; however, its packing matching is significantly worse by an order of magnitude.
Table AssembleFlow against six generative models on COD-Cluster17 17 with 5K, 10K, and all samples. The best results are marked in bold. Baseline Random has the best validity metrics, but they are meaningless since the packing matching is extremely high, remarking that the results collapse. Thus, we mark them in gray.
Packing Matching
Validity
PM (atom)
PM (center)
Collision
Separation
Compactness
Dataset:
COD-Cluster17SSK
Random
54.07 0.42
54.62 0.43
0.31 0.01
99.88 0.01
100.00 0.00
GNN-MD
13.67
0.06
13,80 0.07
27.53 + 0.49
0.22 + 0.11
100.00 0.00
CrystalDDE-VE
15.52 + 148
16.46 0.99
1.20 0.08
0.86 0.86
57.47 7.76
CrystalSDE VP
CrystalFlow-VE VE
14.87 + 7.07
13.08 + 4.51
1.37 007++0.000 0.04
35.70 +:700:0000 0.73
8.40 4.17
CrystalFlow-VP
15.71 2.69
17.10 1.89
1.38
35.43 = 0.88
4.87 1.09
CrystalFlow-LFRR
13.59 + 0.09
13.26 0.09
0.34 0.34+000
97.38
0.10
100.00 0.00
Assemllelow (ours)
7.27 + 0.04
6.13 0.10
0.33 - 0.00
97.64 + 0.36
100.00 0.00
Dataset: COD-Chuster17-108
GNN-MD
13.83 +.06
13,90 = 0,05
27.88 0.49
0.23 0.11
100.00 0,00
DryttalSDE-VE
17.25
17.86 1.11
0.99 0.27
22.99=10.77
34.93 14.99
CrystalSDE-VP
2.220
21.39 1.50
0.53 0.33
52.48 15.44
16.83 18.09
CrystalFlow-Ve VE
16.41
16.71
1.0
1..5
5.47 0.47
CrystalFlow-VV VP
19.39
16.01
33.35
4.23 0.48
CrystaFFlow-LERP
13.54 = 0.03
13,20 : 0.03
0.00
97.32
100.00 0,00
AssembleFlow (ours)
0.03 7.38
6.21 0.05
0.31 0.00
97.73 0.16
99.93 0.05
COD-Cluster77-M
Random
0.07
0.07
0.30 -: 0.00
99.91 0.00
100.00 0.00
GNN-MD
22.30 12.04
14.51 0.82 0.82
24.29 2: 4.558
4.13
98.77 1 1.73
CrystalSDE-VE
17.28 0.73
1892
0.19
12.42
2.51 + 2.37
CrystalSDE-VE VP
18.03 4.566
20.02 3.70
0.55
48.78 - 1.70
6,88 2.82
CrystalFlow-Vy -EE
12.80
15.09 1 0.34
141
35.34 0.28
2.90 1 0.02
CrystalFlow-VP
0.44
13,28 0.48
1.02
33.06 1.31
3.17
CrystalFlow-LEE -ERR
13.61 1 0.00
13.28 0.01
0.34
97.34 1 0.02
99.99 0.01
AssemleeFlww (ours)
6.21
0.22
0.00
32

---

Published as a conference paper at ICLR 2025
H.2 ABLATION STUDIES ON INTERPOLATION ON SO(3)
Empirical results. We consider replacing the SLERP with EMLERP in AssembleFlow, and name it AssembleeFlowEMMLERR We conduct the experiment COD-5000, where we taking the optimal hyperparameters from AssembleFlow.
The results Table As observed, using SLERP is better than EMLERP. We are still running results for COD-10k and COD, and will update the results later.
Table 9: AssembleFlow against six generative models on COD-Clusterl with 5K, 10K, and all samples. The best results marked in bold.
Packing Matching
Validity Separation
PM (atom) 1
PM (center)
Collision
Compactness
Dataset: COD-Chster17-SK
GNN-MD CrystaSDDE-E VE
13.67
13,80 0.07
27.53 + 0.49
0.22 0.11
100.00 0.00
15.52 1.48
16.46 1: 0.99
1.20 +: 0.08
27.17 2 0.86
57.47 1: 7.76
CrystalDEE VP
4.46
CrystalFlow-VE VE
14,87 = 7.07
13.08 4.51
1.37 + 0.04
35.70 0.73
8.40 1: 4.17
CrystalFlow VP
CrystalFlow-LFFR
13.59 0.09
13.26 +-009
0.34 +- +.001
97.38
0.10
100,00 + 0.00
CrystalFlow
7.30
0.04
6.32 1. 0.04
0.37 0.01
93.38
0.54
100,00 1. 0.00
AssemhleFlow (ours)
7.27 0.04
6.13 0.10
0.33 +: 0.00
97.64
0.36
100,00
0.00
COD-Cluster1-10KK
GNN-MD MD
13.83 0.06
13.90 0.05
27.88 1: 0.49
0.23 1. 0.11
100,00 1. 0.00
CrystalSDE-VE
17.25 = 2.46
17.86
1.11
0,99 + 0,27
32.99
10.72
34.93 14.99
CrystaISDEV VP
22.20
3.29
21.39
1.50
0,53 1: 0.35
52.48
15.44
16.83
18.09
CrystalFlowwVy VE
16.41 2.64
16.71
2.35
1,42 + 0.03
33.79
0.51
5.47 0.47
CrystalFlow-Ve VP
19.39 4.37
16.01 1. 3.13
1.44 1: 0.03
33.35
0.55
4,23 1: 0.48
CrystalFlow LERP
13.34 0.03
13.20
0.03
0.32 + 0.00
97.32
0.05
100.00
0.00
AssembleFlow EMLERP
6.46
94.68
AssembleFlow (ours)
7.38 0.03
6.21 + 0.05
0,31 + 0,00
97.73 + 0.16
99.93 + 0.05
Dataset: COD-Cluster 17-A1l
CrystalDDE-VE VE
17.28 = 0.73
18.92 1: 0.03
0.19 1: 0.18
15.47 :: 12,42
2.51 1: 2.37
CrystalSDE VP
CrystalFlow-V1 VE
12.80 1.20
15.09 + 0.344
1.41 + 0.001
35.34 + 0.28
2.90 + 0.02
CrystalFlow VP
13.50
CrystalFlow-LFRR
13.61 0.00
13.28 0.01
0.34 + 0.00
97.34 + 0.02
99.99 + 0.01
AssemhleFlow-EMII EM..FRP
7.28 0.00
6.23 0-001
0.35 + 0.00
93.17 + 0.02
99.98 0.00
AssembleFlow (ours)
7.37
6.21 0.01
+ 0.00
98.15 0.22
99.98 0.00
33


Reviews:
[{"review_id": "Q40lFqbA32", "reviewer": "Reviewer_sWLW", "summary": "The authors present the AssembleFlow framework, which continuously matches weakly correlated structures (like liquid water) to their strongly correlated counterparts, such as crystallized forms like ice. They have carefully designed an underlying theory that leverages advances in flow matching methods to address this problem. AssembleFlow learns SE(3) velocities by decomposing them into translations and rotations. Unlike other methods, they use quaternion representations for rotations, which allows for spherical interpolation. The authors also introduced the idea of reparameterization to directly model SE(3) at time T.\n\nThey performed a comprehensive evaluation by comparing results with other models and with PackMol, a standard simulation-based approach well-established in the field. The authors also computed the distribution of DFT energies for both PackMol and AssembleFlow structures. The method is shown to be 25 times faster than PackMol.", "questions": "Questions:\nEnergy Comparison in Figure 4:\nInstead of comparing the distributions of DFT energies separately for PackMol and AssembleFlow structures, could you plot the energy differences between PackMol and AssembleFlow for the same structures? Generally, interpreting absolute energy values can be challenging.\n\nTypo in Table 2:\nIn Table 2, there appears to be a typo. The separation should be indicated with an upward arrow (↑).", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Strengths:\nPhysics-Based, Theoretically Sound Methodology: The approach includes utilizing quaternions for rotations, spherical interpolation, and using an internal reference frame.\nComprehensive Comparison with Other Methods: The authors provide thorough comparisons with existing models.\nUse of Valuable DFT Energy Calculations: Employing DFT energies as a measure of the stability of the obtained structures adds credibility to their results.", "weaknesses": "Cons:\nPerformance Compared to PackMol: While AssembleFlow shows good results, it does not outperform PackMol. The considerable differences in the energy distributions between PackMol and AssembleFlow structures highlight the advantages of PackMol over AssembleFlow.", "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": "PcntKJ4IQb", "reviewer": "Reviewer_1PjM", "summary": "AssembleFlow introduces a generative model for assembling clusters of rigid molecules. The authors assign inertial frames to each molecule, which reduces the problem to generative modelling on $SE(3)^N$. The authors then solve this with Riemannian flow matching, using linear interpolants for translations and spherical interpolants of quaternions for rotations. The trained model approaches the performance of PackMol but at a much lower computational cost.", "questions": "How are permutations handled? Are the molecules stored in a particular order? Since each molecule has identical structure, flow matching could learn shorter paths if shown examples that have been permutation-aligned by the Hungarian algorithm - see equivariant flow matching [1,2]\n\nI am curious about the relationship between SLERP of quaternions and geodesic interpolation on SO(3) as discussed by FoldFlow. Are they mathematically equivalent?\n\n[1] Klein, L., Krämer, A., & Noé, F. (2024). Equivariant flow matching. Advances in Neural Information Processing Systems, 36.\n\n[2] Song, Y., Gong, J., Xu, M., Cao, Z., Lan, Y., Ermon, S., ... & Ma, W. Y. (2023). Equivariant flow matching with hybrid probability transport. arXiv preprint arXiv:2312.07168.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The work introduces the task of rigid molecular crystallization in the context of deep learning, with significant potential downstream applications. Applying flow matching to spherical interpolation of quaternions is also novel. The paper is clear, well-written, and well-organized, with helpful figures to explain each concept. I appreciated the extensive background provided throughout the text, including on rotation matrices in the appendix.", "weaknesses": "The novelty in this work is mainly in solving a new application domain, rather than pure machine learning developments. As acknowledged by the authors, flow matching models on $SE(3)^N$ were first developed for assembling rigid residues to design protein backbones, such as FoldFlow [1]. The authors claim that inertial frames are a key innovation over models like FoldFlow, since arbitrary molecules do not have a unique way to assign frames, unlike protein residues. However, if the goal is just to uniquely interconvert between atomic coordinates and frames, this can be achieved by storing a canonical pose for each molecule and applying the Kabsch algorithm [2] to align each molecule in a cluster to the canonical pose, obtaining a rotation matrix for each molecule. This generalizes the proposed inertial frames, which specifically use a canonical pose from applying principal component analysis on the point cloud.\n\nThe selected baselines are a somewhat unfair comparison, as none of them use rigid molecules. A simple baseline to accomplish this could just be sampling thousands of random translations and rotations of rigid molecules, and filtering out structures with collisions. This would help emphasize the contribution of this work as a method for *intelligently* assembling rigid frames, not just exploiting molecular rigidity.\n\nSome details of the molecular assembly task are not clear to me. Since a generative model is being used, are multiple samples inferenced before computing metrics? And is the input uncorrelated structure fixed, or is it randomly sampled on every training iteration?\n\nThe fact that packing matching of AssembleFlow and PackMol are similar in Table 2, but that energy distributions are quite different in Figure 4, suggests that packing matching is not that sensitive of a metric. A more sensitive metric could be RMSD, which should be possible to calculate using the Hungarian algorithm. If multiple polymorphs are possible for the same molecule, RMSD could be aggregated using the \"average minimum RMSD\" and \"coverage\" metrics proposed by Ganea et al. [3] for molecular conformer search.\n\nMinor issues:\n- Figure 8 the lower left box content is not centered\n- Table 2 arrow is wrong for separation\n- Line 176 SO(3) must also have det 1\n- Line 285 \"SPLER\"\n- Line 805 \"dynamcis\"\n\n[1] Bose, A. J., Akhound-Sadegh, T., Huguet, G., Fatras, K., Rector-Brooks, J., Liu, C. H., ... & Tong, A. (2023). Se (3)-stochastic flow matching for protein backbone generation. arXiv preprint arXiv:2310.02391.\n\n[2] https://en.wikipedia.org/wiki/Kabsch_algorithm\n\n[3] Ganea, O., Pattanaik, L., Coley, C., Barzilay, R., Jensen, K., Green, W., & Jaakkola, T. (2021). Geomol: Torsional geometric generation of molecular 3d conformer ensembles. Advances in Neural Information Processing Systems, 34, 13757-13769.", "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": "S6TLnpaNpi", "reviewer": "Reviewer_fnyR", "summary": "The paper introduces AssemblyFlow, a flow-matching based generative model to predict the packing of molecules into crystalline form. The method decouples the prediction of crystallization into 2 orthogonal bases, the rotation and the translation of molecular frames. A linear and spherical interpolation scheme are derived for each base, respectively. The score model is parameterized using a SE(3) equivariant network. The model is benchmarked against both generative model and simulation model baselines. Benchmarks and evaluation metrics are developed to evaluate the model performance. AssembleFlow achieves better performance compared with generative model baselines. Compared with simulation baseline PackMol, it achieves comparable performance but is 25 fold faster.", "questions": "- Can the authors explain the difference between their quaternion interpolation and the other SE(3) interpolation scheme like those developed in FrameFlow? Do you expect a difference in performance if you adopt a different SE(3) interpolation scheme? \n- In Fig. 4, it is very clear that PackMol generated structures have a much lower energy compared with AssembleFlow. This indicates a limitation of the generated structures. Can the authors expand on why AssembleFlow generated structures have a much higher energy?\n- The authors mention 2 versions of their SE(3)-equivariant network for inter-molecule modeling. I could not find more results on comparing these 2 versions. Can the authors add more results or point me to the right location?\n- How did the authors predict the initial conformation of the individual molecules? Are those predicted by RDKit? Or do you directly take the conformation from the data?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The paper introduces a rigid frame based flow matching model to generate the structure of organic crystals. The approach is similar to those introduced for frame-based flow matching methods for protein backbone generative, like FrameFlow [1]. Unlike protein, the frame for organic molecules is not well defined. The novelty of this approach is to use the inertial frames of the reference frames for organic molecules. \n- A spherical interpolation scheme is introduced to model the rotation of molecular frames. A quaternion formalism is used to define the interpolation in the SO(3) group. The formalism looks similar to those used for protein frame modeling. It is unclear to me what is the core difference in the author’s approach.\n- A set of new metrics, including packing matching (PM), atomic collision, separation, and compactness are introduced for the crystallization packing problem. The set of metrics are novel and meaningful from the domain perspective. It will benefit the community to further develop methods in the field. \n- The signficant improvements of AssembleFlow compared with other generative model baselines show the advantage of modeling the problem as rigid frames. This advantage is one of the core innovations of the method.\n- AssembleFlow is able to match the performance of MolPack, a simulation baseline. This is quite significant. The 25 fold speed-up compared with MolPack shows the advantage of generative modeling approach.\n\n[1] Yim, Jason, et al. \"Fast protein backbone generation with SE (3) flow matching.\" arXiv preprint arXiv:2310.05297 (2023).", "weaknesses": "- The decoupling of SE(3) into R^3 and SO(3) groups, i.e., rigid flow matching, has been extensively explored in the protein backbone community. The authors need to review more related papers and articulate their innovations in the methodology.\n- As the authors pointed out, numerical stability exists when there is a tie in eigenvalues for the principal axes of inertia. This problem is especially important for symmetric molecules. Can the authors add more details regarding how they address the instability issue? How much will it influence the training stability?\n- The real organic crystals are periodic. However, the authors only model them as non-periodic clusters in this work. It is a foundational limitation of both the dataset and the AssembleFlow model.\n- In real organic crystals, the rigid assumption may not always apply. The molecular conformation may change. Therefore, the rigid frame modeling approach may be limited in predicting the structure of more flexible molecules.", "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": "gZs82y0tkZ", "reviewer": "Reviewer_cXch", "summary": "This paper proposes AssembleFlow, which models the assembly of molecular crystals by assuming a rigid body and operating in the SE(3) space. It learns vector fields on SO(3) and $\\mathbb{R}^3$ separately, using SLERP and LERP, respectively. The reference coordinates systems are constructed with inertial frames.", "questions": "What advantage does the eigendecomposition of the inertial tensor $\\hat{I}$ have over the common approach of using centered $X$?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The paper is written clearly and is easy to understand.\n- The experiment shows good performance.", "weaknesses": "A notable limitation of this work is the claim to be the first to implement rigid generation in SE(3) space for molecular assembly. This claim appears to overlook prior work [1], which closely aligns with the approach presented here, treating molecules as rigid bodies in molecular assembly, separately handling SO(3) and Euclidean space, and applying quaternion representation on SO(3). The differences are minimal, mainly in the methods for constructing rigid frames and in the choice of a normalizing flow over a CNF with flow matching.\n\nAdditionally, the paper emphasizes the decomposition of SE(3) into $\\mathbb{R}^3$ and SO(3) as a unique strength (e.g., “this decomposition also empowers distinct probability paths…, effectively allowing for separate learning”). However, decomposing SE(3) into SO(3) and $\\mathbb{R}^3$ is a widely adopted simplification in SE(3)-based methods ([1,2] among others). To my knowledge, there are no approaches that directly operate on SE(3) without this decomposition, which raises questions about why this common approach is presented as an innovative strength.\n\nGiven these points, I am hesitant to support this work for acceptance in its current form.\n\n-----\n[1] Rigid Body Flows for Sampling Molecular Crystal Structures, Köhler et al., ICML 2023  \n[2] SE(3)-Stochastic Flow Matching for Protein Backbone Generation, Bose et al., ICLR 2024", "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)