Path Gradients after Flow Matching
Authors: Lorenz Vaitl
Leon Klein
Venue: NeurIPS 2025 poster ()
Keywords: Path Gradients
Boltzmann Generators
Continuous Normalizing Flows
Equivariance
Optimal Transport
Flow Matching
Variational Inference
Molecular Dynamics
Primary area: machine_learning_for_sciences
TL;DR: We investigate the benefits of using path gradients to fine-tune CNFs initially trained by Flow Matching, in the setting where a target energy is known.

Abstract: Boltzmann Generators have emerged as a promising machine learning tool for generating samples from equilibrium distributions of molecular systems using Normalizing Flows and importance weighting.
Recently, Flow Matching has helped speed up Continuous Normalizing Flows (CNFs), scale them to more complex molecular systems, and minimize the length of the flow integration trajectories.
 We investigate the benefits of using path gradients to fine-tune CNFs  initially trained by Flow Matching, in the setting where a target energy is known. 
 Our experiments show that this hybrid approach yields up to a threefold increase in sampling efficiency for molecular systems,
 all while using the same model, a similar computational budget and without the need for additional sampling.
 Furthermore, by measuring the length of the flow trajectories 
 during fine-tuning, we show that path gradients largely preserve the learned structure of the flow.

Path Gradients after Flow Matching
Lorenz Vaitl earenz.vait@@out//e
Leon Klein Freie Universität Berlin leon. www.kkeeenfffuaa de
Abstract
Boltzmann Generators have emerged promising machine learning tool for generating samples from equilibrium distributions molecular systems using Normalizing Flows and importance weighting. Recently, Flow Matching has helped speed up Continuous Normalizing Flows (CNFs), scale them to more complex molecular systems, and minimize the length of the flow integration trajectories. We investigate the benefits of using Path Gradients to fine-tune CNFs initially trained by Flow Matching, in setting where the target energy is known. Our experiments show that this hybrid approach yields up to a threefold increase sampling efficiency for molecular systems, all while using the same model, similar computational budget and without the need for additional sampling. Furthermore, by measuring the length of the flow trajectories during fine-tuning, we show that Path Gradients largely preserve the learned structure of the flow.
Introduction
Generative models, ranging from GANs (Goodfellow al. 2014) and VAEs (Kingma Welling 2014) to Normalizing Flows (Rezende & Mohamed 2015 Papamakarios et al. 2019) and Diffusion Models (Ho et al., [2000] Song al. 2001I) have advanced rapidly in recent years, driving progress both in media generation and in scientific applications such as simulation based inference. While scientific workflows often incorporate sppeiii symmetries, they tend to under-exploit crucial resource: the unnormalized target density.
Boltzmann Generators (Noé al. 2019)) are typically trained either via self-sampling (Boyda et al.) 2021 Nicoli et al. 2020 Invernizzi et al. 2022; Midgley el al. 2023), leveraging gradients from the target distribution, using samples from the target without incorporating gradient information (Nicoli et al.
2023 Klein et al. 2023b Draxler et al. 2024| Klein & Noë 2024). However, these approaches each ignore complementary parts of the training signal: either the data its local geometry. Notably, first-order information evaluated at target samples remains underused, despite its potential to improve training. In this work, we close this gap by fine-tuning Continuous Normalizing Flows, initially trained with Flow Matching, using Path Gradients (Roeder et al. [017]] Vaitl (024) on samples from the target distribution. Furthermore, our approach requires computing target gradients only once training sample, avoiding the potentially high cost of repeatedly computing gradients on newly generated samples for self-sampling.
Flow Matching, a method for training CNFs, is based on target samples. It is closely related to diffusion model training and has recently gained traction for its simulation-free training and strong empirical performance, both in generative modeling benchmarks (Lipman et al. 2023 Esser et al. 2024:Jin al., (2024) and in scientific domains (Stark al.) 20224] Jing et al. 2024] Klein Noč 2024). Here, we investigate how incorporating Path Gradients on samples from the target distribution, enhances CNF performance post flow-matching.
Path gradients are low variance gradient estimators, which have strong theoretical guarantees close to the optimum (Roeder al. 2017) and incorporate gradient information from both, the variational
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

and the target distribution Vait et al. 2024). While they have been adopted for training Normalizing
Flows in the field of Lattice Gauge Theory [acchho 2023 Kanwar 2024| Abbott al., 2023
and also variational inference (Agrawal & Domke 2024 Andrade 2023). they remain underused in Binchemistry. In this work, we explore the potential Path Gradient fine-tuning for Boltzmann Generators with promising results. This indicates that even though training with Path Gradients is orders magnitude slower per iteration, its use of additional information and low variance allows to outperform Flow Matching within the same computational constraints in fine-tuning.
We make the following main contributions:
We propose hybrid training using Flow Matching for pre- pre-training and Path Gradients for fine-tuning. We do not require additional samples beyond the original force-labeled data, which were required already for generating the data. In order to ensure fast training, we are the first to optimize CNFs using Path Gradients for the forward KI. and making use of the augmented adjoint method (Vaitl et al. 2022b).
We show for many particle systems as well as alamine dipeptide that this fine-tuning approach can triple importance-samplinn efficiency within the same computational budget.
We investigate Path Gradients' impact on properties trained during Flow Matching, namely flow trajectory length and the MSE Flow Matching loss, and show that these are mostly unaffected by the fine-tuning.
0.2
Flow Matching
Path Gradlents
Hybrid
Start Fine-tuning
15.0
30.0
60.0
Time [s]
Target
Flow Matching 0.067
Path gradients 0.021
Hybrid 0.001
Figure 1: Training a CNF simple 2D Gaussian Mixture Model. Comparison between pure training with Flow Matching, pre-training with Flow Matching and fine-tuning with Path Gradients and pure training with Path Gradients. We that given the same wall-time hybrid training performs best in terms of forward KL divergence. The bottom row shows the target and the final model after training.
2 Method
Generative models have recently garnered significant interest for biochemical applications (Noé et al. 2019 Jumper et al. 2021 Abramson al. (024). In this work, focus the application of enhancing or replacing molecular dynamics simulations.
2.1
Boltzmann Generators
Boltzmann Generators (ooé et al. 2019) aim asymptotically generate unbiased samples from the cquilibrium Boltzmann distribution
1 p exp(-U(())
(1)
Z
with an energy (xx and an unknown normalization factor Z exxp(-U(=)ddrr The Boltzmann distribution describes the probability of states in physical system at thermal equilibrium. Traditional

---

sampling methods, such as Markov Chain Monte Carlo (MCMC) and Molecular Dynamics (MD) simulations, generate samples sequentially. In contrast, Boltzmann Generators are designed to produce independent samples directly from the target distribution. There are many different instances
of Boltzmann Generators, many focus on molecular systems (Dibak et al. [022] Köhler et al. [00T]] Midgley et al.] [023 Ding & Zhang 202 Kim et al. 2024gRizzz et al. 2023] [1mmggnoe et al.
2024
Schönle et al. 2024 Klein & Noč 2024]] Tan et al. 2025). while others are trained to sample
Tan
lattice and other many particle systems (Wirssberger et al. 20200 Ahmad & Cail 2022 Nicoli et al.
2020 2021 Schebek et al. 2024 Abbott et al.] 2023; Kanwar 2024).
Boltzmann Generators employ Normalizing Flow to map samples from a simple base distribution into learned sampling distribution qe. which is trained to approximate the target Boltzmann distribution. We can use Boltzmann Generators to obtain asymptotically unbiased estimators for observables over p(ar) using the sampling distribution go and (self-normalized) importance sampling (Noé et al. 2019: Nicoli et al. 2020).
p(x)
Ep(e) [O(x) = Exo(x)
O
go(x))
for observable O I2 + R. For self-normalized importance sampling, we still have to estimate the normalization constant
p(x)
Z 2= exp(-U = (x)
go (2)
with an MC estimator. This separates Boltzmann Generators from related generative methods, which only generate approximate samples from the Boltzmann distribution (Jing et al. 2022 Abdin & Kim 2023 Klein et al., 2023a Schreiner et al., 2023 Jing et al.) [2024] Lewis et al. 2024).
We can estimate the efficiency of importance sampling using the (relative) effective sampling size (ESS). The ESS
2 ()
L (x) qo
-(Ep
P (x) (x)) 1) qo
ESS == Eggo
(4)
can be estimated either on samples from the model or the target distribution Roughly speaking, able to sample from the target distribution. The estimator
the ESS tells us how efficiently ESS,, based on samples from qo. might fail to detect missing modes in the target distribution, while ESS, requires samples from the target distribution. Both estimators exhibit high variance and might overestimate the performance too few samples are being used, For more information confer (Nicoli et al. 2023)
2.2
Continuous Normalizing Flows
Neural ODEs, introduced by Chen et al. 2018). parameterize a continuous-tima transformation through www.lleettwwedddeeee
vector field By employing the adjoint sensitivity method (Pontryagin] 1987), gradients are computed in backpropagation-likk fashion with constant memory. By computing the ODE of the divergence, we can compute the change in probability, which lets us use Neural ODE as a Normalizing Flow (Chen et al. 20188 Grathwohl et al. 2019).
Using a simple base distribution go (zo) and the transformation
To (ma) = To f' vs (Iz, t)dt,
(5)
we obtain distribution go which we can sample from and compute the probability
(oTo(ro)
log go (To(2a)) = log 40 (2ro - -log
Izo
log go (2o) f Tr
(Dvoler,t) dt.
(6)
Oit.
For training samples from the target distribution, a straightforward way of approximating the target density is by maximum likelihood training, i.e. by minimizing
LML
- Epps) [log (x:))
(7)
3

---

Minimizing Equation 7 is equivalent to minimizing the forward KL divergence up to a constant
KT(plgo) = ptt) p(z) log go))) (0).
(8)
As Köhler et al. (2020) showed, we can construct equivariant CNFs by using an equivariant network as the vector field.
2.3 Flow Matching
Inspired by advances in diffusion models, Flow Matching (Lipman et al. 2023 Albergo et al. 2023
Liu et al. 2022) has emerged as a promising competitor to diffusion models (Esser et al. 2024). Flow Matching enables us to train CNFs in a simulation free manner, i.c. training without the need of the expensive integration of the ODEs (5) and (
The idea behind Flow Matching is to posit a base probability Po and vector field We (x()., which generates Pt from the base density Po. such that the final probability pi equals the target density p. Given these components, Flow Matching aims to minimize the mean squared error between the target vector field ue (2) and the learned vector field vg (2)
(f) = =x-220( [11os(x us (x))1
since pt and i, are not known, it is intractable to compute the objective. Nevertheless, Lipman et al. (2023) showed that we can construct a conditional probability path p(zl22) and corresponding vector field (-12) conditioned on 2 (zo, I),, which has identical gradients to Eq. (9). This yields the conditional Flow Matching objective
LCFM (0) = E0- (1), t) - un (x|=)]1a]
(10)
In this framework, there are different ways of constructing the conditional probability path. We here focus on the parametrization introduced in Tong al. (003). which results in optimal integration paths
= (Ir, 21) and p(2) ス (x),
(11)
ue (x(22) σ zo and Pa(x12) Nallt (1 (1-t) =0,00))
(12)
where T(2g, 21) denotes the optimal transport map between Po and Pl. As this map is again intractable, it is only evaluated per batch during training. We refer to this parametrization as OT FM. In Klein al. (2023b): Song et al. (2023) the authors show that we can extend this to highly symmetric systems, such as many particle e systems and molecules, and still obtain optimal transport paths. The approach enforces the system's symmetries in the OT map by computing the Wasserstein distance over optimally aligned samples with respect to those symmetries. Hereafter, we denote this parameterization as EQ OT FM. For more details on optimal transport Flow Matching refer to Tong al. (2023) and Klein et al. (2023b); Song et al. (2023). 3. Note, that we obtain the original formulation in Lipman et al. (2023) by using p(= = po (To)) (I1). we will refer to as standard FM.
2.4 Path Gradients
Introduced in the context variational auto-encoders (Roeder et al. 2017; Tucker et al. 2019),
Path Gradient estimators have improved performance, especially when applied to normalizing flows (Agrawal et al. 2020; Vaitl et al., 2022a|b Agrawal & Domke 2024). In the field of simulating
Lattice Field Theories, they have become go-to tool for training flows (Bacchio et al., 2023 Kanwar
2024: Abbott et al.) 2023). In these applications, the reverse Kullback-Leibie divergence K T.
minimized by self-sampling, i.e. without existing samples from the target density, but only
samples from fo generated while training.
Path gradients are unbiased gradient estimators which have low variance close to the optimum. We obtain Path Gradients by separating the total derivative of reparametrized gradient for the reverse KL
d
d log
4a (Te())))
K L(gol)) = (aro
(13)
dB
d8
p(To (2x))

---

into two partial derivatives
① 20 (x(x))
i log qo (11)
(xo(
log
(14)
p(x1)
00
00
11 To (200
Path Gradient L(oll)
Score
called the Path Gradient and the score term
Roeder et al. (017)) observed, that the martingale score term vanishes in expectation, but has non-zero variance. Vaitl et al. (2022a) showed that its variance is N la for batch size N and the Fisher Information Ly of the variational distribution qo- The low variance close to the optimum allows Path Gradient estimators to stick-the-landing" (Roeder et al. (017). i.c. having zero variance at the optimum, making it an ideal tool for fine-tuning, Vaitl al. (2024) showed that Path Gradient estimators incorporate additional information about the derivative of both target and variational distribution, opposed to the reparametrized gradient gradient estimator (Kingma & Welling 2014 Mohamed et al., 2020) (see Appendix [A.11.] for a detailed ex explanation). They hypothesized that this leads to less overfitting and generally better fit to the target density (see Vaitt (2024)). These theoretical aspects make Path Gradients a promising suitor for fine-tuning Boltzmann Generators.
Path Gradients given samples Recently, Path Gradient estimators have been proposed for samples from distributions different from the variational density go (Bauer & Mnih (02I). By viewing the forward KL(p1|90 (gg) at t 1 as reverse (L( qo) at = 0,
straightforwardly apply Path Gradients (Vaitl et al. 2024). Here, we assume Po,a to be a (Continuous) Normalizing Flow from P1 as its base density and T to be the parametrized diffeomorphism.
K L(p 46) = (pr( (logp1 (x1) - log (o()))
07 (311)
log)
log (T (x1))
log det.
(15)
Ox
= =log (121)
9 (x1)
Epi()))
[ogpi(z1) log det
loog 90 (o (a1))
Ox
= log Po. (sto)
po (mo)
po.) log loggo(zo)) K (po,ol90)
(16)
This different view lets us compute Path Gradients with the same ease as for the reverse KL, but on given samples from the target distribution.
While Vaiil al. (2022a) and the present work both use Path Gradients, both minimize different losses on different samples. Our work optimizes the Forward KL(p1lgg) with path gradients on existing samples from p. as proposed in Vaitl al. (2024) using the algorithm proposed in Vaitl et al.) (2022a). Vaiil et al. (202aa) minimize the reverse KL(ge via self-sampling, i.c. on samples from the model qo- Training entirely missed, which invalidates all asymptotic guarantees and breaks importance sampling, see e.g. Nicoli al. (2023). This becomes increasingly likely in higher dimensions. Further, the forces for the newly generated samples have to be evaluated, which can increase the cost of training or lead to unstable optimization behavior.
self- samplinm has many drawbacks, mainly: modes of the target p can be
Since the publication Vaidl et al. 2022b), Flow Matching has been established as the de facto training method for CNFss Our work investigates and combines the performance of Flow Matching and Path Gradients. Specifically, we investigate the effect of Path Gradients on the inner workings of the CNF, like e.g. trajectory length, while Vaitl et al. (2022b) simply looked the performance compared to standard self-sampling losses.
In Appendix A. Gradients exhibit zero variance at the optimum, generically, Flow Matching does not. We further
zoom in the properties of the different estimators. We show that while Path
'The score term @ log (x)) is not to be confused with the force term
8 log uatar) which is often called the de
score in the context of diffusion models (Song et al. 202T).
5

---

3.00
Flow Matching
Path Gradients
Hybrid
Start Fine-tuning
2.75
MSE loss
2.50
2.25
2.00
1.75
0.0
15.0
30.0
45.0
60.0
Time [s]
Figure 2: FM loss objective 10 during training on 2D GMM averaged on three runs. Training with Path gradients leaves the MSE largely unchanged.
explicate the theoretical properties of the existing estimators for training via Maximum Likelihood, Path Gradients and Flow Matching.
While the naive calculation Path Gradients requires significantly more compute and memory than Flow Matching, take several steps to obtain constant memory and speed ups in the next section. As we show experimentally, fine-tuning with Path Gradient for a few epochs significantly improves the performance for Continuous Normalizing Flows compared to only using Flow Matching.
3 Path Gradients and Flow Matching
To build an intuition for the behavior of Flow Matching (FM) versus Path Gradient (PG) estimators, we start with a simple 2D Gaussian Mixture Model. We compare three training strategies: 1) FM only, 2) PG only, and 3) a hybrid approach, using FM to quickly approximate the target distribution P1, and subsequently applying PG for fine-tuning. Although slower per training step, PG has strong theoretical guarantees near the optimum.
For Flow Matching, we use the standard formulation proposed in Lipman et al. (023) The dynamics U9 is modeled using a four-layer fully connected neural network with 64 units per layer and ELU activations. Given the simplicity of the model and task, we ignore memory usage in this setup. All experiments are run on CPU and complete in roughly one minute.
As shown Figure FM training rapidly improves initially but soon reaches a plateau with slow improvements after. PG training in contrast, progresses more slowly but eventually reaches and surpasses the FM plateau. The hybrid strategy, i.c. beginning with FM and switching to PG, results in the best performance, combining fast convergence with improved final accuracy. Interestingly. as shown in Figure 2, applying PG after FM has little effect on the MSE loss defined in Equation (10), suggesting that PG fine-tunes the variational distribution go without significantly altering the dynamics vg. Nonetheless, this subtle adjustment leads visibly better samples and improved density matching (cf. Fi)
In the 2D case, we compute the divergence term Equation (6) exactly by directly evaluating the trace of the Jacobian. However, this approach scales quadratically with the number of dimensions and quickly becomes computationally prohibitive in nigher-dimensione settings. A practical solution is Hutchinson's trace estimator (Hutchinson 1989), which provides faster but noisy estimators for the divergence.
Kutta (RK4) scheme with 15 time steps, resulting in 60 function evaluations per integration trajectory. Under these settings, training on a single batch is empirically about 275 times faster with Flow Matching compared to Path Gradients. This discrepancy expected, as the cost of function evaluations differs substantially between the two methods.
For ODE integration, we fourth-order
While these results are promising, the 2D setup allows exact computation of the divergence and may not generalize to higher-dimensional more complex tasks. Moreover, memory usage becomes significant bottleneck in these more realistic scenarios.

