SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Samples generated with a trained Boltzmann generators can be reweighted to the target distribution via Equation (29). Generally, the lower the variance of the log weight distribution, the higher the ESS. Although large outliers might reduce the ESS significantly. We present the log weight distribution for the LJ55 system in Figure 7a, showing that the weight distribution for the BG trained with equivariant OT flow matching has lower variance.
A.9 Approximation methods for equivariant OT flow matching
The approximation given in Equation (16) is in practice performed using the Hungarian algorithm for permutations and the Kabsch algorithm for rotations. However, one could also think of other approximations, such as performing the rotation first or alternating multiple permutations and rotations. We compare different approximation strategies in Figure 7b,c. The baseline reference is computed with an expensive search over the approximation given in Equation (16). Namely, we evaluate Equation (16) for 100 random rotations combined with the global reflection, denoted as  $O_{200}$ , for each sample, i.e.
$$
\hat {c} (x _ {0}, x _ {1}) = \min  _ {o \in O _ {2 0 0} (D)} \tilde {c} (x _ {0}, \rho (o) x _ {1}),
$$
where  $\tilde{c}(x_0, x_1)$  is given by Equation (16). Hence, this is 200 times more expensive than our approach. This baseline should be much closer to the true batch OT solution. Applying our approximation multiple times reduced the transportation cost slightly. Performing the rotations first, lead to inferior results. We observe the same behavior for the other systems. Therefore, we used the approximation given in Equation (16) throughout the experiments, showing that this is sufficient to learn a close approximation of the OT map.
A.10 Different data set sizes
Prior work [17, 18] investigates different training set sizes to train equivariant flows, showing that for simple datasets already quite a few samples are sufficient to achieve good NLL. We test if this also holds for flow matching for the more complex alanine dipeptide system. To this end, we run our flow matching experiments also for smaller data set sizes of 10000 and 1000. We report our findings in Table 5. In agreement with prior work, the NLL becomes worse the smaller the dataset. The same is true for the ESS. However, the sampled integration path lengths are nearly independent on the dataset size.
Figure 7: (a) Log weight distribution for samples generated for the LJ55 system. (b, c) Comparison of different approximation methods to approximate Equation (16) for alanine dipeptide.
Table 5: Comparison of flows trained with different training set sizes. Errors over three runs.
Training type	NLL (↓)	ESS (↑)	Path length (↓)
OT flow matching	Alanine dipeptide - 1000 training samples
-104.97 ± 0.05	0.25 ± 0.04%	10.14 ± 0.05
Eq OT flow matching	-103.38 ± 0.03	0.39 ± 0.08%	9.27 ± 0.02
OT flow matching	Alanine dipeptide - 10000 training samples
-105.16 ± 0.17	0.33 ± 0.07%	10.13 ± 0.02
Eq OT flow matching	-103.36 ± 0.15	0.40 ± 0.06%	9.32 ± 0.01
A.11 Naive flow matching
We provide naïve flow matching, i.e. flow matching without the OT reordering, as an additional baseline. Naïve flow matching results in even longer integration paths, while the ESS and likelihoods are close to the results of OT flow matching, as shown in Table 6. Flow matching with a non-equivariant architecture, e.g. a fully connected neural network, failed for all systems but DW4 and is hence not reported.
B Proofs and derivations
B.1 Equivariant OT flow matching
We show that our equivariant OT flow matching converges to the OT solution, loosely following the theory in [67] and extending it to infinite groups. Let  $G$  be a compact (topological) group that acts on an Euclidean  $n$ -space  $X$  by isometries, i.e.,  $d(g \cdot x, g \cdot y) = d(x, y)$ , where  $d$  is the Euclidean distance. This definition includes the symmetric group  $S_{n}$ , the orthogonal group  $O(n)$ , and all their subgroups. Consider the OT problem of minimizing the cost function  $c(x, y) = d(x, y)^p$ ,  $p \geq 1$  between  $G$ -invariant measures  $\nu_{1}$  and  $\nu_{2}$ .
Let us further define the pseudometric  $\tilde{d}(x,y)^p = \min_{g\in G}d(x,g\cdot y)^p$ . We want to show that any OT map  $T:X\to X$  for cost functions  $c(x,y) = d(x,y)^p$  is also an OT map for  $\tilde{c}(x,y) = \tilde{d}(x,y)^p$ .
Lemma 1. Let  $\nu_{1}$  and  $\nu_{2}$  be  $G$ -invariant measures, then there exists an  $G$ -invariant OT plan  $\pi(g \cdot x, g \cdot y) = \pi(x, y)$ .
Table 6: Results for naïve flow matching. Same hyperparameters as for OT flow matching. Errors over three runs.
Dataset	NLL (↓)	ESS (↑)	Path length (↓)
DW4	1.68 ± 0.01	93.01 ± 0.12%	3.41 ± 0.02
LJ13	-16.10 ± 0.01	57.55 ± 2.20%	3.77 ± 0.01
LJ55	-88.44 ± 0.03	3.15 ± 1.04%	8.66 ± 0.03
Alanine dipeptide	-107.56 ± 0.09	1.42 ± 0.51%	11.00 ± 0.02
Proof. Let  $\pi' \in \Pi(\nu_1, \nu_2)$  be a (not necessarily invariant) OT plan. We average the OT plan over the group
$$
\bar {\pi} (x, y) = \int_ {G} \pi^ {\prime} (g \cdot x, g \cdot y) d \mu (g), \tag {21}
$$
where  $\mu$  is the Haar measure on  $G$ .
An important property of the Haar integral is that
$$
\int_ {G} f (h g) d \mu (g) = \int_ {G} f (g) d \mu (g) \tag {22}
$$
for any measurable function  $f$  and  $h \in G$ . Therefore the average plan is  $G$ -invariant:
$$
\bar {\pi} (h \cdot x, h \cdot y) = \int_ {G} \pi^ {\prime} (h g \cdot x, h g \cdot y) d \mu (g) = \int_ {G} \pi^ {\prime} (g \cdot x, g \cdot y) d \mu (g) = \bar {\pi} (x, y).
$$
Next we compute the marginals
$$
\begin{array}{l} \int_ {X} \bar {\pi} (x, d y) = \int_ {X} \int_ {G} \pi^ {\prime} (g \cdot x, g \cdot d y) d \mu (g) = \int_ {G} \int_ {X} \pi^ {\prime} (g \cdot x, g \cdot d y) d \mu (g) \\ = \int_ {G} \nu_ {1} (g \cdot x) d \mu (g) = \int_ {G} \nu_ {1} (x) d \mu (g) = \nu_ {1} (x) \\ \end{array}
$$
and analogously  $\int_X\bar{\pi} (dx,y) = \nu_2(y)$ . Hence, the average plan is also a coupling,  $\bar{\pi}\in \Pi (\nu_{1},\nu_{2})$
The cost of the average plan is
$$
\begin{array}{l} \int_ {X ^ {2}} c (x, y) \bar {\pi} (d x, d y) \\ = \int_ {X ^ {2}} c (x, y) \int_ {G} \pi^ {\prime} (g \cdot d x, g \cdot d y) d \mu (g) \quad \left(\text {D e f i n i t i o n} \bar {\pi}\right) \\ = \int_ {G} \int_ {X ^ {2}} c (x, y) \pi^ {\prime} (g \cdot d x, g \cdot d y) d \mu (g) \quad (\text {F u b i n i ' s t h e o r e m}) \\ = \int_ {G} \int_ {X ^ {2}} c \left(g ^ {- 1} \cdot x, g ^ {- 1} \cdot y\right) \pi^ {\prime} (d x, d y) d \mu (g) \quad (\text {V o l u m e - p r e s e r v i n g s u b s t i t u t i o n}) \\ = \int_ {G} \int_ {X ^ {2}} c (x, y) \pi^ {\prime} (d x, d y) d \mu (g) \quad (\text {I s o m e t r i c g r o u p a c t i o n}) \\ = \int_ {X ^ {2}} c (x, y) \pi^ {\prime} (d x, d y), \\ \end{array}
$$
and hence the average plan is optimal.
Lemma 2. Let  $T \colon x \mapsto y$  be an  $G$ -equivariant OT map,  $x \in X$ , and  $g^* \coloneqq \operatorname*{arg\min}_{g \in G} c(x, g \cdot T(x))$ . Then  $g^* = \mathrm{id}$ .
Proof.  $T$  maps the orbit of  $x$ , denoted as  $G \cdot x = \{g \cdot x | g \in G\}$ , to the corresponding orbit of  $T(x)$  due to equivariance.
If there exists an  $x^{*} \in X$  such that  $g^{*} \neq \mathrm{id}$  then this property holds true for the entire orbit of  $x^{*}$ .
Now, let us define an  $G$ -equivariant map  $T' \colon x \mapsto y$  such that  $T'(G \cdot x^*) = g^* T(G \cdot x^*)$  and  $T'(x) = T(x)$  everywhere else. However, in this case,  $T$  would not be the OT map, as  $c(x^*, T'(x^*)) < c(x^*, T(x^*))$ . Therefore, it follows that  $g^* = \mathrm{id}$  holds for the  $G$ -equivariant OT map.
Discrete optimal transport Consider a finite number of iid samples from the two  $G$ -invariant distributions  $\nu_{1}, \nu_{2}$ , denoted as  $\{x_{i}\}_{i=1}^{n}$ ,  $\{y_{i}\}_{i=1}^{n}$ , respectively. Let further  $\pi(\nu_{1}, \nu_{2})$  be an  $G$ -invariant OT plan. The discrete optimal transport problem seeks the solution to
$$
\gamma^ {*} = \underset {\gamma \in \mathcal {P}} {\arg \min } \sum_ {i j} \gamma_ {i j} M _ {i j}, \tag {23}
$$
where  $M_{ij} = c(x_i, y_j)$  represents the cost matrix and  $\mathcal{P}$  is the set of probability matrices defined as  $\mathcal{P} = \{\gamma \in (\mathbb{R}^{+})^{n \times n} | \gamma \mathbf{1}_n = \mathbf{1}_n / n, \gamma^T \mathbf{1}_n = \mathbf{1}_n / n\}$ . The total transportation cost  $C$  is given by
$$
C = \min  _ {\gamma \in \mathcal {P}} \sum_ {i j} \gamma_ {i j} M _ {i j}. \tag {24}
$$
Approximately generating sample pairs from the OT plan  $\pi (\nu_{1},\nu_{2})$ , as required for flow matching (Section 3.4), can be achieved by utilizing the probabilities assigned to sample pairs by  $\gamma^{*}$ . Let further
$$
\tilde {\gamma} ^ {*} = \underset {\gamma \in \mathcal {P}} {\arg \min } \sum_ {i j} \gamma_ {i j} \tilde {M} _ {i j} \tag {25}
$$
denote the optimal transport matrix for the cost function  $\tilde{c}(x, y)$  and  $\tilde{C}$  the corresponding transportation cost. Sample pairs  $(x_i, y_j')$  generated according to  $\tilde{\gamma}^*$  are orientated along their orbits to have minimal cost, i.e.
$$
\left(x _ {i}, y _ {j} ^ {\prime}\right) = \left(x _ {i}, g ^ {*} y _ {j}\right), \quad g ^ {*} := \underset {g \in G} {\arg \min } c \left(x _ {i}, g \cdot y _ {j}\right). \tag {26}
$$
Note that the group action  $g$  does not change the probability under the marginals, as  $\nu_{1}$  and  $\nu_{2}$  are both  $G$ -invariant.
Lemma 3. Let  $\gamma^{*}$  and  $\tilde{\gamma}^{*}$  be defined as in Equation (23) and Equation (25), respectively. Then, sample pairs generated based on  $\tilde{\gamma}^{*}$  have a lower average cost than sample pairs generated according to  $\gamma^{*}$ .
Proof. The total transportation cost  $\tilde{C}$  is always less than or equal to  $C$  since the cost function  $\tilde{c}(x_i, y_j) \leq c(x_i, y_j)$  for all  $i$  and  $j$ , by definition. Note that as the number of samples approaches infinity, the inequality between  $\tilde{C}$  and  $C$  becomes an equality.
Note that in practice, samples generated wrt  $\tilde{\gamma}^*$  are a much better approximation, as shown in Section 6. Moreover, by construction, sample pairs generated according to  $\tilde{\gamma}^*$  satisfy  $\arg \min_{g\in G}c(x_i,g\cdot y_j) =$  id, as indicated in Equation (26). This property, which is also true for samples from the  $G$  -invariant OT plan according to Lemma 2, ensures that the generated pairs are correctly matched within their respective orbits. However, sample pairs generated according to  $\gamma^{*}$  do not generally possess this property.
We now combine our findings in the following Theorem.
Theorem 1. Let  $T \colon x \mapsto y$  be an OT map between  $G$ -invariant measures  $\nu_{1}$  and  $\nu_{2}$ , using the cost function  $c$ . Then
$$
c (x, T (x)) = \int_ {G} c (g \cdot x, g \cdot T (x)) d \mu (g) = \min  _ {g \in G} c (x, g \cdot y) \tag {27}
$$
Proof. We prove the parts of Theorem 1 in order.
1. Follows from Lemma 1 and the uniqueness of the OT plan for convex cost functions [68].
2. Follows from Lemma 2.
3. Follows directly from 2., Lemma 2, and
$$
\int_ {X ^ {2}} \min  _ {g \in G} c (x, g \cdot y) \pi (d x, d y) = \int_ {X ^ {2}} c (x, y) \pi (d x, d y). \tag {28}
$$
However, this OT map is not unique.
Hence, we can generate better sample pairs, approximating the OT plan  $\pi (\nu_{1},\nu_{2})$  , by using  $\tilde{c}$  as a cost function instead of  $c$  . Training pairs  $(x_{i},y_{j}^{\prime})$  are orientated along their orbits to have minimal cost (Equation (26)).
These findings provide the motivation for our proposed equivariant flow matching objective, as described in Section 4, as well as the utilization of  $G$ -equivariant normalizing flows, as discussed in Section 5.
B.2 Mean free update for identical particles
The update, described in Section 5, conserves the geometric center if all particles are of the same type.
Proof. The update of layer  $l$  does not alter the geometric center / center mass if all particles are of the same type as  $m_{ij}^{l} = m_{ji}^{l}$  and then
$$
\begin{array}{l} \bar {x} ^ {l + 1} = \sum_ {i} x _ {i} ^ {l + 1} = \sum_ {i} x _ {i} ^ {l} + \sum_ {i, j \neq i} \frac {\left(x _ {i} ^ {l} - x _ {j} ^ {l}\right)}{d _ {i j} + 1} \phi_ {d} (m _ {i j} ^ {l}) \\ = \bar {x} ^ {l} + \sum_ {i, j > i} \frac {\left(x _ {i} ^ {l} - x _ {j} ^ {l}\right) + \left(x _ {j} ^ {l} - x _ {i} ^ {l}\right)}{d _ {i j} + 1} \phi_ {d} (m _ {i j} ^ {l}) \\ = \bar {x} ^ {l} \\ \end{array}
$$
B.3 Selective reweighting for efficient computation of observables
As normalizing flows are exact likelihood models, we can reweight the push-forward distribution to the target distribution of interest. This allows the unbiased evaluation of expectation values of observables  $O$  through importance sampling
$$
\langle O \rangle_ {\mu} = \frac {\mathbb {E} _ {x _ {1} \sim \tilde {p} _ {1} (x _ {1})} [ w (x _ {1}) O (x _ {1}) ]}{\mathbb {E} _ {x _ {1} \sim \tilde {p} _ {1} (x _ {1})} [ w (x _ {1}) ]}, \quad w (x _ {1}) = \frac {\mu (x _ {1})}{\tilde {p} _ {1} (x _ {1})}. \tag {29}
$$
C Technical details
C.1 Code libraries
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 OpenMM[72], ASE [73], and xtb-python [56].
The code will be integrated in the bgflow library https://github.com/noegroup/bgflow.
C.2 Benchmark systems
For the DW4 and LJ13 system, we choose the same parameters as in [17, 18].
DW4 The energy  $U(x)$  for the DW4 system is given by
$$
U ^ {\mathrm {D W}} (x) = \frac {1}{2 \tau} \sum_ {i, j} a \left(d _ {i j} - d _ {0}\right) + b \left(d _ {i j} - d _ {0}\right) ^ {2} + c \left(d _ {i j} - d _ {0}\right) ^ {4}, \tag {30}
$$
where  $d_{ij}$  is the Euclidean distance between particle  $i$  and  $j$ , the parameters are chosen as  $a = 0, b = -4, c = 0.9, d_0 = 4$ , and  $\tau = 1$ , which is a dimensionless temperature factor.
Table 7: Model hyperparameters
Dataset	L	nhidden	Num. of parameters
DW4, LJ13	3	32	22468
LJ55	7	64	204936
alanine dipeptide	5	64	147599
LJ13 The energy  $U(x)$  for the LJ13 system is given by
$$
U ^ {\mathrm {L J}} (x) = \frac {\epsilon}{2 \tau} \left[ \sum_ {i, j} \left(\left(\frac {r _ {m}}{d _ {i j}}\right) ^ {1 2} - 2 \left(\frac {r _ {m}}{d _ {i j}}\right) ^ {6}\right) \right], \tag {31}
$$
with parameters  $r_m = 1$ ,  $\epsilon = 1$ , and  $\tau = 1$ .
LJ55 We choose the same parameters as for the LJ13 system.
Both the LJ13 and the LJ55 system were generated with MCMC with 1000 parallel chains, where each chain is run for 10000 steps after a long burn-in phase of 200000 steps starting from a random generated initial state.
Alanine dipeptide The alanine dipeptide training data at temperature  $T = 300\mathrm{K}$  set is generated through two steps: (i) Firstly, we perform an MD simulation, using the classical Amber ff99SBilden force-field, at  $300\mathrm{K}$  for implicit solvent for a duration of  $1\mathrm{ms}$  [20] using the openMM library.
(ii) Secondly, we relax  $10^{5}$  randomly selected states from the MD simulation for 100 fs each, using the semi-empirical  $GFN2 - xTB$  force-field [56] and the ASE library [73] with a friction constant of 0.5 a.u.
Both simulations use a time step of 1 fs. We create a test set in the same way.
Alanine dipeptide umbrella sampling To accurately estimate the free energy difference, we performed five umbrella sampling simulations along the  $\varphi$  dihedral angle using the semi-empirical  $GFN2 - xTB$  force-field. The simulations were carried out at 25 equispaced  $\varphi$  angles. For each angle, the simulation was initiated by relaxing the initial state for 2000 steps with a high friction value of 0.5 a.u. Subsequently, molecular dynamics (MD) simulations were performed for a total of  $10^{6}$  steps using the  $GFN2 - xTB$  force field, with a friction constant of 0.02 a.u. The system states were recorded every 1000 steps during the second half of the simulation. A timestep of 1 fs was utilized for the simulations.
The datasets are available at https://osf.io/srqg7/?view_only=28deeba0845546fb96d1b2f355db0da5.
Depending on the dataset, different model sizes were used, as reported in Table 7. All neural networks  $\phi_{\alpha}$  have one hidden layer with  $n_{\mathrm{hidden}}$  neurons and  $SiLU$  activation functions. The embedding  $a_{i}$  is given by a single linear layer with  $n_{\mathrm{hidden}}$  neurons.
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.
C.4 Parallel OT batch generation
As Equation (16) needs to be evaluated for each possible sample pair in a given batch, the computation cost of equivariant OT flow matching is quite high. However, a simple way to speed up the equivariant OT training is to generate the batches beforehand or in parallel to the training process. This process is highly parallelizable and can be performed on CPUs, which are in practice usually more available
Table 8: Training schedules
Training type	Batch size	Learning rate	Epochs	Training time
Likelhood [18]	256	5e-4	20	3h
OT flow matching	256	5e-4/5e-5	200/200	0.5h
Equivariant OT flow matching	256	5e-4/5e-5	200/200	0.5h
LJ13
Likelhood [18]	64	5e-4	5	13h
OT flow matching	256	5e-4/5e-5	1000/1000	3h
Equivariant OT flow matching	256	5e-4/5e-5	1000/1000	3h
LJ55
OT flow matching	256	5e-4/5e-5	600/400	17h
Equivariant OT flow matching	256	5e-4/5e-5	600/400	17h
Alanine dipeptide
OT flow matching	256	5e-4/5e-5	1000/1000	6.5h
Equivariant OT flow matching	32	5e-4/5e-5	200/200	25h
(Batch reordering during training)
Table 9: Wall-clock time to reorder a single batch for the two OT flow matching methods. The times for the DW4 and LJ13 system are below 0.01 seconds for OT flow matching and are therefore not reported.
Training type	Batch size	Wall-clock time
DW4
Equivariant OT flow matching	256	3.6s
LJ13
Equivariant OT flow matching	256	4.5s
LJ55
OT flow matching	256	0.01s
Equivariant OT flow matching	256	20.2s
Alanine dipeptide
OT flow matching	256	0.01s
Equivariant OT flow matching	256	22.4s
Equivariant OT flow matching	32	0.4s
than GPUs and also in higher numbers. This also allows for larger batch sizes for the equivariant OT model and comes at little additional cost. Hence, scaling equivariant flow matching to even larger systems should not be an issue. The wall-clock time for the reordering of a single batch on a single CPU is reported in Table 9.
C.5 Biasing target samples
In the case of alanine dipeptide, the transition between negative and positive  $\varphi$  dihedral angles is the slowest process, as depicted in Figure 3d. Since the positive  $\varphi$  state is less probable, we can introduce a bias in our training data to achieve a nearly equal density in both states. This approach aids in obtaining a more precise estimation of the free energy. To ensure a smooth density function, we incorporate weights based on the von Mises distribution  $f_{\mathrm{vM}}$ . The weights  $\omega$  are computed along the  $\varphi$  dihedral angle as
$$
\omega (\varphi) = 1 5 0 \cdot f _ {\mathrm {v M}} (\varphi | \mu = 1, \kappa = 1 0) + 1. \tag {32}
$$
We then draw training samples based on the weighted distribution.
C.6 Effective samples sizes
The effective sample sizes are computed with Kish's equation [74]. We use  $5 \times 10^{5}$  samples for the DW4 and LJ13 system,  $2 \times 10^{5}$  for alanine dipeptide, and  $1 \times 10^{5}$  for the LJ55 system per model.
C.7 Error bars
Error bars in all plots are given by one standard deviation, averaged over 3 runs, if not otherwise indicated. The same applies for all errors reported in the tables. Exceptions are the Mean batch transport cost plots, where we average over 10 batches and the integration error plots, where we average over 100 samples for each number of steps.
C.8 Computing infrastructure
All experiments for the DW4 and LJ13 system were conducted on a GeForce GTX 1080 Ti with 12 GB RAM. The training for alanine dipeptide and the LJ55 system were conducted on a GeForce RTX 3090 with 24 GB RAM. Inference was performed on NVIDIA A100 GPUs with 80GB RAM for alanine dipeptide and the LJ55 system.