FreeFlow: Latent Flow Matching for Free Energy Difference Estimation

Thank you for the detailed review!

---

“The paper [...] doesn't develop an entirely new method, but rather combines existing methodologies into a framework for free energy difference estimation.”

The technical novelty of our work lies not in using a flow to bridge the distributions, but rather in using a latent flow that employs a trans-dimensional change of variables formula to account for dimensionality changes between systems with different numbers of atoms, eliminating the need for dummy atoms and enabling efficient free energy difference estimation in molecular systems. We see our work as an initial and early, yet important step towards enabling free energy estimations through a latent space framework.

“I wonder if there could be benefits to learning a "representation learning" autoencoder in the classical sense—one that can generalize and thus be used for all densities rather than requiring separate autoencoders for each molecular system.”

Thanks for sharing this idea. Exploring a generalizable autoencoder could be an exciting direction for future research. An autoencoder accurate across a set of molecules would save the time to train separate autoencoders and therefore speed up the free energy difference calculations. Nevertheless, this could be challenging as such an autoencoder would need to learn a more complex mapping, requiring a large-scale molecular dataset and costly training. Our current approach focuses on overfitting to maximize accuracy and efficiency for specific pairs of molecules.

“In Figure 4(a), the true target and the estimated target distributions differ significantly. In fact, the true target seems to coincide with the source. Could this be a labeling mistake?”

Thanks for pointing this out. Upon closer examination, we found that this was not a labeling mistake. The observed phenomenon arises because points sampled from different same-dimensional isotropic Gaussian distributions inherently exhibit the same energy distribution when evaluated using their respective energy functions. As a result, our plots were inappropriate to measure the quality of the map. To address this issue, we decided to remove this plot.

Theory (or implications) of the trans-dimensional change of variable

It is indeed true that the map is not lossless, and the change of variables formula is an approximation (for similar work see [1]). This approximation error is minimized to the extent that our data manifold (points we compute the change of variables for) overlaps with the decoder manifold (image of the decoder). We try to minimize this approximation error by overfitting the autoencoders to the data and our results show that this is a promising approach to evaluating free energy differences.

---

[1] Gemici, M. C., Rezende, D., & Mohamed, S. (2016). Normalizing Flows on Riemannian Manifolds (No. arXiv:1611.02304). arXiv. https://doi.org/10.48550/arXiv.1611.02304

Thank you for following up on the discussion. We are happy to provide further clarifications for the questions raised.

"The Jacobian is $J_f(x) = (\frac{\partial f}{\partial x_1}, \frac{\partial f}{\partial x_2})$, so the volume form is simply the square root (the determinant being irrelevant since this is one-dimensional) of $\sum_i \left( \frac{\partial f}{\partial x_i} \right)^2$, correct?"

Yes. There is no Jacobian determinant and the square root of the gradient vector is not the volume change. The change in density can again be computed using Equation 14.

"Given these assumptions, is the equation in Figure 3 valid for all x? This seems counterintuitive, as for generic and $\rho_A$ and $f$, the mapping won't be lossless. For instance, the equation in the top part of Figure 3 holds if is a bijection. "

Yes, the equation in Figure 3 (top) only holds if $f$ is a bijection, and we had intended to show this in the figure.

Since the encoder cannot be a bijection and will be lossy, the change of density is also an approximation. Intuitively, an encoding $z$ collects the probability mass of all $x$ such that $f(x) = z$. This means that we can obtain a meaningful density function over the latent space just using the encoder.

However, if we want to obtain an ambient density from a latent density, the best we can get is a density defined over the fibers $\mathcal{F}(z)$ of the decoder (Equation 13, note that input to $p_X$ is $\mathcal{F}(z)$), where $\mathcal{F}(z) := \{x : f(x) = z\}$.

The decoder change of variables is then an approximation as we are not able to distinguish the densities of the points within the same fiber. This is not an issue if the decoder manifold (decoder’s image) overlaps perfectly with the data manifold, i.e. the autoencoder reconstructs the data samples perfectly, because then we could have a density function defined for each point we will be evaluating it on.

Of course this is hard to achieve in practice. There will then be data points outside the decoder manifold, and the density function we get will not be able to distinguish between them.

"For a surjective $f$, does the reverse equation hold true?"