---

Augmented Adjoint Method In order to have constant memory irrespective of the number of function evaluations, we adapt the augmented adjoint method (Vaitl al., 2022b) for Path Gradient estimators for the forward KI. divergence Equattin T6)
K (P = Exirp1 1p1 d
d
a arr
log
Pu.0 (aro)
0T' (T1)
(17)
40
00
In order to compute the force of the variational distribution
i log Po.a (zoo)
(18)
we use the augmented adjoint method by solving the ODE
Olog (r.(
alog pt.
(vo(1).))
(Ovo(ze;1))
Tr
(19)
dt dr,
dr,
dr.
dz
Oxe
with initial condition
a Plg
2 Olog (x1)
(20)
Ox1
Oz1
from t to (Vaitl et al., 2022b)
Computing Path Gradients requires computing the gradient of the divergence
@vo (12)))))
(21)
Ox.
as well as the ODEs 5 and 19 per integration step. Even for training on single time-step, this is more resource intensive than Flow Matching, which only requires the derivative of the vector field @va(w..()
80 Thus, although the memory demands do not increase with additional integration steps of the ODE, the memory required for Path Gradients greater than with Flow Matching. Since the computational load associated with calculating the terms varies across different architectures, we discuss them individually for every experiment.
For RK4 with 15 integration steps, the ODEs 5 6 and 19 are evaluated in the backward direction (L: 1 + 0) and the ODE and the adjoint method in the forward direction, amounting 300 function evaluations, compared to two (the forward and backward call) with Flow Matching. Thus we expect Flow Matching to be faster than Path Gradients by a factor of least 1500.
Experiments
Motivated by the improvements the toy model, now turn to challenging problems, namely the experiments done in Klein et al. (2023b) and Klein & Noč (2024). To this end, we use the same architecture, a CNF with EGNN (E(n) Equivariant Graph Neural Network) (Satorras et al. 1021 Garcia Satorras et al. 2021) for the vector field. The architecture in Klein & Noë (2024) is mostly the same in Klein 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 & Noč (224)) encodes nearly all atoms differently. Only hydrogens bound to the same atom are treated as indistinguishable. For more details on the model architecture see Appendix A.2 While maintain the architecture and the initial training process with Flow Matching, we change the later training procedure to include Path Gradients. As in Klein et al. (2023b), we investigate optimal transport Flow Matching (OT FM),- and for the second set of experiments on LJ13 equivariant optimal transport Flow Matching (EQ OT FM), and naive Flow Matching (standard FM).
We first evaluate the models on two many particle system with pair-wise Lennard-Jonel interactions with 13 and 55 particles. In addition we investigate Alanine dipeptide (AD2) at T 300K both with classical and cemi-eapirical force field (XTB). In contrast (Klein et al.) 2023b) Klein Noë 2024), we do not bias the training data towards the positive yp state. The bias was originally introduced to make learning problem simpler, as the unlikely state is more prominent (Klein et al. 2023b). This, however, skews the evaluation, since the metrics, namely ESS and Negative log likelihood (NLL), are based on the target distribution, not the biased one. Moreover, it assumes system-specific knowledge of slow- aloyivarying degrees of freedom, information that is typically unavailable for unknown systems.
7

---

ESS-Q
NLL
Trajectory Length
EQ OT
6.8
11.05
OT
-16.10
Standard
0.7
-1615
2.5
0.6
6
Epoch
Epoch
Epoch
Figure 3: Reverse ESS, NLL and trajectory length for Flows trained with standard FM, Optimal Transport and Equivariant Transport during fine-tuning with Path Gradients on LJ13, We can observe that fine-tuning largely leaves the trajectory length unchanged, while substantially improving performance, Mean + sterr over three runs.
Finally, we also investigate the 2AA dataset consisting of dipeptides simulated at T 310K with classical force field (Klein et al.) (222a)) to evaluate the influence of PG in a transferability. For more details the datasets see Appendix A.33 We published code for replicating our experiments In total, we investigate three scenarios:
First, fine-tuning with Path Gradients compared to fine-tuning with OT Flow Matching given the same limited computational resources in Section 4.1
Second, apply Path Gradients on the FM-trained models using unlimited resources to maximize performance after training with Standard FM, OT FM and EQ OT FM as done in Klein et al. (203bb), Section 4.2 In this second case, additionally examine the ODE integration length to understand how PG fine- mne-tuning influences the model.
Finally, we test if the path gradient fine-tuning improves results for transferability in the setting by investigating transferable Boltzmann Generators (Klein & Noë 2024) trained and evaluated on dipeptides, see Section 4.3
4.1 Fine-tuning using the same limited resources
We enforce comparable memory and wall-timm constraints those used for Flow Matching alone by adjusting batch size and employing gradient accumulation. Consequently, our fine-tuning runs with Path Gradients complete less wall-time and occupy similar memory footprint. For full training statistics and implementation details, see Appendix A.44and and Appendix AA.5
Replicating the experiments Klein et al. (2023b), we evaluate how path-gradient fine-tuning impacts model performance. Table presents comparison between models trained with pure OT Flow Matching and those optimized using the hybrid approach of OT FM pre-training and path gradient fine-tuning.
We observe that fine-tuning with Path Gradients improves the ESS metrics, i.c. ESS, and ESS, across most datasets and models, despite their high variance. The hybrid approach reliably improves the NLL for all models and datasets, but one. The only exception is the standard Flow (Klein et al.) 2023b) on AD2 with XTB, where fine-tuning with FM still performs better average. Examining the ESS, we find that both approaches fail to adequately fit the target density. This supports the hypothesis that Path Gradients only effective for fine-tuning when the flow already provides reasonably good approximation of the target. Notably, on LJ55, our hybrid approach nearly triples the ESS, and on AD2 with XTB, doubles in the same training time.
4.2 Effect of Path Gradient fine-tuning on flow trajectory length
We use the provided saved models from Klein et al. 2023b) for IJ13 and investigate the performance and flowed trajectory length during fine-tuning with Path Gradients. We used a batch-size of 256 and trained with Path Gradients for 10 epochs. In Figure 3 we can see that the different pre-trained models have similar NL.Ls and ESS, but differ trajectory lengths. Fine-tuning reliably improves the
wgithub.com/eenz/000/ppoooo hwwppppathhggaaaatraraaaee after---

