SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \mathcal {N} (R \boldsymbol {a} + u; R \boldsymbol {x} + u, \eta^ {2} I) = \frac {1}{Z} \exp \frac {- 1}{2} (R \boldsymbol {a} + u - R \boldsymbol {x} + u) ^ {\top} (\eta^ {- 2} I) (R \boldsymbol {a} + u - R \boldsymbol {x} + u) (18) \\ = \frac {1}{Z} \exp \frac {- 1}{2} (\boldsymbol {a} - \boldsymbol {x}) ^ {\top} R ^ {\top} (\eta^ {- 2} I) R (\boldsymbol {a} - \boldsymbol {x}) \\ = \frac {1}{Z} \exp \frac {- 1}{2} (\boldsymbol {a} - \boldsymbol {x}) ^ {\top} R ^ {\top} R (\eta^ {- 2} I) (\boldsymbol {a} - \boldsymbol {x}) \\ = \mathcal {N} (\boldsymbol {a}; \boldsymbol {x}, \eta^ {2} I) (19) \\ \end{array}
$$
An analogous argument holds for permutations, since any  $\sigma \in S_n$  can be represented as a unitary  $n \times n$  matrix which is simultaneously applied to  $\mathbf{a}$  and  $\mathbf{x}$ .
B.6 Maximum likelihood objective
We train our models from samples by maximising the joint log-density  $\mathbb{E}_{\boldsymbol{x} \sim p(\boldsymbol{x}), \boldsymbol{a} \sim \pi(\boldsymbol{a} | \boldsymbol{x})}[\log q(\boldsymbol{x}, \boldsymbol{a})]$ . This can be shown to be a lower bound on the expected marginal log-density over observations, which is the quantity targeted by models without augmented variables as
$$
\operatorname {E} _ {p (\boldsymbol {x})} [ \log q (\boldsymbol {x}) ] = \operatorname {E} _ {p (\boldsymbol {x})} \left[ \log \operatorname {E} _ {\pi (\boldsymbol {a} | \boldsymbol {x})} \left[ \frac {q (\boldsymbol {x} , \boldsymbol {a})}{\pi (\boldsymbol {a} | \boldsymbol {x})} \right] \right] \geq \operatorname {E} _ {p (\boldsymbol {x}) \pi (\boldsymbol {a} | \boldsymbol {x})} \left[ \log q (\boldsymbol {x}, \boldsymbol {a}) - \log \pi (\boldsymbol {a} | \boldsymbol {x}) \right] \tag {20}
$$
where we drop the latter term from our training objective since it is constant w.r.t. flow parameters. The bound becomes tight when  $q(\pmb{a}|\pmb{x}) = \pi (\pmb{a}|\pmb{x})$ .
B.7 Reverse KL objective
We now show the reverse KL objective involving the joint distribution of observations and augmented variables presented in the main text upper bounds the reverse KL over observations. Using the chain rule of the KL divergence we have
$$
\begin{array}{l} D _ {\mathrm {K L}} (q (\boldsymbol {x}, \boldsymbol {a}) | | p (\boldsymbol {x}) \pi (\boldsymbol {a} | \boldsymbol {x})) = D _ {\mathrm {K L}} (q (\boldsymbol {x}) | | p (\boldsymbol {x})) + \mathbb {E} _ {q (\boldsymbol {x})} D _ {\mathrm {K L}} (q (\boldsymbol {a} | \boldsymbol {x}) | | \pi (\boldsymbol {a} | \boldsymbol {x})) \\ \geq D _ {\mathrm {K L}} \left(q (\boldsymbol {x}) | | p (\boldsymbol {x})\right). \\ \end{array}
$$
The looseness in the bound is  $\mathbb{E}_{q(\pmb{x})}D_{\mathrm{KL}}(q(\pmb{a}|\pmb{x})||\pi (\pmb{a}|\pmb{x}))$  which again vanishes when  $q(\pmb {a}|\pmb {x}) = \pi (\pmb {a}|\pmb {x})$
B.8 Alpha divergence objective
We now show that our joint objective represents an analogous upper bound for the  $\alpha$  divergence when  $\alpha = 2$ , i.e., the objective targeted by FAB [Midgley et al., 2023]. Let  $U(\pmb{x})$  be the unnormalized energy function of interest evaluated at state  $\pmb{x}$ , with  $p(\pmb{x}) \propto \exp(-U(\pmb{x}))$ . We have
$$
\begin{array}{l} D _ {2} (q (\boldsymbol {x}) | | p (\boldsymbol {x})) = \int \frac {p (\boldsymbol {x}) ^ {2}}{q (\boldsymbol {x})} d \boldsymbol {x} \\ \leq \int \frac {p (\boldsymbol {x}) ^ {2}}{q (\boldsymbol {x})} \left(\int \frac {\pi (\boldsymbol {a} | \boldsymbol {x}) ^ {2}}{q (\boldsymbol {a} | \boldsymbol {x})} d \boldsymbol {a}\right) d \boldsymbol {x} \\ = \int \frac {p (\boldsymbol {x}) ^ {2} \pi (\boldsymbol {a} | \boldsymbol {x}) ^ {2}}{q (\boldsymbol {x}) q (\boldsymbol {a} | \boldsymbol {x})} d \boldsymbol {a} d \boldsymbol {x} \\ = D _ {2} \left(q (\boldsymbol {x}, \boldsymbol {a}) | | p (\boldsymbol {x}) \pi (\boldsymbol {a} | \boldsymbol {x})\right) \\ \end{array}
$$
which also becomes tight when  $q(\pmb{a}|\pmb{x}) = \pi(\pmb{a}|\pmb{x})$ . The inequality is true because  $\int \frac{\pi(\pmb{a}|\pmb{x})^2}{q(\pmb{a}|\pmb{x})} \, \mathrm{d}\pmb{a} > 0$ . To see this, take any densities  $\pi$  and  $q$ , and apply Jensen's inequality as
$$
\int \frac {\pi (a) ^ {2}}{q (a)} d a = \mathrm {E} _ {\pi (a)} \frac {\pi (a)}{q (a)} = \mathrm {E} _ {\pi (a)} \left(\frac {q (a)}{\pi (a)}\right) ^ {- 1} \geq \left(\mathrm {E} _ {\pi (a)} \frac {q (a)}{\pi (a)}\right) ^ {- 1} = \int q (a) d a = 1.
$$
B.9 Importance weighted estimator
We define the following estimator  $w \coloneqq \frac{q(\boldsymbol{x},\boldsymbol{a})}{\pi(\boldsymbol{a}|\boldsymbol{x})}$  with  $\boldsymbol{a} \sim \pi(\cdot|x)$ . Assuming that whenever  $\pi(\boldsymbol{a}|\boldsymbol{x}) = 0$  then  $q(\boldsymbol{x},\boldsymbol{a}) = 0$ , we have that this estimator is trivially unbiased:
$$
\begin{array}{l} \mathbb {E} [ w ] = \int \frac {q (\boldsymbol {x} , \boldsymbol {a})}{\pi (\boldsymbol {a} | \boldsymbol {x})} \pi (\boldsymbol {a} | \boldsymbol {x}) d \mu (\boldsymbol {a}) \\ = \int q (\boldsymbol {x}, \boldsymbol {a}) d \mu (\boldsymbol {a}) \\ = q (\boldsymbol {x}) \\ \end{array}
$$
and that is variance is given by
$$
\begin{array}{l} \mathbb {V} [ w ] = \mathbb {E} [ (w - q (\boldsymbol {x})) ^ {2} ] \\ = \mathbb {E} [ w ^ {2} ] - \mathbb {E} [ 2 w q (\boldsymbol {x}) ] + \mathbb {E} [ q (\boldsymbol {x}) ^ {2} ] \\ = \int \frac {q (\boldsymbol {x} , \boldsymbol {a}) ^ {2}}{\pi (\boldsymbol {a} | \boldsymbol {x}) ^ {2}} \pi (\boldsymbol {a} | \boldsymbol {x}) d \mu (\boldsymbol {a}) - 2 q (\boldsymbol {x}) \mathbb {E} [ w ] + q (\boldsymbol {x}) ^ {2} \\ = \int \frac {q (\boldsymbol {x}) q (\boldsymbol {a} | \boldsymbol {x})}{\pi (\boldsymbol {a} | \boldsymbol {x})} q (\boldsymbol {x}, \boldsymbol {a}) d \mu (\boldsymbol {a}) - q (\boldsymbol {x}) ^ {2} \\ \end{array}
$$
thus if  $q(\pmb {a}|\pmb {x}) = \pi (\pmb {a}|\pmb {x}),\mathbb{V}[w] = 0$
The log marginal likelihood may therefore be estimated with
$$
\log q (\boldsymbol {x}) \approx \log \frac {1}{M} \sum_ {m = 1} ^ {M} \frac {q (\boldsymbol {x} , \boldsymbol {a} _ {m})}{\pi (\boldsymbol {a} _ {m} | \boldsymbol {x})} \quad \text {a n d} \quad \boldsymbol {a} _ {m} \sim \pi (\cdot | \boldsymbol {x}), \tag {21}
$$
which becomes accurate as  $\mathbf{M}$  becomes large.
B.10 Effective sample size
For evaluation we use both the reverse and forward effective sample size (ESS), where the former relies on samples from the model, and the latter from samples from the target. For the augmented flow models we report the joint ESS, which lower bounds the marginal ESS. Below, we provide the formulas for the reverse and forward ESS. For simplicity, we define these in terms of the marginal ESS. The joint ESS estimators may be obtained using the same formulas, replacing  $\pmb{x}$  with  $(\pmb{x},\pmb{a})$ .
The reverse effective sample size is given by
$$
n _ {\mathrm {e}, \mathrm {r v}} = \frac {\left(\sum_ {i = 1} ^ {n} w \left(\boldsymbol {x} _ {i}\right)\right) ^ {2}}{\sum_ {i = 1} ^ {n} w \left(\boldsymbol {x} _ {i}\right) ^ {2}} \quad \text {a n d} \quad \boldsymbol {x} _ {i} \sim q (\boldsymbol {x} _ {i}), \tag {22}
$$
where  $w(\pmb{x}_i) = \tilde{p}(\pmb{x}_i) / q(\pmb{x}_i)$  and  $\tilde{p}$  is the unnormalized version of  $p$  [Owen, 2013].
The forward ESS may be estimated with
$$
n _ {\mathrm {e , f w d}} = \frac {n ^ {2}}{Z ^ {- 1} \sum_ {i = 1} ^ {n} w \left(\boldsymbol {x} _ {i}\right)} \quad \text {a n d} \quad \boldsymbol {x} _ {i} \sim p \left(\boldsymbol {x} _ {i}\right), \tag {23}
$$
where the normalizing constant  $Z = \tilde{p} / p$  may be estimated using
$$
Z ^ {- 1} = E _ {p} \left[ \frac {q (\boldsymbol {x})}{\tilde {p} (\boldsymbol {x})} \right] \tag {24}
$$
$$
\approx \frac {1}{n} \sum_ {i = 1} ^ {n} \frac {q \left(\boldsymbol {x} _ {i}\right)}{\tilde {p} \left(\boldsymbol {x} _ {i}\right)} \quad \text {a n d} \quad \boldsymbol {x} _ {i} \sim p \left(\boldsymbol {x} _ {i}\right). \tag {25}
$$
B.11 Conversion of internal coordinate flow density to density on ZeroCoM Cartesian space with global rotation represented
The density of the flow on internal coordinates is not directly comparable to the flows on Cartesian coordinates, as the variables live in a different space (bond distances, angles and torsions instead
of Cartesian coordinates). The internal coordinates, and density may be converted into Cartesian coordinates via passing them through the appropriate bijector and computing the log determinant of the bijector. This is commonly done in existing literature as the target energy typically takes Cartesian coordinates as input [Midgley et al., 2023]. However, this representation is invariant to global rotations of the molecule, making it lower dimensional than the space with which the flows which introduced in this paper operate. Furthermore, this representation has the first atom centered on the origin, where the flows in this paper instead have zero-CoM. To compare the density of the flow on internal coordinates with the flows that operate directly on Cartesian coordinates we have to account for both of these aforementioned factors. We describe how to do this below.
Conversion to density that includes degrees of freedom representing global rotation The rotation-invariant representation of internal coordinates in Cartesian space places the first atom at the origin, the second atom a distance of  $b_{1}$  along the X-axis (i.e. at position  $\{b_{1},0,0\}$ ), and the third atom on the XY plane, at a distance of  $b_{2}$  from the origin at an angle of  $a2$  with respect to the x-axis (i.e. at position  $\{\cos(a_{2})\cdot b_{2},\sin(a_{2})\cdot b_{2},0\}$ ). We define this representation of Cartesian coordinates as  $z\in Z$ , where we drop all coordinates of the first atom (atom 0), the  $y$  and  $z$  coordinates of atom 1, and the  $z$  coordinate of atom 2, all of which are zero, so that  $Z$  contains as many elements as there are degrees of freedom.
To convert this into a space that has degrees of freedom representing global rotations, which we refer to as  $w \in W$ , we simply sample a rotation uniformly over  $SO3$  and apply this to  $w$ .
The rotation has three degrees of freedom, which we represent as follows; First, we apply rotation about the X-axis of  $o_1$  sampled uniformly over  $[-\pi, \pi]$ . Secondly, we rotate about the Z-Axis by sampling  $o_2$  uniformly between  $[-1, 1]$  and rotating by  $\sin^{-1}(o_2)$ . Finally, we rotate about the Y-axis by  $o_3$  sampled uniformly between  $[-\pi, \pi]$ . Together this ensures that the resultant orientation of the molecule in Cartesian space has a uniform density over  $SO3$ . We denote the concatenation of these random variables that parameterize the rotation as  $o \in \mathcal{O}$  where  $\mathcal{O} = \{[-\pi, \pi], [-1, 1], [-\pi, \pi]\}$ .
The density on  $W$  is then given by
$$
q (\boldsymbol {w}) = q (\boldsymbol {z}) q (\boldsymbol {o}) \left| \det  \frac {\partial g (\boldsymbol {z} , \boldsymbol {o})}{\partial \{\boldsymbol {z} , \boldsymbol {o} \}} \right| ^ {- 1}. \tag {26}
$$
where  $g:\mathcal{Z}\times \mathcal{O}\to \mathcal{W}$  is a bijector which applies the rotation parameterized by  $o$  to  $z$ .
It can be shown that the log jacobian determinant of  $g$  is  $2\log (b_{1}) + \log (b_{2}) + \log (\sin (a_{2}))$ . Furthermore,  $\log q(\pmb {o})$  is a constant value of  $-\log 8\pi$ .
Conversion density on Zero-CoM space After applying the above transform we have  $\boldsymbol{w}$  that includes rotations, but is centered on atom 0. Thus, the last step needed is to convert into the Zero-CoM space. As previously, to ensure that our representations have the same number of elements as there are degrees of freedom, we drop the first atom's Cartesian coordinates, as these contain redundant information. We note that for the representation  $\boldsymbol{w} \in \mathcal{W}$  the first atom's coordinates are at the origin, and for the Zero-CoM space  $\boldsymbol{x} \in \tilde{\mathcal{X}}$  they are at  $-\sum_{i=1}^{n-1} x_i$ . The bijector  $h: \mathcal{W} \to \tilde{\mathcal{X}}$  is given by
$$
h (\boldsymbol {w}) = \boldsymbol {w} - \frac {1}{N} \sum_ {i = 1} ^ {n - 1} \boldsymbol {w} _ {i} \tag {27}
$$
where  $N$  is the number of nodes. The log jacobian determinant of  $h$  is simply equal to  $-3\log n$
Overall conversion Thus we can convert the density in  $\mathcal{Z}$  into an equivalent density in  $\tilde{\mathcal{X}}$  using,
$$
\begin{array}{l} \log q (\boldsymbol {x}) = \log q (\boldsymbol {z}) + \log q (\boldsymbol {o}) + \log \left| \det  \frac {\partial g (\boldsymbol {z} , \boldsymbol {o})}{\partial \{\boldsymbol {z} , \boldsymbol {o} \}} \right| ^ {- 1} + \log \left| \det  \frac {\partial h (\boldsymbol {w})}{\partial \boldsymbol {w}} \right| ^ {- 1} \tag {28} \\ = \log q (\boldsymbol {z}) - \log 8 \pi^ {2} - 2 \log (b _ {1}) - \log (b _ {2}) - \log (\sin (a _ {2})) + 3 \log n. \\ \end{array}
$$
C Experiments: Details and Further Results
We provide the details of the hyper-parameters and compute used for each experiment. Additionally, we provide the code https://github.com/lolllcat/se3-augmented-coupling-flows
for replicating our experiments. It also contains the configuration files with the exact hyperparameter for each model and training run.
C.1 Custom Optimizer
For all experiments we use a custom defined optimizer, which improves training stability by skipping updates with very large gradient norms, and clipping updates with a moderately large gradient norms. Our optimizer keeps track of a window the the last 100 gradient norms, and then for a given step: (1) If the gradient norm is more than 20 times greater than the median gradient norm within the window, then skip the gradient step without updating the parameters. (2) Otherwise, if the gradient norm is more than five times the median within the window then clip the gradient norm to five times the median. (3) Use the (possibly clipped) gradients to perform a parameter update using the adam optimizer [Kingma and Ba, 2014]. Lastly, we also include an optional setting which allows for masking of gradient contribution of samples within each batch with very large negative log probability during training, which further improves training stability.
C.2 EGNN implementation
We use our own custom implementation of the EGNN proposed by Satorras et al. [2021b] for our experiments. We found that applying a softmax to the invariant features output from the EGNN helped improve stability. Additionally, we extend Satorras et al. [2021b]'s EGNN to allow for multiple channels of positions/vectors at each node, with interaction across channels. This is needed as the local reference frame  $r$  typically is parameterized by multiple equivariant vectors, and for the case where the multiple sets of augmented variables are used such that there are multiple sets of positions input to the EGNN. We refer to the code for the details of our EGNN implementation.
C.3 Training with samples
C.3.1 DW4, LJ13 and QM9 Positional
Experimental details Below we provide a description of the hyper-parameters used within each experiment. Additionally Tab. 4 provides the run-times for each experiment.
Datasets: For DW4 and LJ13 we use a training and test set of 1000 samples following Satorras et al. [2021a]. We note that this data is biased with a clear difference between the distribution of train and test data (see Fig. 5. 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].
Figure 5: Inter-atomic distances for samples from the train and test data on DW4 and LJ13. Both datasets are biased with visible differences between the train and test sets.
Flow experiment details 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  $\mathcal{N}(\pmb{a};\pmb{x},\eta^2 I)$ , we use  $\eta = 0.1$  for both the base of the flow, and augmented target distribution. For the CARTESIAN-PROJ E-ACF and NON-E-ACF we use RealNVP [Dinh et al., 2017] for the flow transform, while for CARTESIAN-PROJ and VECTOR-PROJ we use a rational quadratic spline transform [Durkan et al., 2019]. For the spline we use 8 bins for all flows. We use 1 set of augmented variables for DW4 and LJ13, and 3 sets
of augmented variables for QM9-positional. For all equivariant GNN's we implemented the GNN proposed by Satorras et al. [2021b] with 3 blocks. All MLP's have 2 layers, with a width of 64 units. For the NON-E-ACF model we use a transformer with 3 layers of self-attention with 8 heads. 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. For QM9 the learning rate is initialised to a value of 0.00006, peaks at 0.0002 after 30 epochs, and then decays back to a final learning rate of 0.00006. With the above learning rate schedule for QM9 the CARTESIAN-PROJ E-ACF was unstable during training which resulted in crashed runs (NaN loss). Hence, it was instead trained with a constant learning rate of 0.00006 at which training was stable. 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. We train the NON-E-ACF for 4 times as many epochs for all problems, which results in roughly the same run time as the SPHERICAL-PROJ E-ACF. For the CARTESIAN-PROJ and SPHERICAL-PROJ E-ACF, we use a coefficient of 10. for the auxiliary loss (Eq. (11)), for weighting it relative to the maximum likelihood loss. Marginal log likelihood in Tab. 1 is calculated using Eq. (21) with 20 samples from the augmented variables  $\mathbf{a}$  per test-set sample of  $\mathbf{x}$ .
Continuous normalizing flow / diffusion experiment details For the E-CNF ML we report the results from Satorras et al. [2021a]. For the E-CNF FM we implement flow matching, with the model's vector field parameterized by the  $\mathrm{E}(\mathfrak{n})$  GNN of Satorras et al. [2021b]. For DW4 and LJ13 we use an EGNN with 3 message passing layers, with a 3-hidden layer MLP with 128 units per layer. For QM9-positional we use an EGNN with 5 message passing layers, with an 4-hidden layer MLP with 256 units per layer. For all problems we use Adam optimizer with a cosine learning rate decay from 1e-4 to 0. Additionally, for QM9, we use EMA with  $\beta = 0.999$  to compute the final training weights which the model is evaluated with. We train for 200, 400, and 1600 epochs on DW4, LJ13 and QM9-positional respectively. For sampling and evaluating log-densities we use the dopri5 ODE solver [Dormand and Prince, 1980] with an rto1 and ato1 of 1e-5. For DW4 and LJ13 we evaluate densities using the exact divergence of the vector field, while for QM9 we approximate the divergence using the Hutchinson trace estimator [Grathwohl et al., 2018] as the exact divergence is computationally expensive.
Our E-CNF-DIFF implementation roughly follows Hoogeboom et al. [2022], but we use the noise schedule as Ho et al. [2020]. We use the same GNN architecture as for our flows, that is the E(n) GNN of Satorras et al. [2021b], to predict the diffusion noise. The network is composed of 3 message passing layers, each of which uses a 3 hidden layer MLP with 256 units per layer. GNN expands its inputs to 256 scalars and 196 vectors per graph node and returns its predictions as a linear combination of the latter 196 vectors. This model has around 90M learnable parameters. We use the AdamW optimiser with a weight decay of  $10^{-4}$  and run it for 100, 200, and 5000 epochs for the DW4, LJ13 and QM9 dataset, respectively. The initial learning rate is set to  $10^{-5}$ . It is linearly increased to  $5 \cdot 10^{-4}$  for the first  $15\%$  of training before being decreased to  $10^{-6}$  with a cosine schedule. Following, Kingma et al. [2021], we sample noise timesteps during such that these are uniformly spaced within each batch, reducing gradient variance. We set the smallest accessible time value to  $10^{-3}$ . For sampling and evaluating log-densities we use the dopri5 ODE solver with a maximum number of evaluations of 1000.
Table 4: Run time for training by maximum likelihood on the DW4, LJ13 and QM9-positional datasets, in hours. All runs used an GeForce RTX 2080Ti GPU. Run times include evaluation that was performed intermittently throughout training. The results are averaged over 3 seeded runs, with the standard error reported as uncertainty.
DW4	LJ13	QM9 positional
E-CNF ML	N/A ± N/A	N/A ± N/A	336 ± N/A
E-CNF-DIFF	N/A ± N/A	N/A ± N/A	3.80 ± 0.00
NON-E-ACF	0.12 ± 0.00	0.54 ± 0.00	20.68 ± 0.09
VECTOR-PROJ E-ACF	0.16 ± 0.00	0.61 ± 0.00	19.98 ± 0.04
CARTESIAN-PROJ E-ACF	0.12 ± 0.00	0.50 ± 0.00	17.13 ± 0.33
SPHERICAL-PROJ E-ACF	0.15 ± 0.00	0.59 ± 0.00	20.61 ± 0.17
Further results Tab. 5 provides the effective sample size of the flows trained by maximum likelihood, with respect to the joint target distribution  $p(\pmb{x}) \pi(\pmb{a}|\pmb{x})$ .
Table 5: Joint effective sample size (\%) evaluated using 10000 samples. The results are averaged over 3 seeded runs, with the standard error reported as uncertainty.
DW4	LJ13
Rev ESS (%)	Fwd ESS (%)	Rev ESS (%)	Fwd ESS (%)
E-CNF FM	5.02 ± 2.22	0.01 ± 0.00	5.37 ± 0.35	0.30 ± 0.18
NON-E-ACF	0.68 ± 0.18	0.06 ± 0.02	0.00 ± 0.00	0.00 ± 0.00
VECTOR-PROJ E-ACF	1.87 ± 0.68	0.16 ± 0.01	1.80 ± 0.17	0.06 ± 0.03
CARTESIAN-PROJ E-ACF	3.08 ± 0.05	0.09 ± 0.03	0.89 ± 0.37	0.00 ± 0.00
SPHERICAL-PROJ E-ACF	3.06 ± 0.60	0.17 ± 0.05	1.50 ± 0.37	0.02 ± 0.00
Varying the number of training samples for DW4 and LJ13 Tab. 6 shows the effect of varying the training dataset size on the performance of the flows. We tune the number of epochs via the validation set; we use 1000, 100, 50 and 35 epochs for DW4 and 1500, 1500, 400 and 100 LJ13 for each dataset size respectively. Additionally we train NON-E-ACF (with data augmentation) for 4 times as many epochs for each experiment, as it is computationally cheaper per epoch and takes a larger number of epochs before it overfits. We see that the performance gap between the non-equivariant and equivariant flow models is larger in the lower data regime.
Table 6: Negative log likelihood results on test set for flows trained by maximum likelihood for the DW4 and LJ13 datasets varying the number of training samples. E-CNF results are from Satorras et al. [2021a]. Best results are emphasized in bold.
# samples	DW4	LJ13
\( 10^2 \)	\( 10^3 \)	\( 10^4 \)	\( 10^5 \)	10	\( 10^2 \)	\( 10^3 \)	\( 10^4 \)
E-CNF	8.31	8.15	7.69	7.48	33.12	30.99	30.56	30.41
NON-E-ACF	10.24	10.07	8.23	7.39	40.21	34.83	33.32	31.94
VECTOR-PROJ E-ACF	8.65	8.69	7.95	7.47	33.64	30.61	30.19	30.09
CARTESIAN-PROJ E-ACF	8.79	8.82	7.93	7.52	36.66	31.72	30.89	30.48
SPHERICAL-PROJ E-ACF	8.66	8.61	7.97	7.38	34.79	31.26	30.33	30.16
C.3.2 Alanine dipeptide
Figure 6: Visualization of the 22-atom molecule alanine dipeptide with its two dihedral angles  $\phi$  and  $\psi$ .
The molecule alanine dipeptide and its dihedral angles  $\phi$  and  $\psi$ , which are used in the Ramachandran plots in Fig. 3, are depicted in Fig. 6. We model it in an implicit solvent at a temperature of  $T = 800\mathrm{K}$ . The data used for maximum likelihood training and testing was generated with a replica exchange MD simulation [Mori and Okamoto, 2010] with 21 systems having temperatures from  $300\mathrm{K}$  until  $1300\mathrm{K}$  with  $50\mathrm{K}$  temperature difference between each of them.
Each ACF has 20 layers with 3 core transformations per layer for the VECTOR-PROJ flow and 1 for the CARTESIAN-PROJ and the SPHERICAL-PROJ flow. The flow trained on internal coordinates has the same architecture as the models in Midgley et al. [2023], i.e. it is a coupling neural spline flow with 12 layers. The models were trained for 50 epochs.
(b) Raw model samples
Figure 7: Ramachandran plot of a SPHERICAL-PROJ E-ACF model trained on alanine dipeptide at  $800\mathrm{K}$ . The samples in Fig. 7c are reweighted with the importance weights. Thereby, the KLD drops from  $1.88 \cdot 10^{-3}$  to  $1.39 \cdot 10^{-3}$ .
(c) Reweighted model samples