In the case of a surjective mapping, the fibers $F(z) = \{x : f(x) = z\}$ will have multiple elements. In this scenario, the density $p_B(f(x))$ in the latent space accounts for the aggregated probability mass of all points $x$ that map to the same $z$. The equation will be exact even though $f$ is lossy, since each latent $z$ will aggregate the probabality mass in $F(z)$.

Thank you for your feedback!

---

Accuracy of the proposed method

While the absolute error of our method is high, we would like to highlight that our work is an early attempt at estimating free energy differences using ML approaches, to alleviate the shortcomings of the classical methods. Our main contribution is that using a latent flow to address the problem of systems with different dimensionalities is a promising direction as an alternative to using dummy atoms. There is certainly more work needed to make the approach applicable and useful in practice.

Additionally, we model our systems using an implicit solvent rather than modeling the solvent atoms explicitly as that would lead to systems with an extremely high dimensionality and thus a harder learning problem. Using an implicit solvent introduces an additional source of error, and that is one the reasons why we focus on correlations rather than absolute errors in our evaluation. Correlations are still useful as free energy differences are often used to compare binding affinities which rely on the relative values between two candidates.

“The cost of the proposed method is relatively large. The training data includes molecular dynamics simulation of the related systems. Thus, when applying the proposed method to new systems, one should perform additional molecule dynamics simulation to collect training data.”

We have clarified in our Introduction that FreeFlow indeed still requires simulation but the amount required is an order of magnitude lower than for alchemical FEP for which one needs to perform simulations for multiple intermediate states connecting the two states.

Thank you for the review! We hope we are able to address the questions raised. Our response is in two parts due to the character limit:

---

“The authors argue that FreeFlow learns more expressive maps between molecular distributions compared to previous normalizing flows solutions. It is not clear that this claim is supported in the experiments.”

We have updated our explanation in the Introduction to reflect that by “expressive” we refer mainly to the less restricted design space enabled by flow matching compared to previous normalizing flows which depended on architectures with tractable Jacobians and fast invertibility. Flow matching and continuous normalizing flows on the other hand can be used with any typical neural network architecture.

“It is not clear how or if avoiding the use of non-physical modifications leads to improved free energy difference estimation. How does FreeFlow compare to methods that use such non-physical modifications?”

Free Energy Perturbation (FEP) which we compare with uses such “non-physical” alchemical intermediates and dummy atoms. A main benefit of our method is that it can be faster because it does not require the intermediate states and because it transforms the distributions in a lower dimensional space.

Comparisons to other baselines such as those described in Figure 2

We have compared FreeFlow with a normalizing flow consisting of coupling layers as in TFEP [2] on our alanine dipeptide task and observed significantly larger estimation errors with estimates on the order of 1,000 kj/mol compared to the reference value of 20.87 kj/mol, as well as unstable likelihood-based training due to the force-field-based potential energies.

In our paper, we compare FreeFlow against alchemical FEP in our main results due to the practical relevance of the method and because there exist works that provide reference ΔF values for the solvent leg (between two unbound ligands) of the thermodynamic cycle.

“Lines 232-233: Could the authors provide some explanation/justification an optimal transport map between ρA and ρB is needed in this setting? What happens if you were to assume independent marginals?”

We thank you for your suggestion and now clarify in Section 3.1 that the primary benefit of approximating an OT map is that the vector field learned with mini-batch OT couplings will have straighter trajectories compared to one learned with independent couplings, and this will reduce the integration error. Specifically for free energy estimation, it has been shown using lower-dimensional systems in [1] that approximating an OT map between ρA and ρB leads to paths with lower free energy compared to linear interpolation, improving the convergence of the estimator.

Auto-Encoders: Validation of over-fitting and reasonable representations & decision on hyper-parameters

We have explained more clearly in Section 3.2 that we train separate autoencoders for each molecule until the mean-squared reconstruction loss on training data converges, since we will be using the training data to estimate ΔF. Additionally for the hyperparameters, we performed validation runs on alanine dipeptide and kept the best-performing ones.

“How do you think the results would change if a different architecture is used, for examples a graph neural network or transformer, which have become de facto architectures for learning molecular representations?”

Thank you for pointing out the significance of comparisons against graph neural networks. We have now added Figure 7 in Appendix A showing that an MLP achieves an almost an order of magnitude lower reconstruction loss than an EGNN with a similar number of parameters. Furthermore, the training time for the MLP is significantly lower (2 min.) compared to the EGNN (17 min.).

