Equivariant flow matching

Authors: Leon Klein, Andreas Krämer, Frank Noe

NeurIPS 2023 | Conference PDF | Archive PDF | Plain Text | LLM Run Details

Reproducibility Variable Result LLM Response
Research Type Experimental We demonstrate the effectiveness of flow matching on rotation and permutation invariant many-particle systems and a small molecule, alanine dipeptide, where for the first time we obtain a Boltzmann generator with significant sampling efficiency without relying on tailored internal coordinate featurization. Our results show that the equivariant flow matching objective yields flows with shorter integration paths, improved sampling efficiency, and higher scalability compared to existing methods.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.
Researcher Affiliation Collaboration Leon Klein Freie Universität Berlin leon.klein@fu-berlin.de Andreas Krämer Freie Universität Berlin andreas.kraemer@fu-berlin.de Frank Noé Microsoft Research AI4Science Freie Universität Berlin Rice University franknoe@microsoft.com
Pseudocode No No explicit pseudocode or algorithm blocks found.
Open Source Code No The code will be integrated in the bgflow library https://github.com/noegroup/bgflow.
Open Datasets Yes The datasets are available at https://osf.io/srqg7/?view_only= 28deeba0845546fb96d1b2f355db0da5.
Dataset Splits No 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.We create a test set in the same way.
Hardware Specification Yes All experiments for the DW4 and LJ13 system were conducted on a Ge Force GTX 1080 Ti with 12 GB RAM. The training for alanine dipeptide and the LJ55 system were conducted on a Ge Force RTX 3090 with 24 GB RAM. Inference was performed on NVIDIA A100 GPUs with 80GB RAM for alanine dipeptide and the LJ55 system.
Software Dependencies No Flow models and training are implemented in Pytorch [69] using the following code libraries: bgflow [8, 17], torchdyn [70], Pot: Python optimal transport [71], and the code corresponding to [18]. The MD simulations are run using Open MM[72], ASE [73], and xtb-python [56].
Experiment Setup Yes We report the used training schedules in Table 8. Note that 5e-4/5e-5 in the second row means that the training was started with a learning rate of 5e-4 for 200 epochs and then continued with a learning rate of 5e-5 for another 200 epochs. All batches were reordered prior to training the model. The only exception is alanine dipeptide trained with equivariant OT flow matching.