FEAT: Free energy Estimators with Adaptive Transport

Thank you for your review. We have addressed your concerns with improved baseline comparisons, detailed statistical analysis, experimental details, and code availability. We are happy to answer any further questions you might have. Should you find our reply satisfactory, we kindly invite you to reconsider your rating.

Responses to Weaknesses

1. In Table 1 results for Neural TI perform poorly: although the text offers some justification, something should be done to make these results a fair comparison, otherwise these results do not have any significant meaning.

Thank you for your suggestion. In our experiment, we tried to keep the network architecture as consistent as possible across different approaches. However, we fully agree that this is not ideal for neural TI, which benefits from task-specific architectural design to ensure boundary conditions and also mitigate the "blindless issue" of score matching losses (e.g., score is insensitive to error in mode weights when modes are separate).

Following your suggestion, we re-evaluated neural TI on GMM-40 and LJ using a better-designed network architecture in the following.

Updated network design for neural TI:

• for GMM, we set the energy network to be

$U(x_t, t) = a_t U_A(x_t) + b_t U_B(x_t) + c_t U_\theta (x_t, t)$

where $a_t = \exp[f_\theta(t) - f_\theta(0)] \cdot (1-t)$, $b_t = \exp[g_\theta(1) - g_\theta(t)] \cdot t$ and $c_t = \exp(h_\theta(t)) \cdot t \cdot (1-t)$, where $f_\theta, g_\theta, h_\theta, U_\theta$ are neural nets and by which we can exactly ensure the boundary condition.

• for LJ, $U_A$ and $U_B$ are more sensitive to noisy $x_t$. Therefore, inspired by [1], we set the energy network to be

$U(x_t, t) = b_t \cdot U_\text{LJ}(x_t, r_t, a_t) + c_t \cdot U_\theta(x_t, t)$

where $r_t$ is the radius parameter in LJ, ranging from 0 to 1. $a_t = \exp(\alpha_\theta(t) )\cdot t \cdot (1-t)$ is a smooth parameter inspired by [1]. $b_t = 1 - \exp(\beta_\theta(t))\cdot t\cdot (1-t)$ and $c_t = \exp(\gamma_\theta(t)) \cdot t \cdot (1-t)$ are scalar to ensure boundary conditions.

Updated neural TI results:

With these updated architectures, we obtained the following updated results for neural TI:

• GMM-40: 0.366 (reference value: 0)
• LJ-55: 235.894 (reference value: 234.77)
• LJ-79: 354.8852 (vs. reference 357.43)

These results show that neural TI can achieve competitive performance when supported by carefully designed network architectures. We will provide updated results on other targets and provide error bars, and we will ensure a fairer discussion on neural TI in our final version.

[1] Máté, B., Fleuret, F. and Bereau, T., 2025. Solvation free energies from neural thermodynamic integration. The Journal of Chemical Physics, 162(12).

2. For certain systems (specifically ALDP), the standard deviations are so large that comparing only the means does not ensure these estimates are trustworthy or reproducible.

We completely agree that both mean and variance are crucial for assessing method reliability. We have conducted extensive additional experiments to address this concern.

We performed 12 runs (6 different model seeds × 2 estimates per model) and provide more statistics:

Their 25, 50 and 75-quantiles are:

• FEAT: 29.367, 29.417, 29.458
• Target FEP: 29.347, 29.422, 29.541

We can observe that FEAT gives more stable results when taking standard deviation and quantiles into account.

3. Code availability

We totally understand your concern. Due to this year’s policy, we are not allowed to directly provide you with links pointing to our code; therefore, we have sent our code link to AC.

Additionally, we will make our code publicly available in the next version of our paper.

Responses to Questions

Q1: Network architecture specifications in Table 3

Sorry for the confusion. "5, 400" refers to 5 layers of neural networks with 400 hidden dimensions; similarly, "4, 64" refers to 4 layers of neural networks with 64 hidden dimensions.

Our method does not necessarily require an H100 to train, smaller GPUs also suffice. We further report the GPU usage of our targets:

Q2: Train/test splitting and validation procedures

We address this from two perspectives:

a) Train/test split: For a given pair of states A & B, we aim to fit a transport between A and B as much as possible, and in theory, we assume that we can sample sufficient configurations in these states. From this perspective, our approach is more akin to stochastic approximation methods from Robbins-Monro [2] and therefore has minimal overfitting issues.

On the other hand, when the data set is provided beforehand and kept fixed, it may not be representative enough, and using the same set for training and testing may introduce bias. To this end, we provide an 8:2 train/test split result on ALDP. The result is not significantly better than that reported in our manuscript, indicating that the dataset we used is sufficient.

b) Hyperparameter tuning: Another question is whether a train/validation split is needed for hyperparameter tuning. The answer is No: in our approach, we tune all hyperparameters based on transport quality by monitoring the energy of transported training samples and the gap between forward and backwards estimators.