We believe this to be because we can overfit on the samples and we have sufficient data, not requiring us to have a more complex geometry-aware model. The EGNN is a more complex function and there’s some optimization error introduced by using it. Furthermore, using a GNN significantly slows down the training procedure (over eight times slower than the MLP) which is detrimental to our approach since our main goal is to speed-up drug screening.

Figure 4.a: Estimated target distribution is far from true target (equidimensional case)

We agree that the figures are not informative and removed them.

Firstly, the first plot showed almost the same energy distribution for both source and target which is correct but not useful for our analysis. This is because points sampled from different same-dimensional isotropic Gaussian distributions exhibit the same energy distribution when evaluated using their respective energy functions. It would be better to plot the points from the source Gaussian under the target Gaussian’s energy function as we do for the mapped and target samples. However, this is not possible for the trans-dimensional case because the source points are of different dimensionality.

We do not have a good explanation for why the energy of the mapped samples in the equi-dimensional case is matching the distribution of the target worse than in the trans-dimensional case. Because of the above and because we are mainly interested in the final free energy difference accuracy, we decided to remove these plots.

“In Figure 5, what is the distribution pairwise distance D(⋅,⋅)?”

We have made it clearer in the Section 4.2 that by D(A,B) we denote the distribution of distances for all pairs (a,b) from A and B; i.e. the set {d(a, b) : a \in A, b \in B}.

“Lines 419-420: How do you observe the distances between B and the mapped sampled M(B) ? In Figure 5, only D(A,B), D(B,B), and D(M(A),B) are reported.”

The “M(B)” in lines 419-420 is indeed a typo and we have now fixed this. Since we map the samples from A with our flow, we only observe D(M(A), B) and not D(M(B), B).

---

[1] Decherchi, S., and A. Cavalli. "Optimal Transport for Free Energy Estimation." The journal of physical chemistry letters 14.6 (2023): 1618-1625.

[2] Wirnsberger, Peter, et al. "Targeted free energy estimation via learned mappings." The Journal of Chemical Physics 153.14 (2020).

We thank the reviewer for the valuable comments and questions!

---

Toy system code to test the method

We added the link to our paper pointing to an example application of FreeFlow to Gaussian distributions https://anonymous.4open.science/r/freeflow-example-26D4

Comparisons with other methods

The primary used methods for FED calculations involve the BAR estimate and Free Energy Perturbation (FEP) which are time-costly because they rely on intermediate transformations between the source and target state for accurate estimates. In our work we compare our estimates against values derived through the common FEP method. Flow-based FEP solutions like FreeFlow optimize the process of FED calculations by reducing the reliance on extensive intermediate transformations which we model through a CNF.

We also trained a normalizing flow consisting of coupling layers as in TFEP [1] for our alanine dipeptide task and obtained significantly less accurate estimates for ΔF. The estimated values were on the order of 1,000 kj/mol compared to the reference value of 20.87 kj/mol. We also observed that the training was highly unstable due to the likelihood training using force-field-based potential energies.

“In Figure 4, the mean of equidimensional gaussians seems to converge to the true mean, but not for the trans dimensional one.”

We have added to the discussion in Section 4.1 that free energy difference estimates obtained from our trans-dimensional change of variables only hold to an approximation, since the encoders are lossy compressors. More specifically, unless we can obtain a zero-loss autoencoder that perfectly recovers the data points, there will be points outside its decoder’s image (the decoder manifold) and particularly for those points the change of variables formula will not hold exactly.

“In Figure 5, the D(B,B) and D(M(A),B) do not completely overlap”

We now clarify in Section 4.2 that the trans-dimensional change of variables formula being an approximation is likely also the reason behind this slight mismatch between D(M(A), B) and D(B, B). However, while we try to learn a map that accurately samples the target distribution, we (and TFEP more generally) can tolerate slightly inaccurate maps since the ultimate goal is to estimate ΔF via importance sampling. In fact, the importance sampling reweighting holds for any map (provided the energy functions are defined everywhere) and results in an unbiased estimator, but our goal is to learn a reasonably accurate map that leads to fast convergence of the ΔF estimate.

---

[1] Wirnsberger, Peter, et al. "Targeted free energy estimation via learned mappings." The Journal of Chemical Physics 153.14 (2020).