We thank the reviewer for the valuable input. We fully agree with every point raised. We brushed over the actual gradient estimators, which is something we will add to the manuscript.

And yes, your guess is correct. The maximum likelihood (ML) and PG estimator minimize the same loss and the estimators only differ by a control variate term. But the PG estimator for the forward KL was actually introduced and discussed in Chapter 4 of [1].

In the following, we first discuss the ML and PG estimators for the forward KL and show benefits of PG over FM for a toy example. This additional discussion is a great addition to the paper, we hope you agree.

The Maximum Likelihood and Path Gradient Estimators.

In short: both the ML and PG estimators are unbiased and consistent, but they differ in variance. For PG, we can give guarantees about their variance, once $q_\theta$ equals to $p$, see e.g. [2, 3,4,5].

Both Maximum Likelihood and the Path Gradient estimators optimize the Forward KL Equation (8). Let’s look at both estimators in detail and compare them. A full derivation can be found in Appendix B.3.2 of [1] & Chapter 4 of [2].

To obtain the estimator $\mathcal G_{ML}$, we first observe that the first term in the KL divergence is constant w.r.t. $\theta$. $KL(p|q_\theta) = E_{p(x_1)} \left[ \log p(x_1) - \log q_\theta(x_1) \right],$ which means that it does not enter in the gradient if we directly estimate the gradients via an MC estimator.

$\frac{d KL(p|q_\theta)}{d \theta} = - E_p(x_1) \left[\frac{d}{d \theta} \log q_\theta(x_1) \right] \approx \mathcal G_{ML} = -\frac{1}{N} \sum_{i=1}^{N} \frac{d }{d \theta} \log q_\theta(x_1^{(i)}), x_1^{(i)} \sim p$.

For the calculation for the estimator, we decompose it fully with

$$q_\theta(x_1) = q_0(T_\theta^{-1}(x_1)) \left| \det \frac{\partial T_\theta^{-1}(x_1)}{\partial x_1} \right|$$

The actual calculation for the ML estimator then is $\mathcal G_{ML} = -\frac 1 N \sum_{i=1}^{N} \frac{d }{d \theta} \left(\log q_0(T_\theta^{-1}(x_1)) + \log \left| \frac{\partial T_\theta^{-1}(x_1)}{\partial x_1}\right| \right) (*).$

For PG, we use Eq. (17) to directly obtain the MC estimator $\mathcal G_{PG}$ $\frac{d}{d \theta} KL(p_{0, \theta}| q_0) = E_{x_1 \sim p_1} \left[\frac{\partial }{\partial x_0} \left(\log \frac{p_{0, \theta}}{q_0}(x_0) \right) \frac{\partial T_\theta^{-1}(x_1)}{\partial \theta} \right] \approx \mathcal G_{PG} = \frac 1 N \sum_{i=1}^N \frac{\partial }{\partial x_0^{(i)} } \left(\log p_{0, \theta}(x_0^{(i)}) - \log q_0(x_0^{(i)}) \right) \frac{\partial T_\theta^{-1}(x_1^{(i)})}{\partial \theta}, x_1^{(i)} \sim p.$ Here $x_0^{(i)}$ is a shorthand for $T_\theta^{-1}(x_1^{(i)})$. If we again decompose the term we expose the full calculation $\mathcal G_{PG} = \frac 1 N \sum_{i=1}^{N} \frac{\partial }{\partial x_0^{(i)} } \left(\log p(T_\theta(x_0^{(i)})) + \log \left| \frac{\partial T_\theta(x_0)}{\partial x_0}\right| - \log q_0(x_0^{(i)}) \right) \frac{\partial T_\theta^{-1}(x_1^{(i)})}{\partial \theta} (**)$

Comparing (*) and (**), we see that the path gradient estimator incorporates $\mathcal G_{PG}$ the gradient information of the target $p$ while $\mathcal G_{ML}$ - and by construction also $\mathcal G_{FM}$ - does not.

Variance of the estimators:

Opposed to $\mathcal G_{ML}$ and $\mathcal G_{FM}$, we have nice guarantees about the variance of the PG estimator $\mathcal G_{PG}$ at the optimum and close to it.

We first recapitulate results from previous work and then, by using a simple example, show that $\mathcal G_{FM}$ does not necessarily exhibit zero variance at the optimum.