---

Table 1:
Comparison between fine-tuning with Optimal Transport Flow Matching and Path Gradients First all models were pre-trained with Optimal Transport Flow Matching like (Klein al. 2023b).
We compare fine-tuning with Flow Matching and Path Gradients. For all experiments, we limited the VRAM and runtime be roughly equivalent. Mean + sterr on three runs. bold: sterrs do not overlap.
Model
Training type
NLL (4)
ESS, %(()
ESS, %(()
LJ13
Only FM
16.09 + 0.03
5.43 82.97
54.36
58.18 + 0.71 82.87 + 0.35
Klein al.
2023b
Hybrid (ours)
16.21 + 0.00
0.40
LJ55
Only FM
88.45 + 0.04
3.74 1.06
2.97 + 0.08
Klein al.
2023b
Hybrid (ours)
89.19 + 0.05
11.04 3.98
13.71 + 3.15
Alanine dipeptide XTB
Only FM
107.89 + 0.07
0.74 = 0.33
0
Klein al.
2023b
Hybrid (ours)
- 107.77 0.18
0.25 + 0.37
0.01
Only FM
-T25.75 - 0.01
3.38
0.50
2.65
KKeein & Noč 2024
Hybrid (ours)
125.82 + 0.02
7.30
1.28
6.59
Alanine dipeptide
Classical
Only FM
-100,14 0.01
4.88 + 0.42
0.26 + 0.19
Klein al.
2023b
Hybrid (ours)
-100.30 0.13
2.86 = 3.05
0.04 004100000 0.06
Only FM
T28.01 0.01
14.42
12.21 1.06
Klein & Noć
2024
Hybrid (ours)
128 26 + 0.02
24.39 6.86
13.89 + 4.91
simulation
Matching
Hybrid
dipeptide Energy
simulation Flow Matching Hybrid
Free energy Het
0.02
-40
-20
20
Energy ksT
Figure 4: Alanine dipeptide results for the TBG model and the classical force field with and without Path Gradient finetuning. Left: Ramachandran plots for the dihedral angel distribution of a reference MD simulation and non reweighted samples from the different TBG models. Right: Corresponding energy distributions of generated samples.
performance of the flows with relative little change the trajectory lengths. For the full statistics Appendix A.7 We observe a similar pattern in the other systems investigated. For alanine dipeptide, for instance, generated samples show significantly improved bond-length and energy distributions, yet the global conformational landscape, captured by the yp and v dihedral-angle distributions in the Ramachandran plot, remains mostly unchanged (see Figure And and Figure
These experiments show that our proposed hybrid approach perfectly suited for maximizing performance while keeping the properties of the model.
4.3 Transferability on dipeptides
Finally, we also investigate whether fine-tuning with Path Gradients improves performance in the transferable setting. To this end we fine-tune Transferable Boltzmann Generators (TBG) (Klein & Noč 2024), which were trained on 200 different dipeptides and evaluate on unseen ones, like in Klein & Noē (2024). The experiments show that fine-tuning with path gradients improves the NLL. and energies for all evaluated unseen test dipeptides. On average, we observe a relative improvement of around 23% in the ESS,, reaching efficiency of 9.79%. For more details see Appendix A.10

---

Discussion
We have presented hybrid training strategy for Boltzmann Generators that uses Flow Matching and Path Gradients. We have shown how to efficiently compute Path Gradients with constant memory and without the need for additional samples. While substantially slower per training step, Path Gradients are a powerful tool for fine-tuning, leveraging first-order information from the energy function an underexplored avenue the scientifie machine learning community. Our results demonstrate that Path Gradients can significantly improve sample quality on the same computational budget as Flow Matching, when the model already reasonably close to the target distribution. Our experiments show that Path Gradients only apply minor changes to the flow trajectory and to the variational distribution qo while still substantially increasing the sampling efficiency and performance.
5.1 Limitations
Path Gradients come with several limitations. First, they require access to a well-defined and differentiable energy function, which restricts their use to domains like molecular modeling and excludes standard tasks such natural image generation. Second, they rely on unbiased training samples. Finally, the method computationally more expensive and tends improve performance only when the model is already close to the target distribution. While Path Gradients do not directly speed up the expensive CNF sampling process, they increase the ESS, thereby reducing the total number of samples needed and, hence, at least partially alleviating the high sampling cost.
5.2 Future Work
Given the similarities between Flow Matching and Diffusion models, extending Path Gradients to diffusion-based frameworks presents exciting and promising direction. The dynamics used in our experiments have also been used for molecular conformation generation using Diffusion models (Hoogeboom et al.) (002). Adoption to their applications would be interesting and suitable candidate for future work.
Furthermore, in model distillation (e.g. (Salimans & Ho 2022 the gradient information of the larger model available, even though there might be no first order information of the original data. Here approach could help to improve the distillation process.
5.3 Broader Impact
This foundational research has no immediate societal impact, but if scalable, could be used to accelerate drug and materials discovery by replacing MD simulations. Potential risks include misuse for biothreat development and the lack convergence guarantees, which may lead to incomplete sampling and misleading conclusions.
Acknowledgements
LV wants to thank Pan Kessel and Shinichi Nakajima for the helpful discussions. LK gratefully acknowledges support by BIFOLD Berlin Institute for the Foundations of Learning and Data.
10

---

References
Abbott, R., Albergo, M., Botev, A., Boyda, D., Cranmer, K., Hackett, D., Kanwar, G., Matthews, A., Racaniere, S., Razavi, A., al. Normalizing flows for lattice gauge theory in arbitrary space-time dimension. ArXiv preprint, abs/2305. 12345, 2023.
Abdin, O. and Kim, P. M. Pepflow: direct conformational sampling from peptide energy landscapes diffusion. bioRxiv, pp. 2023-06, 2023.
through sypernetwork-condttooo
Abramson, J., Adler, J., Dunger, J., Evans, R., Green, T., Pritzel, A., Ronneberger, O., Willmore, L., Ballard, A. J., Bambrick, J., al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, pp. 1-3, 2024.
Agrawal, A. and Domke, J. Disentangling impact of capacity, objective, batchsize, estimator,, and step-size on flow vi. arXiv preprint arXiv:2412.088 2.88844, 2024.
Agrawal, A., Sheldon, D. R., and Domke, J. Advances in Black-Box VI: Normalizing Flows, Importance Weighting, and Optimization. In Advances in Neural Information Pro- cessing Systems 33, 2020. URL https: httcpps:/pproeeedingg htwpeeeediinnnnnuuisss.... wwtp:eeee.p.p..ppppeeeeeeeeeeeee http://////./app////////////////
Abstract.
Ahmad, R. and Cai, W. Free energy calculation of crystalline solids using normalizing flows. Modelling and Simulation Materials Science and Engineering, 30(6))065007, sep 2022. ISSN 0965-0393, 1361-651X. doi: 10. org/article/10 1018 htt.://.00100111111611//////1//1
https::////ooeeceeeee iop
Albergo, M. S., Boffi, N. M., and Vanden-Eijnden, E. Stochastic interpolants: A unifyynn framework for flows and diffusions. ArXiv preprint, abs/2303.08797, 2023, URL https://arxiv, org/ abs/ 08797
Andrade, D. LOFT-Stable Training of Normalizing Flows for Variational Inference. In 6th Workshop on Tractable Probabilistic Modeling, 2023.
Bacchio, S., Kessel, P., Schaefer, S., and Vaitl, I.. Learning trivializing gradient flows for lattice gauge theories. Physical Review D, 107:L051504 2023, doi: www.:///Pyyss....17....... 1504. URL httph:tt//// /tttps://link... aps wwwp//ggggoooo/ooooo htt///ooo111000 1ww.:1113/Phhyeee
107 L051504
Bannwarth, C., Ehlert, S., and Grimme, S. Gfn2-xtb-an an accurate and broadly parametrized self- consistent tight-binding quantum chemical method with multipole electrostatics and density- dependent dispersion contributions. Journal of Chemical Theory and Computation, 1653):1652- 1671, 2019, doi: 10.1021/acs http///cccccc..0111111 176. URL https: http://ddooorrg/o///o/ 10 1020/a0c.s....ee 8b01176 PMID: 30741547.
Bauer, M. and Mnih, A. Generalized Doubly Reparameterized Gradient Estimators. In Proceedings of the 38th International Conference Machine Learning, 2021. http:///rprooeeesnngg html
mlr.
Boyda, D., Kanwar, G., Racanière, S., Rezende, J., Albergo, M. S., Cranmer, K., Hackett, D. C., and Shanahan, P. Sampling using SU (N) gauge equivariant flows. Physical Review D, (7), 2021.
Chen, T. Q., Rubanova, Y., Bettencourt, J., and Duvenaud, D. Neural Ordinary Differential Equations. In Advances in Neural Information Processing Systems 31, 2018.
Dibak, M., Klein, L., Krämer, A., and Noé. Temperature steerable flows and Boltzmann generators. Phys. Rev. Res., 4:L042005, Oct 2022. doi: 10. Ro..o0/ L.0420000.....
Ding, X. and Zhang, B. Computing absolute free energy with deep generative models. Biophysical Journal, 20(3):1954 2021a.
Ding, X. and Zhang, B. Deepbar: A fast and exact method for binding free energy computation. Journal of Physical Chemistry Letters, 12:2509-2515, 2021b. ISSN 19487185. doi: 10. 10.1/s//cs. ipclett.1c00 1c00189.

---

Draxler, F., Sorrenson, P.. Zimmermann, L.. Rousselot, A., and Köthe, U. Free-form flows: Make any architecture normalizing flow. In Dasgupta, S., Mandt, S., and Li, Y. (eds.), Proceedings of The 27th International Conference Proceedings of Machine Learning Research, pp. 2197-2205. PMLR, 02-04 May 2024. URL www.mmwlr...m.........mmm
Artificial Intelligence and Statistics, volume 238 of
https: fttp:://proceedingg mlr. wrw.srrsssse... v23ss//v2/////// wttpev/ae/leee22e2eeeeee html
Esser, P., Kulal, S., Blattmann, A., Entezari, R., Müller, J., Saini, H., Levi, Y., Lorenz, D., Sauer, A., Boesel, F., et al. Scaling rectified flow transformers for high- high-resolution image synthesis. In Forty-first international conference machine learning, 2024.
Garcia Satorras, V., Hoogeboom, E., Fuchs, F., Posner, L,, and Welling, M. E(n) equivariant normalizing flows. In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang. P., and Vaughan, J. W. (eds.). Advances in Neural Information Processing Systems, volume 34, pp. 4181
Curran Associates, Inc., 2021, URL /tt./ttpps: neurips.. www.urpp..cppppe/////0/// cc/paper/2021/filee pdf
Goodfellow, I. J., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley, D.,, Ozair, S., Courville, A. C., and Bengio, Y. Generative Adversarial Nets. In Advances in Neural Information Processing
Systems 27, 2014.
Grathwohl, W., Chen, R. T. Q., Bettencourt, J., Sutskever, and Duvenaud, D. FFJORD: Free-Form Continuous Dynamics for Scalable Reversible Generative Models. In Proceedings of the 7th International Conference on Learning Representations. 2019.
Ho, J., Jain, A., and Abbeel, Denoising Diffusion Probabilistic Models. In Advances in Neural Information Processing Systems 33, 2020.
Hoogeboom, E., Satorras, V. G., Vignac, C., and Welling, M. Equivariant diffusion for molecule generation in 3d. In International conference on machine learning. pp. 8867-8887. PMLR, 2022.
Hutchinson, M. F. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications in Statisiiex-Siulltton and Computation, 18(3): 1059-10699, 1989.
Invernizzi, M., Krämer, A., Clementi, C., and Noé, F. Skipping the replica exchange ladder with normalizing flows. The Journal of Physical Chemistry Letters, 13:11643-11649 2022.
Jin, Y., Sun, Z., Li, N., Xu, K., Jiang, H., Zhuang, N., Huang, Q., Song, Y., Mu, Y., and Lin, Z. Pyramidal flow matching for efficient video generative modeling. arXiv preprint 4rx0:22400.0555, 2024.
Jing, B., Corso, G., Chang, J., Barzilay, R., and Jaakkola, T. Torsional diffusion for molecular conformer generation. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D.,, Cho, K., and Oh, A. (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 24240-24253 Curran Associates, Inc., 2022. URL /ttphttss//prrreeiinnn.. www.wnenrii.... /ctpppps.c/c/pppr///////ee/ 39444442200000007777737777774883 ea94444ffee44444e
.dff
Jing, B., Berger, B., and Jaakkola, т. Alphafold meets flow matching for generating protein ensembles.
arXiv preprint 4rXiv22402.04845 2024.
Jumper. J., Evans, R., Pritzel, A., Green, T., Figurnov, M., Ronneberger, O., Tunyasuvunakool, K., Bates, R., Židek, A., Potapenko, A.. et al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583-58)) 2021.
Kanwar, G. Flow-based sampling for lattice field theories. arXiv preprint arXiv:2401.01299, 2024.
Kim, J. C., Bloore, D., Kapoor, K., Feng, J., Hao, M.-H., and Wang, M. Scalable normalizing flows enable boltzmann generators for macromolecules. International Conference on Learning Representations (ICLR), 2024.
Kingma, D. P. and Welling, M. Auto-Encoding Variational Bayes. In Proceedings of the 2nd International Conference on Learning Representations, 2014. URL http://aax/x/...rrttttt abs/
1312.613 6114
Klein, L. and Noć, F. Transferable boltzmann generators. arXiv preprint arXiv:2406.1442, 2024.
12

---

Klein, L.. Foong, A. Y. K., Fjelde, E., Mlodozeniec, B. K., Brockschmidt, M., Nowozin, S., Noe, F., and Tomioka, R. Timewarp: Transferable acceleration of molecular dynamics by learning time-coarsened dynamics. In Thirty-seventhh Conference on Neural Information Processing Systems, 2023a.
Klein, L.,, Krämer, A., and Noé, F. Equivariant flow matching. Advances in Neural Information Processing Systems, (6:59886-59910, 2023b.
Köhler, J., Klein, L.,, and Noé, F. Equivariant Flows: Exact Likelihood Generative Learning for Symmetric Densities, In Proceedings of the 37th International Conference on Machine Learning, 2020. UtL/tttp://pooeeeanne mlr. wttps//ww.....e.i//////ee... prwsss/111//ohherr/2 html
Köhler, J., Krämer, A., and Noé, F. Smooth Normalizing Flows. In Advances in Neural Information Processing Systems 34, 2021.
Larsen, A. H., Mortensen, J. J., Blomqvist, J., Castelli, E.,, Christensen, R., Dulak, M., Friis, J., Groves, M. N.. Hammer, B.. Hargus, C., Hermes, E. D., Jemnings, P. C.,, Jensen, P. B., Kermode, J., Kitchin, J. R., Kolsbjerg, E. L., Kubal, J., Kaasbjerg, K. Lysgaard, S., Maronsson, J. B., Maxson, T., Olsen, T., Pastewka, L., Peterson, A., Rostgaard, C., Schiotz, J., Schütt, O., Strange, M., Thygesen, K. S., Vegge, T., Vilhelmsen, L., Walter, M., Zeng, Z., and Jacobsen, K. W. The atomic simulation enviromment- a python library for working with atoms. Journal of Physics: Condensed Matter, 29(27):273002, 2017. URL http: /tstt://ssten //stacks lop. org/ 008///9984//222881/ 2-27/a=2700000
Lewis, S., Hempel, T., Jiménezz-Luna, J., Gastegger, M., Xie, Y., Foong, A. Y., Satorras, V. G., Abdin, O., Veeling, B. S., Zaporozhets, I., et al. Scalable emulation of protein equilibrium ensembles with generative deep learning. bioRxiv, pp. 2024-12, 2024.
Lipman, Y., Chen, R. T., Ben-Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling. In Proceedings of the I 11th International Conference on Learning Representations, 2023.
Liu, X., Gong, C., and Liu, Q. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprini arXiv:2209.03003, 2022.
Midgley, L. I., Stimper, V., Simm, N., Schölkopf, B., and Hernández-Lobato, M. Flow Annealed Importance Sampling Bootstrap. In Proceedings of the 11th International Conference on Learning Representations, 2023.
Mohamed, S., Rosca, M., Figurnov, M., and Mnih, A. Monte Carlo Gradient Estimation in Machine Learning. Journal of Machine Learning Research, 21:5183-5244, 2020. URL http://jmlr org org:pappers/221/19-- 346 html
Nicoli, A.. Nakajima, S., Strodthoff, N., Samek, W., Müller, K.-R., and Kessel, Asymptotically unbiased estimation of physical observables with neural samplers. Physical Review E, (2), 2020.
Nicoli, K. A., Anders, C. J., Funcke, L., Hartung, T., Jansen, K., Kessel, P., Nakajima, S., and Stornati, P. Estimation of Thermodynamic Observables in Lattice Field Theories with Deep Generative Models. Physical Review Letters, 126(3):032001, 2021.
Nicoli, K. A., Anders, C. J., Hartung, T., Jansen, K., Kessel, P., and Nakajima, S. Detecting and Mitigating Mode Collapse for Flow-based Sampling of Lattice Field Theories. Physical Review D. 108: 105:1, 2023, doi: 10. 103/PhysRevD.1108 14501. URL pottps://linkt. aps org doi/ 10. Re.00///hyeeeD..... 108. 114501
Noé, F., Olsson, S., Köhler, J., and Wu, H. Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 365(6457))eaaw) 1147, 2019.
Papamakarios, G., Nalisnick, F., Rezende, J., Mohamed, S., and I. Lakshminarayanan, B. Normaliz- http://wxww..rrgbb..a.s.a
ing flows for probabilistic modeling and inference. CoRR, 2019. 1912. 0276211)
Papamakarios, G., Nalisnick, E. T., Rezende, D. J., Mohamed, S., and Lakshminarayanan, B. Normalizing Flows for Probabilistic Modeling and Inference. Journal of Machine Learning Research, 22:57:1-57:64, 2021. URL / //mrr. org.oog///... httpppppppsp////pe 1028.html
13

