SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In this section, we demonstrate the advantages of equivariant OT flow matching over existing training methods using four datasets characterized by invariant energies. We explore three different training objectives: (i) likelihood-based training, (ii) OT flow matching, and (iii) equivariant OT flow matching. For a comprehensive overview of experimental details, including dataset parameters, error bars, learning rate schedules, computing infrastructure, and additional experiments, please refer to Appendix A and Appendix C in the supplementary material. For all experiments, we employ a mean-free Gaussian prior distribution. The number of layers and parameters in the equivariant CNF vary across datasets while remaining consistent within each dataset (see Appendix C). We provide naive flow matching as an additional baseline in Appendix A.11. However, the results are very similar to the ones obtained with OT flow matching, although the integration paths are generally slightly longer.
6.1 DW4 and LJ13
We first evaluate the different loss functions on two many-particle systems, DW4 and LJ13, that were specifically designed for benchmarking equivariant flows as described in [17]. These systems feature pair-wise double-well and Lennard-Jones interactions with 4 and 13 particles, respectively (see Appendix C.2 for more details). While state-of-the-art results have been reported in [18], their evaluations were performed on very small test sets and were biased for the LJ13 system. To provide a fair comparison, we retrain their model using likelihood-based training as well as the two flow matching losses on resampled training and test sets for both systems. Additionally, we demonstrate improved likelihood performance of flow matching on their biased test set in Appendix A.3. The results, presented in Table 1, show that the two flow matching objectives outperform likelihood-based training while being computationally more efficient. The effective sample sizes (ESS) and negative log likelihood (NLL) are comparable for the flow matching runs. However, for the LJ13 system, the equivariant OT flow matching objective significantly reduces the integration path length compared
LJ55 - Sampled path lengths
f)
d)
e)
Figure 2: Results for the LJ55 system (a) Integration paths per particle for OT flow matching (left) and equivariant OT flow matching (right). (b) Mean transport cost (squared distance) for training batches. (c) Integration path length distribution. (d) Integration error for a fixed step size integrator (rk4) with respect to a reference solution generated by an adaptive solver (dropi5). (e) Effective samples vs number of function evaluations, i.e. evaluations of the vector field, for a sampling batch size of 1000. (f) Energy histograms for a flow trained with equivariant OT flow matching.
to other methods due to the large number of possible permutations (refer to Figure 1 for a visual illustration).
6.2 LJ55
The effectiveness of equivariant OT flow matching becomes even more pronounced when training on larger systems, where likelihood training is infeasible (see Appendix A.2). To this end, we investigate a large Lennard-Jones cluster with 55 particles (LJ55).
We observe that the mean batch transport cost of training pairs is about 10 times larger for OT flow matching compared to equivariant OT flow matching (Figure 2b), resulting in curved and twice as long integration paths during inference (Figure 2a,c). However, the integration paths for OT flow matching are shorter than those seen during training, see Appendix A.1 for a more detailed discussion. We compare the integration error caused by using a fixed step integrator (rk4) instead of an adaptive solver (dropi5 [55]). As the integration paths follow nearly straight lines for the equivariant OT flow matching, the so resulting integration error is minimal, while the error is significantly larger for OT flow matching (Figure 2d). Hence, we can use a fixed step integrator, with e.g. 20 steps, for sampling for equivariant OT flow matching, resulting in a three times speed-up over OT flow matching (Figure 2e), emphasizing the importance of accounting for symmetries in the loss for large, highly symmetric systems. Moreover, the equivariant OT flow matching objective outperforms OT flow matching on all evaluation metrics (Table 1). To ensure that the flow samples all states, we compute the energy distribution (Figure 2f) and perform deterministic structure minimization of the samples, similar to [17], in Appendix A.7.
6.3Alanine dipeptide
In our final experiment, we focus on the small molecule alanine dipeptide (Figure 3a) in Cartesian coordinates. The objective is to train a Boltzmann Generator capable of sampling from the equilibrium Boltzmann distribution defined by the semi-empirical  $GFN2 - xTB$  force-field [56]. This semi-empirical potential energy is invariant under permutations of atoms of the same chemical element and global rotations and reflections. However, here we consider the five backbone atoms defining the  $\varphi$  and  $\psi$  dihedral angles as distinguishable to facilitate analysis. Since running MD simulations with semi-empirical force-fields is computationally expensive, we employ a surrogate training set, further challenging the learning task.
Figure 3: Results for the alanine dipeptide system (a) Alanine dipeptide molecule. (b) Mean transport cost (squared distance) for training batches. (c) Integration path length distribution. (d) Ramachandran plot depicting the generated joint marginal distribution over the backbone dihedral angles  $\varphi$  and  $\psi$  after filtering out samples with right-handed chirality and high energies. (e) Energy histograms for samples generated by a flow trained with OT flow matching. (f) Free energy distribution along the slowest transition ( $\varphi$  dihedral angle) computed with umbrella sampling and the equivariant flow.
Alanine dipeptide data set generation The alanine dipeptide training data set is generated through two steps: (i) Firstly, we perform an MD simulation using the classical Amber ff99SBilden forcefield for a duration of 1 ms [20]. (ii) Secondly, we relax  $10^{5}$  randomly selected states from the MD simulation using the semi-empirical  $GFN2-xTB$  force-field for 100 fs each. For more detailed information, refer to Appendix C.2 in the supplementary material. This training data generation is significantly cheaper than performing a long MD simulation with the semi-empirical force field.
Although the Boltzmann generator is trained on a biased training set, we can generate asymptotically unbiased samples from the semi-empirical target distribution by employing reweighting (Appendix B.3) as demonstrated in Figure 3e. In this experiment, OT flow matching outperforms equivariant OT flow matching in terms of effective sample size (ESS) and negative log likelihood (NLL) as shown in Table 1. However, the integration path lengths are still longer for OT flow matching compared to equivariant OT flow matching, as depicted in Figure 3c. Since the energy of alanine dipeptide is invariant under global reflections, the flow generates samples for both chiral states. While the presence of the mirror state reflects the symmetries of the energy, in practical scenarios, molecules typically do not change their chirality spontaneously. Therefore, it may be undesirable to have samples from both chiral states. However, it is straightforward to identify samples of the undesired chirality and apply a mirroring operation to correct them. Alternatively, one may use a SO(3) equivariant flow without reflection equivariance [54] to prevent the generation of mirror states.
Alanine dipeptide - free energy difference The computation of free energy differences is a common challenge in statistical physics as it determines the relative stability of metastable states. In the specific case of alanine dipeptide, the transition between negative and positive  $\varphi$  dihedral angle is the slowest process (Figure 3d), and equilibrating the free energy difference between these two states from MD simulation requires simulating numerous transitions, i.e. millions of consecutive MD steps, which is expensive for the semi-empirical force field. In contrast, we can train a Boltzmann generator from data that is not in global equilibrium, i.e. using our biased training data, which is significantly cheaper to generate. Moreover, as the positive  $\varphi$  state is much less likely, we can even bias our training data to have nearly equal density in both states, which helps compute a more precise free energy estimate (see Appendix C.5). The equivariant Boltzmann generator is trained using the OT flow matching loss, which exhibited slightly better performance for alanine dipeptide. We obtain similar results for the equivariant OT flow matching loss (see Appendix A.4). To obtain an accurate estimation of the free energy difference, five umbrella sampling simulations are conducted along the  $\varphi$  dihedral angle using the semi-empirical force-field (see Appendix C.2). The free energy difference estimated by the Boltzmann generator demonstrates good agreement with the results
Table 2: Dimensionless free energy differences for the slowest transition of alanine dipeptide estimated from various methods. Umbrella sampling yields a converged reference solution. Errors over five runs.
Free energy difference / kB T	MD 5.31	relaxed MD 5.00	Umbrella sampling 4.10 ± 0.26	Boltzmann generator 4.10 ± 0.08
of these simulations (Table 2 and Figure 3f), whereas both the relaxed training data and classical Molecular Dynamics simulation overestimate the free energy difference.
We have introduced a novel flow matching objective for training equivariant continuous normalizing flows on invariant densities, leading to optimal transport integration paths even for small training batch sizes. By leveraging flow matching objectives, we successfully extended the applicability of equivariant flows to significantly larger systems, including the large Lennard-Jones cluster (LJ55). We conducted experiments comparing different training objectives on four symmetric datasets, and our results demonstrate that as the system size increases, the importance of accounting for symmetries within the flow matching objective becomes more pronounced. This highlights the critical role of our proposed flow matching objective in scaling equivariant CNFs to even larger systems.
Another notable contribution of this work is the first successful application of a Boltzmann generator to model alanine dipeptide in Cartesian coordinates. By accurately estimating free energy differences using a semi-empirical force-field, our approach of applying OT flow matching and equivariant OT flow matching to equivariant flows demonstrates its potential for reliable simulations of complex molecular systems.
8 Limitations / Future work
While we did not conduct experiments to demonstrate the transferability of our approach, the architecture and proposed loss function can potentially be used to train transferable models. We leave this for future research. Although training with flow matching is faster and computationally cheaper than likelihood training for CNFs, the inference process still requires the complete integration of the vector field, which can be computationally expensive. However, if the model is trained with the equivariant OT flow matching objective, faster fixed-step integrators can be employed during inference. Our suggested approximation of Equation (15) includes the Hungarian algorithm, which has a computational complexity of  $\mathcal{O}(N^3)$ . To improve efficiency, this step could be replaced by heuristics relying on approximations to the Hungarian algorithm [57]. Moreover, flow matching does not allow for energy based training, as this requires integration similar to NLL training. A potential alternative approach is to initially train a CNF using flow matching with a small set of samples. Subsequent sample generation through the CNF, followed by reweighting to the target distribution, allows these samples to be added iteratively to the training set, similarly as in [27, 37]. Notably, in our experiments, we exclusively examined Gaussian prior distributions. However, flow matching allows to transform arbitrary distributions. Therefore, our equivariant OT flow matching method holds promise for application in scenarios where the prior distribution is sampled through MD at a distinct thermodynamic state, such as a higher temperature or a different level of theory [32, 33]. In these cases, where both the prior and target distributions are close, with samples alignable through rotations and permutations, we expect that the advantages of equivariant OT flow matching will become even more pronounced than what we observed in our experiments.
Building upon the success of training equivariant CNFs for larger systems using our flow matching objective, future work should explore different architectures for the vector field. Promising candidates, such as [58, 59, 60, 61, 62], could be investigated to improve the modeling capabilities of equivariant CNFs. While our focus in this work has been on symmetric physical systems, it is worth noting that equivariant flows have applications in other domains that also exhibit symmetries, such as traffic data generation [63], point cloud and set modeling [64, 65], as well as invariant distributions on arbitrary manifolds [66]. Our equivariant OT flow matching approach can be readily applied to these areas of application without modification.
We gratefully acknowledge support by the Deutsche Forschungsgemeinschaft (SFB1114, Projects No. C03, No. A04, and No. B08), the European Research Council (ERC CoG 772230 "ScaleCell"), the Berlin Mathematics center  $\mathrm{MATH + }$  (AA1-6), and the German Ministry for Education and Research (BIFOLD - Berlin Institute for the Foundations of Learning and Data). We gratefully acknowledge the computing time granted by the Resource Allocation Board and provided on the supercomputer Lise at NHR@ZIB and NHR@Göttingen as part of the NHR infrastructure. We thank Maaike Galama, Michele Invernizzi, Jonas Kohler, and Max Schebek for insightful discussions. Moreover, we would like to thank Lea Zimmermann and Felix Draxler, who tested our code and brought to our attention a critical bug.
[1] Durk P Kingma, Shakir Mohamed, Danilo Jimenez Rezende, and Max Welling. Semi-supervised learning with deep generative models. Advances in neural information processing systems, 27, 2014.
[2] Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real NVP. In International Conference on Learning Representations, 2017.
[3] Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1x1 convolutions. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
[4] Tero Karras, Samuli Laine, Miika Aittala, Janne Hellsten, Jaakko Lehtinen, and Timo Aila. Analyzing and improving the image quality of stylegan. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 8110-8119, 2020.
[5] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems, 30, 2017.
[6] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pages 4171–4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.
[7] Luciano Floridi and Massimo Chiriatti. Gpt-3: Its nature, scope, limits, and consequences. *Minds and Machines*, 30:681–694, 2020.
[8] Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu. Boltzmann generators — sampling equilibrium states of many-body systems with deep learning. Science, 365:eaaw1147, 2019.
[9] John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583-589, 2021.
[10] Joseph L. Watson, David Juergens, Nathaniel R. Bennett, Brian L. Trippe, Jason Yim, Helen E. Eisenach, Woody Ahern, Andrew J. Borst, Robert J. Ragotte, Lukas F. Milles, Basile I. M. Wicky, Nikita Hanikel, Samuel J. Pellock, Alexis Courbet, William Sheffler, Jue Wang, Preetham Venkatesh, Isaac Sappington, Susana Vázquez Torres, Anna Lauko, Valentin De Bortoli, Emile Mathieu, Regina Barzilay, Tommi S. Jaakkola, Frank DiMaio, Minkyung Baek, and David Baker. Broadly applicable and accurate protein design by integrating structure prediction networks and diffusion generative models. bioRxiv, 2022.
[11] Esteban G Tabak, Eric Vanden-Eijnden, et al. Density estimation by dual ascent of the log-likelihood. Communications in Mathematical Sciences, 8(1):217–233, 2010.
[12] Esteban G Tabak and Cristina V Turner. A family of nonparametric density estimation algorithms. Communications on Pure and Applied Mathematics, 66(2):145-164, 2013.
Figure 4: Path lengths distributions for models trained with OT flow matching. a) LJ13 system b) Alanine dipeptide system c) LJ55 system.
A Additional experiments
A.1 Equivariant flows can achieve better paths than they are trained for
We make an interesting observation regarding the performance of equivariant flows trained with the OT flow matching objective. As shown in Figure 4a-c, these flows generate integration paths that are significantly shorter compared to their training counterparts. While they do not learn the OT path and exhibit curved trajectories with notable directional changes (as seen in Figure 1d and Figure 2a), an interesting finding is the absence of path crossings. Any permutation of the samples would increase the distance between the latent sample and the push-forward sample. To quantify the length of these paths, we compute the Euclidean distance between the latent sample and the push-forward sample. The relatively short path lengths can be attributed, in part, to the architecture of the flow model. As the flow is equivariant with respect to permutations, it prevents the crossing of particle paths and forces the particles to change the direction instead, as an average vector field is learned.
A.2 Memory requirement
As the system size increases, the memory requirements for likelihood training become prohibitively large. In the case of the LJ55 system, even a batch size of 5 necessitates more than 24 GB of memory, making effective training infeasible. In contrast, utilizing flow matching enables training the same model with a batch size of 256 while requiring less than 12 GB of memory.
For the alanine dipeptide system, likelihood training with a batch size of 8 demands more than 12 GB of memory, and a batch size of 16 exceeds 24 GB. On the other hand, employing flow matching allows training the same model with a batch size of 256, consuming less than 3 GB of memory.
The same holds for energy based training, as we need to evaluate Equation (3) and Equation (4) as well during training. Therefore, we opt to solely utilize flow matching for these two datasets.
A.3 Comparison with prior work on equivariant flows
Both [17] and [18] evaluate their models on the DW4 and LJ13 systems. However, [17] do not provide the numerical values for the negative log likelihood (NLL), only presenting them visually. On the other hand, [18] use a small test set for both experiments. Additionally, their test set for LJ13 is biased as it originates from the burn-in phase of the sampling.
Although the test sets for both systems are insufficient, we still compare our OT flow matching and equivariant OT flow matching to these previous works in Table 3. It should be noted that the value range differs significantly from the values reported in Table 1, as the likelihood in this case includes the normalization constant of the prior distribution.
A.4 Alanine dipeptide - free energy difference
Here we present additional results for the free energy computation of alanine dipeptide, as discussed in Section 6.3. The effective sample size for the trained models is  $0.50 \pm 0.13\%$ , which is lower due to training on the biased dataset (compare to Table 1). We illustrate the Ramachandran plot in
Table 3: Comparison of different training methods for the DW4 and (biased) LJ13 dataset as used in [18] for  $10^{5}$  and  $10^{4}$  training points, respectively.
Dataset	Likelihood [18]	OT flow matching	Equivariant OT flow matching
DW4 NLL (↓)	7.48 ± 0.05	7.21 ± 0.03	7.18 ± 0.02
LJ13 NLL (↓)	30.41 ± 0.16	30.34 ± 0.09	30.30 ± 0.03
Figure 5: Alanine dipeptide - free energy experiments. The model is trained on the biased dataset. (a) Ramachandran plot depicting the generated joint marginal distribution over the backbone dihedral angles  $\varphi$  and  $\psi$  after filtering out samples with right-handed chirality and high energies. (b) Energy histograms for samples generated by a flow trained with OT flow matching. (c) Free energy distribution along the slowest transition ( $\varphi$  dihedral angle) computed with umbrella sampling, the equivariant flow, and a relaxed MD trajectory.
Figure 5a, the energy distribution in Figure 5b, and the free energy projections for Umbrella sampling, model samples, and relaxed MD samples in Figure 5c.
We obtain similar results with equivariant OT flow matching, as shown in Table 4.
A.5 Alanine dipeptide integration paths
We compare the integration paths of the models trained on the alanine dipeptide dataset in Figure 6a,b. Despite the higher number of atoms in the alanine dipeptide system (22 atoms), we observe a comparable difference in the integration paths as observed in the LJ13 system. This similarity arises from the presence of different particle types in the alanine dipeptide, resulting in a similar number of possible permutations compared to the LJ13 system.
A.6 Integration error
The integration error, which arises from the use of a fixed-step integration method (rk4) instead of an adaptive solver (dopri5), is quantified by calculating the sum of absolute differences in sample positions and changes in log probability (dlogp). The discrepancies in position are depicted in Figure 2d, while the difference in log probability change are presented in Figure 6c. Both of these are crucial in ensuring the generation of unbiased samples that accurately represent the target Boltzmann distribution of interest.
A.7 Structure minimization
Following a similar approach as the deterministic structure minimization conducted for the LJ13 system in [17], we extend our investigation to the sampled minima of the LJ55 system. We minimize the energy of samples generated by a model trained with equivariant OT flow matching, as well as samples obtained from the test data. The resulting structures represent numerous local minima within the potential energy distribution.
In Figure 6d, we observe a good agreement between the distributions obtained from both methods. Furthermore, it is noteworthy that both approaches yield samples corresponding to the global minimum. The energies of the samples at  $300\mathrm{K}$  are considerably higher than those of the minimized structures.
Table 4: Dimensionless free energy differences for the slowest transition of alanine dipeptide estimated with a Boltzmann Generator trained with the equivariant OT flow matching objective. Umbrella sampling yields a converged reference solution. Errors over five runs.
Free energy difference / kB T	Umbrella sampling 4.10 ± 0.26	Equivariant OT flow matching 3.95 ± 0.19
Figure 6: (a,b) Integration paths for alanine dipeptide for models trained with OT flow matching and equivariant OT flow matching, respectively. The color code for the atoms is as usual: gray - H, black - C, blue - N, red - O. (c) Integration error for a fixed step integrator for the LJ55 system. (d) Minimized samples from a model trained on the LJ55 system with equivariant OT flow matching.
A.8 Log weight distribution