BoltzNCE: Learning likelihoods for Boltzmann Generation with Stochastic Interpolants and Noise Contrastive Estimation

We thank the reviewer for their time and their helpful comments. We are glad that the reviewer found our paper novel and well written. We would also like to mention that in addition to results on alanine dipeptide, we now include new experiments demonstrating transferability For transferability we train on a set of dipeptides and infer on dipeptide systems not present in the training set. With a bit of fine-tuning of the EBM on the unseen dipeptide, we obtain accurate free energy estimates and low energy/torus W-2 scores. Visually our energy histograms have nearly perfect overlap and free energy projections have great agreement with the test set.

Imaging and other experiments:

I would like to see experimental evaluation on a simple imaging dataset.

We thank the reviewer for their suggestion on solving the inverse problems on images. We definitely think this could be an interesting direction of research and experiment/algorithm design with an EBM deserves further thought. However, for the rebuttal period, we had to focus on other requested experiments on molecular systems (transferability and scalability) as efficiently solving the Boltzmann distribution is the primary purpose of this work. Further details on this can be seen in our reply to reviewer 3 vcq2.

InfoNCE objective:

The INFO NCE loss is not clear...

Thank you for your comments and questions, we have improved the writeup for the InfoNCE objective in our manuscript. The InfoNCE objective is an orthogonal objective function that maximizes the conditional likelihood $\rho_\theta(t|x_t)$.