In general, the variance of the Path Gradient estimators $\mathcal G_{PG}$ is bounded by the squared Lipschitz constant of the term $\left(\log p_{0, \theta}(x_0) - \log q_0(x_0) \right) = \left(\log p(x_1) - \log q_\theta(x_1) \right)$ ([6] Section 5.3.2).

Thus, if the target density $p$ is not well approximated by $q_\theta$, the variance of the gradient estimator can be large and training with PG might not be beneficial. If $\log p$ and $\log q_\theta$ are close in the sense that the Lipschitz constant of their difference is small, we can assume path gradient estimators to be helpful.

Gradient estimators at the optimum

In the case of perfect approximation, i.e. $q_\theta(x) = p(x) \forall x \in \mathcal X$, the following statements about the ML estimator and the Path Gradient estimators are known. The gradients for path gradient estimators are deterministically $0$, i.e. $E[\mathcal G_{PG}] = 0, Var[\mathcal G_{PG}] = 0$, while for ML the variance is generically nonzero $E[\mathcal G_{ML}] = 0, Var[\mathcal G_{ML}] = \frac 1 N \mathcal I_\theta$. Where $\mathcal I_\theta$ is the Fisher Information Matrix of $p_{0, \theta}$ [2].

Why is that?

Already the review [7] notes in Section 2.3.3 that the duality of the KL divergence means that fitting the model $q_\theta$ to the target $p$ using ML is equivalent to fitting $p_{0,\theta}$ to the base $q_0$ under the reverse KL . This means that the results from [2] directly hold. We only adapted the notation.

Variance of $\mathcal G_{FM}$ for a toy example

We do not have a term for the variance of $\mathcal G_{FM}$ for general settings, but we can show that it does not deterministically vanish like for $\mathcal G_{PG}$ via a simple example. This means better behavior for PG than for FM at the optimum, which we also verified empirically.

Our assumptions aim to simplify the example as much as possible.

• First, assume the standard loss for Flow Matching.
• Further, assume the two densities to be the same $D$-dimensional Normal distribution $q_0 = p = \mathcal N(0,I)$.
• Finally, assume the CNF is the identity parametrized by a single parameter $\theta$, i.e. $v_\theta = \theta = 0$.

In this example $q_\theta(x_1) = q_0(x_1) I$ approximates the target density $p$ perfectly and the optimal gradient estimators is $0$. Yet, in this setting the variance of the estimator $\mathcal G_{FM}$ is non-zero: $Var[\mathcal G_{FM}] =\frac 8 {ND},$ preventing the model from staying at the optimum during training.

Proof:

The gradient estimator for the FM loss Eq (10) is $\mathcal G_{FM} = \frac 1 N \sum_i \frac{\partial }{\partial \theta} \left(|| v_\theta - ( x_1^{(i)} - {\tilde x}_0^{(i)}) ||^2\right),$ where ${{\tilde x}_0^{(i)}, x_1^{(i)} \sim \mathcal N(0,I)}$. (Note that here ${\tilde x}_0^{(i)}$ and $x_1^{(i)}$ are independently sampled. Before $x_0^{(i)}$ was the transformed sample).

First, we break down the terms