---

Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Kopf. A., Yang, E., Vetito, Z., Raison, M., Tejani, A.. Chilamkurthy, S., Steiner, B., Fang. I.., Bai, J., and Chintala, S. Pytorch: An imperative style, high-performance deep learning library. In Advances Neural Information Processing Systems 32, pp. 4024-8035. Curran Associates, Inc., 2019. URL
meur pttps/s....pppp..ppppppe////// deep-leernninggee.eeeeeeeeeenn-e 1ibrary......
1a15-pytorcc-aae an- an-imperative aetb--ytord--a--mmeraiiiiii pdf
Poli, M., Massaroli, S., Yamashita, A., Asama, H., Park, J., and Ermon, S. Torchdyn: Implicit models and neural numerical methods in pytorch. In Neural Information Processing Systems, Workshop on Physical Reasoning and Inductive Biases for the Real World, volume 2021.
Pontryagin, L. S. Mathematical theory of optimal processes. Routledge, 1987.
Rezende, D. and Mohamed, Variational inference with normalizing flows. International conference on machine learning, pp. 1530-1538. PMLR, 2015.
Rizzi, A., Carloni, P., and Parrinello, M. Multimap targeted free energy estimation, 2023.
Roeder, G., Wu, Y., and Duvenaud, D. Sticking the Landing: Simple, Lower-Varianee Gradient Estimators for Variational Inference. In Advances in Neural Information Pro- cessing Systems 30, 2017. URL https: h/tppee:s//rooeeinggs www.ggg...riis... www.urps.ccpppperr//.aa/ pccppapper2017/hash (9116ffffar7/el15444444444444440
Abstract. html
Salimans, T. and Ho, J. Progressive distillation for fast sampling of diffusion models. arXiv preprint
arXiv:2202.00512. 2022.
Satorras, V. G., Hoogeboom, E., and Welling, M. E (n) equivariant graph neural networks. In International conference on machine learning, pp. 9323-9332. PMLR, 2021.
Schebek, M., Invernizzi, M., Noé, F., and Rogal, J. Efficient mapping of phase diagrams with conditional boltzmann generators. Machine Learning: Science and Technology. 2024.
Schönle, C., Gabrié, M., Lelièvre, T., and Stoltz, G. Sampling metastable systems using collective variables and jarzynski-crooks paths. arXiv preprint arXiv:2405.18160, 2024.
Schreiner, M., Winther, O., and Olsson, S. Implicit transfer operator learning: Multiple time- resolution models for molecular dynamics. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https: fttt:::///oenreviiete wttp://www.eeetetuornriiieee.eee iicrrummmm==dxdx> JiuA2
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. Score-Based Gener- ative Modeling through Stochastic Differential Equations. In Proceedings of the 9th International Conference on Learning Representationss 2021.
Song, Y., Gong, J., Xu, M., Cao, Z., Lan, Y., Ermon, S., Zhou, H., and Ma, W.-Y. Equivariant flow matching with hybrid probability transport for 3d molecule generation. In Thirry-seventh Conference on Neural Information Processing Systems, 2023. URL https::/ net/ forummidd
Stark. H., Jing, B., Barzilay, R., and Jaakkola, T. Harmonic ceff-conditioned flow matching for joint multi-ligand docking and binding site design. In Forty-first International Conference on Machine Learning, 2024.
Tamagnone, S., Laio, A., and Gabrié, M. Coarse-grained molecular dynamics with normalizing flows. Journal of Chemical Theory and Computation, 20 (1188)77796-8855 2024.
Tan, C. B., Bose, A. J., Lin, C., Klein, L.,, Bronstein, M. M., and Tong, A. Scalable equilibrium sampling with sequential boltzmann generators. arXiv preprint arXiv:2502. 18462, 2025.
Tong, A., Malkin, N., Huguet, G., Zhang, Y., Rector-Brooks, J., Fatras, K., Wolf, G., and Ben- gio, Y. Conditional flow matching: Simulation-free dynamic optimal transport. arXiv preprint 2302.03300 2(3), 2023.
14

---

Tucker, G., Lawson, D., Gu, S.. and Maddison, C. J. Doubly Reparameterized Gradient Estimators for Monte Carlo Objectives. In Proceedings of the 7th International Conference on Learning Representations, 2019.
Vaitl, L. Path Gradient Estimators for Normalizing Flows. PhD thesis, Technische Universität Berlin, 2024.
Vaitl, L., Nicoli, K. A., Nakajima, S., and Kessel, P. Gradients should stay on path: better estimators of the reverse-and forward KL divergence for normalizing flows. Machine Learning: Science and Technology, 3(4):045006, 2022a.
Vaitl, L., Nicoli, K. A., Nakajima, S., and Kessel, P. Path-Gradient Estimators for Continuous Normalizing Flows. In Proceedings of the 39th International Conference Machine Learning. mlr.
2022b. URL https://proceeddii http:///ia.......ppee/////////aa (tress//v122/aatt1222 html
Vaitl, I.., Winkler, L.., Richter, L.,, and Kessel, P. Fast and unified path gradient estimators for nor- malizing flows. In to be presented 12th International Conference on Learning Representations, 2024.
Wirnsbegee, P. Ballard, A. J., Papamakarios, G.. Abercrombie, S., Racanière, S., Pritzel, A., Jimenez Rezende, D., and Blundell, C. Targeted free energy estimation via learned mappings. The Journal of Chemical Physics, 153(14):1441 12, 2020.
15

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: In Section 3 we present how to efficiently do fine-tuning in constant memory
and Section we show both, improved performance and minimal changes to the trajectory length.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
In Section 5.1 we discuss the limitations of the presented approach.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g. if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
16