$\rho_\theta(t|x) = \frac{\rho_\theta(t,x_t)}{\rho_\theta(x_t)} = \frac{\rho_\theta(t,x_t)}{\int\rho_\theta(t',x_t)dt'} = \frac{\exp(E_\theta(t,x_t))}{\int \exp(E_\theta(t',x_t))dt'}$

where $E_\theta$ is the energy based model. This leads us to minimize the negative log likelihood of the conditional:

$\mathcal{L_{\text{NLL}}}(\theta) = E_{x_t}\left[-\log \frac{\exp(E_\theta(t,x_t))}{\int \exp(E_\theta(t',x_t))dt'}\right]$

The intractable integral in the denominator is approximated by appropriately sampling a set of negative time points $t'$ resulting in the InfoNCE loss:

$\mathcal{L_{\text{InfoNCE}}}(\theta) = E_{x_t}\left[ -\log \frac{\exp(E_\theta(t,x_t))} {\sum_{t' \in \{\tilde{t}\} \cup t } {\exp(E_\theta(t',x_t))}}\right]$

please add some more explanation on that. why is it reasonable to sample negative points around t with a small variance?

In practice, on the interpolant process, only time points $t'$ close to $t$ will have a non-negligible likelihood $\rho_\theta(t', x_t)$, since $t$ also determines the amount of noise added to the clean sample. Therefore, it is reasonable to sample $t'$ using a distribution with small variance centered around $t$.

While the paper makes a very good experimental point, that both the NCE and the SM loss are needed, I would also like some (heuristic) explanation why that is.

Thank you for your question, we expand on it further below:

The score matching loss minimizes the Fisher divergence between the model and data distributions, and therefore also implicitly minimizes the KL divergence. However, as demonstrated in our 2D experiments, while the score matching objective effectively identifies the modes of the data distribution, it does not reliably capture their relative weights. The InfoNCE loss, being orthogonal to score matching, serves as a complementary regularizer that helps the model learn more accurate relative weighting between modes.

Distillation of FM model as a baseline

As a baseline one could also consider distilling the FM model, i.e., train a model which has similar likelihoods/odes but with much fewer steps

Thank you for the question. Most FM distillation methods we are familiar with fall under the category of consistency models. Due to their algorithmic design, these methods do not support likelihood estimation. The only method we are aware of that potentially reduces the number of integration steps is rectified flows, however, even this approach does not address the challenge of scaling inference with dimensionality. Furthermore, since our FM models are already trained with mini batch OT, rectified flow distillation may not reduce the transport cost by much.

Citations

I would also like to see ['citations'] when introducing the methods

Thank you for your comment. We have now added citations throughout the Methods and Results sections, including references related to stochastic interpolants, vector field and score matching objectives, endpoint parameterization, optimal transport, GVP, Graphormer, and more. Additionally, we have included more citations to several seminal works on score matching and flow matching in the main text.

Free Energy

please explain for me (not necessarily in the paper), why the free energy of a certain angle is such a good proxy for likelihood evals.

We believe the free energy projection of a system is a relevant baseline for the following two reasons:

• It represents a high dimensional integral of probability: The equation of a free energy projection along a reaction coordinate is given by:

  $$F(r') = -K_BT \ln \rho(r')$$

  where $\rho(r')$ is the probability density of observing the system at position $r'$ along the coordinate,

  $$\rho(r') = \frac{1}{Z} \int_x \delta(r(x) - r') e^{(-U(x)/K_BT)} \, dx$$

  In principle, if we solve the above integral along all points of the reaction coordinate, we solve a ($D-1$) dimensional integral, where $x \in R^D$. Thus, this serves as a good metric on how well the model matches the ground truth $R^D$ distribtuion.

  For dipeptides the process of going from the negative $\varphi$ angle to the positive $\varphi$ angle is considered a slow process as there are regions of high energy/low probability in between the two. Therefore, this serves as an ideal reaction coordinate to study dipeptide systems.

• Domain Relevance: Applied studies in the biophysical/ biochemical/ structural biology domain work on elucidating the free energy projection along a reaction coordinate. This helps the researchers identify the relative stability of different modes along the coordinate as well as the rate of reaction along the coordinate.

Other Questions

In the contributions you write "We show that the learned likelihoods can be used to reweight conformations to match the Boltzmann distribution. To the best of our knowledge, this is the first method to recover the Boltzmann distribution without requiring exact likelihood computations." I do not understand this, there have been a lot of papers recovering Boltzmann dists using approximate likelihoods and CNFs?

Thank you for your question. We do want to emphasize here that, at the time of submission, we were only aware of Boltzmann generator methods that compute likelihoods using the change of variable equation (both discrete and continuous versions). For example, CNFs use the continuous change of variable equation to compute likelihoods. We call these exact likelihoods methods. As of now, we are only aware of only one other (concurrent) work [1] that does not use exact likelihoods. We will revise our claim accordingly and clarify that our method also demonstrates strong performance in transferabile settings.

Does training the EBM on the side deteriorate the sampling (i.e, W_2) performance of the FM method?

Thank you for your question, the EBM does not interfere with the training of the FM model and is only trained to match the likelihoods of samples generated by the FM model. Therefore training it on the side does not affect the performance of the FM model.

[1] T. Akhound-Sadegh et al., “Progressive Inference-Time Annealing of Diffusion Models for Sampling from Boltzmann Densities,” arXiv.org, 2025. https://arxiv.org/abs/2506.16471 (accessed Jul. 30, 2025).

We thank reviewer vcq2 for their insights and are glad that they find our work well motivated. We are happy to see that the reviewer agrees with us that our EBM method achieves performance comparable to exact-likelihood normalizing flows while offering orders-of-magnitude faster inference.

2. Additional Experiments

As suggested by the reviewers, we conducted additional experiments to evaluate the transferability of our method. Specifically, we train our model on 200 dipeptides and test it on 7 unseen dipeptides, closely following the experimental setup described in [1]. Similar to our ADP experiments, we bias the training set to ensure adequate coverage of all relevant modes along the $\varphi$ angle.

To establish a molecular dynamics baseline, we run 1 μs simulations and compute E-W2, T-W2, and $\Delta F$ estimation errors with respect to the test set provided in [1]. We also train the model from [1] on our biased dataset. However, inference with this model requires several days of computation, so we were unable to obtain results for it within the rebuttal period. We will include these results in the camera-ready version.

During inference, we fine-tune our pre-trained EBM using samples generated for each unseen dipeptide. Additionally, we exclude the top 0.2% of importance weights during reweighting, following the approach introduced in [2], which significantly improves the stability of the reweighting process. Our results show that the method successfully transfers across systems while offering a substantial inference-time advantage over MD and ECNF baselines. The performance of our method compared to the MD baseline across the 7 dipeptides is summarized in the table below.

As shown in the table above, our method achieves a 6× speedup over conventional molecular dynamics while producing W2 and $\Delta F$ metrics that closely match the ground truth. Visually, the reweighted energy histograms show near-perfect overlap, and the predicted free energy projections align closely with the ground truth. These results support the conclusion that our method generalizes well and is transferable for inference beyond the training data.

Reweighting Improves Results

"reporting energy W2, torus W2, etc for both the proposal and the reweighted samples would help support the effectiveness of the EBM"

This is a very interesting idea and provides a great opportunity to show case the strength of our EBM model. We report energy W2 and torus W2 for both the proposal and reweighted samples in the table below. In all cases, we see that reweighting significantly improves generated distributions. The EBM (BoltzNCE Vector field) is able to efficiently reweight samples from the GVP vector field model and beats the ECNF baseline on these metrics

3. Questions

It is not immediately clear to me how the statement made on line 54 relates to the density chasm problem.

Thank you for the question. The density chasm problem arises when the samples drawn from the 2 distributions are easily distinguishable. Such a case results in vanishing gradients[3] under the NCE objective. Therefore to narrow the gap between the distributions that are compared in the contrastive loss we employ annealing through the use of stohcastic interpolant and ensure that the negative samples are in closely proximity to the positive sample.

Algorithm 1 in appendix B states the EBM uses data samples but the method states the samples at t=0 comes from the generative model.

Good catch! We have fixed it to reflect that we train the EBM on generated samples.

I am curious to see how does the generative model perform if the score and vector field are parameterized using $\nabla E_\theta$ from the EBM?

This is an interesting point that is worth exploring further. Our method takes advantage of a good sampler (Flow Matching) and a good likelihood estimator (EBM). We leverage both these properties to get the best of both worlds. For the scope of this rebuttal, however, we mainly focused on experiments related to transferability and leave sampling with $\nabla E_\theta$ to future work.

In conclusion, we thank the reviewer for their great suggestions of extra experiments. We believe that the experiments they proposed have truly strengthend our paper and we hope that they see our novel experiments on transferability as even stronger proof of our method. We also show that our EBM is well correlated with likelihoods from the Jacobian trace integral with orders of magnitude speed up. With this, we hope the reviewer considers improving their score on our paper.

[1]Leon Klein and Frank Noé. Transferable boltzmann generators. arXiv preprint arXiv:2406.14426, 2024.

[2]Tan, Charlie B., Avishek Joey Bose, Chen Lin, Leon Klein, Michael M. Bronstein, and Alexander Tong. "Scalable equilibrium sampling with sequential boltzmann generators." arXiv preprint arXiv:2502.18462 (2025).

[3]Liu, B., Rosenfeld, E., Ravikumar, P., & Risteski, A. (2021). Analyzing and improving the optimization landscape of noise-contrastive estimation. arXiv preprint arXiv:2110.11271.

We thank the reviewer for their feedback and we are happy to see they agree that our method provides orders of magnitude speed up with competitive accuracy.

1. Experimental Results

The experiments are somewhat weak, as they only cover the alanine dipeptide

We agree that additional experimental validation would benefit our paper and provide stronger proof that our method works. To this order, we show that our method is transferable by training on 200 dipeptides and testing on 7 unseen dipeptides. With a bit of fine tuning on the unseen dipeptide, we achieve good performance while attaining a significant time advantage over MD and benchmark ECNF methods. Visually our energy histograms have nearly perfect overlap and free energy projections have great agreement with the test set. For more details on transferability, refer to our reply to reviewer vcq2. These results are discussed in depth in our reply to reviewer vcq2 (reviewer 3).

when samples are generated with one model, but weighted by another, it must be ensured that both models agree on the generative probability

We agree! To showcase this, we calculated Spearman correlation coefficients between the EBM output and exact likelihood as calculated by the Jacobian trace integral for the alanine dipeptide syteme. Our final paper version will include the scatterplots but we report just the quantative value below.

As shown in Table 1, the EBM output exhibits strong agreement with the flow-matching likelihoods. Moreover, we found that our EBM outperforms the Hutchinson estimator, achieving better correlation while being over two orders of magnitude faster at inference on 10,000 samples. Interestingly, increasing the number of estimator calls in the Hutchinson method does not improve its correlation. Due to the use of an adaptive step-size ODE solver, 4 and 8 estimator calls are actually faster than 2 calls.

2. Formulas and equations

lacks derivations of the formulas, nor have sections 3 and 4 any citations

Thank you for the comments. We have now included multiple new citations throughout the methods and results sections so that the reader can follow where some parts of the method are from. Some of the main papers we cite heavily include:

• For stochastic interpolant equations

  • Nanye Ma, Mark Goldstein, Michael S Albergo, Nicholas M Boffi, Eric Vanden-Eijnden, and Saining Xie. Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers. In European Conference on Computer Vision, pages 23–40. Springer, 2024.368

• For Boltzmann Generator formulations and equations

  • Leon Klein, Andreas Krämer, and Frank Noé. Equivariant flow matching. Advances in Neural Information Processing Systems, 36:59886–59910, 2023

• For Endpoint Objective

  • Bowen Jing, Bonnie Berger, and Tommi Jaakkola. Alphafold meets flow matching for generating protein ensembles. arXiv preprint arXiv:2402.04845, 2024

• For mini-batch OT

  • Alexander Tong, Kilian Fatras, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. arXiv preprint arXiv:2302.00482, 2023.459

We have also included citations to other seminal score matching and flow matching works.

3. Coupling concerns

The coupling here[in the EBM] must be the optimal coupling resulting from Schrödinger bridge matching

Thank you for your comment. To test your hypothesis, we trained versions of the EBM model with both OT and independant coupling. We found that the models converged in both cases and we saw only minor changes in model performance. We report the performance of both methods in the table below. With this we conclude that OT coupling may not be necessary for model convergence.

The authors mention briefly (lines 183-184) that they approximate the coupling by minibatch optimal transport via the Hungarian algorithm, but they fail to discuss to what degree this approximation fulfills the requirements for the coupling, nor what these requirements are in the first place.

Thank you for your comment. As shown in [1], minibatch OT approximates entropic regularize OT coupling of the two distributions. The hungarian algorithm is of similar flavor and is mainly used to for scalability to larger batch sizes.

4. Generator and Emulator Clarifications

it is hard to understand exactly which sub-algorithms according to which formulas are included or excluded in the different model variants

We apologize for the confusion, we added the following descriptions more explicitly in the tables and figures in our newest version in a similar manner as the table below.

We have now also added a description of the umbrella sampling method in our paper.

It is also unclear why reweighting was not included in the experiment in figure 4 and table 1

Figure 4 and Table 1 aimed to compare the performance of the emulators not the generators. ie. we wanted to show the impact of the choice of objective functions without the influence of reweighting. Furthermore, our results showed that a better emulator does not neccesarily create a better generator due of the design constraints that generators have to follow. However, we agree that this section may confuse readers and detract from the main point of the paper and have moved this discussion to the appendix. We have, however, now added additional W-2 results, before and after reweighting on the Boltzmann Generators comparison table. See further detail on this in our response to reviewer vcq2.

How other methods fail to estimate the free energy in figure 5

Thank you for the question. We have provided free energy projections for the other methods in the supplementary material. We hypothesize that the other methods have slightly more inaccurate free energy projection predictions due to inaccurate likelihood estimates that can occur due to error propogation during ODE integration. This is especially prevelant when one tries to integrate likelihoods on CNF models trained with the endpoint objective.

Overall we thank the reviewer for their valuable comments, we hope that they see our novel experiments on transferability provides even stronger proof that our method unlocks the potential for a truely boltzmann generator. We also address concerns regarding couplings for the EBM with emperical results and provide additional clarification on methods and equations. With this, we hope the reviewer considers improving their score on our paper.

[1] Tong, Alexander, et al. "Simulation-free schrodinger bridges via score and flow matching." arXiv preprint arXiv:2307.03672 (2023).

We thank the reviewer for their thoughtful review and questions, and we are pleased that they found our approach practical for scaling to larger systems. We would also like to mention that in addition to results on alanine dipeptide, we now include new experiments demonstrating transferability. For transferability we train on a set of 200 dipeptides and infer on 7 dipeptide systems not present in the training set and show accurate free energy estimates and low energy/torus w-2 scores. Visually our energy histograms have nearly perfect overlap and free energy projections have great agreement with the test set. For more details on transferability, refer to our reply to reviewer vcq2.

The method requires training an additional EBM, whereas other approaches typically train only a single CNF model.

We acknowledge that our method requires training two models instead of one. However, we believe that this is a key strength, rather than a limitation. Flow matching models are good for sampling but hard for likelihood evaluations. EBMs on the other hand provide easy access to likelihoods but are not great for sampling. By combining the two types of model we leverage the best of both worlds.

Moreover, our method design significantly reduces inference-time costs, as evaluating the likelihood of a sample requires only a forward pass through the EBM. This is beneficial as it reduces the constraints on the emulator/neural sampler as it gets rid of the requirement of invertibility for Boltzmann generation. Furthermore, in transferability scenarios, i.e. inferring on systems not present in the training set, the EBM requires only a bit of fine tuning for significant inference time gains. The two model approach is also advantageous when scaling to higher dimensional systems, where training an additional neural network is less computationally expensive than calculating Jacobian trace integrals.

The architecture choices are complex, yet architectural ablations are missing. A discussion is needed to clarify that performance gains are not merely due to increased parameter counts.

We acknowledge that our flow matching architectures are slightly more complex than the ECNF model (to account for multiple vector features), however, the GVP models our study has significantly fewer parameters (108,933) than the benchmark ECNF (147,599). We now report the parameter counts in our updated paper.

The endpoint models diverge near t = 0, requiring manual truncation of the integration. This instability reduces the effectiveness of likelihood-based sampling.

We agree with the reviewer that the endpoint model is not well-suited for likelihood-based sampling. However, we do not view this as a limitation of our method. Our approach leverages an EBM to estimate likelihoods, thereby avoiding the computational burden of the continuous change-of-variable equation. A key insight from our method is that introducing a secondary EBM for likelihood estimation relaxes the design constraints on the Emulator models. Since efficient likelihood evaluation is no longer required for Boltzmann generation, the Emulator can be trained using the endpoint objective. We demonstrate this in our experiments on alanine dipeptide and emphasize that this is the central takeaway from our use of the endpoint objective.

If the Boltzmann emulator is trained on biased data with missing modes, the EBM is unable to recover those missing modes.

We agree with the reviewer regarding the issue of missing modes; however, this is a limitation common to most generative models and not specific to our method. The emulator needs to be able to sample sufficiently from the modes of the dataset for effective reweighting with the EBM. Notably, the authors of the original ECNF paper[1] also reported needing to bias the data toward underrepresented modes in order to obtain accurate free energy differences. Our experimental design follows the same approach.

Is it possible to just use EBM for sampling, since it already provide a score of the model density.

This is an interesting idea worth exploring further. However, we emphasize that the sample distribution must approximately match the likelihoods predicted by the EBM to achieve the Boltzmann. In practice, achieving this with an EBM typically would require running Langevin dynamics for large amounts of time, which may be computationally infeasible.

In your dipeptide experiments, how does the model perform if the emulator training dataset are more biased?

We biased our dataset to account for the underrepresented mode along the positive $\varphi$ angle. This step is necessary to obtain accurate free energy differences method. We adopted the same method for biasing as employed in the original ECNF work[1] to ensure a fair and equitable comparison between the two methods. Again, we acknowledge the need for our Emulator model to be able to cover all the modes, but this is a limitation that is common to most generative models and is discussed in previous work as well[1].

[1]Leon Klein, Andreas Krämer, and Frank Noé. Equivariant flow matching. Advances in Neural Information Processing Systems, 36:59886–59910, 2023.

Overall, we thank the reviewer for their thoughtful comments and hope they will consider increasing their score in light of the new experimental results and our responses to their concerns.