SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The E-CNF and the E-CNF-DIFF model have  $\mathrm{E}(\mathfrak{n})$  GNN with 4 message passing layers, each using a 3-layered MLP with 128 hidden units each, as a vector field and score function, respectively. Both were trained for 20 epochs.
Since the version of the VECTOR-PROJ flow as well as the E-CNF and E-CNF-DIFF models are parity equivariant, the model learns to generate both chiral forms of alanine dipeptide [Midgley et al., 2023]. For evaluating the model, we filtered the samples from the model to only include those that correspond to the L-form, which is the one almost exclusively found in nature.
To plot and compute the KLD of the Ramachandran plots, we drew  $10^{7}$  samples from each model. Since computing the likelihood of the E-CNF and the E-CNF-DIFF model is so expensive, we evaluated these models only on  $10\%$  of the test set, i.e.  $10^{6}$  datapoints.
Limiting the number of samples To test how well E-ACFs generalize thanks to their equivariance, we trained our three variants of E-ACFs and an NON-E-ACFon subsets of the training dataset with various sizes. The results are given in Table 7. For all training dataset sizes the equivariant architectures clearly outperform the non-equivariant counterpart. Notably, the NON-E-ACF trained on the full dataset performs worse than the VECTOR-PROJ E-ACF trained on  $10\%$  of the data and the CARTESIAN-PROJ and SPHERICAL-PROJ E-ACF trained on just  $1\%$  of the data.
Table 7: NLL of the models trained on subsets of the training dataset with different sizes. The best performing model, i.e. the one with the lowerst NLL, is marked in bold for each training dataset size.
Training dataset size	102	103	104	105	106
NON-E-ACF	-28.8	-76.8	-120.4	-164.3	-177.8
VECTOR-PROJ E-ACF	-31.0	-111.4	-169.5	-183.6	-188.5
CARTESIAN-PROJ E-ACF	-65.1	-155.7	-185.9	-188.1	-188.6
SPHERICAL-PROJ E-ACF	-53.1	-119.3	-183.6	-188.3	-188.6
Reweighting Given the high effective sample sizes we obtained for our models (see Tab. 2), we also attempted to reweight the samples. As already observed by Midgley et al. [2023], for a large number of samples, i.e.  $10^{7}$ , there were a few high outlier importance weights, which spoil the importance weight distribution. Hence, we clipped the highest  $0.2\%$  of the importance weights as in [Midgley et al., 2023].
The results for our SPHERICAL-PROJ E-ACF model are shown in Fig. 7. The reweighted Ramachandran plot is significantly closer to the ground truth in low density regions. The KLD is reduced from  $1.88 \cdot 10^{-3}$  to  $1.39 \cdot 10^{-3}$ .
Alanine dipeptide at  $T = 300\mathrm{K}$  Although our experiments were done on alanine dipeptide at  $T = 800\mathrm{K}$ , we also tried  $T = 300\mathrm{K}$ . We trained our SPHERICAL-PROJ E-ACF model on the same
(b) Raw model samples
Figure 8: Ramachandran plot of a SPHERICAL-PROJ E-ACF model trained on alanine dipeptide at  $300\mathrm{K}$ . The raw and reweighted results are shown. Through reweighting, the KLD drops from  $2.79 \cdot 10^{-3}$  to  $1.97 \cdot 10^{-3}$ .
dataset used in [Midgley et al., 2023]. As shown in Fig. 8, we got good performance. Even without reweighting our model has a KLD of  $2.79 \cdot 10^{-3}$ . When doing reweighting, it decreases to  $1.97 \cdot 10^{-3}$ . For comparison, the flows trained on internal coordinates presented in [Midgley et al., 2023] had a KLD of  $(7.57 \pm 3.80) \cdot 10^{-3}$ , and reweighting worsened the results given the low ESS. However, these models had about 6 times fewer parameters than our model.
C.4 Energy-based training
For both the DW4 and LJ13 experiments we use an identical flow architecture and optimizer setup to the flows trained by maximum likelihood with samples App. C.3.1. We use a batch size of 128 and 1024 for DW4 and LJ13 respectively. We train for 20000 and 14000 epochs for DW4 and LJ13 for the E-ACF. For the NON-E-ACF we train for 60000 and 56000 epochs for DW4 and LJ13 respectively, which corresponds to roughly the same training time as the SPHERICAL-PROJ E-ACF because the NON-E-ACF is faster per iteration. Reverse ESS is estimated with 10,000 samples from the flow. Forward ESS is estimated with the test sets.
Hyper-parameters for FAB: We set the minimum replay buffer size to 64 batches and maximum size to 128 batches. For the importance weight adjustment for re-weighting the buffer samples to account for the latest flow parameters, we clip to a maximum of 10. For AIS we use HMC with 5 leapfrog steps, and 2 and 8 intermediate distributions for DW4 and LJ13 respectively. We tune the step size of HMC to target an acceptance probability of  $65\%$ . We perform 8 buffer sampling and parameter update steps per AIS forward pass.
Table 8: Run time for energy-based training with FAB on the DW4 and LJ13 problems. DW4 was run using a GeForce RTX 2080Ti GPU and LJ13 with a V-4 TPU. 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
NON-E-ACF	8.9 ± 0.0	38.1 ± 0.0
VECTOR-PROJ E-ACF	7.5 ± 0.0	45.0 ± 0.1
CARTESIAN-PROJ E-ACF	6.3 ± 0.0	42.7 ± 0.0
SPHERICAL-PROJ E-ACF	7.3 ± 0.0	44.7 ± 0.0