\mathcal G_{FM}= \frac{1}{N} \sum_{i=1}^{N} \frac{\partial }{\partial \theta} \frac{1}{D} \sum_{d=1}^{D}\left(v_\theta^d +\tilde x_0^{(i),d} - x_1^{(i),d} \right)^2 = \frac{1}{ND} \sum_{i=1}^N \sum_{d=1}^D 2 \left(v_{\theta}^d + \tilde x_0^{(i),d} - x_1^{(i),d}\right) \frac{\partial v_{\theta}^d}{\partial \theta}.$$ Because $v_\theta$ is parametrized by $\theta$, we set in $\frac{\partial v_\theta^d}{\partial \theta}=1$ and $v_\theta^d=0$ and separate the terms $$\mathcal G_{FM}= \frac {1} {ND} \sum_i \sum_d 2 ( \tilde x_{0}^{(i),d} - x_{1}^{(i),d}) = \frac 2 {ND} (\sum_d \sum_i \tilde x_{0}^{(i),d} - \sum_d \sum_i x_{1}^{(i),d}). $$ We can compute the distribution of $\mathcal G_{FM}$ by using the property $\sum_{i=1}^N a_i \sim \mathcal N(0, N \sigma_a^2)$ for $a \sim \mathcal N(0, \sigma_a^2)$. So $\sum_d \sum_i \tilde x_{0}^{(i),d} \sim N(0, D N)$ and $\sum_d \sum_i \tilde x_{0}^{(i),d} - \sum_d \sum_i x_{1}^{(i),d} \sim \mathcal N(0, 2 D N).$ The expectation of the estimator $\mathcal G_{FM}$ then is 0 and the variance is: $Var[\mathcal G_{FM}] = \frac 4 {N^2D^2} Var[(\sum_d \sum_i \tilde x_{0}^{(i),d} - \sum_d\sum_i x_{1}^{(i),d})] = \frac 4 {N^2D^2} 2DN = \frac 8 {ND}$. Interestingly, the derivative $\mathcal G_{FM}$ is invariant to re-ordering because the sums of $x_1$ and $\tilde x_0$ are invariant to permutation. The result thus also holds for OT-FM. [1] L. Vaitl et al., “Fast and unified path gradient estimators for normalizing flows” ICLR 2024 [2] L. Vaitl “Path Gradient Estimators for Normalizing Flows" PhD Thesis [3] G. Roeder et al., "Sticking the landing: Simple, lower-variance gradient estimators for variational inference" NeurIPS (2017) [4] G.Tucker et al. "Doubly reparameterized gradient estimators for monte carlo objectives" ICLR 2019 [5] L. Vaitl et al. "Path-gradient estimators for continuous normalizing flows" ICML, 2022 [6] S. Mohamed et al. "Monte carlo gradient estimation in machine learning" JMLR 2020 [7] G. Papamakarios et al. "Normalizing flows for probabilistic modeling and inference" JMLR 2021 # Other concerns - X axis, figure caption, L206, too strong statement, L226 second set of experiments: Thanks for the input. We will fix these in the final manuscript. - Isn't "applying path gradients at the end" the same as fine-tuning? True, in the paper the difference is not very clear. In the first setting we used limited memory, compute and time constraints to compare FM and the hybrid approach. This is to show that the improvements do not only come from additional training. The second setting does not limit computational resources or time and training with PG starts after full training as done in Klein et al. 2023b. This was done to see how much PG can help training. We will clarify this in the final manuscript. - Difference between the models labeled Klein et al 2023b and Klein & Noe 2024. We agree with the reviewer that the difference should have been explained better in the paper. The architecture in Klein & Noe 2024 is mostly the same as in Klein et al. 2023b, with the main difference of the encoding of the atoms. While the model in Klein et al. 2023b treats the same atom type as indistinguishable, the model in Klein & Noe 2024 encodes nearly all atoms differently. Only hydrogens bound to the same atom are treated as indistinguishable. We will add description in the final version. - Paper Formatting We have not changed the template

We thank the reviewer for their thoughtful comments and for highlighting the clarity and structure of the paper, as well as the practical advantages of our hybrid approach. We hope that our responses to the raised concerns help clarify our contributions and demonstrate the significance and robustness of the method.

Weaknesses:

"The quantitative comparison in Table 1 is limited to comparing pre-trained OT-FM models with FM fine-tuning or with Path Gradient fine-tuning. Additional comparisons of a hybrid training using other base models would be helpful."

We agree with the reviewer that additional comparisons would be helpful. So for the rebuttal we did additional experiments.

We evaluated our hybrid approach on the Dipeptide setting from Transferable Boltzmann generators [3]. This training set consists of 200 dipeptides and the evaluation is done on unseen dipeptides. Our approach yielded an improvement in the NLL on all evaluated molecules with an average improvement of 0.20, which is in line with our experiments on AD2. For additional information see our rebuttal for reviewer Fj6i.

Still, we want to note that in Section 4.2 we evaluated fine-tuning with EQ OT FM and Standard FM on LJ13 in addition to the comparisons in Table 1. For the experiments of Table 1 we stuck to OT FM because the ESS and NLL were similar for the three FM variants and Klein & Noe 2024 used OT FM.

