SE(3) Equivariant Augmented Coupling Flows

Authors: Laurence Midgley, Vincent Stimper, Javier Antorán, Emile Mathieu, Bernhard Schölkopf, José Miguel Hernández-Lobato

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

Reproducibility Variable Result LLM Response
Research Type Experimental When trained on the DW4, LJ13, and QM9-positional datasets, our flow is competitive with equivariant continuous normalizing flows and diffusion models, while allowing sampling more than an order of magnitude faster. [...] 4 Experiments
Researcher Affiliation Academia Laurence I. Midgley University of Cambridge lim24@cam.ac.uk Vincent Stimper Max Planck Institute for Intelligent Systems University of Cambridge vs488@cam.ac.uk Javier Antorán University of Cambridge ja666@cam.ac.uk Emile Mathieu University of Cambridge ebm32@cam.ac.uk Bernhard Schölkopf Max Planck Institute for Intelligent Systems bs@tue.mpg.de José Miguel Hernández-Lobato University of Cambridge jmh233@cam.ac.uk
Pseudocode Yes Algorithm 1: Flow block f [...] Algorithm 2: Joint density evaluation
Open Source Code Yes Code: https://github.com/lollcat/se3-augmented-coupling-flows
Open Datasets Yes For DW4 and LJ13 we use a training and test set of 1000 samples following Satorras et al. [2021a]. [...] For QM9-positional we use a training, validation and test set of 13831, 2501 and 1,813 samples respectively also following Satorras et al. [2021a].
Dataset Splits Yes For QM9-positional we use a training, validation and test set of 13831, 2501 and 1,813 samples respectively also following Satorras et al. [2021a].
Hardware Specification Yes Importantly, sampling and density evaluation of the E-ACF on an A100 GPU takes roughly 0.01 seconds. [...] All runs used an Ge Force RTX 2080Ti GPU . [...] DW4 was run using a Ge Force RTX 2080Ti GPU and LJ13 with a V-4 TPU.
Software Dependencies No The paper mentions software like "Adam optimizer", "dopri5 ODE solver", and custom EGNN implementation, but does not provide specific version numbers for any of these.
Experiment Setup Yes All flow models using 12 blocks (see Alg. 1 for the definition of a block). For the noise scale of the augmented variables sampled from N(a; x, η2I), we use η = 0.1 for both the base of the flow, and augmented target distribution. [...] We use a cosine learning rate schedule, that is initialised to a value of 0.00002, peaks at 0.0002 after 30 epochs, and then decays back to a final learning rate of 0.00002 for DW4 and LJ13. [...] We use a batch size of 32 for all experiments, and 100, 400 and 800 epochs of the datasets for DW4, LJ13 and QM9-positional respectively.