---

Answer: [Yes]
Justification: The cross-references for the formulae for Path Gradients on samples in Section 2.4 well as the adapted augmented adjoint state method Equation (9) cross- referenced and numbered. The same holds true for equations on EQ-OT- OT- and standard
FM in Section 2.3
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proof' provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide all the necessary information in Sections A.S
Further we recreate previously published experiments. Additionally we also provide code and checkpoints for the trained models.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurlPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the partieular way they provide for reproducibiliity In the case of closed-source models, it may be that access to the model is limited in
17

---

some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: We provide code and trained models in the supplementary material. We further describe where to download the datasets in Appendix A.3
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
guidddss/ggide////eeeeeoesssssee CoodeSummeiieeeee ssi onPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: Appendix A.3 and Appendix A.7 specify data splits, hyperparameters, the finding of hyperparameters and optimizer choice
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer:
[Yes]
Justification:
We present standard errors on three repetitions for our experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
18

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπτог rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: In Section 1 and the Appendix we specify the main CPU or GPU, the compute time as well the VRAM usage, where is relevant.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics /tt http://wwwpp..ccouuuui.i/ccctcaa
Answer: Yes
Justification: We conform the Code of Ethics and have communicated the impact of the work in Section 5.3
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., s if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We do this in Section
Guidelines:
The answer NA means that there is no societal impact of the work performed.
19

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Does not apply
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best. faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We mention the licenses for the data, code, and models in Appendix A.3 and Appendix A.13
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
20

---

If assets are released, the license, copyright information, and terms of use in the thcode has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssit
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
Code is documented.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Does not apply.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: Does not apply
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
21

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: Only used for writing/editing
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy bettts://neurre hctps////ooeeeecccc//2///////// for what should or should be described.
22

---

Flow Matching
Hybrid Path Grads
-0.064
-0.048
-0.032
-0.016
0.000
0.016
0.032
0.048
0.064
0.080
p. (log p log 98)
Figure 5: Loss in space after training as done in Figure
|
A
Appendix
A.1 2D GMM experiment
The 2D GMM is Gaussian Mixture model with 4 equally weighted Gaussians
N
c + 0.01
d 44000.0)
0
(22)
with a,, by. C, d; ~ (1 1). 2, 43. The training set consists of 2000 samples, the (Llg)))) and MSE loss are evaluated on 2048 samples. We used Adam with default parameters and Ir=le-2 for pure FM/PG and pre-training with FM and Ir=5e-3 for finetuning with PG. All experiments on an Intel 17-1165G7 CPU and ignoring the validation, training completes in roughly 45 seconds.
Figure 5 shows the weighted difference
p(x) log p(z) go (x)
(23)
after training.
A.2 Architecture
We use the same model architecture as introduced in Klein et al. (023b); Klein Noë (2024). We
here summarize it briefly, closely following the presentation in Klein & Noč (2024).
The underlying normalizing flow model for Boltzmann Generator is a CNE The corresponding vector field vp(L,x) is parametrized by an O(D)- and S(N) -equivarrri graph neural network (EGNN) Garcia Satorras et al. (2021) ); [Satorras et al. (021). The vector field vo (zr 4) consists of consecutive EGNN layers. The position of the i-th particle x; updated according to the following
23

---

set of equations:
N. - a.);
mly = de (1e 11 dtt)
(24)
(Ir x²)
att o +
dij + Dd (m:)))
(25)
jfi
h = on (k),
m4 Σ
Sm Emm(mi))ms
(26)
574
N
08 (2 ((c,,t)) t)2 = 27 2
2
[ -2))
(27)
where the o, represent different neural networks, d,y is the Euclidean distance between particle and I is the time, a; is embedding for each particle.
The proposed model architecture in Klein et al. (2023b) uses for alanine dipeptide distinct encodings a, for all backbone atoms and the atom types for all other atoms. In contrast, the model architecture in Klein & Noč (2024) uses distinct encodings for all atoms, except for Hydrogens bond to the same Carbon atom. We refer to this model as TBG, which stands for transferable Boltzmann Generator, even though we do not deploy it in transferable way this work.
For more details see Klein et al. (2023b) Klein & Noč 20024).
A.3 Datasets
Here we provide more details on the investigated datasets.
We use the same training and test splits as defined in Klein et al. (2023b); Klein & Noć (2024).
Lennard-Jones systems
The energy (xx for the Lennard-Jones systems is given by
U(x)
Σ
(28)
dig
where di is the distance between particle and j. The authors of Klein et al. (2023b)
(CC BY 4.0) made the datasets available here: httpss//osf 1o/ srgg7//view, only
the following number of test samples: LJ13: ESS;; 5 x 10;; ESS,, NLL: 5 x 10a.
For computing the metries
LJ55: E x 105; ESSp, NLL: 1 x 10
Alanine dipeptide The classical alanine dipeptide dataset was generated with an MD simulation, using the classical Amber ff99SBildn force-field 300K for implicit solvent for duration of 1 ms Köhler et al. (202T) with the openMM library. The datasets is available as part of the public bgmol www.///nogrou/bggnoo... (MIT licence) repository here: /tttps://gitii http:///gggaabbb..coooooroooo//p
The alanine dipeptide dataset with the semi empirical force-field, was generated by relaxing 105 randomly selected states from the classical MD simulation. The relaxation was performed with the semi-empirical GFN2-x7B force-field for 100 fs each, using (Bannwarth et al. 2019) and the ASE library (Larsen al., 2017) with friction constant of 0.5 a.u. The test set created
in the same way. The authors of Klein et al.] (2023b) made the relaxed alanine dipeptide with the semi empirical force field available here (CC BY 4.0): httts:://oo io/ wrggg//viewwonnnn
28deeba0345546fb96d..... b2235aaaabaa
For AD2-XTB, we precompute the target forces of the samples, re-use them during PG training. For computing the metrics we use the following number of test samples:
E 2 x 1 ESS,,, NLL: 1 105.
Dipetide dataset (2AA) The dipeptide dataset was introduced by Klein et al. (2023a) and is The dataset consists
available at https: htttps//huuggtt http://hggggggaaceeee/ httgggiiffaae../daatteeeeeeeeeee of classical MD trajectories of dipeptides at 310K. There are 200 trajectories in the train set, simulated for 50ns each and 100 each in the validation and test set, simulated for 1 each.
24

---

Alanine dipeptide - Free energy prajection
Alanine dipeptide - Free energy projection
Bond length distribution: C.N
Free energyyeyy Raa
Plye
Free energy /Kat,
15
Mundhigg rygend
listrt)
Hyped tyer "MAeighted
--c liyhrid reweighted
10
22
-1
12
14 Distance /
13
3.5
V
Figure 6: Alanine dipeptide results for the TBG model and the classical force field with and without Path Gradient finetuning. Left and middle: Free energy projection for the yp and nb dihedral angles, respectively. Right: Bond-length distribution for a Carbon Nitrogen bond.
A.4 Finetuning experiments
We used Adam with learning rate of le-4 for PG and only the training set provided. For each of the training points, we require access to the force of the target. We use Hutchinson's 's estimator for estimating the trace of the Jacobian. Importantly, in order to keep the memory constant, we avoid saving checkpoints. For our experiments we used A100 GPUs.
For LJ13, we used a batch-size of 64 instead of 256, which uses about 90% of the original memory (FM B, PG 1.. For fine-tuning we used 2 epochs for Path Gradients, which took 117 min, instead of 135 min for 1000 epochs with Flow Matching. For LJ55, we use the same batch-sizes resulting 13.20 B for Flow Matching and 14.76G for Path Gradients. We fine-tuned with Path Gradients for 1 epoch, taking 394 minutes instead of 440 min for 400 epochs with Flow Matching.
measuring the ESS and neg-loglikelihood on the target density, we removed the reweighting from the Alanine Dipeptide data. If we changed the target density to its reweighted version and computed the gradients of the weighting w.r.t. the samples, we could also straightforwardly apply Path Gradients to the reweighted distribution. Further, fine-tuning on AD2 with Path Gradients showed some instabilities during training. To fight these, we employed gradient clipping to norm of I and gradient accumulation to batch-size of around 1000 in batches of 50. The resulting fine-tuning with Path Gradients uses 2.44GB VRAM and 131 min, compared 200 minutes with 2.95G B with Flow Matching and batch-size 256.
Because
We used the number of samples like (Klein al. 2023b) for the ESS, but note that higher number might have been beneficial for more reliable estimates.
For the experiments, we did not run exhaustive hyperparameter tuning. The Batch-sizes were set to have a similar memory footprint. The hybrid PG learning rate was set to the middle between the initial and FM-fine-tuning one after some preliminary experiments.
A.5 Additional results for alanine dipeptide
Here we present additional alanine dipeptide results for the TBG model Figure We observe that the &o and vp dihedral-angle projections change little after trthh-gradiet fine-tuning. Moreover, more reflected by a higher ESS (see Table In contrast, all bond-length
reweighting to the target distribution via Equation (2) succeeds in both cases, though efficient post-finettuinn, distributions align much more closely with the target after fine-tuning, exemplified here by a Carbon Nitrogen bond in Figure 6
A.6 Additional results for alanine dipeptide- classical with more memory
For classical AD2, we further investigated the performance of fine-tuning with PGs when allowed larger memory footprint. We fine-tuned with batch-size 1024 for two epochs, which resulted in 22GB VRAM usage and runtime 92 minutes. In Table 2 we can see that this yielded a further improvement over Table [
25

---

Table 2: Additional experiment for fine-tuning TBG (Klein & Noč 2024) with Path Gradients on
classical Alanine Dipeptide. First all models were pre-trained with Optimal Transport Flow Matching like in (Klein et al. 2023b). We compare fine-tuning with Flow Matching and Path Gradients. Here we did not limit the memory usage and trained PG with batch-size 1024 for 2 epochs. mean sterr over 3 runs.
Training type
NLL (4)
ESS, %(t)
ESS, %(1)
Only FM
128.01 + 0.01
14.42 2.44
12.21 1.06
Hybrid (ours) unlimited
128.33 = 0.04
29.47 + 2.24
19.57 + 4.79
U13 EQ OT Flow Matching
U13 OT Flow Matching
y13 standard Fow Matching
Optimal burspo map Hatciing Hyond
Optimol toroport mop Marthing Hydrl)
optnaalr map
Flon Matching Hybeld
04
62
Integration trajectory length
23 25
13 21 Integration trajectory length
Integration trajectory length
Figure 7: Integration trajectory lengths for the Lennard-Jones system with 13 particles. Compared are different Flow Matching models as well as the same models after fine-tuning with Path Gradients (Hybrid).
A.7 Additional results for experiments on flow trajectory length
Fine-tuning on the 13-particle Lennard-Jones system used an A100 for 10 epochs (152 min/run) starting from the Klein et al. (2023b) checkpoints, with training and evaluation 5 x 10 samples. As shown in Table path- gradient fine-tuning boosts NLL and ESS across the board while leaving flow trajectory lengths nearly unchanged (see Figure 7). One of the nine runs diverged during training, was discarded and repeated.
Table 3: Performance metrics before and after fine- tot-ttnning for different methods for the LJ13 system.
mean + sterr over 3 runs.
Method
Fine- tuneituning
Trajectory Length
NLL
ESS, (%)
ESS, (%
before after
3.76 + 0.00
16600 + 0.01
60.45 + 0.89
40.23 +: 20.11
Standard
3.86 + 0.00
-16.22 0.00
85.78 + 0.69
85.93 + 0.57
before after
2.80 + 0.01
-16.01
57.16 + 0.88
57.07 0.54
Optimal Transport
2.78 + 0.04
-16.22
85.91 + 0.08
85.78 + 0.02
Equivariant Op-
before
2.11 0.00
-16.04 0.00
58.11 1.93
19.77 12.36
timal Transport
after
2.19 0.00
-16.23 0.00
87.77 + 0.24
58.40 29.20
A.8 Comparison with maximum likelihood training
We here present comparisons with maximum likelihood training baseline IJ13. To this end we compare to Garcia Satorras et al. (2021), which nearly the same architecture as in Klein et al. (2023)), but the model trained via maximum likelihood training instead of flow matching. The results are shown in Table
A.9 Additional energy histograms
We show additional energy histograms of the systems investigated in Section in Figure 8 The average energy for every investigated system is smaller and closer to the target energy distribution after PG finetuning.
26

---

Method
NLL
ESS, %
ML training (Garcia Satorras et al. 2021
15.83 + 0.07
39.78 + 6.19
Only FM (Klein et al. (023b)
16.09 0.03
54.36 5.43
Hybrid approach (Ours)
- 116.21 + 0.00
82.97 + 0.40
Table Results for the Lennard-Jone, system with 13 particles (1.113).
a)
U13 Energy distribution
b)
1J55 Energy distribution
Classical Energy Distribution
0.06
0.05
0.05
Natching Hybrid
Natching Hybrid
Natching Hybrid
0.04
0.020
0.03
0.015
0.03
0.010
0.02
0.01
0.00
50 -40
0.00
Energy
Energy
Energy In KgT
d)
Classical Energy Distribution AT
Classical Energy Distribution GN
f)
Classical Energy Distribution IM
0.07
MD smulation
How Matching
0.06
How Matching llbrd
HD simulation
0.05
Hlow Matching Hybrd
MD siulacion
llyrid
0.05
0.04
0.03
0.03
0.02
0.02
0.02
0.01
0.01
0.01
0.00
300
280
260
240
220
-00
0.00
-380
-350
340
-320
-300
-290
0.00
188
-160
-140
-120
100
Energy K=T
Energy In kgT
Energy in Rg?
g)
Classical Energy Distribution RL
h)
Classical Energy Distribution
i)
Classical Energy Distribution
Fow Matching Hybrid
MD simuet on
0.05
HD simuletio
Matching
0.06
simutation Matching
0.04
0.05
0.04
0.03
0.02
0.03
0.02
0.01
0.01
0.00
-620
-600
-580
-560
-540
-520
0.00
-360
-340
-220
500
-260
-260
-380
-360
-320 -300
Energy in KgT
Energy in Ko
Energy in Rg7
Figure 8: Energy histograms for the a) LJ13 systems, b) the LJ55 system, and c-i) Different didpeptides from the testset for the transferable Boltzmann Generator (TBG).
A.10 Results for transferable Boltzmann Generators (TBG) evaluated on dipeptides
We applied the hybrid approach to transferable Boltzmann Generators (TBG) on dipeptides as introduced in Klein & Noe (2024). We again fine-tune the pretrained Boltzmann Generator from Klein Noč 2024) with path gradients. Training took 6 days on an A100 with learning rate 0.00001. This transferable Boltzmann Generator is trained on subset of all possible dipeptides and evaluated on unseen ones (for more details see Appendix A.3).. Due to the expensive evaluation, we evaluated the model 16 test dipeptides and chose the same subset as in Klein & (2024). Our experiments show that fine-tuning with path gradients improves the NLL and energies for all evaluated unseen test dipeptides with an average improvement -10.93 from -100.72 NLL and also shows average improvements for the ESS, of around 23% to an efficiency of 9.79% as shown in Table 5 An example energy histogram is shown Figure g
A.11 About the different estimators
In the following, first discuss the Maximum Likelihood (ML) and Path Gradient (PG) estimators for the forward KL and show benefits of PG over FM for a toy example.
A.11.1 The Maximum Likelihood and Path Gradient Estimators
In short: both the MI. and PG estimators are unbiased and consistent, but they differ in variance. For PG, we can give guarantees about their variance, once qo equals to p, see e.g. (Vaitl 2024 Roeder et al. 2017 Tucker et al. 2019 Vaitl et al. 2022b).
27