The ability to optimize the flow trajectory length, as in OT-FM and EQOT-FM, is unique to FM approaches. These recent developments are of growing interest in the community. Our work seeks to understand their behavior under hybrid training. Other base models, like other CNFs or Coupling Type Flows were already evaluated in previous work [2,3,4] and showed improved performance with Path Gradients.

A hybrid approach applied to diffusion models and their Probability Flow ODE is a very promising direction for future work, but merits a thorough investigation.

"An important baseline that is missing in the experiments is the method proposed in Satorras et al. [1]."

We agree with the reviewer that the model in [1] is a good baseline and for the LJ13 we will gladly incorporate the baseline.

For the more challenging datasets, we therefore only compared the models proposed in Klein et al. 2023 [6] and Klein et al. 2024 [7], which outperform Satorras et al. [1].

"The authors suggest that Path Gradients are suitable for fine-tuning only when the flow is already close enough to the target distribution. However, this claim is not supported in the experimental setup. How does the final model perform when fine-tuning with Path Gradient is performed when the flow is not close to the target?"

We agree with the reviewer that Path Gradients (PG) can also improve training beyond the fine-tuning regime. Prior work (e.g., [2,3,4]) has demonstrated strong performance when training models exclusively with PG, outperforming other KL-based objectives (including Maximum Likelihood training).

Yet, our preference for a hybrid approach over exclusively PG is motivated by:

• Time efficiency: FM is significantly faster per step than PG and ML.

• Theoretical guarantees: PG has convergence guarantees only near the optimum (see Rebuttal to Reviewer DmNE for a thorough analysis).

• Numerical stability: Like ML training, PG involves Jacobians and their derivatives which suffer from numerical instabilities, whereas FM is often more stable.

Especially the time efficiency was shown in our experiments:

• In our GMM experiment (Section 3), we directly compare PG, FM, and the hybrid approach. FM trains faster initially, but PG eventually surpasses FM in final performance. The hybrid strategy outperforms PG alone by combining FM’s rapid early progress with PG’s refinement.
• On the complex AD2-XTB (Table 1), we observe that, given limited time, FM outperforms fine-tuning with PG when using a less expressive model [3]. This supports our point: when the model is underfitting or far from the target distribution, PG fine-tuning provides limited benefit.

We will revise the manuscript to reflect this more nuanced perspective and avoid overstating the limitations of PG.

"One notable weakness of the proposed approach (already addressed in the limitations section) is that it relies on the availability of a well-defined and differentiable energy function over the data domain, restricting its applicability to other domains."

We agree with the reviewer. That said, the setting still covers a wide range of practically important domains like natural sciences – including drug discovery, material discovery or protein folding (e.g., AlphaFold) – and can even be applied to model distillation, like used in StableDiffusion3.

[1] Garcia Satorras, Victor, et al. "E (n) equivariant normalizing flows." Advances in Neural Information Processing Systems 34 (2021): 4181-4192.

Questions:

"What is the performance of the hybrid approach using an increased runtime?"

For LJ13 we show the results with increased runtime in Table 3. For AD2-classical we give results for TBG [7] in Table 2. Summarizing, we were able to obtain the following results

Additionally, we ran the Hybrid approach on LJ55 once for 3 PG epochs instead of one and got the following performance NLL: -89.38, $ESS_q$ 25.44 %.

"Why do the values of ESS-Q and NLL in Figure 3 of the OT + Path Gradient fine-tuning baseline on LJ13 not match the results reported in Table 1?"

The settings are different:

• In Table 1. we used a pretrained model (1000 epochs FM with lr 5e-4) which we fine-tuned with PG, limited batch-size (64) and time.
• For the results in Figure 3, we used a model which was “fully trained” with FM (1000 epochs FM with lr 5e-4 and 1000 epochs with lr 5e-5) and finetuned with PG and batch-size 256. The larger batchsize means that an epoch had less training-steps.

"What is the runtime of the hybrid approach fine-tuning (Table 1)?"

We list the runtimes in A.4.

"While the authors suggest that they show how to compute Path Gradients with constant memory and without relying on additional samples, this was already addressed in prior work [2]. Can you clarify the novelty to Vaitl et al. [2]?"

There are two major differences:

• 1.) While Vaitl et al. [2] and our work both use PG, they minimize a different loss on different samples. Our work optimizes the Forward KL $KL(p|q_\theta)$ with PG on existing samples from $p$, as proposed in [3] using the algorithm proposed in [2]. Vaitl et al. [2] minimize the reverse KL divergence $KL(q_\theta|p)$ via self-sampling, i.e. on samples from the model $q_\theta$. This has many drawbacks, mainly: modes of the target can be entirely missed, which invalidates all asymptotic guarantees and breaks importance sampling, see e.g. [5]. This becomes increasingly likely in higher dimensions.

• 2.) Since the publication [2], Flow Matching has been established as the de facto training method for CNFs. Our work investigates and combines the performance of Flow Matching and Path Gradients. Specifically, we investigate the effect of PG on the inner workings of the CNF, like e.g. trajectory length, while [2] simply looked at the performance compared to standard self-sampling losses.

That being said, forward KL based training can straightforwardly be combined with self-sampling based approaches (like [2]) to artificially increase the number of samples.

[2] Vaitl, Lorenz, et al. "Path-gradient estimators for continuous normalizing flows." International conference on machine learning. ICML, 2022.

[3] L. Vaitl et al., “Fast and unified path gradient estimators for normalizing flows”, ICLR 2024

[4] A. Agrawal et al., "Advances in black-box VI: Normalizing flows, importance weighting, and optimization." NeurIPS2020

[5] Detecting and Mitigating Mode-Collapse for Flow-based Sampling of Lattice Field Theories, Nicoli, Kim A. et al. PRD 23

[6] Klein, Leon et al. "Equivariant flow matching." NeurIPS (2023).

[7] Klein, Leon, and Frank Noé. "Transferable boltzmann generators." NeurIPS (2024)

We thank the reviewer for their thoughtful and encouraging feedback. We appreciate them pointing out the practical relevance of the problem, the strength of the experimental evidence, and the clarity of the presentation.

Weaknesses:

"The experiments are performed on well-established but relatively small molecular systems (Lennard-Jones clusters and Alanine dipeptide). While these are excellent for controlled comparisons, demonstrating the method's effectiveness on a larger and more complex system (e.g., a larger protein) would significantly bolster the paper's impact."

We agree that additional and more challenging systems improve the impact of the paper. So we additionally evaluated the performance of our approach on the Dipeptides dataset in (Klein & Noe, 2024) to test the transferability to unseen systems. This is currently the largest system Boltzmann Generators have shown transferability for.

Again, we fine-tune the available pretrained Boltzmann Generator from (Klein & Noe 2024) with path gradients. This transferable Boltzmann Generator is trained on a subset of all possible dipeptides and evaluated on unseen ones. Our experiments show that fine-tuning with path gradients improves the NLL and energies for all evaluated unseen test dipeptides and also shows average improvements for the ESS from $7.96$ to $9.79$.

The full results can be found in our rebuttal to reviewer Fj6i.

Questions:

"The hybrid strategy requires switching from Flow Matching pre-training to Path Gradient fine-tuning. Your work appears to do this after a fixed number of epochs. Is there a more principled way or a clear heuristic to determine the optimal time to make this switch? For instance, should one monitor the FM loss for a plateau? A generalizable strategy for this would increase the practical utility of your method."

This is a good idea. In principle, monitoring the FM loss plateau is a reasonable heuristic and inexpensive to implement. However, our primary focus was understanding and comparing the behavior of PG, FM, and their combination—rather than on optimizing the switching criterion itself. For the experiments in Section 4.1, we simply used half the available training time for FM and PG each, for Section 4.2 we used the training regime in (Klein et al. 2023b) for pretraining with FM.

While designing a more principled or generalizable strategy is an interesting direction, its effectiveness can depend on the dataset, model, and optimization dynamics. We therefore leave this to future work.

"The performance gains are very impressive on the tested systems. How do you expect this hybrid approach to scale to much larger biomolecules? In such systems, force calculations are more expensive and the energy landscapes are far more complex. Do you hypothesize that the relative benefit of a few epochs of PG fine-tuning will increase or decrease in that regime?"

We thank the reviewer for this insightful question. We believe our hybrid approach has two key advantages that support its scalability to larger systems:

• Path Gradients typically scale well with dimensionality. They have successfully been applied to high-dimensional systems to simulate Quantum Field Theories (>10^3 dimensions [1,2] and even >10^5 dimensions [3,4]). Further, our additional experiments on the unseen dipeptides show a positive correlation of 0.53 between the NLL gain per dimension and the number of particles. This is a good indicator that fine-tuning with PG scales well to higher dimensions for molecules as well.
• A practical strength of our method is that force evaluations are required only once. In many real-world applications, these forces are already available from prior simulations, like in MD or MCMC sampling where forces are required repeatedly. This “recycling” makes our method especially attractive when force evaluations are expensive, as is the case for large biomolecular systems.

However, the effectiveness of Path Gradients remains limited by the expressiveness of the model. If the energy landscape is too complex to be approximated well, we do not expect fine-tuning with PG to resolve this fundamental mismatch. We thus expect the benefits of Path Gradients to extend to larger molecular systems - as long as the model’s capacity is sufficient to capture the complex energy landscape of the target distribution.

[1] Vaitl, Lorenz, et al. "Path-gradient estimators for continuous normalizing flows." International conference on machine learning. ICML, 2022.

[2] Bacchio, Simone, et al. "Learning trivializing gradient flows for lattice gauge theories." Physical Review D 107.5 (2023): L051504

[3] Abbott, Ryan, et al. "Normalizing flows for lattice gauge theory in arbitrary space-time dimension." arXiv preprint arXiv:2305.02402 (2023).

[4] L. Vaitl “Path Gradient Estimators for Normalizing Flows", PhD Thesis

We thank the reviewer for their positive feedback. We're glad the clarity and practical value of our method came through.

Weaknesses

"It seems that the main benefit of path-gradient fine-tuning is that it improves the efficiency of sampling high quality samples at inference time of a flow model that is already quite good. This is not really a bottleneck in scaling Boltzmann generators since it is quite cheap to sample the diffusion model once its trained. For Lennard Jones and Alanine dipeptide, the energy is pretty cheap to evaluate and so importance sampling is straight-forward. It’s hard for me to see how this technique would scale up to harder problems we are unable to approach with flow/diffusion samplers (such as Chignolin in [1]), if it requires the flow model to already be quite good and unbiased."

We agree that path-gradient fine-tuning improves the efficiency of the Boltzmann Generators. However, we disagree that this does not solve the major bottleneck of sampling speed.

The inference with Boltzmann Generators based on CNFs is expensive if one aims to do importance sampling. The major bottleneck is usually not the energy evaluation, which can be cheap, but the integration of the trace of the Jacobian (eq. (6)). Hence, having access to a better Boltzmann Generator, such as one fine-tuned with path-gradients, allows faster sampling as less model evaluations are required per effective sample.

Even when avoiding correction via importance sampling, the samples are significantly closer to the reference (see Figure 4 right) and the target is approximated better. So, while our training approach is not a panaceum for samplers, in settings where flow models are already useful, we showed that we can obtain significant gains in performance and soften the bottleneck for importance sampling.

"Requiring unbiased data seems like a major downside to this method, as this assumes you should already have a way to generate true samples via MCMC or otherwise. I think a more relevant situation would be when you have some biased data (due to limitations of another approach) and want to improve sampling towards the desired boltzmann, even at the cost of additional evaluation (see [1])."

We appreciate the reviewer’s point, relying on unbiased samples is indeed a limitation. However, we’d like to clarify a few key aspects, to put this downside into perspective:

• In a setting where you have access to the derivative of the energy, you can always apply MCMC or MD. The settings with Boltzmann Generators, as well as the approach by (Tan et al. 2025) [1] or protein folding fall in this regime. The issue with MCMC or MD is that they are not easily parallelizable and exhibit a high autocorrelation. Deep generative models simply aim to speed up the process. Sample-efficient training methods, like Path Gradients, are of great value in this setting.
• In practice, a small bias in the training data is tolerable. For example, in our AD2 and TBG experiments, the training data likely contained mild sampling bias due to finite-length MD runs, yet the models performed well after PG fine-tuning. In additional experiments on dipeptides that we did for this rebuttal, we also saw an improved performance even though the “dipeptides, which were simulated each with a classical force field for 50 ns and, therefore, may not have reached convergence” (Klein & Noe 2024).
• In these additional experiments we further investigated the transferability to unseen Dipeptides and saw that PG was able to generalize to unseen systems, further relaxing this downside. In the transferable setting, the model needs only equilibrium samples from the training systems, no data is required for the unseen test systems to which it transfers. Below we summarized the additional experiments.

