Efficient Generation of Molecular Clusters with Dual-Scale Equivariant Flow Matching

Amorphous molecular solids are disordered organic systems with several applications in the field of organic optoelectronics. Their mechanical flexibility and solution processability make them attractive alternatives to inorganic semiconductors [22, 14, 19]. Different molecules tend to pack differently in amorphous solids, which lead to different network extensivity, electronic properties, and transport rates [13]. The influence of packing structure on optoelectronic properties has led to significant interest in computationally simulating these materials with molecular dynamics (MD) simulations [13, 6] to enhance our understanding of the correlations between structural design and corresponding packing and optoelectronic properties. This can inform better design strategies for the next generation of organic electronics devices.

Having the ability to generate amorphous solids in a computationally efficient manner can enable faster sampling from the multi-molecule Boltzmann distribution, as well as property-guided generation of optimally packed configurations. There have been several works developing efficient methods to sample from Boltzmann distributions or their emulated (non re-weighted) versions [17, 11]. Scaling up to larger systems (such as amorphous clusters) is however computationally expensive if generation is to be performed at the all-atom (AA) level. In the field of protein generation, coarse-grained (CG) representations are used to collapse amino acids into single beads, which drastically reduces system size while still retaining higher-level structural features. There have been works that generate proteins in their coarse-grained representations, and other works that back-map coarse-grained proteins into their all-atom structure [12, 4, 24, 10].

In this paper, we develop a unified dual-scale flow matching method that improves accuracy of all-atom flow matching methods by separating training and inference into different stages for coarse-grained and all-atom resolutions. Moreover, we show that by performing a majority of the inference compute at the coarse-grained resolution, one can reduce the time taken for inference integration without sacrificing accuracy, paving the way for more efficient generation of larger systems. As a demonstrative example, in this work, we used a dataset of amorphously packed clusters consisting of 5 Y6 [25] molecules that we obtained from MD simulations. As the coarse-grained representation for our experiments, we pre-defined a mapping scheme that collapses the 505 atoms system into 65 beads (see Figure 1(a)). With gains in efficiency obtained from this method, we plan to scale up to larger systems in the future that are more representative of thin films in devices, as well as test other choices of coarse-graining mappings.

Conditional Flow Matching (CFM). This is a new paradigm of training continuous normalizing flows (CNFs) efficiently in a simulation free manner [3, 2, 15]. An ODE parameterized by a learnable vector field $v_{\theta}(\mathbf{x},t)$ transforms an easy-to-sample prior distribution $p_{0}$, into the more complex target distribution $p_{1}$. Parameters $\theta$ are learned through the regression objective

We chose the probability path definition to be $p_{t}(\mathbf{x}|\mathbf{z})=\mathcal{N}(\mathbf{x}|t\mathbf{x_{1}}+(1-t)\mathbf{x_{0}},\sigma^{2})$ , which gives $u(\mathbf{x}|\mathbf{z})=\mathbf{x_{1}}-\mathbf{x_{0}}$ as the conditional vector field. Similar to work by Stark et al. [20], we trained $v_{\theta}$ to predict $\mathbf{x_{1}}$ rather than $\mathbf{x_{1}}-\mathbf{x_{0}}$, and parameterized $v_{\theta}$ by SE(3) equivariant refinement tensor field network (TFN) [23] layers. We present a comparison against other parameterizations in Section A.1.

Harmonic Prior. Besides the standard gaussian prior, another common choice for prior in molecule generation tasks is the harmonic prior. The harmonic prior was introduced in Jing et al. [9] and used within the flow-matching framework by Stark et al. [20] as a more chemically plausible prior. The prior distribution $p_{0}(\mathbf{x_{0}})$ is interpreted as a Boltzmann distribution with a quadratic potential energy function, i.e., $p_{0}(\mathbf{x_{0}})\propto e^{-E(\mathbf{x})}=e^{-\mathbf{x}^{T}\mathbf{H}\mathbf{x}}$, where $\mathbf{H}$ is chosen such that $E(\mathbf{x})=\Sigma_{i,j\in\mathcal{E}}||\mathbf{x}_{i}-\mathbf{x}_{j}||^{2}$. We need to have the bond connectivity information $\mathcal{E}$ a-priori to sample from the harmonic prior.

We evaluated generated samples on 3 metrics (shown in Table 1). To evaluate distributional matching capabilities, we measured Jensen Shannon Divergence (JSD) between distributions of generated and MD distributions for bond lengths and angles. We evaluated efficiency through the average time taken per integration step during inference on one NVIDIA A100 GPU. We noticed that the dual-scale flow matching method improved upon a single-scale flow matching method (with Gaussian/Harmonic priors) by 15-25% on bond length and angle JSDs, while also decreasing the inference time by $\sim$85%.

We also tested the influence of the ratio of integration steps taken in coarse-grained and all-atom resolutions, on the 3 metrics to see how much inference time can be saved without sacrificing accuracy (shown in Figure 1(b)). We denote the ratio of number of CG and AA integration steps by CG:AA. It is interesting to see that even for large CG:AA ratios, we significantly decreased inference time with negligible sacrifice to JSD performance.

The preliminary results are encouraging, opening up several avenues of further exploration for the future. We plan to scale to larger systems ($\sim$200 Y6 molecules) which is more representative of active layers in devices, and also test our approach on different chemistries and system sizes. Secondly, while we pre-defined a coarse-graining mapping scheme in this paper, we plan to test other schemes in the future to identify correlations between the chosen schemes and performance. Finally, we plan on expanding to other metrics that are also typically used in Boltzmann generators and conformer generation such as effective sample size (ESS), average minimum RMSD (AMR) and coverage [11, 7, 8].

In this work, we developed a dual-scale flow matching method that improves upon single-scale flow matching-based samplers in accuracy as well as inference efficiency, as demonstrated on a dataset of amorphous Y6 clusters. Moreover, it is interesting to see that we can push the ratio of CG:AA during inference integration to large values without sacrificing the prediction accuracy, while boosting inference speed. While we tested on clusters of size 5 in this paper, we plan to expand to larger clusters in the future whose packing properties are more representative of active layers in optoelectronic devices. We also plan on testing the influence of coarse-graining mapping scheme on performance since we tested a single pre-defined scheme in this work.