---

Dipeptide
NLL (Before)
NLL (After PG)
ESS, (Before) in %
ESS, (After PG) in %
KS
100.822
1100999
6.0
4.6
AT
-75.49
-75.61
20.8
21.5
GN
-65.45
-65.55
19.2
25.6
LW
1148.91
-149.22
0.4
3.6
NY
118.24
-118.47
9.5
10.5
IM
100.80
-10.001
3.1
4.6
TD
-8.09
-81.19
4.1
11.6
HT
-00.006
-03.227
0.2
5.5
KG
-88.60
-88.79
5.0
7.4
NF
116.455
116.667
3.7
12.2
RL
133.77
136.12
1.3
1.5
ET
-89.94
-90.08
6.3
4.1
AC
-63.90
-63.91
31.8
33.3
GP
-71.63
-71.77
10.5
8.0
KQ
119.36
-119.68
4.3
2.0
RV
126.28
-12..62
1.3
0.6
Average 100.773
-100.93
7.96
9.79
Table 5: Comparison of NLL and ESS before and after PG for different dipeptides from the testset.
Both Maximum Likelihood and the Path Gradient estimators optimize the Forward KL Equation Let's 's look at both estimators in detail and compare them. A full derivation can be found in Appendix B.3.2 Vaiti et al. (024)) Chapter 4 of Vaiti (2024).
To obtain the estimator GML, we first observe that the first term in the KI. divergence constant w.r.t. 0.
L(plg) = Ep(z)) p(<1) - log go(x1))
(29)
which means that it does not enter in the gradient if we directly estimate the gradients via MC estimator.
dK L(plgo)
[ d log (o()))
Ex(())
d0
N
9ML =
log (z(xx))) (1)), 2(0) , pp
(30)
d0
For the exact calculation of the estimator, we decompose it fully with
0T, (x1) dx
go (0 qo (Tg (z1))
det
(31)
The actual calculation for the ML estimator then is
1 d N
(log (105
81 0 (x11) Or
GML
Σ
( (21)) log
(32)
do
For PG, we use Eq. 17 directly obtain the MC estimator GPG
OT, (x1))
d0 LL = Enao
log po.a (xo) 00
GPG
N
1
(8)
(10gpo. log go(=x)))
(1)
OT,
g
00
1
(x
(2)
(6)
p.
(33)
Here (1) is shorthand for To (zr(0) (0)) we again decompose the term we expose the full calculation
GPG
1
N
(o + log
(To(
log log go(=)"))))
(11)
01,
(0
(4)
(34)
N
Oxo (
axo
00
28

---

Comparing 32 and 34
we see that the path gradient estimator incorporates GPG the gradient log while GMI. and by construction also GFM does not.
information of the target p-
8mp
A.11.2 Variance of the estimators
Opposed to GML and GFM. we have nice guarantees about the variance of the PG estimator GPG at the optimum and close it.
We first recapitulate results from previous work and then, by using a simple example, show that GrM does not necessarily exhibit variance at the optimum.
In general, the variance the Path Gradient estimators GPG is bounded by the squared Lipschitz constant of the term
o log (o (logp(z1) log go(x)))
(35)
see Mohamed et al. (2020) Section 5.3.2.
Thus, if the target density is not well approximated by qo, the variance of the gradient estimator can be large and training with PG might not be beneficial. If log p and log qo are close the sense that the Lipschitz constant of their difference is small, can assume path gradient estimators to be helpful.
Gradient estimators at the optimum In the case of perfect approximation, i.e. go(x) p(z)vr E 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. B[GrG] = O, Var[Gro] = 0, while for ML the variance is generically nonzero E[GMr]] = 0, Var(9ur.) N Io. Where Ia is the Fisher Information Matrix of Po.e (Vaitl et al. 2022a).
Why is that? Already the review by Paaamakariiss et al. (202T) notes in Section 2.3.3 that the duality of the KL. divergence means that fitting the model qo to the target using ML is equivalent to fitting Po, to the base 90 under the reverse KL. This means that the results from Vaitl (2024) directly hold. We only adapted the notation.
A.11.3 Variance of GFM for a toy example
We do not have term for the variance of GrM for general settings, but we can show that it does not deterministically vanish like for GPG via simple example. This means better behavior for PG than for FM at the optimum, which we also verified empirically. Our assumptions aim simplify the example much as possible.
First, assume the standard loss for Flow Matching.
Further, assume the densities to be the same D-dimensional Normal distribution = p = (0, I).
Finally, assume the CNF is the identity parametrized by a single parameter 0, i.e. 0 = 0.
In this example (x(2)) = qu(21) I approximates the target density p perfectly and the optimal gradient estimator is O. Yet. in this setting the variance of the estimator GrM is non-zero:
8
Var[Gra]
(36)
ND
preventing the model from staying at the optimum during training.
Proof: The gradient estimator for the FM loss Eq 10 is
GFM
Σ
(lle
(a (1)
(1) )1))
(37)
N
20
where = x1 10,
NNte that here z(") and (=()) are independently sampled. Before x0( was the transformed sample
29

---

First, we break down the terms
NV
à
1
D
9PM
1o,d + = - [1dd
N
00
D
i=1
d=1
rug.d
vo.d + 1/..d F1,e)
(38)
ND
08
d=1
Because is parametrized by 0, we set in
Deg. 80
and V0,d = 0 and separate the terms
9PM ND
ΣΣ> 2
1.
(1)
2
ΣΣ
(1) O,d
(39)
ND
We can compute the distribution of GFM by using the property
Σ a4 (O0, for ~((,
(40)
So the sum over dimensions and samples follows the normal distribution with variance DN:
ΣΣ
(0) o,d ~ NO,
(41)
and the difference has twice its variance
ΣΣ I
ΣΣ क
= N(0,2DN)
(42)
The expectation of the estimator GFM then is and the variance is simply
Var[g
Var N2
(6) o,d
ΣΣ
+220 N2
2DN = ND
(43)
d
Interestingly, the derivative GFM is invariant to re-ordering because the sums of 2ı and To invariant to permutation. The result thus also holds for OT-FM. This property is due to the simple assumption, the vector field vo is independent of z,,
A.12 Pseudocode for forward KL path gradients via augmented adjoint
Algorithm Augmented Adjoint Dynamics
1: function FORWARD(FF xr, V log
2: i. div +- black box dynamics(t, 2T)
3: Vlog t gradient, [-V log 9 + div]
4: return Vog g, -iiv
5: end function
Algorithm 2 Pathwise Gradient Estimator
1: function PATIWSSEGRAADIENNTSSTTTTAC.C prior, target, flow)
2:
logpi <- - target.energy(())
3: log P1 gradient,, (logP1)
Integrate using Augmented Adjoint state method To, V log Po,9, log det flo.iinteratekuuAAjioonntzz V P1) inverse=True) inve
4:
5:
log qo +- - priorr.nerry(zo))
6:
P Compute gradient of loss w.r.t. sample wo 7: V., L < 1 (V log Po.0 log 4o) & Backpropagate using standard Adjoint state method N
log qo + gradient,, (log 9o)
8:
path gradients +- gradient, (To detach( V ()))
9: end function
30

---

A.13 Code libraries
We primarily use the following code libraries: PyTorch (BSD-3) Paszke et al. 2019), bgflow (MIT
license) (Noé et al.) 2019: Köhler et al. 2020). torchdyn (Apache License 2.0) (Poli et al. 2021). Additionally, we use the code from (Garcia Satorras al. (021) (MIT license) for EGNNs, well as the code from (Klein et al. 2023b) (MIT license) and (Klein & Noë 2024) (MIT license) for models and dataset evaluations.
31


Reviews:
[{"review_id": "PauVgWYzya", "reviewer": "Reviewer_Fj6i", "summary": "This work proposes a method to perform fine-tuning of a pretrained CNF model sampling a Boltzmann distribution via path-gradients (PG). Because computing PG only requires the data x_1 and unnormalized energy gradients \\nabla_{x_1} E(x_1), (which are supposedly already computed to generate training data), PG requires no additional compute in terms of energy evaluations. By fine-tuning with PG, the authors show a significant increase in the effective sample size (ESS) on a number of well-studied energy functions when compared to flow-matching or path-gradients alone.", "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?\n\n- 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.", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n- The method is easy to understand and performing path-gradients yields a large boost in effective sample size (ESS).\n\n- Unlike importance sampling or relaxation, Path-gradients require no additional energy evaluation and preserves the target Boltzmann distribution, given that the model is a CNF trained on unbiased data.\n\nWeaknesses:\n\n- 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.\n\n- 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]). \n\n[1] Tan et al 2025 \"Scalable Equilibrium Sampling with Sequential Boltzmann Generators\"", "was_revised": true, "final_justification": "My concerns were addressed and they provided some interesting additional experiments on transferable energies. I think its a good paper.", "rebuttal": "We thank the reviewer for their positive feedback. We're glad the clarity and practical value of our method came through.\n## Weaknesses\n#### *\"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.\"*\n\nWe 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. \n\nThe 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. \n\n\nEven when avoiding correction via importance sampling, the samples are significantly closer to the reference (see Figure 4 right) and the target is approximated better.\nSo, 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.\n\n\n#### *\"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]).\"*\n\nWe 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:\n- In a setting where you have access to the derivative of the energy, you can always apply MCMC or MD. \nThe settings with Boltzmann Generators, as well as the approach by (Tan et al. 2025) [1] or protein folding fall in this regime.\nThe 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.\nSample-efficient training methods, like Path Gradients, are of great value in this setting.\n- In practice, a small bias in the training data is tolerable. \nFor 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).\n- 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.\nBelow we summarized the additional experiments.\n\n[1] Tan et al 2025 \"Scalable Equilibrium Sampling with Sequential Boltzmann Generators\"\n## Questions:\n#### *\"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?\"*\n\nIndeed, Flow Matching (FM) and Path Gradients (PG) optimize different objectives: \n\nFM 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$. \nSince 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$.\n\nSo for a minimal FM loss the KL is also minimized, while the reverse does not hold.\nOur 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$.\n\n#### *\"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.\"*\n\nWe 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]).\n\nPGs 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.\n\nAs 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. \n\nFurther, 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. \n\n\n\n\n\n## Additional Dipeptide Experiments:\nWe 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:\n\n| Dipeptide | NLL (Before) | NLL (After PG) | ESS (Before) | ESS (After PG) |\n|-----------|--------------|----------------|---------------|----------------|\n| AC        | -63.90       | **-63.91**     | 31.8          | **33.3**       |\n| AT        | -75.49       | **-75.61**     | 20.8          | **21.5**       |\n| ET        | -89.94       | **-90.08**     | **6.3**       | 4.1            |\n| GN        | -65.45       | **-65.55**     | 19.2          | **25.6**       |\n| GP        | -71.63       | **-71.77**     | **10.5**      | 8.0            |\n| HT        | -103.06      | **-103.27**    | 0.2           | **5.5**        |\n| IM        | -106.80      | **-107.01**    | 3.1           | **4.6**        |\n| KG        | -88.60       | **-88.79**     | 5.0           | **7.4**        |\n| KQ        | -119.36      | **-119.68**    | **4.3**       | 2.0            |\n| KS        | -100.82      | **-100.99**    | **6.0**           | 4.6            |\n| LW        | -148.91      | **-149.22**    | 0.4           | **3.6**        |\n| NF        | -116.45      | **-116.67**    | 3.7           | **12.2**       |\n| NY        | -118.24      | **-118.47**    | 9.5           | **10.5**       |\n| RL        | -135.71      | **-136.12**    | 1.3      | **1.5**            |\n| RV        | -126.28      | **-126.62**    | **1.3**       | 0.6            |\n| TD        | -81.09       | **-81.19**     | 4.1           | **11.6**       |\n|           |              |                |               |                |\n| Average   | -100.73      | **-100.93**    | 7.96          | **9.79**       |\n\n[2] Vaitl, Lorenz, et al. \"Path-gradient estimators for continuous normalizing flows.\" International conference on machine learning. ICML, 2022.\n\n[3] Bacchio, Simone, et al. \"Learning trivializing gradient flows for lattice gauge theories.\" Physical Review D 107.5 (2023): L051504\n\n[4] Abbott, Ryan, et al. \"Applications of flow models to the generation of correlated lattice QCD ensembles.\" Physical Review D 109.9 (2024): 094514.\n\n[5] Abbott, Ryan, et al. \"Normalizing flows for lattice gauge theory in arbitrary space-time dimension.\" arXiv preprint arXiv:2305.02402 (2023).\n\n[6]  L. Vaitl “Path Gradient Estimators for Normalizing Flows\", PhD Thesis\n\n[7] L. Vaitl et al., “Fast and unified path gradient estimators for normalizing flows”, ICLR 2024"}, {"review_id": "AN4lwcKH7w", "reviewer": "Reviewer_F87i", "summary": "This paper proposes a hybrid training strategy for Continuous Normalizing Flows (CNFs) used as Boltzmann Generators for molecular systems. The authors identify that Flow Matching (FM) provides a computationally efficient way to pre-train a model using target samples, but it ignores valuable information from the target energy function. Conversely, training methods that use energy gradients, like Path Gradients (PG), are often computationally intensive. The proposed solution is to first pre-train a CNF with Flow Matching to rapidly obtain a good initial model, and then fine-tune this model using Path Gradients. This hybrid approach is shown to significantly improve the sampling efficiency of the resulting model—by up to a factor of three—within a similar computational budget as pure Flow Matching, all without requiring additional data. Furthermore, the work demonstrates that this fine-tuning process refines the model's accuracy while largely preserving the efficient, short integration paths learned during the Flow Matching phase.", "questions": "1.  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.\n2.  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?", "limitations": "Yes. The authors provide a dedicated and candid \"Limitations\" section 5.1. They correctly identify the main constraints of the method: it requires a differentiable energy function, it is computationally more expensive per step than Flow Matching, and it performs best as a fine-tuner when the model is already a reasonable approximation of the target distribution. This is a transparent and sufficient discussion of the work's scope.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\n* The paper addresses the important and practical problem of how to best leverage all available information (both data samples and their energy gradients) for training Boltzmann Generators. The proposed hybrid approach is a simple, novel, and highly effective strategy that provides a clear recipe for improving upon existing state-of-the-art methods.\n* The technical quality of the work is high. The claims are supported by strong and clear experimental evidence on several standard molecular dynamics benchmarks. \n* The paper is very well-written and easy to follow. It provides a concise yet thorough background on the relevant concepts (CNFs, Flow Matching, Path Gradients) before clearly motivating and presenting the proposed hybrid method. The experiments are well-designed and the results are presented in a clear and compelling manner.\n\n**Weaknesses**\n\n* 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.", "was_revised": true, "final_justification": "I had a full discussion with the author and I kept my score.", "rebuttal": "We thank the reviewer for their thoughtful and encouraging feedback. \nWe appreciate them pointing out the practical relevance of the problem, the strength of the experimental evidence, and the clarity of the presentation.\n\n## Weaknesses:\n#### *\"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.\"*\n\nWe agree that additional and more challenging systems improve the impact of the paper.\nSo 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. \n\nAgain, 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$. \n\nThe full results can be found in our rebuttal to reviewer Fj6i.\n## Questions:\n#### *\"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.\"*\n\nThis 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.\nFor 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.\n\n 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.\n\n#### *\"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?\"*\n\nWe thank the reviewer for this insightful question. We believe our hybrid approach has two key advantages that support its scalability to larger systems: \n- *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. \n- 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.\n\nHowever, 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. \nWe 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.\n\n[1] Vaitl, Lorenz, et al. \"Path-gradient estimators for continuous normalizing flows.\" International conference on machine learning. ICML, 2022.\n\n[2] Bacchio, Simone, et al. \"Learning trivializing gradient flows for lattice gauge theories.\" Physical Review D 107.5 (2023): L051504\n\n[3]  Abbott, Ryan, et al. \"Normalizing flows for lattice gauge theory in arbitrary space-time dimension.\" arXiv preprint arXiv:2305.02402 (2023).\n\n[4]  L. Vaitl “Path Gradient Estimators for Normalizing Flows\", PhD Thesis"}, {"review_id": "n3tEY83IzB", "reviewer": "Reviewer_P3GD", "summary": "In this paper, the authors propose using Flow Matching for pre-training and Path-Gradient during fine-tuning, in the case where the target energy is well-defined. The authors show that the proposed hybrid approach can significantly increase the sampling efficiency given the same computational budget on many particle systems. They further show that fine-tuning with Path Gradients does not affect the properties of the base model, while simultaneously improving the performance.", "questions": "- What is the performance of the hybrid approach using an increased runtime?\n- 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? \n- What is the runtime of the hybrid approach fine-tuning (Table 1)?\n- 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]?\n\n[2] Vaitl, Lorenz, et al. \"Path-gradient estimators for continuous normalizing flows.\" International conference on machine learning. PMLR, 2022.", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths:\n\n- The paper is well-structured, clearly written, and easy to follow.\n- The proposed hybrid approach relying on Path Gradients shows an improvement in the performance and the sampling efficiency, using the same model and without relying on additional samples. \n\n\nWeaknesses:\n\n- 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. \n- An important baseline that is missing in the experiments is the method proposed in Satorras et al. [1].\n- 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?\n- 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. \n\n[1] Garcia Satorras, Victor, et al. \"E (n) equivariant normalizing flows.\" Advances in Neural Information Processing Systems 34 (2021): 4181-4192.", "was_revised": true, "final_justification": "The authors have addressed my concerns. Many thanks!", "rebuttal": "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.\n## Weaknesses:\n#### *\"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.\"*\n\nWe agree with the reviewer that additional comparisons would be helpful. So for the rebuttal we did additional experiments.\n\nWe evaluated our hybrid approach on the Dipeptide setting from Transferable Boltzmann generators [3].\nThis training set consists of 200 dipeptides and the evaluation is done on unseen dipeptides. \nOur 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.\nFor additional information see our rebuttal for reviewer Fj6i.\n\nStill, 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.\n\n 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. \n\nA hybrid approach applied to diffusion models and their Probability Flow ODE is a very promising direction for future work, but merits a thorough investigation.\n\n#### \"*An important baseline that is missing in the experiments is the method proposed in Satorras et al. [1].*\"\n\nWe agree with the reviewer that the model in [1] is a good baseline and for the LJ13 we will gladly incorporate the baseline.\n\n| Method | $\\text{NLL}$ | $\\text{ESS}_q$ |\n|--------|--------------|----------------|\n| Satorras et al. [1]    | $-15.83 \\pm 0.07$ | $39.78 \\pm 6.19$ |\n| Only FM [6]    | $-16.09 \\pm 0.03$      | $54.36 \\pm 5.43$         |\n| Hybrid approach (Ours)   | **$\\mathbf{-16.21 \\pm 0.00}$** | **$\\mathbf{82.97 \\pm 0.40}$** |\nNote, that the model architectures for [1] and [6] are mostly the same. Klein et al. 2023 [6] already outperforms the model from [1] by using FM training instead of Maximum Likelihood (ML) training. \n\nFor 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].\n\n#### \"*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?*\"\n\nWe 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).\n\nYet, our preference for a hybrid approach over exclusively PG is motivated by:\n- **Time efficiency**: FM is significantly faster per step than PG and ML.\n\n\n- **Theoretical guarantees**: PG has convergence guarantees only near the optimum (see Rebuttal to Reviewer DmNE for a thorough analysis).\n\n\n- **Numerical stability**: Like ML training, PG involves Jacobians and their derivatives which suffer from numerical instabilities, whereas FM is often more stable.\n\nEspecially the time efficiency was shown in our experiments:\n- 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.\n- 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.\n\nWe will revise the manuscript to reflect this more nuanced perspective and avoid overstating the limitations of PG.\n\n#### \"*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.*\"\n\nWe 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. \n\n[1] Garcia Satorras, Victor, et al. \"E (n) equivariant normalizing flows.\" Advances in Neural Information Processing Systems 34 (2021): 4181-4192.\n\n## Questions:\n#### \"*What is the performance of the hybrid approach using an increased runtime?*\"\n\nFor LJ13 we show the results with increased runtime in Table 3. For AD2-classical we give results for TBG [7] in Table 2. \nSummarizing, we were able to obtain the following results\n| System         | Time      | Method           | NLL                  | ESS_q              | \n|----------------|-----------|------------------|----------------------|--------------------|\n| LJ13           | limited   | OT FM + PG       | -16.21 ± 0.00        | 82.97 ± 0.40       | \n| LJ13           | unlimited | EQ OT FM + PG    | **-16.23 ± 0.00**    | **87.77 ± 0.24**   | \n|                |           |                  |                      |                    |              \n| AD2 classical  | limited   | OT FM + PG       | -128.26 ± 0.02       | 24.39 ± 6.86       | \n| AD2 classical  | unlimited | OT FM + PG       | **-128.33 ± 0.04**   | **29.47 ± 2.24**   |\n\nAdditionally, we ran the Hybrid approach on LJ55 once for 3 PG epochs instead of one and got \nthe following performance NLL: -89.38, $ESS_q$ 25.44 %. \n\n\n#### *\"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?\"*\n\nThe settings are different: \n- 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. \n- 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.\n\n#### *\"What is the runtime of the hybrid approach fine-tuning (Table 1)?\"*\n\nWe list the runtimes in A.4.\n| LJ13  | LJ55  | AD2   |\n|-------|-------|-------|\n| 252 min | 834 min | 331 min |\nAll done on A100.\n\n#### *\"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]?\"*\n\nThere are two major differences: \n- 1.) While Vaitl et al. [2] and our work both use PG, they minimize a different loss on different samples.\nOur 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].\nVaitl 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.\n\n- 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.\n\nThat being said, forward KL based training can straightforwardly be combined with self-sampling based approaches (like [2]) to artificially increase the number of samples.\n\n[2] Vaitl, Lorenz, et al. \"Path-gradient estimators for continuous normalizing flows.\" International conference on machine learning. ICML, 2022.\n\n[3] L. Vaitl et al., “Fast and unified path gradient estimators for normalizing flows”, ICLR 2024\n\n[4] A. Agrawal et al., \"Advances in black-box VI: Normalizing flows, importance weighting, and optimization.\" NeurIPS2020\n\n[5] Detecting and Mitigating Mode-Collapse for Flow-based Sampling of Lattice Field Theories, Nicoli, Kim A. et al. PRD 23\n\n[6] Klein, Leon et al. \"Equivariant flow matching.\" NeurIPS (2023).\n\n[7] Klein, Leon, and Frank Noé. \"Transferable boltzmann generators.\" NeurIPS (2024)"}, {"review_id": "pyWjw4HMsL", "reviewer": "Reviewer_DmNE", "summary": "The paper presents a method for fine-tuning trained CNFs by using gradient information at the target samples in the context of learning Boltzmann generators (with access to samples). The method is shown to improve KL, ESS, and energy based metrics when compared to further FM training with the same computational cost, despite incurring significantly greater computational cost per training step.", "questions": "Actionable items are discussed in Strengths And Weaknesses.", "limitations": "Yes.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 1, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strengths**\nThe paper's principle motivation is that the training of Boltzmann generators should effective leverage energy and samples, a point which should be commended and pursued in the literature. The experiments indeed show across-the-board improvements on top of CNFs from prior works as a result of the energy fine-tuning procedure pursued here. The experiments are done with proper training time controls, and there is careful accounting and explanation of the NFE cost of the method.\n\n\n**Most serious concerns**\nTo this reviewer familiar with CNFs and Boltzmann generators but not the path gradient method, it is not immediately clear how the path gradient loss here differs from the standard maximum likelihood loss with CNFs. The two cited path gradient papers consider the problem of minimizing reverse KLs, as written in Eqs 13, 14. Instead, the authors recast likelihood maximization $KL(p_1, q_\\theta)$ into a reverse KL by considering $T^{-1} \\theta(x)$ with $x$ drawn from target samples to be the parameterized approximation of the fixed prior $q_0$. This is a significant enough departure that a substantial discussion about the properties of this new objective is merited. It leads to some clunky notation (which should have a standalone short explanation), but more importantly, since it appears identical to an objective that does not use gradient information, it is not clear where the target gradient information is coming from. At the very least, the flipped variant of Eq 14 actually optimized in practice should be written out and dissected. My best guess is that the term inside the expectation is somehow not equal to $\\frac{\\partial}{\\partial \\theta}\\log\\frac{p_1(x_1)}{q_{1,\\theta}(x_1)}$ but differs by a control variate that contains $\\nabla \\log p_1(x_1)$. If the authors can provide a satisfactory explanation in the rebuttal (and promise include it in the supplement for publication), I am open to raising the score.\n\n**Other concerns**\n- Why is there an x axis for Fig 2 but not Fig 1?\n- Please label figure captions with the task / system being sampled.\n- L206 memory -> compute\n- The difference between the two evaluation settings is not very clear. Isn't \"applying path gradients at the end\" the same as fine-tuning?\n- The text should state the difference between the models labeled Klein et al 2023b and Klein & Noe 2024. \n- L226 Standard, OT, and EQ OT FM are all mentioned, but only one row per system is shown in Table 1. Please clarify. Is L226 referencing only the 2nd set of experiments?\n- \"We expect path gradients to become SOTA for BGs and generative models in science\" - this statement is far too strong.", "was_revised": true, "final_justification": "The authors have addressed my primary concern regarding clarity and have committed to including detailed explanations and background in the revision.", "rebuttal": "We thank the reviewer for the valuable input. \nWe fully agree with every point raised. \nWe brushed over the actual gradient estimators, which is something we will add to the manuscript. \n\nAnd 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.\nBut the PG estimator for the forward KL was actually introduced and discussed in Chapter 4 of [1].\n\nIn 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.\nThis additional discussion is a great addition to the paper, we hope you agree. \n\n# The Maximum Likelihood and Path Gradient Estimators.\nIn short: both the ML and PG estimators are unbiased and consistent, but they differ in variance.\nFor PG, we can give guarantees about their variance, once $q_\\theta$ equals to $p$, see e.g. [2, 3,4,5].\n\nBoth Maximum Likelihood and the Path Gradient estimators optimize the Forward KL Equation (8).\nLet’s look at both estimators in detail and compare them.\nA full derivation can be found in Appendix B.3.2 of [1] & Chapter 4 of [2].\n\nTo obtain the estimator $\\mathcal G_{ML}$, we first observe that the first term in the KL divergence is constant w.r.t. $\\theta$.\n$$KL(p|q_\\theta) = E_{p(x_1)} \\left[ \\log p(x_1) - \\log q_\\theta(x_1) \\right],$$\nwhich means that it does not enter in the gradient if we directly estimate the gradients via an MC estimator.\n\n$\\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$.\n\nFor the calculation for the estimator, we decompose it fully with\n$$\nq_\\theta(x_1) = q_0(T_\\theta^{-1}(x_1)) \\left| \\det \\frac{\\partial T_\\theta^{-1}(x_1)}{\\partial x_1} \\right|\n$$\n\nThe actual calculation for the ML estimator then is\n$$\\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) (*).$$\n\nFor PG, we use Eq. (17) to directly obtain the MC estimator $\\mathcal G_{PG}$\n$$\\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. $$\nHere $x_0^{(i)}$ is a shorthand for $T_\\theta^{-1}(x_1^{(i)})$.\nIf we again decompose the term we expose the full calculation\n$$\\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} (**)$$\n\nComparing (*) 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.\n# Variance of the estimators:\nOpposed 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.\n\nWe 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.\n\nIn general, the variance of the Path Gradient estimators $\\mathcal G_{PG}$ is bounded by the squared Lipschitz constant of the term\n$\\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).\n\nThus, 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.\n## Gradient estimators at the optimum\nIn 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.\nThe gradients for path gradient estimators are deterministically $0$, i.e.\n$E[\\mathcal G_{PG}] = 0, Var[\\mathcal G_{PG}] = 0$,\nwhile for ML the variance is generically nonzero $E[\\mathcal G_{ML}] = 0, Var[\\mathcal G_{ML}] = \\frac 1 N \\mathcal I_\\theta$.\nWhere $\\mathcal I_\\theta$ is the Fisher Information Matrix of $p_{0, \\theta}$ [2].\n### Why is that?\nAlready 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 .\nThis means that the results from [2] directly hold. We only adapted the notation.\n\n## Variance of $\\mathcal G_{FM}$ for a toy example\nWe 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.\nThis means better behavior for PG than for FM at the optimum, which we also verified empirically.\n\nOur assumptions aim to simplify the example as much as possible.\n- First, assume the standard loss for Flow Matching.\n- Further, assume the two densities to be the same $D$-dimensional Normal distribution $q_0 = p = \\mathcal N(0,I)$.\n- Finally, assume the CNF is the identity parametrized by a single parameter $\\theta$, i.e. $v_\\theta = \\theta = 0$.\n\nIn this example $q_\\theta(x_1) = q_0(x_1) I$ approximates the target density $p$ perfectly and the optimal gradient estimators is $0$.\nYet, in this setting the variance of the estimator $\\mathcal G_{FM}$ is non-zero: \n$$Var[\\mathcal G_{FM}] =\\frac 8 {ND},$$ \npreventing the model from staying at the optimum during training.\n### Proof:\nThe gradient estimator for the FM loss Eq (10) is\n$$\\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)}$.\n*(Note that here ${\\tilde x}_0^{(i)}$ and $x_1^{(i)}$ are independently sampled. Before $x_0^{(i)}$ was the transformed sample)*.\n\nFirst, we break down the terms\n$$\n\\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}.$$\n\nBecause $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\n$$\\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}). $$\n\nWe can compute the distribution of $\\mathcal G_{FM}$ by using the property\n$\\sum_{i=1}^N a_i \\sim \\mathcal N(0, N \\sigma_a^2)$ for $a \\sim \\mathcal N(0, \\sigma_a^2)$.  \nSo \n$\\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).$\n\nThe expectation of the estimator $\\mathcal G_{FM}$ then is 0 and the variance is:\n$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}$.\n\nInterestingly, 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.\n \n[1] L. Vaitl et al., “Fast and unified path gradient estimators for normalizing flows” ICLR 2024\n\n[2] L. Vaitl “Path Gradient Estimators for Normalizing Flows\" PhD Thesis\n\n[3] G. Roeder et al., \"Sticking the landing: Simple, lower-variance gradient estimators for variational inference\" NeurIPS (2017)\n\n[4] G.Tucker et al. \"Doubly reparameterized gradient estimators for monte carlo objectives\" ICLR 2019\n\n[5] L. Vaitl et al. \"Path-gradient estimators for continuous normalizing flows\" ICML, 2022\n\n[6] S. Mohamed et al. \"Monte carlo gradient estimation in machine learning\" JMLR 2020\n\n[7] G. Papamakarios et al. \"Normalizing flows for probabilistic modeling and inference\" JMLR 2021\n# Other concerns\n- X axis, figure caption, L206, too strong statement, L226 second set of experiments:\n\nThanks for the input. We will fix these in the final manuscript.\n- Isn't \"applying path gradients at the end\" the same as fine-tuning?\n\nTrue, 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.\nWe will clarify this in the final manuscript.\n- Difference between the models labeled Klein et al 2023b and Klein & Noe 2024.\n\nWe 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.\n\n- Paper Formatting\n\nWe have not changed the template"}]

Author rebuttal:

Meta-review: The authors present a method to augment Boltzmann Generators using the output of MD simulations, requiring only an initial energy and gradients of the energy along a sampled path. They demonstrate good results on model systems like alanine dipeptide. While the reviewers raised concerns about scalability and applicability in situations where a good energy function was not available, the authors managed to win many of the reviewers over in the rebuttal period. There was a consensus that the paper should be accepted, which I endorse.
Decision: Accept (poster)