Learning Escorted Protocols For Multistate Free-Energy Estimation
Authors: Lars Holdijk
Nithishwer Mouroug Anand
Michael M. Bronstein
Max Welling
Venue: ICLR 2026 Poster (2026)
Keywords: Free-Energy
Jarzynski Equality
Crooks Fluctuation Theorem
Non-Equilibrium
Transport
Thermodynamics
Stochastic Thermodynamics
Flow Matching
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: Conditional Flow Networks and Density Matching for Multistate Escorted Free-Energy Estimation

Abstract: Estimating relative free energy differences between multiple thermodynamic states lies at the core of numerous problems in computational biochemistry. Traditional estimators, such as Free Energy Perturbation and its non-equilibrium counterpart based on the Jarzynski equality, rely on defining a switching protocol between thermodynamic states and computing the free energy difference from the work performed during this process. In this work, we present a method for learning such switching protocols within the class of escorted protocols that combine deterministic and stochastic steps. For this purpose, we use Conditional Flow Matching, and  introduce Conditional Density Matching (CDM)  for the purpose of estimating the change in Free-Energy. We further reduce the variance in the multistate setting by coupling multiple flows between thermodynamic states into a Flow Graph, enforcing estimator consistency across different transition paths.

Published as a conference paper at ICLR 2026
LEARNING ESCORTED PROTOCOLS FOR MULTISTATE ENERGYY ESTIMATION
Lars Holdijk
Nithishwer Mouroug Anand
Department of Computer Science
Department of Biochemistry
University of Oxford
University of Oxford
Michael Bronstein
Max Welling
Department of Computer Science University of Oxford, AITHYRA
AMLAB
University of Amsterdam, CuspAI
ABSTRACT
Estimating relative free energy differences between multiple thermodynamic states lies at the core of numerous problems in computational biochemistry. Traditional estimators, such as Free Energy Perturbation and its non-equilibrium counterpart based on the Jarzynski equality, rely on defining a switching protocol between thermodynamic states and computing the free energy difference from the work performed during this process. In this work, we present method for learning such switching protocols within the class of escorted protocols, which combine deterministic and stochastic steps. For this purpose, we use Conditional Flow Matching and introduce Conditional Density Matching (CDM) to estimate changes in free energy. We further reduce the variance in the multi-state setting by cou- pling multiple flows between thermodynamic states into flow graph of escorted protocols, enforcing estimator consistency across different transition paths.
INTRODUCTION
In recent years there has been significant interest in the application of machine learning to problems in computational biochemistry, such as protein folding (Jumper et al., 2021; Abramson et al., 2024; Bose et al., 2024) and molecular conformer generation (6ómez-Bombarell et al., 2018) for drug discovery (Wan et al., 2022) and materials design (Merchant et al., 2023). While initial approaches focused generating single static examples of systems of interest, recent efforts have shifted towards generating the full dynamic ensemble of the system (Noé et al., 2019; Holdijk et al., 2023; Tan et al., 2025; Akhound-Sadege
al., 2024). This shift towards the entire dynamical ensemble has opened up the possibility of using these methods in computational biochemistry beyond what generating static samples allows. Notable among these is the problem of estimating free-energy differences between different thermodynamic states, a crucial aspect of many tasks in this field, such as binding affinity prediction (Mobley and Gilson, 2017) and other components of the drug discovery pipeline (Cournia and Chipot, 2024). Traditional estimators such multiwindow Free-Energy Perturbation (FEP) (Zwanzig, 1955; Wang et al., 2015) and its non -quuiiiiiirr variants, defined through the Jarzynski equality (Jarzynski, 1997). among the most commonly used for this purpose.
In this work, we specifically consider the class of Escorted Non-EQuilibrium (E-NEQ) estimators (Vaikuntanathan and Jarzynski, 2008; 2011) based the aforementioned Jarzynski equality. These estimators combine dynamics that preserve given stationary distribution p(z, t). such as Langevin dynamics, with deterministic escorting vector b(z, field to transition samples between thermody- namic states. To minimize variance in the finite-sample limit, the preserved stationary distribution and the escorting vector field must together satisfy the continuity equation @p(a.t) v (p(z, t)b(2, t)) 0 (Zhong and WeWeese, 2024). Traditionally, finding such pair is non-trivial.
To overcome this barrier, we propose learning the stationary distribution and escorting vector field jointly by extending the Conditional Flow Matching (CFM) framework (Lipman et al., 2023) with an additional Conditional Density Matching (CDM) objective (Sec. 3). We hypothesize that learning the stationary distribution and driving force in a unified framework will allow them to compensate for each other's small errors, leading to more accurate estimator than using only vector fields,

---

Published as a conference paper at ICLR 2026
AF A-C
c
AF A-A-
AFB
to
B
Figure I: (a) Illustration processes transforming between two thermodynamic states (blue and red) of the same molecular system. Using our proposed Conditional Density Matching framework we learn approximate intermediate densities (grey) along which the learned flow (grey dotted lines) transports samples, The top stochastic line shows how stochastic processes purely targeting the intermediate distributions will not fully reach the target state, while the middle process shows that, due integration error, deterministic processes may deviate from true path. Combining stochastic and deterministic processes allows two approaches to compensate, resulting in more accurate sampling. connecting node (A) to obtain the free-energy AF difference between the nodes not connected.
(b) Illustration of combining two free-energy estimates with a central
such as done in methods such as Targeted Free-Energy Perturbation (TFEP) (Jarzynski, 2002) as illustrated in Fig. 1. To further improve the applicability of the framework, we propose two important practical considerations in the form of Lie-Trotte splitting to reduce the computational cost of the work calculation (Sec. 4.1) and a flow graph construction to reduce the combinatorial complexity of learning protocols for multi-state free-energy estimation (Sec. 4.2).
The remainder of our paper organized as follows to introduce the core contributions of our work:
Sec. introduces the theory of non-equilibrium free-energy estimation, focusing on topics such as the Jarzynski equality, escorted switching protocols, and bi-directional estimators.
Sec. then relates escorted switching protocols with the Flow Matching framework and introduces one of our core contributions in the form of conditional Density Matching. Collectively these form the learning objectives for our proposed Escorted Non- -on--QQilibrium (E-NEQ) estimator.
To improve the efficiency of multistate free-energy estimation, Sec. 4.1 proposes a splitting scheme and Sec. 4.2 proposes the concept of Escorted Protocol Flow Graphs combine multiple trained escorted switching protocols to reduce training requirements.
In Sec. 5 the E-NEQ estimator is experimentally validated using the well-known Alanine Dipeptide (ADP) system. ADP is important benchmark due to its multistate free-energy surface.
2
BACKGROUND
As stated, in this work we are interested in estimating the relative free-energy differences AFA between pairs of thermodynamic states A and B. Each thermodynamic state is associated with potential energy function UA R3N + I, which defines the distribution of possible microscopic states x E RaN through the Boltzmann distribution
/ exp (-30 dr.
(*(x)= ZA exp (=33(((())),
ZA =
JRAN
Generally, 3 = Tx with ky representing the Boltzmann constant and T the temperature of the system in Kelvin. Using the partition function ZA. the equilibrium free-energy of state A can be defined as FA- - In A and their relative free-energy difference
AF FB FA = - Zn
ZA
As exemplified by the definition of the partition function, estimating free energies by directly integrating over the entire space 3N is prohibitively expensive. As such, various fields have developed methods to estimate this quantity more efficiently.
2.1 EON-LIQUIM
ESTIMATION
We will specifically focus on the class of alchemical Eon-Euuiiibrium (NEQ) free-energy estimation methods, which are subset of the class of path-based free-energy estimation methods. In bash- free-energy estimation, a central object of in the switching protocol U UA-BB.
2

---

Published as a conference paper at ICLR 2026
Definition 2.1 (Switching Protocol). Given two thermodynamic states A and B with poten-
tial energy functions UA and UB. switching protocol UA-B R3N x [0, tyl 7 R is a time-dependent potential energy function with boundary conditions U UA- 0) (2) and UA-n(I, L4) = (n)
Using such a switching protocol U BA a system x initially equilibrated with respect to the potential energy function UA can be driven from state A to state B according to the time-dependent dynamics
L,p(a, t), where L, is a time-dependent forward (Fokker-PlanckKLiouv/i the target distribution p(a;; L) xx exp (-3U A- (ar, ()) is stationary for the frozen-time dynamies, i.e. Cep(., t) = O. Notably, it does not have to be the case that the system is in equilibrium at all times, i.c. p(z, t) = p(z, does not need hold for all
dep(2,
p(z, 0) PA(x).
(3)
generator, specified so that for all I
Under such dynamics defined by the protocol UA-B, the second law of thermodynamics tells us that, with thermodynamic work defined as
" OU AUA-n((r dt, ("
W WA-B(X)=
(4)
Öt
the average work performed on the system along trajectories x generated by the dynamics in Eq. (3) an upper bound on the free-energy difference between the two states
AFA-B S (WA-B(x))U_-s-
(5)
where the shorthand (DUA-- denotes the ensemble average over trajectories x generated by the
dynamics in Eq. (3). Crucially, equality only holds here for quasi-static (infinitely slow, ły = oo) processes where the system remains in equilibrium at all times, i.c p(a, t) p(x, t) for all f.
2.1.1 JARZYNSKI EQUALITY
For any such switching protocols, Jarzynski (Jarzynski, 1997) showed the remarkable result that instead of restricting to quasi-static processes to obtain an equality for the thermodynamic work we can consider the ensemble average of the exponential of the work.
Theorem 2.2 (Jarzynski Equality (JE) (Jarzynski, 1997)). Given switching protocol UA-B and dime-dependent dynamics with frozen time stationary distribution p(æ, t) x for all t E [0,t,l as in Eq. (3), we have that along trajectories generated by the dynamics
(z: t))
--a xx(x)
BAFA 18,
UA-B
Important to note here is that the class of dynamical processes considered in the JE only require stationarity. The JE does not depend ergodicity and is therefore valid for wide range of dynamics including both time-dependent stochastic dynamics such as underdamped and overdamped Langevin well deterministic dynamics, e.g. Hamiltonian. Extra care is needed when momenta considered (e.g., underdamped Langevin dynamics). In this case, a state-dependent Hamiltonian, rather than the potential energy, should be used in the definition of work in Eq. (4).
Using standard Monte Carlo integration with trajectories x,, generated according to Eq. (3), the JE provides a consistent but generally biased estimator of the A.FA- in the finite-sample setting:
N
AF R AF AF = -8 In (-8Wa-n(xn)
JE Estimator Variance Both the bias and the variance of this estimator grow with the excess work Wee = (W)u.-R - AFA-B (Geiger and Dellago, 2010; Gore et al., 2003), which for every time t lower bounded by the Kullback-Leinn
divergence between the two distributions:
W 2 8- DKL
t)), Vt e e [0,
(8)
As such, intuitively, while the JE is valid for any switching protocol the variance and the bias of the estimator are determined by how much the instantaneous distribution p(z, t) lags behind the stationary distribution p(a, t) at every time l defined by the switching protocol U-
3

---

Published as a conference paper at ICLR 2026
2.1.2 ESCORTED JARZYNSKI EQUALITY
To reduce the amount of lag between the instantaneous distribution p(z, 1) and the target distribution p(z, 1) Vaikuntanathan and Jarzynski (2008) introduced the concept of escorted switching protocols.
Definition 2.3 (Escorted switching protocol). Given two thermodynamic states A and B with potential energy functions UA and UB. an escorted switching protocol (UA-B. b) consists of uimedeepenentt potential UA-B R3N x [0, t - and a time deeeddepndeett vector field b: RRN x [0. ty] + 3N with boundary conditions U (A+B(Z 0) (z) and UA-B(I,ty) = (n(r) and bar, 0) b( (z) = 0.
Given an escorted switching protocol, the escorted dynamics driving system from state A to state B adds deterministic advection based on the time-dependent vector field b
0,p(x, t) = t) t) = Lip(x,t) t) V ()(xx,) t)),
Ú(z,) 0) PA
(9)
Here L, denotes the forward operator that satisfies frozen time stationarity Lep(,, () 0 for p(x,, () x Sxp(-BUA-nn(C 1)), as defined Eq. (3). Under escorted dynamics, Vaikuntanathan and Jarzynski (2008) showed that with an alternative definition of the work, the JE still holds.
Theorem 2.4 (Escorted Jarzynski Equality (E-JE) Vaikuntanathan and Jarzynski (2008)). Given an escorted switching protocol (UA-BB.b) b) and lime-dependeri dynamics as in Eq. (9), we have that along trajectories generated by the dynamics
(e (e-AA-a(())
E
BAPAA-B
(10)
(UA-1..)
where the escorted work Wr WA-BB(x) along trajectory x is defined as
W WA-B(x)= = 1," (D.UA-D(MM,)) b(r,.)) (UA-B(FL))- B- MZZ de. (11)
Notably, because is not involved in enforcing that the dynamics have stationary distribution p(x, 1) x exp (-p f)) can be freely chosen to minimize the amount of lag and with that the variance of the escorted estimator AF = -3- In( En e --aa .x(x)) Notably among the choices of b are the case where O, which is equivalent the non-escorted case, as well as the following optimal choice of b as shown by Zhong et al. (2023):
Theorem 2.5 (Optimal Escorted Switching Protocol (Zhong et al., 2023)). 1fbx 1) and p(x,, 1) & exp (=sUA-u-((, ()) collectively solve the continuity equation O.p + V (pb) = 0, then, for every trajectory x generated by the dynamics Eq. (9), have WA-B( = AFA-B-
Crucially, what this shows is that if we learn b and U B collectively to solve the continuity equation, while maintaining the boundary conditions U (z, 0) = U (xx) and UA, t,) (2),, then single trajectory x generated by the dynamics in Eq. (9) suffices to estimate the free-energy difference A-AB In the remainder this work we will explore how we can learn these components using Conditional Flow Matching and our proposed Conditional Density Matching.
2.1.3 B1-DIRECTIONAL SAMPLING AND THE CROOKS IDENTITY
In addition to introducing the escorting velocity field b. a second approach to reduce the variance of our estimator of the free-energy difference is to use bi-directional estimators. Instead of only evolving samples from the initial thermodynamic state A, it is equally straightforward to start from state B and define a reverse protocol UB-A with boundary conditions UB-AA(2: (z. 0) = UB (x) and A ty) (x) In the context of non-equilibrium free-energy estimation, this idea underpins the Crooks Fluctuation Theorem (CFT) (Crooks, 1998):
PA-B(WW (W)
(-W) exp[3 (W- AFA+B)].
(12)
Here PA- (W) is the probability of observing work W for a process starting from 2o ~ PA and evolved under the forward dynamics, while PB+A(-W)) is the probability of observing work - W
4

---

Published as a conference paper at ICLR 2026
under the reme-reversed protocol. This fluctuation theorem holds both for the non-escorted case and the escorted case with WA-B(x) (x) replacing (xx).
For the escorted case, the time-reversed protocol (UB-AA, b) is given by Un-a( () Ua-B(z,t, t) and b(z, t) = b(x, ts t). Consequently the time-reversed escorted dynamics driving the system from state B to state A are given by
O,pn(2,1) = (,)n( (,,((,,))=_,n((,,1 ()
v (b(2,, (n(x,))) DR(Z,0) - (B()).
(13)
Here Ê, denotes the reverse operator of a non-escorted dynamics that satisfies frozen time stationarity L,p(z,)) t) = for p t) = p(z,t, t, t) x axp(-BUB-AC t)). In the case where dynamics with momenta are considered, the momenta have to be reversed in the reverse generator Ê.
Crucially, the Crooks Fluctuation Theorem expresses AF B as the solution single- single-parameter problem. Rather than directly estimating AFA-B via the JE, which often suffers from high variance and finite sample bias due to its exponential average, the CFT can be inverted to solve for AFa-B as the parameter that best fits given set of observed forward and reverse work for a given protocol.
The remainder of this work will only consider using the CFT instead of the JE. Practically, this implemented by the Bennett Acceptance Ratio (BAR) estimator (Bennett, 1976).
2.2 RELATED WORK
Traditional Free-Energy Estimation In addition to the traditional alchemical non-equilibrium free-energy estimation approaches introduced in the previous section, large collection of other approaches have been proposed. We briefly outline the core lines of work here. Closely related to the approaches considered in this work are alchemical equilibrium approaches such as (iterative) Free- Energy Perturbation (FEP) (Zwanzig, 1955) and Thermodynamic Integration (TI) (Kirkwood, 1935), which depend on the Zwanzig equation (Zwanzig, 1955). Next are the non-alchemical path-based approaches such transition path sampling (Bolhuis et al., 2002) and nudged elastic band sampling (Henkelman et al., 2000). Non-alchemical approaches rely on frameworks such as Transition State Theory (Truhlar et al., 1996) and the Arrhenius/Eyring equation (Eyring, 1935).
Neural Free-Energy Estimation Methods Within the family of alchemical free-energy estimation approaches. most proposed neural approaches focus on the Targeted Free-Energy Perturbation method (TFEP) (Jarzynski, 2002). The TFEP method can be considered a special case of the escorted switching protocol where only the deterministic vector field b is used to define an invertible mapping. TFEP has been studied within the machine learning context in Wirnsberger et al. (2020); Rizzi al. (2023; 2021): Erdogan al. (2024); Zhao and Wang (2023), which all use variants of MLE- trained discrete normalizing flows or Flow Matching. Within the same alchemical family, Máté et al. (2024; 2025) propose a neural version of Thermodynamic Integration. Lastly, most closely related to methods proposed here is the work by He et al. (2025), which also parameterizes an escorted switching protocol. Crucially, their proposed method FEAT uses two different protocols for the forward and backward process between two states and does not consider the multistate setting.
LEARNING ESCORTED PROTOCOLS USING FLOW AND DENSITY MATCHING
Having discussed the background of (Escortedd)) Non-Equilibrium free-energy estimators, we now propose a new method for parameterising and learning the escorted switching protocol (U% A- my b²) such that they collectively solve the continuity equation, and adhere to the boundary conditions UA (z, 0) = (z(z) and (Z, t) = Uz(z). We propose to parameterise both components by learning by using standard Conditional Flow Matching (CFM) (Lipman al., 2023) and U.? using an extension of CFM, which we call Conditional Density Matching.
3.1 LEARNING By USING CONDITIONAL FLOW MATCHING
Conditional Flow Matching (CFM) is general framework for learning a vector field of that drives samples from one arbitrary distribution Po to another p1 along a set of time-dependent intermediate distributions pt. As it can generally be assumed that there is no access to the ground truth vector field it or samples from the intermediate distributions Pt beyond the initial and final distributions, CFM approaches this by considering conditional time-dependent distribution Pr (z 2) generated by conditional vector field 2 (z, 2) and a coupling distribution q(2). A common choice is to have q(2) defined as Optimal Transport coupling (Tong al., 2024) enforce in to follow straight paths.
5

---

Published as a conference paper at ICLR 2026
Using this coupling of conditional and marginal vector fields, Lipman et al. (2023) showed that if we learn the vector field v (2t, t) by regressing on the conditional vector field By (2 t 2)) using
LCrM = En~Uni(0, =-)(()) (x)p(()))))) ( - (x 2) [)]
(14)
then this is equivalent to regressing on the vector field or (24) directly.
Crucially, given paired conditional vector field and conditional distribution, under minor regu- larity conditions, for any choice of coupling distribution q(2) the marginal vector field ey (x())) = Ne [e(dr]s)p
p((2,) 7a(=) and marginal distribution pt (2+) (pr(2el2)) are shown to jointly solve the continuity equation (Tong et al., 2024). As such, if we set po = PA P1 PB. ty (x, s) b (z, t) under the rescaling t sty, and define p. p(æ, t) x exp (- 1-- - t)). then CFM provides a valid approach to finding the escorting vector field 6 for our escorted switching protocol.
3.2 LEARNING U. USING CONDITIONAL DENSITY MATCHING
This still leaves us with the problem of learning the time-dependent potential (z, t). For this purpose, we employ a similar trick to that used by CFM to learn the time-dependent potential tr,, x exp (- a t)) from the conditional time-dependent distribution pt (z+ 2) and consider the following extended version of the Maximum Likelihood Estimation (MLE) objective
LDM (,), 2. (xx)) | log (zx, (),,
(15)
which adds an additional time dependence. To highlight the similarity with Flow Matching, we denote this object as the Density Matching (DM) objective.
Due to LDM being defined as an expectation over the unknown distribution p((2z), learning U% - (z) directly using this objective is not possible. However, similar to the CFM objective, we can equivalently express another maximum likelihood objective using the conditional distribution (a(e 2) instead. We denote this as the Conditional Density Matching (CDM) objective:
LCDM = Ee--Un((..) (). a~<(()) 2x.pp)))) [- log (Ix, f)]
(16)
Similar to the CFM objectives, the DM and CDM objectives have equivalent gradients, VoLDM VoLCDM. and thus MLE using the conditional distribution will result in the same learned marginal
distribution p as MLE using the marginal distribution.
In summary, combining the CFM objective reviewed above with the proposed CDM objective, we thus learn pair of escorting vector field (xx, I) and potential - (1r, i) that collective solve the continuity equation. When these two components are collectively used as the escorted switching protocol in an E-NEQ estimator, this results in a low- variance estimate of AF
4 EFFICIENT MULTI-STATE
ESTIMATION
With the previous section proposing a method for learning the escorted protocol (U%=.. b²). we now consider some practical restrictions of E-NEQ estimators. Namely, (i) the work calculation becoming prohibitively expensive because a very small global time step di being required to deal with the numerical instability of the stationary distribution preserving dynamics, and (ii) in the multi-state setting the number of escorted switching protocols grows exponentially with the number of states. We will address both of these issues in the following two sections. In the appendix we further discuss three more standard, but important, practical considerations.
4.1 EFFICIENT WORK CALCULATION BY LIE-TROTTER SPLITTING
Molecular Dynamics simulation known to quickly become numerically unstable when using large time-steps due to the potential energy including sharply peaked components such as Lennard-Jones potentials. As such, when simulating the dynamics of our escorted protocol in Eq. (9), we are required use very small global time-step dí. Unfortunately, as a result of this, in the work calculation we therefore must evaluate the divergence many times, which can be computationally very expensive for any system of considerable size. Ideally, we would therefore like to decouple the divergence calculation from the unstable stationary distribution preserving dynamics.
For this purpose we propose to use -iottrotter (Trotter, 1959) splitting to split the combined dynamics into two separate steps; first step using b(z, t) and then a step using the stationarity- preserving distribution preserving dynamies. Given our escorted dynamics La La & where

---

Published as a conference paper at ICLR 2026
& - (b(x, t),(x,t)) t)) is the escort transport, Lie-Trotee Trotter step of size h dt frozen at time t approximates the full evolution by composing the subflows as
h&r s e hhr eh:.
(17)
Under such split dynamics, simplified version of the Escorted Jarzynski Equality holds.
Theorem 4.1 (Split Escorted Jarzynski Equality). Given an escorted switching protocol (UA-B,b) b) and split escorted time-dependent dynamics e gbLt ss e e where La is the stationarit--rrseevvng operator (i.e. Lep(r,t) = 0 with p(z, 1) & e e-BUs- and Exp- = -V. (b(,,,, t)p) is the escort transport, consider a vie-Trotter time grid tv = kh with N h = y and the split update
d' = IK + b(Ok, hx Ek+1 Kt, (a'e)) )
(18)
where K2 is any transition kernel that preserves p(a, 1) at frozenn time L. 1f pA(x), then
(e (e-Baa-s (UA-B.b.N)
= -BAFA-FF O(h),
(19)
where the split escorted work defined by
N-1
@U @UAABB(Zk tx). - 3-1 2..... b(øx, tx) bw)) h.
W W)=
(20)
Ōt
k=0
In particular, (19) holds exactly the limit h - 0, recovering the Escorted Jarzynski Equality.
Proof. Consider the escort step as a deterministic kernel Lr, (x/(xr) 62+b(2,teM (x') and the stationarity prreseriing step as kernel Ki, with stationary distribution p(zz, tx) such that the com- posed update is the kernel Pt. = Kr, Lt.- Then the proof follows from the discrete-time E-JE in Vaikuntanathan and Jarzynski (2011).
While this still requires us to use the global time discretization well as for the calculation of the work, it allows for smaller internal time discretization within the stationarity- saatoonartty-peservingg kernel K2 independent of the escort step. This split can be used to significantly reduce the number of divergence evaluations V. b(ar, t) needed in the work calculation.
dt for each individual substep as
For reference, operator splitting is common topic in MD simulation (Frenkel and Smit). It is, for example, at the core of the BAOAB splitting for underdamped Langevin dynamics and the velocity Verlet integrator for Hamiltonian Monte Carlo (Leimkuhler and Matthews, 2013; Swope et al., 1982).
4.2 EFFICIENT MULTI-STATE FE ESTIMATION USING FLOW-GRAPH BASED MBAR
With the combined CFM and CDM approach learn the escorted switching protocol (U2- bb) between the two thermodynamic states A and B, we now consider how to extend this to the case of estimating the relative free-energy difference (AF->) namic states. Within the context of the fruggdiicoovery pipeline, lead optimization for example often involves large assays of ligands for which we need assess their binding affinity (Wang et al., 2015).
between collection of K thermody-
Naively, using single reference state h we can train K 1 escorted switching protocols [(02 b }.K b to calculate {AF, ) 1==1,==1 using the E-JE. Considering that Free-Energy differences are a state function, we can then obtain all not-directly connected estimates as AFri - -A + A.Fr-j- While this would result in a consistent estimate for all pairs of states, it is generally understood that to obtain accurate multi-state free-energy is best to obtain all [Wi+g) individually and use the self-consistent Multistate Bennett Acceptance Ratio (MBAR) (Shirts and Chodera, 2008) to obtain (AFF,,) Following this line of thinking,
a straightforward extension of the CFM/CDM framework would therefore be to learn a U.Ly and escorting vector fields b for each pair of states and j individually and use them obtain W W;-, - according to Eq. (4). However, for set of K thermodynamic states, this would require K K(K 1)/2 models, which would quickly become infeasible.
Instead, we propose to construct an Escorted Protocol Flow Graph.
7

---

Published as a conference paper at ICLR 2026
Definition 4.2 (Escorted Protocol Flow Graph (EPFG)), Given collection of K thermodynamic states and K escorted switching protocols {U b - with central node k Escorted Protocol Flow Graph (EPFG) is constructed by considering the K thermodynamic states to be the nodes and the escorted switching protocols {(Uk-s, bb= between all states and j represent edges. Here all edges not connected the central node k are given by
(1
k0 (z, by 2t), 22-2bb (1, 26)) ifo << of
(U)-)(2,t,,b
(21)
(0 k- (x, (1, 21 1y), (r, (1, if if S si
Using the EPFG and the escorted protocols defined as its edges, including the constructed concate- nated ones, we can now obtain [ ا-ررا1) for all combinations of states and while only requiring the training of K - 1 protocols. This enables more accurate free-energy estimates using the MBAR estimator compared to using the pair based approach while minimizing the training costs.
5 EXPERIMENTS
We evaluate the feasibility of learning the thermodynamic flow network of coupled (J8 and escorting field b using the well-known alanine dipeptide (ADP) system. While limited in size, ADP shows similar challenges to larger systems such as high-energy barriers and solvent effects, Furthermore, as ADP in solvent has six distinct metastable states o,, aDD 8, C5 and a' as visualized in Fig. 2. These states are known to exhibit significant differences con- formational free energy it serves as an well-understood benchmark for the multi-state setting.
100
Figure 2: Umbrella sampling FES reconstruc- tion and training samples for ADP.
Data Generation For each metastable state, we generated 10,000 samples using a harmonic flat-bottom constraint on the torsion angles defined by the metastable state boundaries, specified in the appendix. Training samples for each state are visualized in Fig. 2. We observe that the training samples cover each state but do not represent the correct ratios observed in full equilibrium sampling.
Baselines To obtain a baseline value to validate our method against, we performed an extensive Umbrella Sampling (US) estimate using large number of windows and the same setup as used for data generation. Additionally, we compare our method against a neural version of Targeted Free-Energy Perturbation (TFEP). For our comparison we will use the vector field 6d learned using Conditional Flow Matching to implement the TFEP method. As such, the only difference between TFEP and E-NEQ can come from the inclusion of the stationary distribution preserving dynamics.
Model Details We used the same model definitions and training setup for all learned escorted protocols. The escorting vector field 6 implemented using the SE(3)-equivariant graph neural network (Satorras et al., 2021) 1 with an additional learnable time-embedding component to make it time-dependent (Tan et al., 2025). We use optimal transport coupling q (Tong et al., 2024) to enforce that the escorting vector field and the time-dependent density not only collectively solve the continuity equation but also align with the dynamical optimal transport problem (Benamou and Brenier, 2000). While not discussed in depth in this work, there is close connection between the dynamical optimal transport problem and the amount of dissipated work in escorted switching protocols (Zhong et al., 2023; Zhong and De Weese, 2024). We leave exploring this interplay further for future work.
The potential U is parameterised as the negative log probability of a discrete normalizing flow with conditional affine coupling layers and similar learnable time-embedding component. We chose to use a discrete-time normalising flow here, instead of a more general energy-based model, to minimize the complexity of maximum- likelhoood training.
Integration Details For the implementation of the escorted dynamics in Eq. (9), we use the following setup. The escorting vector field is implemented in all experiments using a Runge-Kutta integrator, while we consider three different options for the stationary preserving component of the split dynamics: underdamped Langevin, overdamped Langevin, and Hamiltonian Monte Carlo, all of which use Hatropolis-Hasting (MH) correction to preserve the stationary distribution in finite time. The split operators E and L. are iteratively applied for 100 time steps, during which the work is calculated.
8

---

Published as a conference paper at ICLR 2026
Table 1: Quantitative results of the estimated AF between the central or state of the escorted flow graph the directly connected states, as as the Mean Absolute Error (MAE) over all pairs states including those not directly connected. The & column denotes the type of integrator used: Underdamped Langevin (UL), Overdamped Langevin (OL) or Hamiltonian Monte Carlo (HMC). EPFG denotes whether the estimates were obtained using pairwise summation (EPFG=N), using the escorted protocol flow graph (EPFG=Y).
Method
C
EPFG
op
C5
a'
MAE
US
7.42 0.16
12.07
0.40
-1.11 0.03
1.37 0.05
6.55 0.13
TFEP
0.005
12.39
0.06
+.00
2.39 # 0.04
5.78
1.17
E-NEQ
UL
N
7.93 #0113
12182
0.11
1.47 #.221
5.06
0.93
(ours)
UL
Y
0.12
12.54
0.78
1.31
0.06
0.88
OL
0.07
12.86 0.08
0.26
1.78
5.59
0.96
OL
Y
8.06 +.06
12.09 0.06
0.04
1.69 0.05
0.04
0.59
HMC
7.96
0.17
0.02
1.57 0.22
0.05
0.99
HMC
Y
7.33
12.89 0.16
0.50 0.09
1.83 0.20
5.94
0.95
5.1 RESULTS
In Tab. 1, we report the main quantitative results of learning the escorting vector field 6 and corre- sponding potential U for use the E-NEQ estima- tor, compared the TFEP approach. We note that the & state was chosen as the central node when constructing the flow graph. As such, all reported free-energy differences are relative to an-
Estimated (k//ml))
10
5
E-NEQ improves over TFEP for the same amount of divergence calculations Across all three integration approaches, we find that E-NEQ with learned escorted dynamics outperforms TFEP. This holds for both the directly reported AF values us- ing the central energy differences across pairs states, with state a noticeable exception. In almost all cases the found free-energy differences closely resemble the US bascline with /3 as notable exception.
E-NEQ (R'=0,91, r=0.9
-5
TFEP -0900. r-0.95)
0.0
2.5
7.5 US AF (k]/mol)
5.0
10.0
12.5
state and the MAE of the free-
Figure 3: Correlation of estimated ener-ynergy dif- ferences from E-NEQ and TFEP against the Um- hrella Sampling MBAR baseline.
further study of the correlation between the estimated free-energy differences and the true (as reported by US) free-energy differences similarly shows that the E-NEQ estimator outperforms TFEP. As highlighted in Fig. 3, both E-NEQ and TFEP show strong correlations, with R² and Pearson correlation coefficients of 0.91 and 0.96 for E-NEQ and 0.90 and 0.95 for TFEP. Both methods firmly within kcal/mol of the true free-energy differences, and most cases even within kJ/mol. This is well within the generally considered acceptable tolerances for free-energy estimation.
Concatenating Protocols improves Multistate Free-Energy Estimation Comparing the results in table 1 between using pairwise consistent estima- tion between all states using only directly connected states (EPFG=N) and using a EPFG with concate- nated protocols (EPFG=Y), we find that the latter in almost all cases more accurate. Notably, it is not only the MAE that improves, but also the estimated free-energy differences between the states directly connected to the reference a state, despite using the same trajectories.
m
Fig. 4 we visualise E-NEQ trajectories starting from the 3 state towards the az state following concatenated protocol. The trajectories show the con- catenation of the characteristically straight paths of OT-CFM starting from the 3 and visiting the state before finally transitioning the ar state.
Figure 4: Trajectories E-NEQ estimator us- ing a concatenated protocol starting from the / state towards the Yo state using the underdamped Langevin integrator.
9

---

Published as a conference paper at ICLR 2026
Table 2: Number of successful trajectories starting in each state for the different integration schemes. The maximum number of trajectories is 50,000, 10,000 to each other state.
f
dt
acc. rate
OR
a
QD
C5
a'
Underdamped Langevin
le-4
0,94
49459
48871
48383
49183
48328
49419 49539
Overdamped Langevin
lc-8
0,91
49535
49420
49265
49414
49172
Hamiltonian Monte Carlo
le-4
0,96
44056
46892
38892
40145
32477
47582
Stationary Preserving Dynamics comparison Lastly, we observe that, within the E-NEQ estima- tor with learned dynamics, underdamped and overdamped Langevin seem to be more stable than Hamiltonian Monte Carlo. Specifically, found that over longer transitions, occurs during the concatenation of protocols for multistate free-energy estimation, the HMC approach was more likely to diverge, resulting in discarded samples and ultimately higher finite sample bias. The discrepancy between the number of successful trajectories starting in each state between the different integrators reported in Tab. 2. MH correction does not help in this case due the instability primarily caused by the deterministic vector field putting the samples in high-energy regions of the learned potential. As exemplified by the extremely small step size for overdamped Langevin dynamics, our proposed Lie- Trotter splitting of the work calculation is a necessary step to make the computation feasible within an acceptable compute budget.
DISCUSSION
In this work we have proposed a method to learn the escorting switching protocol (U B-: bo) to construct an E-NEQ estimator with minimal variance. For this purpose we considered the framework of Conditional Flow Matching to learn 6c and proposed an extension named Conditional Density Matching to learn U - Furthermore, we considered two practical considerations in the form of Lie- Trotter splitting to reduce the computational cost of the work calculation and flow graph construction to reduce the combinatorial complexity of learning multiple protocols in the multi state setting. In our experimental evaluation using the ADP system, which has six well-defined metastable states, we have shown our proposed method for learning the escorted protocols to be effective.
L Limitations and future work While ADP is fitting benchmark for our method due to its multiple metastable states, similar challenges to larger systems, and generally a complexity level similar in size to studied in other related work (Rizzi et al., 2021: 2023; Máté al., 2024; Erdogan et al., 2024; He et al., 2025), is still a toy problem within the context of practical applications. We therefore emphasize the need to focus on scalability in future work. Preliminary experiments have shown that to achieve this emphasis should be placed on accurately learning the correct time-dependent potential 0% U- For this purpose we believe that recent advances in scaling discrete time normalising flows is an important avenue to explore (Rehman et al., 2025; Zhai et al., 2024).
In addition to this, we would like note that in this work we choose to align the proposed method as closely as possible to the presented theory in Sec. 2 to provide well grounded first exploration of neural escorted free-energy estimation. As such, we believe that by lifting some of these strict constraints performance could be improved. For example, removing the Eetropolis-Hasting step that used to ensure that stationary preserved under discretized dynamies, and using adaptive ODE solvers instead of Kanga-Kutter to set the lie-trotter splitting could therefore be beneficial. Both of these are common practice for standard molecular dynamic simulation. Furthermore, extensions of the current framework through hard constraints for enforcing the continuity equation instead of just encoding it in the joint CFM and CDM learning objectives, or incorporating the EPFG already during the training process could be beneficial. Lastly, we note that while we consider Flow Matching as the base for our E-NEQ estimation other generative modelling approachs in the form of Schrodinger Bridge marching and recent advances in Flow-Maps should also be considered in future work.
ACKNOWLEDGMENTS
LH is supported by the EPSRC Centre for Doctoral Training in Autonomous Intelligent Machines and Systems (EP/S024050/1), MB and LH are partially supported by the EPSRC Turing AI World- Leading Research Fellowship No. EP/X040062/1 and EPSRC AI Hub No. EP/Y028872/1. NMA is supported by an MRC Enterprise (Industrial CASE) Studentship (MSD2021 1538327) at the University of Oxford.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew J. Ballard, Joshua Bambrick, Sebastian W. Bodenstein, David A. Evans, Chia-Chun Hung. Michael O O'Neill, David Reiman, Kathryn Tunyasuvunakool, Zachary Wu, Akvilè Zemgulyté, Eirini Arvaniti, Charles Beattie, Ottavia Bertolli, Alex Bridgland, Alexey Cherepanov, Miles Congreve, Alexander Cowen-Rivers, Andrew Cowie, Michael Fig- umov, Fabian B. Fuchs, Hannah Gladman, Rishub Jain, Yousuf A. Khan, Caroline M. R. Low, Kuba Perlin, Anna Potapenko, Pascal Savy, Sukhdeep Singh, Adrian Stecula, Ashok Thillaisun- daram, Catherine Tong. Sergei Yakneen, Ellen D. Zhong, Michal Zielinski, Augustin Zídek, Victor Bapst, Pushmeet Kohli, Max Jaderberg, Demis Hassabis, and John M. Jumper. Accurate structure prediction of biomolecular interactions with AlphaFold 3. Nature, pages 1-3, May 2024. ISSN 1476-4687. doi: 10. 101011041116-00000 084-188887664447777
Tara Akhound-Sadegh Jarrid Rector-Brooks, Avishek Joey Bose, Sarthak Mittal, Pablo Lemos, Cheng-Hao Liu, Marcin Sendera, Siamak Ravanbakhsh, Gauthier Gidel, Yoshua Bengio, Nikolay Malkin, and Alexander Tong. Iterated Denoising Energy Matching for Sampling from Boltzmann Densities, February 2024.
Jean-David Benamou and Yann Brenier. A computational fluid mechanics solution to the Monge- Kantorovich mass transfer problem. Numerische Mathematik, 84(3)337--39,, January 2000. ISSN 0945-3245. doi: 10.1007/s0021 10050002.
Charles H Bennett. Efficient estimation of free energy differences from monte carlo data. Journal of Computational Physics, 22(2):245 268, 1976.
Peter G. Bolhuis, David Chandler, Christoph Dellago, and Phillip L.. Geissler. Transition Path Sampling: Throwing Ropes Over Rough Mountain Passes, in the Dark. Annual Review of Physical Chemistry, (3(1))291-3318 October 2002. ISSN 0066-426X, 1545-1593, doi: 10.1146/annurev htt.eemm....0231..1111.
Avishek Joey Bose, Tara Akhound-Sadegh Guillaume Huguet, Kilian Fatras, Jarrid Rector-Brooks, Cheng-Haa Hao Liu, Andrei Cristian Nica. Maksym Korablyov, Michael Bronstein. and Alexander Tong. Sto3)-Stcchastic Flow Matching for Protein Backbone Generation, April 2024.
Ricky T. Q. Chen and Yaron Lipman. Flow Matching on General Geometries, February 2024.
Zoe Cournia and Christophe Chipot. Applications of Free-Energy Calculations to Biomolecular Processes. A Collection. Journal of Chemical Information and Modeling, 64(7):2129-213) April 2024. ISSN 1549-9596. doi: 10.102 1ac.///asciii.4400000 4c00349.
Gavin E. Crooks. Nonequilibrium Measurements of Free Energy Differences for Microscopically Reversible Markovian Systems. Journal of Statistical Physics, 1085):141--1487, March 1998.
ISSN 1572-9613. doi: 10. (00102//A:102320211999999
Ege Erdogan, Radoslav Ralev, Mika Rebensburg, Céline Marquet, Leon Klein, and Hannes Stark. FreeFlow: Latent Flow Matching for Free Energy Difference Estimation. October 2024.
Henry Eyring. The Activated Complex in Chemical Reactions. Journal of Chemical Physics, 107-115, February 1935. ISSN 0021-9606. doi: 10. 10.1066/1.1749644,
Daan Frenkel and Berend Smit. Understanding Molecular Simulation: From Algorithms to Applica- lions.
Philipp Geiger and Christoph Dellago. Optimum protocol for fast-switching free-energy calculations. Physical Review E. 81(2):02117, February 2010. ISSN 1539-3755, 1550-2376. doi: 10.1 103/ PhysRevE.81 81.021227.
Rafael Gómez-Bombarelli, Jennifer N. Wei, David Duvenaud, José Miguel Hernández-Looatto Benjamín Sánchez- Lenghliig,, Dennis Sheberla, Jorge eguilera-Iparragunne Timothy D. Hirzel, Ryan P. Adams, and Alán Aspuru-Guzik, Automatic Chemical Design Using Data-Driven Continuous Representation of Molecules. ACS Central Science. 4(2):268-22 276, February 2018.
ISSN 2374-7943 doi: 10. 2at110011/ccccn /acc/ccesssccc.70005522

---

Published as a conference paper at ICLR 2026
Jeff Gore, Felix Ritort, and Carlos Bustamante. Bias and error in estimates of equilibrium free- energy differences from nonequilibrium measurements. Proceedings of the National Academy of Sciences, 120022)::544-12566 October 2003. ISSN 0027-8424, 1091-6490. doi: 10. 10.1003/pnas. 1635159100.
Jiajun He, Yuanqi Du, Francisco Vargas, Yuanqing Wang, Carla P Gomes, José Miguel Hernández- Lobato, and Eric Vanden -ijnden... Feat: Free energy estimators with adaptive transport. arXiv preprint arXiv:2504.11555 2025.
Graeme Henkelman, Blas P. Uberuaga, and Hannes Jónsson. A climbing image nudged elastic band method for finding saddle points and minimum energy paths. The Journal of Chemical Physics, (13(22)990119990)) December 2000. ISSN 0021-9606, 1089-7690. doi: 10.1063/1.132961
Lars Holdijk, Yuanqi Du, Ferry Hooft, Priyank Jaini, Berend Ensing, and Max Welling. Stochastic optimal control for collective variable free sampling of molecular transition paths. Advances in Neural Information Processing Systems, 36:79540-79556 2023.
C. Jarzynski. Nonequilibrium Equality for Free Energy Differences. Physical Review Letters, 78(14): 2690-26993 April 1997. ISSN 0031-9007, 1079-7114. doi: 10.1 10.100//yysReeeettt8..666
C. Jarzynski. Targeted free energy perturbatioon. Physical Review E, 65(4):046122, April 2002. ISSN 1063-65 1095-3787. doi: 10.100/PhyyRRcE.6650666222
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, Alex Bridgland, Clemens Meyer, Simon A. A. Kohl, Andrew J. Ballard, Andrew Cowie, Bernardino Romera- Paredes, Stanislav Nikolov, Rishub Jain, Jonas Adler, Trevor Back, Stig Petersen, David Reiman, Ellen Clancy, Michal Zielinski, Martin Steinegger, Michalina Pacholska, Tamas Berghammer, Sebastian Bodenstein, David Silver, Oriol Vinyals, Andrew W. Senior, Koray Kavukcuoglu, Push- meet Kohli, and Demis Hassabis. Highly accurate protein structure prediction with AlphaFold. Nature, 596(7873):583-588)) August 2021. ISSN 1476-4687. doi: 10. 10.1008/<41886-02 00044421---000838888 021-03811-11111
John G. Kirkwood. Statistical Mechanics of Fluid Mixtures. The Journal of Chemical Physics, 3(5): 300-313, May 1935. ISSN 0021-9606. doi: 10. 1063/1.1749633
Benedict Leimkuhler and Charles Matthews. Rational Construction Stochastic Numerical Methods for Molecular Sampling. Applied Mathematics Research eXpress, 2013(1):34 56, January 2013. ISSN 1687-1200. doi: 10. 10.1093/amrx/ab0000
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow Matching for Generative Modeling, February 2023.
Bálint Máté, François Fleuret, and Tristan Bereau. Neural Thermodynamie Integration: Free Energies from Energy-based Diffusion Models, June 2024.
Bálint Máté, François Fleuret, and Tristan Bereau. Solvation free energies from neural thermodynamic integration. The Journal of Chemical Physics, 162((2), 2025.
Amil Merchant, Simon Batzner, Samuel S. Schoenholz, Muratahan Aykol, Gowoon Cheon, and Ekin Dogus Cubuk. Scaling deep learning for materials discovery. Nature, 824(7990):80-85, December 2023. ISSN 1476-4687. doi: 10.1038/s41588 00088//41586002--07735533
023-06735-9,
David L Mobley and Michael K Gilson. Predicting Binding Free Energies: Frontiers and Benchmarks. 2017.
Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu. Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 365(6457) 1e44 2019.
Danyal Rehman, Oscar Davis, Jiarui Lu, Jian Tang, Michael Bronstein, Yoshua Bengio, Alexander Tong, and Avishek Joey Bose. Fort: Forward-only regression training of normalizing flows. arXiv preprint arXiv:2506. 01 2025.
12

---

Published as a conference paper at ICLR 2026
Andrea Rizzi, Paolo Carloni, and Michele Parrinello. Targeted Free Energy Perturbation Revisited: Accurate Free Energies from Mapped Reference Potentials. The Journal of Physical Chemistry Letters. (2(39):9449-9455 October 2021. ISSN 1948-7185, 1948-7185. doi: 10.1021/acs.jpcker 1c02135.
Andrea Rizzi, Paolo Carloni, and Michele Parrinello. Multimap targeted free energy estimation. Proceedings of the National Academy of Sciences, 120(46))e23043388120 Novembeeer 0027-8424, 1091-6490.
November 2023. ISSN
doi: 10.1073//pas.2304338122..
Victor Garcia Satorras, Emiel Hoogeboom, Fabian B. Fuchs, Ingmar Posner, and Max Welling, E(n) Equivariant Normalizing Flows. arXiv:2105, 09016 [physics, stat], June 2021.
Michael R. Shirts and John D. Chodera. Statistically optimal analysis of samples from multiple equilibrium states. The Journal Chemical Physics, 129(12):124105, September 2008. ISSN 0021-9606, 1089-7690, doi: 10.1063/1. 00166//.29781777
William C. Swope, Hans C. Andersen, Peter H. Berens, and Kent R. Wilson. A computer simulation method for the calculation of equilibrium constants for the formation of physical clusters of molecules: Application to small water clusters. The Journal of Chemical Physies, 76(1):637-644
January 1982. ISSN 0021-9606, 1089-7690. doi: 10.1063/1.442711
Charlie B. Tan, Avishek Joey Bose, Chen Lin, Leon Klein, Michael M. Bronstein, and Alexander Tong. Scalable Equilibrium Sampling with Sequential Boltzmann Generators, February 2025.
Alexander Tong, Kilian Fatras, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector- Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport, March 2024.
G.M. Torrie and J.P. Valleau. Nonphysical sampling distributions Monte Carlo free-energy estimation: Umbrella sampling. Journal of Computational Physics, 23(2):187-199, February 1977. ISSN 00219991. doi: 10. (220116/002-999911(19911111
Hale F Trotter. On the product of semi-groups of operators. Proceedings of the American Mathemati- cal Society, 10(4)5545-551, 1959.
Donald G. Truhlar, Bruce C. Garrett, and Stephen J. Klippenstein. Current Status of Transition- State Theory. The Journal of Physical Chemistry, 100(31): January 1996. ISSN 0022-3654. doi: 10.1021/jp9537484
Suriyanarayanan Vaikuntanathan and Christopher Jarzynski. Escorted Free Energy Simulations: Improving Convergence by Reducing Dissipation. Physical Review Letters, 100619)) May 2008. ISSN 0031-9007, (079-71144 doi: 10. 100110//PyysReeLett.00..... 100.10000ee
Suriyanarayanan Vaikuntanathan and Christopher Jarzynski. Escorted free energy simulations. The Journal of Chemical Physics, 134(5):054107, February 2011. ISSN 0021-9606, 1089-7690. doi: 10.1063/1 15.44677..344469. .3544679.
Jiff Vymētal and Jiří Vondrášek. Metadynamics As Tool for Mapping the Conformational and Free-Energy Space of Peptides The Alanine Dipeptide Case Study. The Journal of Physical Chemistry B, 11((16):563-5642 April 2010. ISSN 1520-6106. doi: 10. 10.1021/jp1009500
Fangping Wan, Daphne Kontogiorgos- Heintz, and Cesar de la models for peptide design. Digital Discovery, 1(3):195-25) 208, June 2022. ISSN 2635-098X. doi:
Deep generative
10.1039/DIDD00024A.
Lingle Wang. Yujie Wu, Yuqing Deng. Byungchan Kim, Levi Pierce, Goran Krilov, Dmitry Lupyan, Shaughnessy Robinson, Markus K. Dahlgren, Jeremy Greenwood, Donna L. Romero, Craig Masse, Jennifer L. Knight, Thomas Steinbrecher, Thijs Beuming, Wolfgang Damm, Ed Harder, Woody Sherman, Mark Brewer, Ron Wester, Mark Murcko, Leah Frye, Ramy Farid, Teng Lin, David L. Mobley, William L. Jorgensen, Bruce Berne, Richard A. Friesner, and Robert Abel. Accurate and Reliable Prediction of Relative Ligand Binding Potency in Prospective Drug Discovery by Way of a Modern Free-Energy Calculation Protocol and Force Field. Journal of the American Chemical Society, 137(7):2695-2700 February 2015. ISSN 0002-7863. doi: 10. 10.101//j5112555
13

---

Published as a conference paper at ICLR 2026
Peter Wirnsberger, Andrew J. Ballard, George Papamakarios, Stuart Abercrombie, Sébastien Racanière, Alexander Pritzel, Danilo Jimenez Rezende, and Charles Blundell. Targeted free energy estimation via learned mappings. The Journal of Chemical Physics, 153(14):144112, October 2020. ISSN 0021-9606, 1089-7690. doi: 10. 00.1063/5.0018003
Shuangfei Zhai, Ruixiang Zhang, Preetum Nakkiran, David Berthelot, Jiatao Gu, Huangjie Zheng, Tianrong Chen, Miguel Angel Bautista, Navdeep Jaitly, and Josh Susskind. Normalizing flows capable generative models. arXiv preprint arXiv:24122.06329, 2024.
Lu Zhao and Lei Wang. Bounding free energy difference with flow matching. Chinese Physics Letters, 12020,,, December 2023. ISSN 0256-307X, 1741-3540. doi: 10.1088/0256-007X 20222/120201
Adrianne Zhong and Michael R. DeWeese. Beyond Linear Response: Equivalence between Thermo- dynamic Geometry and Optimal Transport, April 2024.
Adrianne Zhong, Benjamin Kuznets-Speck, and Michael R. DeWeese. Time-Asymmerric Fluctuation Theorem and Efficient Free Energy Estimation, December 2023.
Robert W Zwanzig. High-temperature equation of state perturbation method. ii polar gases. The Journal of Chemical Physics, 23(10):191--192,, 1955.
A APPENDIX
A.I BENNETT ACCEPTANCE RATIO (BAR) AND MULTISTATE BENNETT ACCEPTANCE RATIO (MBAR)
In this section we briefly introduce the Bennett Acceptance Ratio (BAR) and Multistate Bennett Acceptance Ratio (MBAR) estimators introduced in the main text. First, the BAR estimator be interpreted as a maximum likelihood solution for the logistic model that is implied by the Crooks Fluctuation Theorem (CFT), Given collection of forward x and backward x trajectories and their corresponding work values WA Wa-n(xxf and WB-A( respectively, BAR solves:
+ exp (WABB (x,) AF) -) APD)
1 + exp (B AF) - (22)
where e Na and is equal to 0 for all our experiments. To find AF we can now use standard root-finding algorithms. The approximate solution for FF found using this estimator is known to have minimal variance among estimators that use both forward and backward trajectories (Bennett, 1976).
N
Adjusted to the multistate setting, the MBAR estimate is given by the set of equations (Shirts and Chodera, 2008):
N
exp (-3Wa-,(x
Fx = -- In
(23)
K EE exp (Fi-8Wi-p
3W: (Wi+)((x))
where W (x(+(x( is the work done on the n-th trajectory moving from state i to state j. Note that here we obtain absolute free-energies Fi with respect the reference state k. Thus, obtain the free-energy differences between all pairs of states, we can use the following equation:
AF F, Fg.
(24)
A.2 VALIDITY OF METHOD
In Sec. 3 we showed approach leam the potential U and vector field 5 that can collectively be used as an escorted protocol for free-energy estimation. For this purpose we aligned the objective of learning the vector field with standard Conditional Flow Matching, and introduced extension
14

---

Published as a conference paper at ICLR 2026
named Conditional Density Matching to learn the potential density. In this section provide more formal derivation to show that this forms a valid method.
For this purpose first give formal derivation of Conditional Density Matching, and subsequently show how jointly using CFM and CDM provides a valid escorted protocol.
A.2.1 CONDITIONAL DENSITY MATCHING
With q(2) the coupling distribution and (z 2) the conditional probability path generated by conditional vector field (x 2), the marginal probability path is defined as follows:
pe(xt) (a(x))- (atx)=-Daaee pa(
=) ((=) dz.
(25)
Our objective is to parameterize density pl (or equivalently, a potential U. via pf exp(-BU()) and train it to approximate Pt using Maximum Likelihood Estimation. We call the approach to directly do this Density Matching,
LDM Ez-Uni(O, 2..p. (w.) [- log p (I,, t)]
(26)
but as discussed in the main body of the paper, this is infeasible as we are not able to directly sample from pt- This is where our alternative learning method, Conditional Density Matching (CDM), comes in. CDM defined using the conditional probability path and the coupling distribution, which
we have direct access to:
LCDM (m,,,)(,,1)) (mg()) (empp(ll) [- log 1 (ate, Q)]
(27)
Crucially, as mentioned in the main paper but not further elaborated on, the DM and CDM objectives have equivalent gradients, VoLDM = VoLCDM- Similar to the argument used for CFM, we can thus use the CDM loss for optimization and consequently also optimize for the DM loss.
We now formally show that these gradients are the same.
Proposition A.1 (CDM Gradient Equivalence). The Density Matching and Conditional Density Matching objectives,
LDM Eu-.Ui((011 logy" (oo log (Iz. (xr.
(28)
CCDM
(29)
have identical gradients with respect 6:
VoLDM = V LCDM-
(30)
Proof. We show the stronger statement that LDM = LCDM- Starting from LCDM:
LCDM Uni(0,(, 2n -- log (zr,t)
(31)
[/ po(xe 22)(() (- log (x), dz dr
E,-Uni(a,1)
(32)
[/ (f pelore 2) =))(() dz (- log (xx, 1)) dra
Be-Uni(O,))
(33)
Em-uni(o.1) / pa(aas) (- -logog (ogg((xx,))) dra
(34)
EL lologllog log (z,,
(35)
LDM-
(36)
Following this, now directly have that VoLDM VoLCDM.
What is crucial here is that we can do this at the same time as the learning of the velocity field bo, At every epoch, we sample (I,)
q(2), t ~ Uni(O, 1), and consequently use our definition of the conditional probability path to obtain H, ~ (x, 22), which can now be used to learn both b
using the CFM loss and nd using the CDM loss.
15

---

Published as a conference paper at ICLR 2026
A.2.2 JOINTLY USING CFM AND CDM PROVIDES VALID ESCORTED PROTOCOL
We now focus on how CFM and CDM can be used to jointly train the escorting vector field 6 via CFM and the time-dependent potential U via CDM to produce a pair (U%_,) bo) that provides valid escorted protocol as defined Thm. 2.3. We first recall the key criteria for such a protocol:
1. Boundary conditions on U: U A- (z, 0) = (x( and UA-B(II = -n(())
2. Boundary conditions on b: b(z, 0) b(z
Notably, if these two boundary conditions are satisfied, any choice of (UA- b) works as valid escorted switching protocol. A hard requirement on how the two are related only comes in when we want to achieve zero variance in the estimator, as discussed in Thm. 2.5. In this case we add the following requirement:
3. Continuity equation: exp(-3UA-B(2,0))
0pp(z,t) + V ((a,, 0, where p(r,t) &
As the two boundary conditions can trivially be satisfied, we now show that the CFM and CDM learning frameworks collectively target the continuity condition. For this purpose we first restate well-known result from Conditional Flow Matching:
Theorem A.2 (Conditional continuity equation implies marginal continuity equation). Let q(2) be coupling distribution over pairs z = (æo, x1) with zo ~ P and # ~ pB, and let the pair
{pe(aa) 2). (a(; 2)} be conditional probability path and conditional vector field such that for each z the conditional continuity equation
Oepe(a 2) + (pe(ax 2) (e 1=)
(37)
is satisfied. Then with the marginal probability path and marginal vector field defined as
Pe(xe) (Po( (2))
(38)
(x(x)) uz (1 pe 22)
(39)
(((
q(x)
the pair (pe(xe), (x())) collectively solves the marginal continuity equation
dipc(æt) + (pr(xx) ve(xz)) 0.
(40)
Proof. See, for example, Tong et al. (2024).
Note here that it is crucial that the conditional probability path (z,, z) and conditional vector field ve (2t 2) are chosen consistently, i... such that they jointly satisfy the conditional continuity equation (37). A standard choice is the optimal transport (OT) conditional path with linear interpolation It (1 t) + 1. for z = (mo, 21), which gives
(= 2) = (a (1 t).o + txr,, o²).
(x 2) x To.
(41)
where a > 0 is small fixed variance. This pair satisfies (37) by construction, as the conditional density simply transported along the constant velocity field 0 = I1 Zo.
We can now state the validity of the joint CFM/CDM escorted protocol as direct consequence of the above Theorem A.2 and Theorem A.I:
Corollary A.3 (Validity of joint CFM/CDM escorted protocol). Under the same conditions as the above theorem, let o6 minimise Lcrm and p" minimise LCDM- Then, under the identification 60(x, 1) 1 1 1/tz) and pº (x, t) xx exp(-3U* (x, ()), the escorted protocol (U bo* satisfies the conditions of Theorem 2.5.
Proof. By the standard CFM result, 80" (x,, (xx) recovers the marginal vector field. By Thm. A.1 and the optimality of maximum likelihood estimation, pa (#x, = px(2)) recovers the marginal density. By the above theorem, this pair jointly solves the continuity equation 8,pt + V
(p, v,) = 0. The time rescaling preserves the continuity equation, giving 0,p + V (p" 5o = which is exactly the condition of Theorem 2.5.
16

---

Published as a conference paper at ICLR 2026
The above analysis shows that, with appropriately chosen conditional probability path and condi- tional vector field that collectively solve the conditional continuity equation, we recover an optimal escorted switching protocol.
A.3 FULL ALGORITHMIC OVERVIEW
Below we briefly provide the full algorithmic overview for multistate free-energy estimation using E-NEQ. Code is available at: http:/ghhhubooom//ssoooooii/////
A.3.1 TRAINING
Generating Training Samples. Given set of K thermodynamic states of interest and their corresponding potential energies {U%
] k= generate dataset of N₂ equilibrium samples
{x W. from each state & using standard molecular dynamics simulations.
2. Learning K-1 Escorted Protocols. Select one of the thermodynamic states as the reference state & and learn the K escorted protocols { {U R-o by (-))):ss Jiyk using Conditional Flow Matching and Conditional Density Matching as proposed in Sec.
3. Constructing the Flow Graph. Construct the fully connected Escorted Protocol Flow Graph by concatenating the learned protocols to define edges between all pairs of states, as defined Eq. (21).
A.3.2 OBTAINING A
1. Running the Escorted Protocols. For each pair of states (i.j), run fixed-length non- equilibrium trajectories using the escorted dynamics in Eq. (9) (forward) and Eq. (13) (reverse), and compute the escorted work for each trajectory.
2. Estimating Free-Energy Differences. Use the collected forward and reverse work values )
(W-+)) 1,1=1 with the MBAR estimator obtain a self-consistent estimate of the free- energy differences between all pairs of states.
A.3.3 PRACTICAL DETAILS
So far we have assumed each thermodynamic state to have its own potential energy function Us. In practice, this is only the case for small number of free-energy estimation problems, such as solvation free-energy. For most other forms free-energy, such as conformational free-energy which we consider in experimental evaluation, the states have the same potential energy U but correspond to different regions of the phase-space. In this case, a state-restricted potential energy function U, has to be defined as:
E Ω
U.(x) = [ x $ Ω4
(42)
where Ω2 is the region of space that is considered to be part of thermodynamic state i. Restricting potential energies to specific regions of space is a common technique in Molecular Dynamics simulations (Torrie and Vallcau, 1977).
Equivariances Molecular systems are in general invariant to rigid-body transformations, such as rotations and translations. Including these transformations in architecture design of 6d and U is thus desirable and can significantly improve the performance of the method. However, in the case of Flow Matching extra care needs to be taken to ensure that the target vector field v also takes these transformations into account. This is formalised in the framework of Riemannian Flow Matching (Chen and Lipman, 2024; Bose al., 2024).
Thermodyname Irrelevant Degrees of Freedom In addition to the reduction in complexity that can be achieved by considering the symmetry structure of the system, an additional inductive bias can be introduced by considering the degrees of freedom of the system that do not contribute to the free-energy difference between the states interest. For example, in the case Alanine Dipeptide (ADP) number of Carbon atoms each connected to a single other heavy atom and have their remaining valences satisfied by hydrogen atoms. While these Hydrogen atoms considerably fluctuate during the simulation and cause significant spikes in the potential energy, they generally have
17

---

Published as a conference paper at ICLR 2026
uniform contribution to the free-energy across all the states of interest. As such, we do not have to include them in the representation the system used for U and b".
Note that removing degrees of freedom that are considered thermodynamic irrelevant has to be done with care, as incorrect removal can have detrimental effect on the performance of the method. For example, while the atoms in the ADP system attached to terminal Carbon atoms are considered thermodynamic irrelevant, the same not true for the atoms attached to central heavy atoms.
A.4 EXPERIMENTAL SETUP
A.4.1 MOLECULAR DYNAMICS SIMULATION DETAILS
The simulation was run at 300 K using a standard Langevin middle integrator with a BAOAB splitting scheme and time step of 2 fs. Samples were saved every 1000 steps. The Amber 14-alndx force-field was used with the GBn2 implicit solvent model. The same force field and implicit solvent model used the U and UB potentials in the Umbrella Sampling simulations."
A.4.2 METASTABLE STATE DEFINITIONS
For the Alanine Dipeptide (ADP) system, we use the following definition of the metastable states based on the & and & dihedral angles:
0n = { - 120 < S 0, -100 S 1 S 90)
(43)
3 = -120 - o
(44)
C5 f {-180 م -220
< Φ --10) - -200000000
180)
U {120 5 Φ 180, -180 S S -110) U [120 < < 180)
(45)
a = {-180 < -120, -110 < 90) U {120 S o 180, 110 - 90)
(46)
a1 (0 S $ 120, -90
90 < < 90}
(47)
= {0 Φ < 120, 180 S Φ - -90 0 < Φ < 120,90 90 S - 180}
(48)
adapted from (Vymētal and Vondrášek, 2010)
We have visualised the 5 metastable states in figure 2.
A.5 STANDARD JARZYNSKI EQUALITY RESULTS
As an additional baseline to compare our E-NEQ method against we consider using only the potential 7 learned with the proposed density matching method the switching protocol in the Jarzynski Equality. We only consider overdamped Langevin dynamics, and compare 3 different number of interpolation steps: 500, 5,000 and 50,000. The results are presented in Tab. 3.
Even at 50,000 steps, JE does not reach the accuracy of E-NEQ at 100 steps (MAE 0.93), highlighting the benefit of the learned escorting field.
Table Extended Baseline quantitative results the estimated AF between central au of escorted flow graph and the directly connected states, as well as the Mean Absolute Error (MAE) over all pairs of states including those not directly connected.
Method
steps
ap
C5
a'
MAF
US
7.42 + 0.16
12.07 0.40
.1.11 + 0.03
1.37 a 0.05
6.55 0.13
TEEP
8.60 = 0.05
12.39 0.06
0.04
2.39
0.08
1.17
JE
500
0.75
1.62 0.57
2.13
-2525
4.25
6.16
JE.
5,000
-0.21 + 0.79
16.83 0.94
7.88 0.46
4.29
4.95
7.73
JE
50,000
3.27 0.61
13.82
3.07 3.0741.4
-.39
4.58
3.92
18


Reviews:
[{"review_id": "tzfsxxSn4x", "reviewer": "Reviewer_k7x6", "summary": "This paper proposes a method for learning escorted switching protocols to estimate free-energy differences between multiple thermodynamic states. The approach combines Conditional Flow Matching (CFM) to learn a deterministic escorting vector field and introduces Conditional Density Matching (CDM) to learn a time-dependent potential. The authors also propose Lie-Trotter splitting to reduce computational cost and a flow graph construction to scale to multiple states. Experiments on alanine dipeptide (ADP) show improvements over Targeted Free-Energy Perturbation (TFEP) in terms of estimation accuracy.", "questions": "- Why is the Metropolis–Hastings correction necessary here, and how is it made computationally feasible in high-dimensional settings?\n- How are the Langevin dynamics discretized? What step size, friction, and stochastic gradient estimation schemes are used?", "limitations": "", "rating": 2, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "- The paper tackles an important problem in computational biochemistry: efficient and accurate free-energy estimation across multiple states.\n\n- The idea of jointly learning the escorting vector field and the time-dependent potential via CFM and CDM is a meaningful attempt to reduce variance in free-energy estimators.\n\n- The use of Lie-Trotter splitting and flow graph construction addresses practical challenges in computational cost and scalability.\n\n- Experimental results on ADP show that the proposed E-NEQ estimator outperforms TFEP in several settings.", "weaknesses": "Overall, it is hard for me to follow this paper, and I think there are many related but not well discussed in this work. But I would be open to adjusting my score if the authors can address my concerns.\n\n- It is very hard for me to identify what the actual contribution of this paper is.\n  - If I understand the work correctly, Sections 2 and 3 both appear to be preliminary material rather than novel methods. \n  - Section 4 supposedly presents the proposed algorithm, but the description is vague and lacks a clear algorithmic definition. I would recommend a pseudo-code to present the algorithm clearly.\n  - It is unclear to me what exactly is being optimized or learned, and what the goal of this paper is to address.\n- The methodology is vague to me.\n  - Section 3 mentions flow matching but does not clearly define the forward and backward processes, how the reverse flow is parameterized, or how training is performed.\n  - Section 4 seems to involve a neural network trained with ideas similar to parallel tempering, but no pseudo-code or mathematical formulation is provided.\n  - Section 5 discusses Overdamped & Underdamped Langevin and Hamiltonian Monte Carlo, which makes the algorithm a bit clearer. But it is unclear to show the continuous-time SDEs, their discretization schemes, or how stochastic gradients are estimated from data.\n  - The use of Metropolis–Hastings correction seems questionable given the high computational cost in high-dimensional generative models, yet the paper does not justify this design choice.\n- There is no dedicated related work section, and the paper does not discuss how it differs from prior research in replica exchange or parallel tempering methods. To me, the approach appears very similar in spirit to parallel tempering/replica exchange MCMC methods [1–4], which also sample across multiple states to balance exploration and exploitation. The paper does not clarify whether the proposed method offers any theoretical or empirical advantage over the related work.\n- There are some undefined notations or some presentation issues. If I missed some, please let me know where they are. For example, I did not find some important quantities such as $\\alpha_R$, $\\alpha_L$, $\\alpha_D$, $\\beta$, $C_5$, and $\\alpha'$ in Table 2, and variables $\\phi$ and $\\psi$ in Figure 4. This makes the results hard to interpret for readers unfamiliar with Alanine Dipeptide.\n\nReferences:\n\n[1] Accelerating Nonconvex Learning via Replica Exchange Langevin Diffusion. ICLR 2019\n\n[2] Non-convex Learning via Replica Exchange Stochastic Gradient MCMC. ICML 2020\n\n[3] Spectral Gap of Replica Exchange Langevin Diffusion on Mixture Distributions, Stochastic Processes and their Applications, 2022 \n\n[4] Non-reversible Parallel Tempering: A Scalable Highly Parallel MCMC Scheme, Journal of the Royal Statistical Society Series B: Statistical Methodology, 2022", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AupseoV4wg", "reviewer": "Reviewer_6Qux", "summary": "The authors propose a neural formulation of the Escorted Jarzynski Equality (E-JE) to estimate free-energy differences with reduced variance. The authors reinterpret the E-JE framework using neural potentials $U_\\theta$ and escort fields $b_\\phi$, trained via conditional flow matching to approximate the optimal nonequilibrium protocol. The paper is well written, with clear theoretical exposition and solid mathematical grounding, providing an intuitive bridge between nonequilibrium statistical mechanics and neural flow–based modeling. Experiments on the alanine dipeptide system demonstrate the feasibility of the approach and compare its performance with classical methods such as Umbrella Sampling and TFEP.", "questions": "1. **JE baseline and variance reduction**\nSince the standard Jarzynski Equality (JE) corresponds to the case ($b$ = 0), could you include this baseline to quantify the actual variance reduction achieved by the learned escort field?", "limitations": "", "rating": 4, "confidence": 2, "soundness": 1, "presentation": 2, "contribution": 2, "strengths": "1. Provides clear and pedagogical background on the Jarzynski Equality (JE) and its escorted extension, making a technically challenging topic highly accessible.\n2. Builds on a solid theoretical foundation, with mathematically consistent derivations and well-motivated neural extensions.\n3. The formulation elegantly connects non-equilibrium statistical mechanics with modern neural flow modeling, offering a promising bridge between physics-based and data-driven approaches.", "weaknesses": "**Limited discussion of related work and positioning**\nThe paper presents a neural formulation of the Escorted Jarzynski Equality (E-JE), but does not clearly situate it within the broader landscape of neural flow, Schrödinger bridge, or optimal control–based free-energy estimation methods. Moreover, the boundary between background explanation and related work discussion is somewhat blurred, making it difficult to discern which parts of the formulation are newly contributed by this paper versus prior theoretical developments. A clearer distinction between foundational background and the paper’s novel contributions would help readers better appreciate the originality of the work.\n\n**Lack of verification for continuity-equation consistency**\nA central theoretical condition of the proposed method is the continuity equation\n$$\\partial_t p_\\theta + \\nabla \\cdot(p_\\theta b_\\phi)=0,$$\nensuring that $U_\\theta$ and $b_\\phi$ form a consistent escorted protocol. However, the paper does not quantitatively evaluate or visualize how well this condition is satisfied after training.\n\n**Missing JE ($b$ = 0) baseline — variance reduction unquantified**\nThe main purpose of E-JE is to reduce variance relative to the standard JE, but the JE ($b$ = 0) baseline is entirely omitted from experiments.\nAlthough both yield identical $\\Delta F$ values, without reporting $\\mathrm{Var}(e^{-\\beta W})$ it is impossible to demonstrate the actual benefit of the learned escort field.\n\n**Limited experimental diversity and missing ablations**\nExperiments are restricted to a single system (alanine dipeptide) and comparisons are limited to classical methods such as TFEP and umbrella sampling. There is no evaluation against recent neural baselines or component-wise ablations.\n\n**Insufficient explanation for separating $b_\\phi$ and $U_\\theta$**\nIt is unclear why $b_\\phi$ (flow field) and $U_\\theta$ (potential) must be trained as independent networks rather than being coupled through a gradient or consistency constraint. A brief discussion on the physical or algorithmic motivation for this design choice (e.g., Helmholtz decomposition or non-conservative flow representation) would improve clarity.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VIFudb2dsd", "reviewer": "Reviewer_mX26", "summary": "The paper learns escorted non-equilibrium protocols for free-energy estimation by pairing Conditional Flow Matching for the escort field with a “Conditional Density Matching” objective for the time-dependent potential so the pair approximately satisfies the continuity equation. It also adds two pragmatic pieces: Lie–Trotter splitting to make work evaluation cheaper and an Escorted Protocol Flow Graph (EPFG) so K states only need K–1 learned protocols but still feed MBAR with all pairwise work values. On alanine dipeptide, the approach improves accuracy over TFEP and benefits from EPFG concatenation.", "questions": "- How is the density matching (DM) in Equation 15 different from maximum likelihood estimation? (I’m confused by the terminology—why call it “Density Matching” when Eq. (15) is a negative log-likelihood over \\((x,t)\\)?) In line 263, the authors say “MLE”—that’s maximum likelihood estimation, right (still need to define it!)? \n- The method learns both $b_\\phi(x,t)$ and $U_\\theta(x,t)$ to “collectively” solve the continuity equation. Is there an identifiability or degeneracy issue—i.e., multiple $(b,U)$ pairs giving similar likelihood but very different work statistics? What regularization or constraints help?\n- In the Lie–Trotter split, how large can the time step be before the discrete E-NEQ estimator’s bias shows up meaningfully in MBAR? Any diagnostic or adaptivity you recommend during training/evaluation? \n- EPFG builds all pairwise work values from K–1 learned edges. Do we lose anything versus learning each pair directly (e.g., variance concentration on hard pairs), and could the graph topology (choice of hub k) be learned or adapted?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Overall, well-written paper.\n- Jointly learn the escort vector field and the time-dependent potential (CFM + CDM) so the learned pair targets the continuity equation, which is exactly what E-NEQ needs.\n- Lie–Trotter splitting reduces the number of divergence evaluations in the work computation—practically important for MD-style forces.\n- Train only $K–1$ protocols yet still populate MBAR with all $W_{i\\to j}$, cutting training cost while improving multi-state consistency.", "weaknesses": "- The empirical validation is limited to ADP; it’s not obvious how the training and stability behave on larger biomolecular benchmarks. \n- Some integrator choices remain touchy on long concatenations.\n- The CDM/DM construction is essentially an MLE on a time-indexed density $p_\\theta(x,t)$; the paper introduces new terminology (“Density Matching”) but the core estimator is maximum likelihood on $(x,t)$ samples (later replaced by conditional sampling). That raises questions about identifiability (many $(b_\\phi,U_\\theta)$ pairs can satisfy the continuity equation approximately) and about how model misspecification in $p_\\theta(x,t)$ propagates into work-estimate bias.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Dj9c4APXBC", "reviewer": "Reviewer_X5Rb", "summary": "A paper proposing an original solution to an important computational physics problem, namely computing free energy changes in non-equilibrium thermodynamics. The proposed solution uses flow matching techniques in an innovative way to suggest a more efficient/ accurate algorithm to sample Monte Carlo trajectories over which to estimate such free energy changes.", "questions": "- Is it conceivable that your algorithmic improvements to learning the escorting protocol could be extended to other common ML scenarios where distributions are to be matched, e.g. diffusion models?\n- Your approach is a practical solution to minimising the bias/ variance of the Monte-Carlo Jaczinsky estimator, does it come with guarantees/ special scenarios in which it could turn out to be exact?\n- I didn't get much rationale for the choice of flow architecture etc, which potentially could be a factor in determining efficiency gains. Did you just take an off-the-shelf approach or are there mileage in optimising that side?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "I really enjoyed reading this paper, albeit as a somewhat knowledgeable outsider. The authors make a good effort to provide a comprehensive background accessible to non-specialists (although I wonder how well it fares with non-physicists). The idea of using ML to minimise the Jaczinsky lower bound is elegant and the proposed improvements in terms of adopting larger time steps/ extending to multi-state estimation are potentially important for the community. The empirical validation is well carried out albeit not extensive (similar in that to physics papers)", "weaknesses": "The main weakness to me is how much this paper could appeal outside of the computational physics community, which is but a (small) strand in the ICLR community. Some opportunities to broaden the appeal of the paper are listed below in the questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)