These validation metrics eliminate the need for separate validation sets.

[2] Robbins, H. and Monro, S., 1951. A stochastic approximation method. The annals of mathematical statistics, pp.400-407.

Thank you for your detailed and insightful review. We are delighted that you found our work exceptionally well-organized, original, effective, and significant. We now address your concerns with new experiments on larger molecular systems and clarifications on experimental details:

Responses to Weaknesses and Questions

Q1 & Weakness: Experimental results on more complex molecules

Thank you for your great suggestions. Following your recommendation, we now provide results on significantly larger systems:

Alanine Tetrapeptide Results:

• System size: 43 atoms (129 dimensions)
• FEAT estimate: 65.51 ± 1.51 kcal/mol
• Reference value (MBAR with 20 intermediate distributions): 64.09 kcal/mol

Chignolin Results:

• System size: 175 atoms (525 dimensions)
• For this larger system, we found that it is more effective to decompose the estimation task: instead of directly learning a transport map between the vacuum and solvent states, we estimate the free energy using FEAT in two stages—from a Gaussian to vacuum, and from a Gaussian to solvent—and then take the difference between the two estimates.
• FEAT estimate: 190.784 +- 0.414 kcal/mol (due to time limit, we did not train 3 times, and this std is obtained by doing 3 inferences with one model.)
• Reference value (MBAR with 20 intermediate distributions): 190.887 kcal/mol

These results demonstrate FEAT's applicability to larger, more realistic systems. We will update these results in our camera-ready version.

Q2: Data pair in stochastic interpolant training

This is an excellent question. The FEAT framework is general and supports multiple pairing strategies, for example:

• OT: Using optimal transport (OT) plans for structured correspondence
• Random: Independent sampling from each endpoint distribution

In our experiments on simpler systems such as GMM-40 and LJ-55, we use randomly selected pairs. For more complex targets, we sample batches randomly but apply mini-batch OT to determine better pairings, which helps to stabilize training and improve network fitting. We found that omitting the OT step leads to slower convergence and reduced performance, as random pairing can lead to less "straight" vector fields, making training more challenging.

In table 3, we have a hyperparameter called "OT pair" to reflect if we use this approach in our training. We will provide a brief discription and discussion in our camera-ready version.

Q3: Figure 2 equation reference typo

Thank you for catching it. We will fix in our camera-ready version.

Thank you for your constructive review. We are pleased that you found our work theoretically sound, well-organized, effective, and novel. We have addressed your concerns with further experiments and clarifications. We are committed to addressing any further concerns you may have. Should you find our reply satisfactory, we kindly invite you to reconsider your rating.

Responses to Weaknesses

W1: Comparison with DeepBAR as an additional baseline

Thank you for this excellent suggestion! However, FEAT is orthogonal to DeepBAR rather than directly comparable.

DeepBAR estimates free energy differences by first estimating absolute free energies for each system using generative models (e.g., NF or CNF), then taking their difference. FEAT can actually improve DeepBAR by providing better transport learning for the absolute free energies.

We provide results on GMM and ALDP-S using DeepBAR with CNF and FEAT below:

DeepBAR-FEAT outperforms DeepBAR-CNF, especially on ALDP-S. Interestingly, on GMM-100, both achieve better performances compared to standard Target FEP w. FM and FEAT, likely because learning transport from Gaussian to GMM is easier than directly learning between two complex GMMs.

We will include this discussion and comparison in our camera-ready version.

W2: Additional baselines from related works

Since DeepBAR did not release their datasets, we evaluated on a comparable system: Chignolin (175 atoms, 525 dimensions) in vacuum vs. implicit solvent:

• Reference value (MBAR): 190.887 kcal/mol
• DeepBAR-FEAT result: 190.784 +- 0.414 kcal/mol (due to time limit, we did not train 3 times, and this std is obtained by doing 3 inferences with one model.)

This showcased our approach's scalability.

Responses to Questions

Q1: Can FEAT be generalized to non-Boltzmann sampling?

Yes. FEAT can handle any distribution admitting a differentiable density $p(x)$. The Boltzmann form is not required: our framework applies to any probability distribution satisfying these mild regularity conditions.

Q2: When deriving Eq. 36, do you assume perfect control?

No. Equation 36 relies entirely on importance sampling (lines 619-624) and does not require perfect control. This is a key advantage of FEAT: intermediate samples need not remain at equilibrium.

However, we do need the control to be "reasonable" with sufficient overlap between forward and backward processes. Poor control can increase estimation error, but this is easily detected through:

• Quality of transported samples
• Gap between ELBO and EUBO
• Discrepancy between forward and backward estimators

Q3: How do you estimate $\dot I_t$ for learning $v_t$?