[1] Tan et al 2025 "Scalable Equilibrium Sampling with Sequential Boltzmann Generators"

Questions:

"The path gradient is based on the backwards KL, but the flow-matching marginal vector-field is not minimizing the same objective. Yet it is claimed that the structure of the flow isnt changed very much. Do you have intuition about what happens to the marginal of the path in flow matching?"

Indeed, Flow Matching (FM) and Path Gradients (PG) optimize different objectives:

FM minimizes a trajectory-level MSE to match a reference trajectory $u_t$, while PG minimizes the KL divergence between the induced distributions $q_\theta$ at time $t=1$ and $p$. Since different trajectories can lead to the same transformation (e.g., via time reparametrization and early stopping in $t$), the KL is agnostic to the actual trajectory for $t \in [0,1)$, i.e. it depends only on the endpoint of the transformation at $t=1$.

So for a minimal FM loss the KL is also minimized, while the reverse does not hold. Our experiments suggest that when PG is used after FM pretraining, the variational distribution $q_\theta$ is already close to the target distribution $p$, and PG fine-tuning adjusts the trajectory only slightly without deviating too much from the reference trajectory $u_t$.

"Similar to my concern above: If this is considered to be a SOTA method for Boltzmann generators, how do path-gradient scale on larger systems which current diffusion samplers have trouble with (larger peptides like Chignolin)? why not compare to other strategies for fine-tuning / reweighting such as “Scalable Equilibrium Sampling with Sequential Boltzmann Generators”. These methods do not even require unbiased data, but at the cost of additional energy evaluations."

We agree with the reviewer that scalability is an open question. However, scaling to larger systems often comes with high computational costs. In this work, we aim to show how to improve current Boltzmann Generators with path-gradient fine-tuning and focus on the architecture trained with flow matching. However, path-gradient fine-tuning is also applicable to a large class of Normalizing Flows architectures (see [7]).

PGs have been demonstrated to perform well in high dimensional settings e.g. in O(10^3) dimensions in [2,3] or even in O(10^5) in [4,5,6] in quantum field theory. Still, Chapter 6 of [6] suggests that a limiting factor is often the complexity of the target density compared to the modeling power of the model.

As the reviewer mentions, the method in [1] requires additional energy evaluations compared to traditional BGs, which can be a bottleneck depending on the energy function. Importantly, for the path-gradient training, no additional energy evaluations are required during training, as the forces of the training sample are used.

Further, note that PG-training can readily be combined with the non-equilibrium sampling scheme from Proposition 2 in [1], since the former is a training scheme and the latter is a test-time sampling scheme.

Additional Dipeptide Experiments:

We applied our hybrid approach to transferable Boltzmann Generators on dipeptides as introduced in Klein et al. 2024. We again fine-tune the pretrained Boltzmann Generator from Klein et al. 2024 with path gradients. Training took 6 days on an A100. This transferable Boltzmann Generator is trained on a subset of all possible dipeptides and evaluated on unseen ones. Our experiments show that fine-tuning with path gradients improves the NLL and energies for all evaluated unseen test dipeptides and also shows average improvements for the ESS. We will include these findings in the final version. For more detailed results, see table below:

[2] Vaitl, Lorenz, et al. "Path-gradient estimators for continuous normalizing flows." International conference on machine learning. ICML, 2022.

[3] Bacchio, Simone, et al. "Learning trivializing gradient flows for lattice gauge theories." Physical Review D 107.5 (2023): L051504

[4] Abbott, Ryan, et al. "Applications of flow models to the generation of correlated lattice QCD ensembles." Physical Review D 109.9 (2024): 094514.

[5] Abbott, Ryan, et al. "Normalizing flows for lattice gauge theory in arbitrary space-time dimension." arXiv preprint arXiv:2305.02402 (2023).

[6] L. Vaitl “Path Gradient Estimators for Normalizing Flows", PhD Thesis

[7] L. Vaitl et al., “Fast and unified path gradient estimators for normalizing flows”, ICLR 2024