SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\operatorname {S h i f t C o M} (\tilde {\boldsymbol {x}}, \boldsymbol {a}) \triangleq (\tilde {\boldsymbol {x}} - \bar {\boldsymbol {a}}, \boldsymbol {a} - \bar {\boldsymbol {a}}) \quad \text {w i t h} \quad \bar {\boldsymbol {a}} \triangleq \frac {1}{n} \sum_ {i = 1} ^ {n} [ \boldsymbol {a} ] ^ {i}. \tag {5}
$$
This operation is invertible, with inverse ShiftCoM  $(\tilde{a}, x)$ , and has unit Jacobian determinant.
Algorithm 1: Flow block f	Algorithm 2: Joint density evaluation
Inputs: Zero-CoM observation x̂, augmented variable a, Coupling transforms F1, F2(x, x̂) ← ShiftCoM(x̂, a)(x̂, x̂) ← FM(1) ∘ ... ∘ F1(1)(x̂, x̂)(a, x̂) ← ShiftCoM(x̂, x̂)(a, x̂) ← FM(2) ∘ ... ∘ F1(2)(a, x̂)	Inputs: (x, a) ~ p, base density q0, (f(l))l=1(x̂(0), a(0)) ← (x - x̂, a - x̂)logdet ← 0for l = 1, ..., L do(x̂(l), a(l)) ← f(l)(x̂(l-1), a(l-1))logdet ← logdet + |∂f(l)(x̂(l-1), a(l-1))/∂(x̂(l-1), a(l-1))|
Output: x̂, a	Output: q0(x̂(L), a(L)) + logdet
Putting the building blocks together Our flow transform is built as a sequence of  $L$  blocks. Each block, described in Alg. 1, consists of two equivariant coupling layers, see Fig. 1. Our observations  $\tilde{\boldsymbol{x}}\in \tilde{\mathcal{X}}$  are lifted onto  $\mathcal{X}$  with ShiftCoM, they are transformed with  $M$  core transformations  $\left(\mathcal{F}_i^{(1)}\right)_{i = 1}^M$ , and ShiftCoM is applied one more time to map the observations back to the zero-CoM hyperplane. After this, our augmented variables  $\pmb{a}$  are transformed with  $\left(\mathcal{F}_i^{(2)}\right)_{i = 1}^M$ .
Joint density evaluation  $q(\pmb{x}, \pmb{a})$  is performed with Alg. 2. We first subtract the observed variables' CoM from both the observed and augmented variables. We then apply our  $L$  flow transform blocks before evaluating the transformed variables' density under our base distribution  $q_{0}$ , which is described next. The log determinant of the core flow transform,  $f$ , has a contribution from the projection, transform in the invariant space, and inverse projection (see App. B.3 for details).
3.2  $\operatorname{SE}(3) \times S_n$  Invariant base distribution
Again, we assume  $\mathcal{A} = \mathcal{X}$ , i.e.  $k = 1$ , with the generalization given in App. B.4. Our invariant choice of base distribution is  $q_{0}(\boldsymbol{x}, \boldsymbol{a}) = \tilde{\mathcal{N}}(\boldsymbol{x}; 0, I) \mathcal{N}(\boldsymbol{a}; \boldsymbol{x}, \eta^{2}I)$  where  $\boldsymbol{x} \in \mathbb{R}^{3n}$  and  $\boldsymbol{a} \in \mathbb{R}^{3n}$  refer to  $\boldsymbol{x}$  and  $\boldsymbol{a}$  flattened into vectors,  $\eta^{2}$  is a hyperparameter and we denote Gaussian distributions on  $\tilde{\mathcal{X}}$  as  $\tilde{\mathcal{N}}$  [Satorras et al., 2021a, Yim et al., 2023] with density
$$
\tilde {\mathcal {N}} (\boldsymbol {x}; 0, I) = (2 \pi) ^ {- 3 (n - 1) / 2} \exp \left(- \frac {1}{2} \| \tilde {\boldsymbol {x}} \| _ {2} ^ {2}\right). \tag {6}
$$
We sample from it by first sampling from a standard Gaussian  $\mathcal{N}(0, I)$  and then removing the CoM. On the other hand, the distribution for  $\pmb{a}$  is supported on  $\mathcal{A}$  which includes non-zero CoM points. It is centered on  $\pmb{x}$ , yielding joint invariance to translations. The isotropic nature of  $q_{0}$  makes its density invariant to rotations, reflections, and permutations [Satorras et al., 2021a, Yim et al., 2023].
3.3 Training and likelihood evaluation
In this section, we discuss learning and density evaluation with augmented variables.
Invariant augmented target distribution We assume the density of our observations  $p$  is  $\mathrm{SE}(3)\times S_n$  invariant. Our target for augmented variables is  $\pi (\pmb {a}|\pmb {x}) = \mathcal{N}(\pmb {a};\pmb {x},\eta^2 I)$ , where  $\eta^2$  matches the variance of the base Gaussian density over  $\pmb{a}$ . This satisfies joint invariance  $p(g\cdot x)\pi (g\cdot a|g\cdot x) = p(x)\pi (a|x)$  for any  $g\in \mathrm{SE}(3)\times S_n$ , as shown in App. B.5.
Learning from samples When data samples  $x \sim p$  are available, we train our flow parameters by maximizing the joint likelihood, which is a lower bound on the marginal log-likelihood over observations up to a fixed constant
$$
\mathbb {E} _ {\boldsymbol {x} \sim p (\boldsymbol {x}), \boldsymbol {a} \sim \pi (\boldsymbol {a} | \boldsymbol {x})} [ \log q (\boldsymbol {x}, \boldsymbol {a}) ] \leq \mathbb {E} _ {p (x)} [ \log q (\boldsymbol {x}) ] + C. \tag {7}
$$
Learning from energy When samples are not available but we can query the unnormalized energy of a state  $U(\pmb{x})$ , with  $p(\pmb{x}) \propto \exp(-U(\pmb{x}))$ , we can minimize the joint reverse KL divergence. By
the chain rule of the KL divergence, this upper bounds the KL between marginals
$$
D _ {\mathrm {K L}} \left(q (\boldsymbol {x}, \boldsymbol {a}) \mid p (\boldsymbol {x}) \pi (\boldsymbol {a} | \boldsymbol {x}))\right) \geq D _ {\mathrm {K L}} \left(q (\boldsymbol {x}) \mid p (\boldsymbol {x})\right). \tag {8}
$$
However, the reverse KL encourages mode-seeking [Minka, 2005] which may result in the model failing to characterize the full set of meta-stable molecular states. Therefore, we instead use flow annealed importance sampling bootstrap (FAB) [Midgley et al., 2023], which targets the mass covering  $\alpha$ -divergence with  $\alpha = 2$ . In particular, we minimize the  $\alpha$ -divergence over the joint which leads to an upper bound on the divergence of the marginals
$$
D _ {2} \left(q (\boldsymbol {x}, \boldsymbol {a}) \mid \mid p (\boldsymbol {x}) \pi (\boldsymbol {a} \mid \boldsymbol {x})\right) \triangleq \int \frac {p (\boldsymbol {x}) ^ {2} \pi (\boldsymbol {a} \mid \boldsymbol {x}) ^ {2}}{q (\boldsymbol {x} , \boldsymbol {a})} d \boldsymbol {a} d \boldsymbol {x} \geq \int \frac {p (\boldsymbol {x}) ^ {2}}{q (\boldsymbol {x})} d \boldsymbol {x} \triangleq D _ {2} \left(q (\boldsymbol {x}) \mid \mid p (\boldsymbol {x})\right). \tag {9}
$$
To compute unbiased expectations with the augmented flow we rely on the estimator  $\mathbb{E}_{p(\boldsymbol{x})}[f(\boldsymbol{x})] = \mathbb{E}_{q(\boldsymbol{x},\boldsymbol{a})}[w(\boldsymbol{x},\boldsymbol{a})f(\boldsymbol{x})]$  where  $w(\boldsymbol{x},\boldsymbol{a}) = p(\boldsymbol{x})\pi (\boldsymbol{a}|\boldsymbol{x}) / q(\boldsymbol{x},\boldsymbol{a})$ . Minimizing the joint  $\alpha$ -divergence with  $\alpha = 2$  corresponds to minimizing the variance in the joint importance sampling weights  $w(\boldsymbol{x},\boldsymbol{a})$ , which allows for the aforementioned expectation to be approximated accurately.
Evaluating densities To evaluate the marginal density of observations we use the importance weighed estimator  $q(\pmb{x}) = \mathbb{E}_{\pmb{a} \sim \pi(\cdot|\pmb{x})} \left[ \frac{q(\pmb{x}, \pmb{a})}{\pi(\pmb{a}|\pmb{x})} \right]$ , noting that  $\pi$  is Gaussian and thus supported everywhere. The estimator variance vanishes when  $q(\pmb{a}|\pmb{x}) = \pi(\pmb{a}|\pmb{x})$ , as shown in App. B.9.
4.1 Training with samples: DW4, LJ13 and QM9 positional
First, we consider 3 problems that involve only positional information, with no additional features such as atom type or connectivity. Thus, the target densities are fully permutation invariant. The first two of these, namely DW4 and LJ13, are toy problems from Kohler et al. [2020], where samples are obtained by running MCMC on the 4 particle double well energy function (DW4) and 13 particles Leonard Jones energy function (LJ13) respectively. The third problem, i.e. QM9 positional [Satorras et al., 2021a], selects the subset of molecules with 19 atoms from the commonly used QM9 dataset [Ramakrishnan et al., 2014] and discards their node features.
For our model, Equivariant Augmented Coupling Flow (E-ACF), we consider all projection types (VECTOR-PROJ, CARTESIAN-PROJ, SPHERICAL-PROJ) and compare them to: (1) NON-E-ACF: An augmented flow that is not rotation equivariant but is translation and permutation equivariant, as in [Klein et al., 2023a]. This model uses the same structure as the E-ACF but replaces the EGNN with a transformer which acts directly on atom positions, without any projection. We train NON-E-ACF with data-augmentation whereby we apply a random rotation to each sample within each training batch. (2) E-CNF ML: The SE(3) equivariant continuous normalizing flow from Satorras et al. [2021a] trained by maximum likelihood. (3) E-CNF FM: An SE(3) equivariant continuous normalizing flow trained via flow matching [Lipman et al., 2023, Klein et al., 2023b]. (4) E-CNF-DIFF: An SE(3) equivariant diffusion model [Hoogeboom et al., 2022] evaluated as a continuous normalizing flow. All equivariant generative models use the SE(3) GNN architecture proposed by Satorras et al. [2021b]. The CARTESIAN-PROJ exhibited numerical instability on QM9-positional causing runs to crash early. To prevent these crashes it was trained at a lower learning rate than the other E-ACF models. App. C.3.1 provides a detailed description of these experiments.
On Tab. 1, we see that on DW4 and LJ13, E-CNF FM performs best, while E-ACF is competitive with E-CNF ML and E-CNF-DIFF. We note that both DW4 and LJ13 have biased datasets (see App. C.3.1). On QM9-positional, the E-ACF is competitive with E-CNF-DIFF and E-CNF FM, while the under-trained E-CNF ML from Satorras et al. [2021a] performs poorly. We expect that with further tuning E-CNF-DIFF and E-CNF FM would match the performance of SPHERICAL-PROJ on QM9-positional. The NON-E-ACF performs much worse than the E-ACF, despite being trained for more epochs, demonstrating the utility of in-build equivariance. Furthermore, Fig. 2 shows that the distribution of inter-atomic distances of samples from our flow matches training data well. Importantly, sampling and density evaluation of the E-ACF on an A100 GPU takes roughly 0.01 seconds. For the CNF trained by flow matching (E-CNF ML) and score matching (E-CNF-DIFF), sampling takes on average 0.2, and 5 seconds respectively. Thus, the E-ACF is faster for sampling than the CNF by more than an order of magnitude.
Table 1: Negative log-likelihood results for flows trained by maximum likelihood on DW4, LJ13 and QM9-positional. E-CNF ML results are from Satorras et al. [2021a]. Best results are emphasized in bold. The results are averaged over 3 seeded runs, with the standard error reported as uncertainty.
DW4	LJ13	QM9 positional
E-CNF ML	8.15 ± N/A	30.56 ± N/A	-70.2 ± N/A
E-CNF FM	6.31 ± 0.05	28.81 ± 0.34	-143.52 ± 0.26
E-CNF-DIFF	8.01 ± 0.03	31.02 ± 0.12	-158.30 ± 0.15
NON-E-ACF	10.07 ± 0.03	33.32 ± 0.34	-76.76 ± 1.77
VECTOR-PROJ E-ACF	8.69 ± 0.03	30.19 ± 0.12	-152.23 ± 6.44
CARTESIAN-PROJ E-ACF	8.82 ± 0.08	30.89 ± 0.09	-138.62 ± 0.74
SPHERICAL-PROJ E-ACF	8.61 ± 0.05	30.33 ± 0.16	-165.71 ± 1.35
Figure 2: Inter-atomic distances for samples from the train-data and SPHERICAL-PROJ E-ACF.
4.2 Training with samples:Alanine dipeptide
Next, we approximate the Boltzmann distribution of alanine dipeptide in an implicit solvent at temperature  $T = 800\mathrm{K}$ . We train the models via maximum likelihood on samples generated by a replica exchange MD simulation [Mori and Okamoto, 2010], which serve as a ground truth. Our generated dataset consists of  $10^{6}$  samples in the training and validation set as well as  $10^{7}$  samples in the test set. Besides our E-ACF using the different projection schemes that we introduced in Sec. 3.1, we train the non-SO(3) equivariant flow (NON-E-ACF) with data augmentation similarly to the previous experiments. Moreover, we train a flow on internal coordinates as in Midgley et al. [2023]. The joint effective sample size (ESS) is reported for the E-ACF models which is a lower bound of the marginal ESS (see Eq. (9)). For the CNFs, the Hutchinson trace estimator is used for the log density [Grathwohl et al., 2018]. This results in a biased estimate of the ESS, which may therefore be spurious. Further details about model architectures, training and evaluation are given in App. C.3.2.
Table 2: Alanine dipeptide results. KLD is the empirical KLD of the Ramachandran plots (see Fig. 3). Forward ESS is estimated with the test set. Reverse ESS is estimated with  $10^{5}$  model samples. Results are averaged over 3 seeded runs, with the standard error reported as uncertainty.
KLD	NLL	Rev ESS (%)	Fwd ESS (%)
Flow on internal coordinates	(2.01 ± 0.04) · 10-3	-190.15 ± 0.02	1.61 ± 0.03	-
E-CNF FM	(3.83 ± 0.18) · 10-2	-190.20 ± 0.09	(3.1 ± 0.2) · 10-4	(2.5 ± 0.8) · 10-128
E-CNF-DIFF	(8.86 ± 0.49) · 10-3	-188.31 ± 0.01	(8.1 ± 1.1) · 10-4	(5.1 ± 4.1) · 10-237
NON-E-ACF	(1.66 ± 0.01) · 10-1	-184.57 ± 0.35	0.14 ± 0.07	(5.5 ± 4.5) · 10-30
VECTOR-PROJ E-ACF	(6.15 ± 1.21) · 10-3	-188.56 ± 0.01	19.4 ± 13.4	(9.4 ± 7.7) · 10-7
CARTESIAN-PROJ E-ACF	(3.46 ± 0.28) · 10-3	-188.59 ± 0.00	52.5 ± 3.2	(9.7 ± 7.9) · 10-9
SPHERICAL-PROJ E-ACF	(2.55 ± 0.29) · 10-3	-188.57 ± 0.00	48.4 ± 7.2	(5.0 ± 4.1) · 10-14
The results are shown in Fig. 3 and Tab. 2. All variants of our E-ACF clearly outperform the NON-E-ACF, as the Kullback Leibler divergence (KLD) of the Ramachandran plots is significantly lower and the NLL as well as the reverse and forward ESS, see App. B.10, are higher. The flow trained on internal coordinates is only marginally better regarding the NLL and KLD than our best models, i.e.
Figure 3: Ramachandran plots, i.e. marginal distribution of the dihedral angles  $\phi$  and  $\psi$  (see App. C.3.2), obtained with MD (ground truth) and various models.
the CARTESIAN-PROJ and SPHERICAL-PROJ E-ACF, while we outperform it considering the ESS. Note that the flow on internal coordinates explicitly models the angles  $\phi$  and  $\psi$ , while the E-ACF operates on the underlying Cartesian coordinates. The E-ACFs outperform the E-CNF-DIFF model in all metrics, but the E-CNF trained with flow matching has a slightly lower NLL while having a significantly higher KLD on the Ramachandran plot. This could be due to a better performance on other marginals or on some outlier data points in the test set. The forward ESS is very low for all models, which suggests that the models do not cover some regions in the target distribution, and that the reverse ESS is spurious. Alternatively, this may be from numerical instabilities in the models. To the best of our knowledge, our models are the first to learn the full Boltzmann distribution of a molecule purely on Cartesian coordinates while being competitive with a flow trained on internal coordinates. Additionally, we train the best E-ACF model at  $T = 300\mathrm{K}$  where it also performs well (see App. C.3.2).
4.3 Energy-based training: DW4 and LJ13
Lastly, we demonstrate that our proposed flow can be trained on the DW4 and LJ13 problems using only the target's unnormalized density with the FAB algorithm [Midgley et al., 2023]. The annealed importance sampling procedure within FAB requires sampling from the flow and evaluating its density multiple times. This is used within the training loop of FAB making it significantly more expensive per parameter update than training by maximum likelihood. Given that sampling and density evaluation with CNFs is very expensive, training them with FAB is intractable. Thus, we only report results for our flow, as well as for the NON-E-ACF. We train the NON-E-ACF for more iterations than the E-ACF, such that the training times are similar, given that the NON-E-ACF is faster per iteration. App. C.4 provides further details on the experimental setup.
Table 3: Results for training by energy with FAB. Best results are emphasized in bold. Results are averaged over 3 seeded runs, with the standard error reported as uncertainty. Reverse ESS is estimated with  $10^{4}$  samples from the flow. Forward ESS is estimated with the test sets.
Rev ESS (%)	DW4 Fwd ESS (%)	NLL	Rev ESS (%)	LJ13 Fwd ESS (%)	NLL
NON-E-ACF	35.94 ± 2.63	5.45 ± 4.10	7.38 ± 0.01	5.38 ± 3.66	4.14 ± 3.10	33.22 ± 0.96
VECTOR-PROJ E-ACF	84.29 ± 0.41	83.39 ± 0.79	7.11 ± 0.00	59.60 ± 1.13	65.20 ± 1.61	30.33 ± 0.03
CARTESIAN-PROJ E-ACF	82.44 ± 0.50	80.08 ± 0.64	7.13 ± 0.00	60.68 ± 0.41	65.54 ± 0.37	30.34 ± 0.01
SPHERICAL-PROJ E-ACF	80.44 ± 0.88	81.46 ± 0.95	7.14 ± 0.00	62.09 ± 0.76	66.13 ± 0.11	30.21 ± 0.02
Tab. 3 shows that the E-ACF trained with FAB successfully approximates the target Boltzmann distributions, with reasonably high joint ESS, and NLL comparable to the flows trained by maximum likelihood. Additionally, the ESS may be improved further by combining the trained flow with AIS, this is shown in App. C.4. In both problems the NON-E-ACF performs worse, both in terms of ESS and NLL. All models trained by maximum likelihood have a much lower ESS (see App. C.3.1). This is expected, as unlike the  $\alpha = 2$ -divergence loss used in FAB, the maximum likelihood objective does not explicitly encourage minimizing importance weight variance. Furthermore, the flows trained by maximum likelihood use a relatively small, biased training set, which therefore limits their quality.
Augmented flows have been used for improving expressiveness of Boltzmann generators [Köhler et al., 2023b,c]; however, these models were not equivariant. Klein et al. [2023a] proposed an augmented normalizing flow architecture to provide conditional proposal distributions for MD simulations and use a coupling scheme similar to ours. However, this model only achieves translation and permutation equivariance and the authors make their flow approximately rotation invariant through data augmentation. In our experiments, we found data augmentation to perform significantly worse than intrinsic invariance. In principle, Klein et al. [2023a]'s model could be made fully invariant by substituting in our flow's projection-based coupling transform.
An alternative to our equivariant flow and equivariant CNFs are the equivariant residual flows proposed in Bose et al. [2022]. Alas, residual flows require fixed-point iteration for training and inversion. This is expensive and may interact poorly with energy-based training methods such as FAB [Midgley et al., 2023] which require fast exact sampling and densities. Furthermore, Bose et al. [2022] found that the spectral normalization required for residual flows did not interact well with the equivariant CNN architecture in their experiments.
There has been recent progress in improving the sampling speed of CNFs/diffusion models [Tong et al., 2023, Song et al., 2023] and on using these models for sampling from unnormalized densities [Vargas et al., 2023, Zhang and Chen, 2022, Zhang et al., 2023]. Thus, in the future CNF/diffusion models trained by energy may prove to be competitive with discrete-time flow based methods.
The strategy of projection into a local reference frame to enforce equivariance has been successfully employed in existing literature, specifically for protein backbone generation [Jumper et al., 2021, Yim et al., 2023]. Here we have focused on modelling the full set of Cartesian coordinates of a molecule, but an interesting avenue for future work is to extend our framework to other domains, such as modelling rigid bodies, which has applications to protein backbone generation [Jumper et al., 2021, Yim et al., 2023] and many-body systems [Kohler et al., 2023c].
Limitations Although our flow is significantly faster than alternatives such as CNFs, the expensive EGNN forward pass required in each layer of the flow makes it more computationally expensive than flows on internal coordinates. Additionally, we found our flow to be less numerically stable than flows on internal coordinates, which we mitigate via adjustments to the loss, optimizer and neural network (see App. B.3, App. C.1, App. C.2). Our implementation uses the E(3) equivariant EGNN proposed by Satorras et al. [2021a]. However, recently there have been large efforts towards developing more expressive, efficient and stable EGNNs architectures [Fuchs et al., 2020, Batatia et al., 2022, Musaelian et al., 2023, Liao and Smidt, 2023]. Incorporating these into our flow may improve performance, efficiency and stability. This would be especially useful for energy-based training, where the efficiency of the flow is a critical factor.
We have proposed an SE(3) equivariant augmented coupling flow that achieves similar performance to CNFs when trained by maximum likelihood, while allowing for faster sampling and density evaluation by more than an order of magnitude. Furthermore, we showed that our flow can be trained as a Boltzmann generator using only the target's unnormalized density, on problems where internal coordinates are inadequate due to permutation symmetries, and doing so with a CNF is computationally intractable. It is possible to extend our model to learn the Boltzmann distribution of diverse molecules, by conditioning on their molecular graph, which we hope to explore in the future.
We thank Gábor Csányi and his group for the helpful discussions. Laurence Midgley acknowledges support from Google's TPU Research Cloud (TRC) program and from the EPSRC through the Syntech PhD program. Vincent Stimper acknowledges the Max Planck Computing and Data Facilities for providing extensive computing resources and support. Javier Antorán acknowledges support from Microsoft Research, through its PhD Scholarship Programme, and from the EPSRC. José Miguel Hernández-Lobato acknowledges support from a Turing AI Fellowship under grant EP/V023756/1. José Miguel Hernández-Lobato and Emile Mathieu are supported by an EPSRC Prosperity Partnership EP/T005386/1 between Microsoft Research and the University of Cambridge. This work has been performed using resources provided by the Cambridge Tier-2 system operated by the University of Cambridge Research Computing Service (http://www.hpc.cam.ac.uk) funded by an EPSRC Tier-2 capital grant. It was also supported by the German Federal Ministry of Education and Research (BMBF): Tübingen AI Center, FKZ: 01IS18039B; and by the Machine Learning Cluster of Excellence, EXC number 2064/1 - Project number 390727645.