This appears to be a misunderstanding. The derivative of the interpolant $\dot I_t$ can be calculated exactly from our Eq.17: it does not need to be estimated or learned. Instead, we use it to learn $v_t(x) = \mathbb{E}[\dot I_t| I_t = x]$ (Eq.18; Eq. 2.10 in [1] ), which is exactly the minimizer of the loss in Eq.19 and is not noisy.

[1] Michael S. Albergo, et al. Stochastic Interpolants: A Unifying Framework for Flows and Diffusions

Q4: Comparison between averaged forward-backward estimators vs. min-variance estimator

This is a great suggestion! In fact, the minimum-variance estimator (blue curve) is theoretically guaranteed to perform no worse than averaged forward-backward estimators (red curves), as it's derived to minimize estimation error.

To see this, we provide an empirical verification on ALDP-S:

• After 2000 iters: Blue = 29.30 +- 0.36; Avg red = 28.86 +- 1.19
• After 5000 iters: Blue = 29.54 +- 0.16; Avg red = 29.56 +- 0.73

Q6: Can FEAT utilize intermediate distributions if available?

Yes. We can learn multiple FEATs in sequence: Given intermediate distributions $U_2...U_{N-1}$ between $A$ (energy $U_1$) and $B$ (energy $U_N$), we learn FEAT between each neighboring pair $U_i$ and $U_{i+1}$. The total $\Delta F$ is the sum of individual estimates. This approach may reduce individual network burden compared to learning a single $A\to B$ transport.

Q7: How do you control overfitting?

We answer this from two perspectives:

a) Train/test split: For a given pair of state A & B, we aim to fit a transport between A and B as much as possible, and in theory we assume that we can sample sufficient configurations in both states. From this perspective, our approach is more akin to stochastic approximation methods from Robbins-Monro [2] and therefore has minimal overfitting issues.

On the other hand, when the data set is provided beforehand and kept fixed, it may not be representative enough, and using the same set for training and testing may introduce bias. To this end, we provide an 8:2 train/test split result on ALDP. The result is not significantly better than that reported in our manuscript, indicating that the dataset we used is sufficient.

b) Hyperparam tuning: Another question is whether a train/valid split is needed for hyperparam tuning. The answer is No: in our approach, we tune all hyperparameters based on transport quality by monitoring the energy of transported training samples, or the gap between forward and backward estimators. These validation metrics eliminate the need for separate validation sets.

[2] Robbins, H. and Monro, S., 1951. A stochastic approximation method. The annals of mathematical statistics, pp.400-407.

Q8: Reference computation times for Tab 2 to highlight the cost of sampling from intermediate distributions

For ALDP, generating samples for each target requires ~1 day on our machine. Collecting all the targets for MBAR can hence take between 1-10 days, depending on whether the simulations are run in parallel.

In contrast, our training process takes only 1-2 hours, which is significantly faster than sampling from all intermediate densities. We will include these rough numbers and the corresponding discussion in our camera-ready version.

Q10 & Q11: Robustness w.r.t. discretization steps and sample size

Following your suggestion, we conducted robustness analysis on GMM:

FEAT

Target FEP

These results highlight FEAT's greater robustness toward fewer steps/samples, compared to Target FEP w flow matching. The results obtained with fewer steps (larger discretisation error) also reflect our discussion in lines 238-239.

Q12: Target score matching with different time intervals

Following your suggestions, we evaluate FEAT with $t \sim U(0, 0.25)$ and $U(0.75, 1)$ in TSM on ALDP-S. It achieves $29.37 \pm 0.02$ across 3 independent runs. Compared to the choice with $U(0, 0.5)$ and $U(0.5, 1)$ ($29.38\pm 0.04$), we do not see obvious difference.

This is because TSM helps the score learn better at boundary. As long as we sample $t$ close to the boundary, TSM can provide useful signal. Also, different from neural TI that fully relies on the learned energy, FEAT rely on non-equilibrium transports with both the score and the vector field, leading to a higher robustness towards the TSM choice.

Q13: Can FEAT be applied to different systems without retraining?

Yes. If we train the network using datasets from multiple targets and condition the network on the target parameters, the resulting model becomes transferable, allowing us to apply FEAT to similar systems without retraining.

We showcase this on GMM-40 with different scaling factor: concretely, we scale the state B with scaling $0.5, 0.7, 0.9, 1.1, 1.3, 1.5$ when training. Our network will also take the scaling as an input. After training, we evaluate $\Delta F$ on unseen scaling factors and report the error and std across 3 runs in the following table:

These results show that the transferable FEAT model achieves good accuracy across a range of unseen targets. Also, as expected, interpolation yields better performance than extrapolation.

Q14: Performance without data preprocessing

Omitting the data processing step significantly complicates training, resulting in larger bias and variance. Following your suggestion, we provide an ablation study on ALDP-S. Without pre-processing:

• FEAT: $37.40 \pm 0.38$
• Target FEP w FM: $44.2 \pm 4.0$.

Neither method attains a reasonable value, while FEAT exhibits a slightly better performance.

However, we emphasize that this is not a weakness of our proposed framework. Rather, it reflects a general challenge in generative modeling. We also expect that using higher-capacity architectures (e.g., transformers) could help mitigate this issue.

Formatting and typos (Weakness 3/4 and Q5/9):

Thank you for pointing these typos out! We will fix them in our camera-ready version.

Broader Impact is not counted to the page limit per instruction, so our main text is within 9 pages.

Thank you for your continued participation and thoughtful feedback.

I see that for the Chignolin system, you report results only for the two-stage FEAT. What are the results for the one-stage FEAT?

We found learning a map from one complex system to another system is harder than mapping from a simple system to only one complex system, This requires longer time to converge and also potentially better network with higher capacity. In our attempts, we obtained 314.2 +- 26.9 ( = 186.9 +- 16.0 kcal/mol vs reference value 190.8 kcal/mol) using one-stage FEAT (obtained by 3 runs with one model). It is not crazily off, but it is worse than two-stage FEAT. We will include a discussion in our final version.

You mention that performance does not change much when using target score matching with different time intervals. Could this robustness to the choice of time intervals be a property unique to relatively simple systems? For example, does this robustness hold for the Chignolin system?

• In fact, we did not observe a performance difference in our case on both GMM and ALDP, which ablates two systems with fairly different complexities.
• Whilst we agree with the reviewer that more ablation of the Uniform interval is indeed interesting, we would like to iterate that the used settings, which were inherited from previous works, are not the primary focus/contribution of our work.
• The Chignolin experiment is expensive to run, and we will not be able to thoroughly test this during the discussion period. For completeness, we are happy to attempt an ablation of this in the final version, as it may be useful for the wider community. This is, however, tangential to our contribution.

Thank you again for your time, and we are happy to address any further concerns you may have.

Thank you for your valuable review. We are pleased that you found the method in our paper effective and novel. We have addressed your concerns with new experiments on larger systems and detailed clarifications about our experimental setup. We are committed to addressing any remaining concerns you may have. Should you find our reply satisfactory, we kindly invite you to reconsider your rating.

Responses to Questions

Q1: Could the authors include more real-world case studies (beyond toy systems) for free energy prediction?

Thank you for this important suggestion. We provide FEAT's results on two larger systems: Alanine Tetrapeptide (polypeptide with 4 amino acids) and Chignolin (protein with 10 amino acids).

Alanine Tetrapeptide Results:

• System size: 43 atoms (129 dimensions)
• FEAT estimate: 65.51 ± 1.51 kcal/mol
• Reference value (MBAR with 20 intermediate distributions): 64.09 kcal/mol

Chignolin Results:

• System size: 175 atoms (525 dimensions)
• For this larger system, we found that it is more effective to decompose the estimation task: instead of directly learning a transport map between the vacuum and solvent states, we estimate the free energy using FEAT in two stages—from a Gaussian to vacuum, and from a Gaussian to solvent—and then take the difference between the two estimates.
• FEAT estimate: 190.784 +- 0.414 kcal/mol (due to time limit, we did not train 3 times, and this std is obtained by doing 3 inferences with one model.)
• Reference value (MBAR with 20 intermediate distributions): 190.887 kcal/mol

These results demonstrate FEAT's applicability to larger, more realistic systems. We will update these results in our camera-ready version.

Q2: Could the authors provide details on the training sets used for the alanine dipeptide (2-state), Lennard-Jones (LJ), and ALDP-S systems?

We included comprehensive details about each system in Appendix G.1, including state definitions and sample collection procedures. Could you elaborate on which specific details are missing? We will ensure any additional details you suggest are included in our camera-ready version.

Clarification on System Differences:

Two-state alanine dipeptide (ALDP-T):

• We split the Boltzmann distribution for alanine dipeptide in vacuum into two metastable state distributions
• State A: blue shaded region in Fig 2(a)
• State B: pink shaded region in Fig 2(a)
• Physical meaning: Free energy difference when alanine dipeptide transitions between metastable conformational states

ALDP-S (Solvation):

• State A: entire Boltzmann distribution for alanine dipeptide in vacuum
• State B: entire Boltzmann distribution for alanine dipeptide in implicit solvent
• Physical meaning: Solvation free energy difference between vacuum and implicit solvent environments

Q3: What are the limitations regarding the applicability of the FEAT method to explicit solvent systems or larger (e.g., biomolecular) systems?

Explicit solvent has been a long-standing challenge for simulating biomolecules with machine learning forcefields as explicit solvent molecules are often too large to fit into GPU memory. It is a challenge of independent interest and is orthongoal to what we proposed in our framework. We hence only consider the implicit solvent in our study.