Energy Matching: Unifying Flow Matching and Energy-Based Models for Generative Modeling
Authors: Michal Balcerak
Tamaz Amiranashvili
Antonio Terpin
Suprosanna Shit
Lea Bogensperger
Sebastian Kaltenbach
Petros Koumoutsakos
Bjoern Menze
Venue: NeurIPS 2025 poster ()
Keywords: Generative Modeling
Energy-Based Models
Flow Matching
Optimal Transport
Inverse Problems
Likelihood Modeling
Primary area: deep_learning
TL;DR: Energy Matching unifies flow matching and energy-based models in a single time-independent scalar field, enabling efficient transport from noise to data while retaining explicit likelihood information for flexible, high-quality generation.

Abstract: Current state-of-the-art generative models map noise to data distributions by matching flows or scores. A key limitation of these models is their inability to readily integrate available partial observations and additional priors. In contrast, energy-based models (EBMs) address this by incorporating corresponding scalar energy terms. Here, we propose Energy Matching, a framework that endows flow-based approaches with the flexibility of EBMs. Far from the data manifold, samples move from noise to data along irrotational, optimal transport paths. As they approach the data manifold, an entropic energy term guides the system into a Boltzmann equilibrium distribution, explicitly capturing the underlying likelihood structure of the data. We parameterize these dynamics with a single time-independent scalar field, which serves as both a powerful generator and a flexible prior for effective regularization of inverse problems. The present method substantially outperforms existing EBMs on CIFAR-10 and ImageNet generation in terms of fidelity, while retaining simulation-free training of transport-based approaches away from the data manifold. Furthermore, we leverage the flexibility of the method to introduce an interaction energy that supports the exploration of diverse modes, which we demonstrate in a controlled protein generation setting. This approach learns a scalar potential energy, without time conditioning, auxiliary generators, or additional networks, marking a significant departure from recent EBM methods. We believe this simplified yet rigorous formulation significantly advances EBMs capabilities and paves the way for their wider adoption in generative modeling in diverse domains.

Energy Matching: Unifying Flow Matching and Anergy-Bassee Models for Generative Modeling
Michal Balcerak University of Zurich michal. balcerak@uzh.cc ch
Tamaz Amiranashvili University of Zurich Technical University of Munich
Antonio Terpin ETH Zurich
Suprosanna Shit University of Zurich
Lea Bogensperger
Sebastian Kaltenbach
University of Zurich
Harvard University
Petros Koumoutsakos
Bjoern Menze University of Zurich
Harvard University
Abstract
Current state-of-the-art generative models map noise to data distributions by match- ing flows or scores. A key limitation of these models is their inabiliity to readily integrate available partial observations and additional priors. In contrast, energy- based models (EBMs) address this by incorporating corresponding scalar energy terms. Here, we propose Energy Matching, framework that endows flow-based approaches with the flexibility of EBMs. Far from the data manifold, samples move from noise to data along irrotational, optimal transport paths. As they approach the data manifold. equilibrium distribution, explicitly capturing the underlying likelihood structure the data. We parameterize these dynamics with single dime-independent scalar field, which serves as both a powerful generator and flexible prior for effective regularization of inverse problems. The present method substantially outperforms existing EBMs on CIFAR-10 and ImageNet generation terms of fidelity, while retaining simulation-fiee free training of transport- based approaches away from the data manifold. Furthermore, we leverage the flexibility of the method to introduce interaction energy that supports the exploration of diverse modes, which we demonstrate in a controlled protein generation setting. This approach learns a scalar potential energy, without time conditioning, auxiliary generators, or additional net- works, marking significant departure from recent EBM methods. We believe this simplified yet rigorous formulation significantly advances EBMs capabilities and paves the way for their wider adoption in generative modeling in diverse domains.
entropic energy term guides the system into Boltzmann
Introduction
Generative models learn to map from a simple, easy-to-sample distribution, such as a Gaussian, to desired data distribution. They do so by approximating the optimal transport (OT) map-such as in flow matching [Lipman et al. 2023 Liu et al.) 2023 Albergo and Vanden- 2023 or through iterative noising and denoising schemes, such as diffusion models [Ho et al., 2020 Song
et al. 2021]. In addition to being highly effective in sample generation, diffusiion- and flow-base
models have also been used as priors to regularize poorly posed inverse problems [Chung et al. 2023
Mardani et al. [204]] nen-Hamu et al. [224]. However, these models do not explicitly capture the unconditional data score and instead model the score of smoothed manifolds different noise levels.
The measurement likelihood, on the other hand, is not tractable on these noised manifolds. As
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Action Matching
OT-Flow Matching
Energy -aased Model
Energy Matching (Ours)
80(2, 1) E
(z t) E Rd
Ex( E
Vo(a) E R
Figure 1: Trajectories (green lines) of samples traveling from a noise distribution (black dots; here, Gaussian mixture model) to data distribution (blue dots; here, two moons in TTng all 2023
under four different methods: Action Matching [Neklyudov et al., [002]] Flow Matching (OT- CFM
[Tog et al. 2023]. EBMs trained via contrastive divergence [Hinton 2002 and proposed Energy Matching. We highlight several individual trajectories red to illustrate their distinct behaviors. Both Action Matching and Flow Matching learn time-dependent transports and are not trained for traversing the data manifold. Conversely, EBMs and Energy Matching are driven by time- inneinndeeendent fields that can be iterated indefinitely. allowing trajectories to navigate across modes. While samples from EBMs often require additional steps to equilibrate (see, e.g., the visible mode collapses that slow down sampling from the data manifold), Energy Matching directs samples toward the data distribution in "straight" paths, without hindering the exploration of the data manifold.
result, existing approaches repeatedly shuttle between noised and data distributions, leading to crude approximations of complex, intractable terms Daras et al. [2024). For example, DPS [Chung et al., 2023 approximates an intractable integral using single sample. More recently, D-Flow (Ben-Hamui
2024 optimizes initial noise by differentiating through the simulated trajectory. To the best of our knowledge, these models lack a direct way navigate the data manifold in search of the optimal solution without repeatedly transitioning between noised and data distributions.
et al.
EBMs [Hopfield] 1982 Hinton 2002 LeCun et al. 2006 provide alternative approach for approximating the data distribution by learning a scalar-valued function E(x) that specifies an unnormalized density p(z) exp (-E(x)). Rather than explicitly mapping noise samples onto the data manifold, EBMs assign low energies to regions high data concentration and high energy elsewhere. This defines a Boltzmann distribution from which one can sample, for example, via Langevin sampling. In doing so, EBMs explicitly retain the likelihood information likelihood information possibly together with additional priors simply by adding their energy terms (Du and Mordatch 2019). Moreover, direct examination local curvature the data allioolddallwwss the computation of local intrinsic dimension (LID) (an important proxy for data complexity) whereas diffusion models only approximate such curvature in the proximity of noise samples.
Exx). This then be used in conditional generation (e.g.. to solve inverse problems),
Despite the theoretical elegance of using a single, time-independent scalar energy, practical EBMs have historically suffered from poor generation quality, falling short of the performance of diffusion or flow matching models. Traditional methods [Song and Kingma 2021] for training EBMss such contrastive divergence via Markov chain Monte Carlo (MCMC) local acore-based approaches [Song and Ermon 2019). often fail to adequately explore the energy landscape in sigh--imensional spaces, leading to instabilities and mode collapse. Consequently, many methods resort time- conditioned ensembles [Gao et al. 2021]. hierarchical latent ensembles (Cui and Han 2024], or combine EBMs with separate generator networks trained in cooperation [Guo et al. 2023 Zhang
jet al. 2024 Yoon et al. 2024), thereby requiring significantly higher parameter counts and training
complexity.

---

Contributions. In this work, we propose Energy Matching, two-regime training strategy that combines the strengths of EBMs and flow matching; see Figure !
When samples lie far from the data manifold, they efficiently transported toward the data, Once near the data manifold, the flow transitions into Langevin steps governed by an internal energy component, enabling precise exploration of the Boltzmann- like density well around the data distribution. This straightforward approach produces a time- time-indepenaee independent scalar energy field whose gradient both accelerates sampling and shapes the final density well-via contrastive objective that directly learns the score at the data manifold yet remains efficient and stable to train. Empirically, our method significantly outperforms existing EBMs on both CIFAR-I and ImageNet generation in terms of fidelity, and compares favorably to flow-matching and diffusion models-withe auxiliary generators or time-dependent EBM ensembles.
Our proposed process complements the advantages of flow matching with an explicit likelihood modeling, enabling traversal of the data manifold without repeatedly shuffling between noise and data distributions. This simplifies both inverse problem solving and controlled generations under prior. In addition, to encourage diverse exploration of the data distribution, we showcase how repulsive interaction energies can be easily and effectively incorporated, with an application to conditional protein generation. Finally, we also showcase how analyzing the learned energy reveals insight on the LID of the data with fewer approximations than diffusion models
"Code repository: http:///gttbbb.cca//aaalllrrr//s
2 Energy matching
In this section, we show how a scalar potential V(x) can simultaneously provide an opaimal-transport- like flow from noise data while also yielding Boltzmann distribution that explicitly captures the unnormalized log-likelihood of the data.
(JKO) scheme. The starting point of our approach is the JKO scheme [Jordan et al. 1998). which is the basis of the success of numerous recent generative models [Xu al. [023 Terpin al. 2024 Choi al. [024]. The JKO scheme describes the discrete-time evolution of probability distribution P: along energy peergy-inimizzing trajectories in the Wasserstein space,
The Mindaan-inndrleeherr
P++A: arg min W3 pc)
Vo(c))dp(s) + +o(l) / p(x)
log log p(x)dx.
24t
Transpor Cost
Potential Energy
Internal Fnergy (-Entropy)
Here, 9 denotes the learnable parameters of the scalar potential (z(z, and s(t) is temperature- parameter tuning the entropic term. The transport cost is given by the Wasserstein distance,
min / (ET(
W3 (o,
|| vl y11 dr,
where [',p, pr) the set of couplings between and Pt. i.e., the set probability distributions on Rd Rd with marginals p and Pe. Here, d the dimensionality the data. Henceforth, we call OT coupling any " that yields the minimum in (2). When = (id, T) P, i.e., it is the pushforward of the map + (æ, T(r)) for some function 7, we say that an OT map from p to Pt-
Differently from most literature, we consider E(t) to be dependent on time and study the behavior of Equation (1) as co. To fix the ideas, consider, for instance, linear scheduling:
0,
0SI<T'
e(1) = Emax I-pe
T" TS1,,
Emax:
12
First-order optimality conditions. Following Terpin et al. 2024 we analyze ( at each time /
via its first-order optimality conditions [Lanzetti et al. 2024| 2025 These conditions characterize the properties of the desired solution and thus represent the optimization goal:
3

---

At x v) + V. Vz(x)) + =t)VV log Pc+ = 0,
(x. E supp(?)
(4)
where Yr is OT plan between the distributions Pt+ and pt and supp(?)) is the support of Y:- That is. this condition has to hold for all pairs of points in the support of pr+ar and pr that are coupled by OT. Intuitively, analyzing 4 provides us with two key insights:
1. For times i <' (l) 0 and (4 becomes
1 c - v) + + Vo(x 0 (r, E supp(?))-
(5)
Al
That is, the system is in an OT, flow-like, regime. Near the data manifold, which we aim at modeling with the equilibrium distribution Peg of
1
Pt+Ar 2 Poq and, thus, for t > 1, x xe for all (I, v) E supp(?r). Then, we can simplify as
Emax Vx log Peg (x) V (o(
(1) x exp
Va(a)
Emax
Thus, the equilibrium distribution is described by an EBM, exp (-E(x)), with E(x)
(z(2)
Gmax
Our approach in a nutshell. Combining the two insights above, we propose a generative framework that combines OT and EBMs learn fime-independett scalar potential Vo(x) whose Boltzmann distribution.
Vo(z)
Peq (z) exp
(6)
Emax
matches Pdata- To transport samples efficiently from noise Po Peq Pdatas we use two regimes:
Away from the data manifold: E =z O. The flow is deterministic and OT-like, allowing rapid movement across large distances in sample space.
Near the data manifold: E Emax- Samples diffuse into a stable Boltzmann distribution, properly covering all data modes.
By combining the long-range transport capability of flows with the local density modeling flexibility of EBMs, we achieve tractable sampling and explicitly encode the unnormalized log-likelihood -Ve(x)/Emax of the underlying data distribution; see Figure I
2.1 Training objectives
In practice, we balance the two objectives by initially training V, exclusively with the optimal- transport-like objective (8 = 0, see Section 2.1.I), ensuring a stable and consistent generation of high-quality negative samples for the contrastive phase. Subsequently, we jointly optimize both the transport-based and contrastive divergence objectives, progressively increasing the effective temperature to E = Emax as samples approach the data manifold (i.c., the equilibrium distribution); see Section 2.1
2.1.1 Flow-like objective for
We begin by constructing a global velocity field v that carries noise samples {ro} to data samples {Iaata) } with minimal detours. For this, we consider geodesics in the Wasserstein space (Ambrosio et al., 2008). Practically, we compute the OT coupling 7 between two uniform empirical probability distributions, one supported on mini- bbtch of the data, and one supported on a of noise samples with the same cardinality. These samples are drawn from an aasy-to-sample distribution; in our case, a Gaussian. Since the probability distributions are uniform and empirical with the same number of samples, a transport map T is guaranteed to exist [Ambrosio et al. 20088.
Remark 2.1 (OT solver), Depending on the method used to compute the OT coupling, an explicit OT map may or may not be obtained. Similarly, the number of noise samples differs from the mini-bateh size B. the resulting OT coupling generally will not correspond to a map. In this case, one can adapt the algorithm by defining a threshold wen and considering all pairs (Iaaras xo) for which the coupling value satisfies x² (xdata, so) > weh- In our experiments, we used the POT solver (Flamury) al.) yy02Ty and did not observe benefus fromfssfoom from using a sample size different from B, consistent with previous approaches [Tong et al. 2023

---

Then, for each data point Edata we define the interpolation =( (2T(td + ta which a point along the geodesic. The velocity each I, ddata T (Faata) (i.e., the samples move from the noise to the data distribution at constant speed) and, in this regime, we would like to have
-V Re Tdata -T (Edata). For this, we define the loss:
Lot E
#daaeD (0,,")
(aVo(1)) + Edata Tzaata) )12]
This objective be interpreted a flw-matching formulation under the assumption that the velocity field is both time-independen and given by the gradient of scalar potential, thereby imposing an irrotational condition. This aligns naturally with OT, which also yields irrotational velocity field-any any rotational component would add unnecessary distance to the flow and thus inflate the transport cost without benefit. Our experimental evidence adds to the recent study by (Sun) et al. 2025]. in which the authors observed that inne-independent velocity fields can, under certain conditions, outperform dependepn noise-condiiooned fields sample generation.
Algorithm Phase I (warm-up).
1: Initialize model parameters 0
2: for iteration n = 0, .........
3: Sample mini-batch (#data)) )
Data samples
4: Sample mini-batch {z0,> ~ ~N(07)
b Random Gaussian samples
5: T < OTsolver([Taaua.
(k0,b})
& Compute OT map
Sample {ta 1 (10,T")
Typically r' b Interpolation along geodesics
warm-up
7: Set interpolations It, + -1- 1b) T(raata.) + to Idata.b
8: Lor(0) + > / [V Edata, - 2 (I
P Loss function
9:
0 + 0 - aVo LoT
Di Gradient update with learning rate
10:
end for
11:
return 0
Trained
2.1.2 Contrastive objective LCD
Near the data manifold, Vz(x) is refined so that (x) x exp ( distribution. We adopt the contrastive divergence loss described in EBMs Hinton 2002
(=)
Emax) matches the data
Vo(z)
(ā)
LCD Rarroaa
E (xweg(pe))
Emax
Emax
where ē are "negative" samples of the equilibrium distribution induced by Vg. We approximate these samples using an MCMC Langevin chain (Welling and Tch! 2011]. We split the initialization for negative samples: half begin at real data, and half begin at the noise distribution. This way, V/ Va(x) forms well-defined basins around high-density regions while also shaping regions away from the manifold, correcting the generation. The sg(-) indicates stop- stop-gradient operator, which ensures gradients do not back-propas propaatte through the sampling procedure.
2.1.3 Dual objective and implementation notes
To balance the deterministic flow-like regime (where Re 0) away from the data manifold and the stochastic Boltzmann regime (where z Emax) near equilibrium, adopt the linear temperature schedule described in (3). We introduce a dataset-specific hyperparameter ÅCD to stabilize the contrastive objective by appropriately weighting Len relative to LOT. The resulting algorithm is described in detail in Algorithm and Algorithm Since A benefiis benefits from high-quality negatives, begin with Algorithm [ (and, thus, with Lot only) ensure sufficient mixing of noise-initialized negatives.
Given the trained models, define sampling time Ta Although convergence to the equilibrium distribution guaranteed only as т; > oo, we empirically observe that sample quality (measured with Fréchet inception distance (FID)) plateaus by T, 3.25 on CIFAR-I0: see Section A.2 The sampling procedure, which optionally includes conditional and interaction terms, detailed in Algorithm In practice, we implement training using explicit Euler-Heun predictor-come correcto scheme, while for simplicity the algorithms illustrate only explicit updates. Additionally, the constant factor 1 1/Emax in LCD is absorbed into ACD.
updates and sampling with an
Section A.T discusses how the landscape of Vi evolves across these two phases. Hyperparameters for each dataset, along with intuitions guide their selection for new datasets, provided in Section
5

---

Algorithm 2 Phase 2 (main training).
1: 0 <- Opretriined
Initialize from Algorithm
2: for iteration
0, L.....dd do Lor < Use lines from Algorithm
4: Initialize negative samples {= 8
/ from noise and/or data
Negative samples
m = 0.1..... ........ Maaaggvnn do 6: for b = to B do
7:
e(m)
Emax,
if initialized from data (Optional)
(((mi)) from 3
otherwise
8:
Sample 76 NV T)
9:
(o)+1)
(m)) Fneg.
AIV. Vogco) (1 (crimeg)) )
(m) 7p
Langevin dynamics step
10:
for
11:
end for
LCD Ex=1 Vo(Eaata.o) (2 (8f
12:
Contrastive divergence loss
13:
L(0) + LoT +ADD LCD
14:
Update 0 - 0 a voo(())
D Gradient descent step
15:
end for
return 0
Trained
Table 1: FIDI score comparison for unconditional CIFAR-10 generation (lower is better). Unless otherwise specified, we use results for solvers that most closely match our setup (325 fixed-step Euler-Heun [Butcher] 2016)). indicates reproduced methods, while all other entries reflect the best reported results. EGC in its unconditional version has been reported in Zhu et al. 2024
Learning Unnormalized Likelihnod
Learning Transport/Seure Along Trujectories
Ensembles: Diffusion (ome or many) EBMs
Dilfusion Models
Hicrarchical EBM Diffusion CCui and Han
2024
8.93
DDPM
Ho ot 2020
6,45
(uoo
162M parums, steps) et
Cooperarie DKL TeMM params)) Zhu et al
2024
4.31
NCSN++ 1107M params, 1000 steps)
Song et
21Z
2.45
Cooperative DRL-large (145M paramsy Zha ct
2024
3.68
Energy-based Models
Flow-based Models
ImprovedCD Du et ]02 CLEL-large (32M paramss Lee et al.
25.
Action Matching Neklyudov et al 2023 Flow-matching Lipman al AuS
10.07
2023
8.61
6.35
Energy Matching (50M params,
3.34
OT-CFM" (37M params]
2023
4.04
Applications
In this section, we demonstrate the effectiveness and versatility of our proposed Energy Matching approach across three applications: (i) unconditional generation (ii) inverse problems, and (iii LID estimation. The model architecture and all the training details are reported in Section D
3.1 Unconditional generation
We compare four classes of generative models: (1) Diffusion models, which deliver state-of-the-art quality but typically require many sampling steps; (2) Flow-based methods, which learn OT paths for more efficient sampling with fewer steps; (3) EBMs, which directly model the log-density as scalar field, offering flexibility for inverse problems and constraints but sometimes at the expense of sample quality; and (4) Ensembles (Diffusion with one or many EBMs), which combine diffusion's 's robust sampling with elements of EBM flexibility but can become large and complex train. Our approach, Energy Matching, offers a simple (a single time-independent scalar field) yet powerful EBBM-based framework. We evaluate our approach on CIFAR-10 Krizhevsky and Hinton! 2009) and ImageNet32x32 [Deng al.) 2009 Chrabaszez et al. 2017 datasets, reporting FID scores in Table and Table 2 respectively. Our method outperforms state-of-the-ard EBMs, reducing the FID score by more than 50%.

---

Table 2: FID, score comparison for unconditional ImageNet 32x32 generation (lower is better). Unless otherwise specified, results for solvers that most closely match our setup (300 fixed-step Euler-Hewl [Butcher] 2016).
Learning Unnormalized Data Likelihood
Learning Transuot/Sccre Along Noised Trajectories
Ensembles: Diffusion + (ome or many) EBMs
Ditfusion Models
Cooperative DRL (40M panams) Zhu et al. 2024
9.35
DDPM++ r62M params, 7000 steps) Kim et 2021
8.42
Energy-based Models
Flow-based Models
ImprovedCDD al.
mawchinthhn Lipman
2023
params)
CLE--base Lee al.
(7M params)
22.16 16
CLEL-large Lee et
2023
(32M params))
15:47
Eoergy Matching (304
Ours)
6.64
3.2 Inverse problems
In many practical applications, we interested in recovering some data æ from noisy measurements y generated by operator A, y A(x) + wr where A((,, v2 In this setting, the posterior distribution x given y is
exp exp(-Eo(x)) p(rly)) x exp -(w)))
A(x))1
x p((x)
x p(ylax)
where Ex) is an energy function which one can learn from the data, an EBM. Because we want to sample x given measurement y. this reconstruction task often referred to an inverse problem. Here, |yy - Ax| encodes the measurement fidelity with Ç controlling the balance between this fidelity term and the prior. We obtain the prior term Ex(x)
Va(()) by training Vo(x) via Energy Matching.
Smax
Samples from this posterior can be drawn by starting from random sample (0)) ~ A/O, I) and following Langevin update. We detail the algorithm for generating solutions to inverse problems in Algorithm 3 (which also incorporates additional interaction energy W (z, x²) between generated samples). We demonstrate our model's capabilities qualitatively through controlled inpainting task and quantitatively via protein inverse design benchmark. Specific hyperparameters are detailed in Section
Algorithm 3 Unconditioonal/condiiona sampling with optional interaction energy
1: for m = 1 to M do
2: Initialize xm from noise and/or data
Initialize each chain
3: end for
<- [Ta /AI
Number Langevin steps for sampling time T.
5: for n 0. 1, I,...,N 1 do
6: for = 1,2.... M do
Prior data fidelity interaction
7:
c(N)
Emax,
if initialized from data (Optional)
e(nAt) from
otherwise
8:
U (x(
(n)
+ Vo (2m)) (n)
+ I A(z(n)))1 /12 + g(()) Engo W (win) (n) xx
9: Sample
(m) NN(O I)
Im (n+1) Im
At V.U (2i + V2e(") (m)
Langevin dynamics step
end for
12: end for
13: return (Im (N)
b Final samples
Controlled inpainting. Suppose we want to recover two images from a masked image while encouraging diverse reconstructions. EBMs allow this by introducing an additional interaction energy, B(#1 =2))1 where B has ones in the region of interest (focusing diversity there) and zeros elsewhere, and a is a hyperparameter controlling the interaction's strength. Specifically, we define p(#1, x2 y) x p(11 3) p(z2
W(11, x22)
(app-WWw,, x2)). which gives high probability to pairs (21, that lie far apart in the specified region B. This encourages exploring the edges of the
7

---

posterior rather than just its modes, and with suitable W, samples shift toward rare events without needing many draws. To illustrate the interaction term's advantages for diverse reconstruction, we apply our method to a CelebA [Liu et al., 2015] 64 x 64 inpainting task. As shown in Figure 2 we start from partially observed (masked) face and aim to reconstruct two distinct high-fidelity completions.
p x2ly) xe p( (p)
p x(11))2211 x p ())((22)))
BB (01 022) m2)
exp
σ²
Figure 2: Controlled inpainting for diverse reconstructions. On the left is the masked face. On the right are two reconstructions: the top pair without the interaction term and the bottom pair with The interaction term applies in the solid red square (where B has ones), and the measurement matrix is the dotted blue square (zeros inside, ones outside). By encouraging 21 and 22 to differ in the target region, the interaction yields wider range of completions while preserving fidelity.
Protein inverse design. In Figure demonstrate method's performance the inverse
design problem of generating Adeno-Associate Virus (AAV) capsid protein segments (Bryant et al.) 2021 Given desired functional property (fitnes))heree here defined the predicted viral packaging efficiency normalized between 0 and 1- goal is to design novel protein sequences satisfying this target condition. Beyond achieving high fitness, practical inverse design requires generating diverse candidate sequences to ensure robustness in the downstream experimental validation [Jain et al. 2022). We evaluate two benchmark splits (medium and hard), which correspond to subsets the original AAV dataset differing in baseline fitness distributions and required mutational distance from known tigh-performing variants Kirjner et 2024). Leveraging the latent-space representation of VLGPO (Bogensperger et al. 2025]. we employ our inference-time tunable repulsion term, allowing explicit control over the diversity of the designed proteins. This enables a flexible trade-off between fitness and diversity, resulting in high fitness scores alongside substantially improved sequence diversity. See Section B for experimental details and dataset descriptions.
Langevin sampler with an
3.3 Local intrinsic dimension estimation
Real-world datasets, despite displaying high number of variables, can often be represented by lower-dimensional manifolds concept referred to as the manifold hypothesis [Fefferman et al.) 2016). The dimension of such a manifold is called the intrinsic dimension. Estimating the LID at given point reveals its effective degrees of freedom or directions of variation, offering insight into data complexity and adversarial vulnerabilities. We defer the precise definition to Section C
Diffusion-based approaches. Recent work leverages pretrained diffusion models to estimate the LID Kamkari et al. 2024 Stanczuk et al. [0244] by examining the learned score function. However, since these models do not learn the score at the data manifold
1), their estimates become unreliable there. Consequently, current methods rely on approximations, for instance by evaluating the score in the proximity the data manifold (t. = lo), where computations remain sufficiently reliable.

---

AAV Medium
AAV Hard
10
15
Eg dWIS
ChAS
Energy Matching w/ repulsiio
Se-dWJS
Flow-hased methods
8
AdaL.ead
Score-based methods
Diversity?
Diversityf
Other non-likellbood
10
GWG
Energy Matching
6
Atial sead
VI.GPO
Energy Matching
4
GGS
5
GGS VLGPO
LatProtRL
LatProtRL
0.4
0.5
0.6
0.7
0.3
0.4
0.5
0.6
0.7
Fitnesst
Fitness
Figure 3: Fitness diversity trade-off for protein inverse design on the AAV Medium (left) and Hard (right) benchmarks. We compare Energy Matching method (blue), with diversity explicitly controlled by repulsion strength parameter (xx =²), against leading flow-based (purple), score-based (orange), and other non-likelihood methods (black). Fitness measures how well generated sequences satisfy the target property (predicted viral packaging efficiency), while diversity quantifies the average Levenshtein distance between sequences in each generated batch.
Spearman's 's correlation
MNIST
CIFAR-10
ESS Johnsson et al. 2014
0.444
0.326
FLIPD Kamkari
2024
0.837
0.819
NB Stanczuk et al. 2024
0.864
0.894
Energy Matching (Ours)
0.877
0.901
Table 3: Spearman's 's correlation coefficients of LID estimates with PNG compression rate, Bench- marks results reported in [Kamkari al., 2024].
Hessian-based LID Estimation. Unlike diffusion models, EBMs explicitly parametrize the relative data likelihood. This explicit parametrization enables efficient analysis of the curvature of the underlying data manifold in this example, estimating the LID. To this end, we compute the Hessian at a given data point and perform its spectral decomposition. We define near-zero eigenvalues as those whose absolute values lie within a small threshold (in our experiments, we set for MNIST [Deng 2012] and reflects the number of flat directions and thus reveals the local dimension. As shown in Table 3 the LID estimates we obtain exhibit stronger correlations with PNG compression size (evaluated on 4096 images) using Spearman's 's correlation. Figure ooffers qualitative illustrations. Our EBM-based approach compares favorably to diffusion-basee methods, as it relies on fewer approximations by performing computations exactly the data manifold rather than merely in its vicinity.
matrix
for CIFAR-10) 10). The count of near-zero eigenvalues
Hessiinn spectrum
Hessian spectrum
LID: 21
LID: 136
Eigesvvate
1,000
a
Elppavaite
L,00
500
1003
100
900
500
$00
003
Index
Index
Figure 4: Qualitative results for LID estimation using the Hessian spectrum of Vo(x). Left: Spectrum for low-LID image. Right: Spectrum for high-LID image. The cigenvalues quantify curvature along principal directions (eigenvectors). marked in red) indicates locally "flat" regions, revealing the LID. Intuitively, higher image complexity often corresponds to a higher LID.
degenerate spectrum (many near-zero eigenvalues,
PNG is lossless compression scheme specialized for images and can provide useful guidance when no LID ground truth is available Kamkari al. 22044.

---

Conclusion and limitations
Contributions. We introduced generative framework, Energy Matching, that reconciles the advantages of EBMs and OT flow matching models for simulation-free likelihood estimation and efficient high-fidelity generation. Specifically, it:
Learns a time-independent scalar potential energy whose gradient drives rapid high-fidelity sumpling-surpassing state-of-the-art energy-based models-while also forming Boltzmann-likk density well suitable for controlled generation. All without auxiliary generators.
Offers efficient sampling from target data distributions on par with the state-of-the-art, while learning the score at the data manifold with manageable trainable parameters overhead.
Offers simulation-free principled likelihood estimation framework for solving inverse prob- lems-where additional priors can be easily introduced and enables the estimation of data point's LID with fewer approximations than score-based methods.
Limitations, First, our method requires an additional gradient computation with respect to the input, which can increase GPU memory usage (c.g., by 20-40%), particularly during training. Second, when estimating the LID (Section 3.3) for very high-dimensional datasets, computing the
full Hessian spectrum may be impractical due to its computational complexity of O(d)); in such cases, spectrum methods such as random projections or iterative solvers can be employed instead.
Outlook. Contrary to widespread belief, we demonstrated that iimeiindependent irrotational meth- ods for generative flows are highly effective and offer an exciting direction for future research. Our Energy Matching approach has the potential yield novel insights into controlled generation and inverse problems for cancer research Weidner et al. [2024), Balcerak et al. [2024]. molecules and proteins Wu al. [2022], Bilodeau et [2022]. computational fluid dynamics Gao et al.) (2024),
Shysheya al. M Molinaro al. [2024], and other fields where precise control over generated samples and effective integration of priors or constraints are crucial. Moreover, Energy Matching aligns naturally with recent generative Al trends toward scaling inference for new capabilities (Zhang
cet al. 2025 Ma et al. 2025 further broadening its potential impact across scientific and engineering
domains.
Acknowledgments and Disclosure of Funding
This research was supported by the Helmut Horten Foundation and the European Cooperation in Science and Technology (COST).
References
Charu C. Aggarwal, Alexander Hinneburg, and Daniel A. Keim. On the surprising behavior of distance metrics in high dimensional space. International Conference on Database Theory (ICDT), 2001.
Michael Albergo and Eric Vanden- Building normalizing flows with stochastic interpolants. In International Conference on Learning Representations (ICLR), 2023.
Ambrosio, N. Gigli, and Savaré. Gradient flows: metric spaces and in the space of probability measures. Lectures in Mathematics ETH Zürich. Birkhäuser Basel, 2008.
Jason Ansel, Edward Yang, Horace He, Natalia Gimelshein, Animesh Jain, Michael Voznesensky, Bin Faster machine learning through dynamic python bytecode transformation and graph compilation. In International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS), 2024.
Bao, Peter Bell, David Berard, Evgeni Burovski, et Pytorch 2:
Michal Balcerak, Tamaz Amiranashvili, Andreas Wagner, Jonas Weidner, Petr Karnakov, Johannes Paetzold, Ivan Ezhov, Petros Koumoutsakos, Benedikt Wiestler, et al. Physics-regularieed multi- modal image assimilation for brain tumor localization. In Advances in Neural Information Processing Systems (NeurIPS), 2024.
10

---

Heli Ben-Hanne Hamu, Omri Puny, Itai Gat, Brian Karrer, Uriel Singer, and Yaron Lipman. D-flow: differentiating through flows for controlled generation. In International Conference on Machine Learning (ICML), 2024.
Camille Bilodeau, Wengong Jin, Tommi Jaakkola, Regina Barzilay, and Klavs F Jensen. Generative models for molecular discovery: Recent advances and challenges. Wiley Interdisciplinary Reviews: Computational Molecular Science, 12(5):e 2022.
Lea Bogensperger, Dominik Narnhofer, Ahmed Allam, Konrad Schindler, and Michael Krauthammer. A variational perspective generative protein fitness optimization. In International Conference on Machine Learning (ICMI.), 2025.
David Brookes, Hahnbeom Park, and Jennifer Listgarten. Conditioning by adaptive sampling for robust design. In International conference on machine learning (ICML), 2019.
Drew H. Bryant et al. Deep diversification of an aav capsid protein by machine learning. Nature Biotechnology, 39:691-6888 1-69, 2021.
John C. Butcher. Numerical Methods for Ordinary Differential Equations. John Wiley & Sons, 3rd edition, 2016.
Jaemoo Choi, Jaewoong Choi, and Myungjoo Kang. Scalable wasserstein gradient flow for generative modeling through unbalanced optimal transport. In International Conference on Machine Learning (ICML), 2024.
Patryk Chrabaszez, Ilya Loshchilov, and Frank Hutter. A downsampled variant of imagenet as an alternative to the cifar datasets. arXiv preprint arXiv:1707.088// 2017.
Hyungjin Chung, Jeongsol Kim, Michael T Mccann, Mare L Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. In International Conference on Learning Representations (ICLR), 2023.
Jiali Cui and Tian Han. Learning latent space hierarchical ebm diffusion models. In International Conference on Machine Learning (ICML), 2024.
Giannis Daras, Hyungjin Chung, Chieh-Hsin Lai, Yuki Mitsufuji, Jong Chul Ye, Peyman Milanfar, Alexandros G Dimakis, and Mauricio Delbracio. A survey on diffusion models for inverse problems. arXiv preprint arXi:22410.00083 2024.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In IEEE/CVF Computer Vision and Pattern Recognition (CVPR), 2009.
Li Deng. The mnist database of handwritten digit images for machine learning research. IEEE Signal Processing Magazine, (9(6))1411144 2012. 2012.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, G Heigold, S Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations (ICML), 2020.
Yilun Du and Igor Mordatch. Implicit generation and generalization in energy-based models. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
Yilun Du, Shuang Li, Joshua Tenenbaum, and Igor Mordatch. Improved contrastive divergence training of energy based models. In International Conference on Machine Learning (ICML), 2021.
Charles Fefferman, Sanjoy Mitter, and Hariharan Narayanan. Testing the manifold hypothesis. Journal of the American Mathematical Society, 2944):933-1049 2016.
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, et al. POT: Python Optimal Transport. Journal of Machine Learning Research, 22(78):1-8, 2021.

---

Han Gao, Sebastian Kaltenbach, and Petros Koumoutsakos. Generative learning for forecasting the dynamics of high-dimensional complex systems. Nature Communications, 15(1):8904, 2024.
Ruiqi Gao, Yang Song, Ben Poole, Ying Nian Wu, and Diederik P Kingma. Learning energy-based models by diffusion recovery likelihood. In International Conference Learning Representations (ICLR), 2021.
Will Grathwohl, Kevin Swersky, Milad Hashemi, David Duvenaud, and Chris Maddison. Oops took gradient: Scalable sampling for discrete distributions. In International Conference on Machine Learning (ICML), 2021.
Qiushan Guo, Chuofan Ma, Yi Jiang, Zehuan Yuan, Yizhou Yu, and Ping Luo. Egc: Image generation and classification via a diffusion energy-based model. In /EEE/CVF International Conference on Computer Vision (ICCV), 2023.
Geoffrey E Hinton. Training products of experts by minimizing contrastive divergence. Neural computation, 14(8):1771-1800, 2002.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
J. Hopfield. Neural networks and physical systems with emergent collective computational abilities. Proceedings of the National Academy of Sciences, 79(8):2554-2558, 1982.
Aapo Hyvärinen. Connections between score matching, contrastive divergence, and pseudolikelihodd for continuous-valuee data. Neural Computation, 18(8):1527-1550, 2006.
Zarif Ikram, Dianbo Liu, and M Saifur Rahman. Antibody sequence optimization with gradient- guided discrete walk-jump sampling. ICLR 2024 Workshop on Generative and Experimental Perspectives for Biomolecular Design, 2024.
Moksh Jain et al. Biological sequence design with gflownets. In International Conference on Machine Learning (ICML), 2022.
Kerstin Johnsson, Charlotte Soneson, and Magnus Fontes. Low bias local intrinsic dimension estimation from expected simplex skewness. IEEE transactions on pattern analysis and machine intelligence, 37(1):196-202 2014.
Richard Jordan, David Kinderlehrer. and Felix Otto. The variational formulation of the fokker-planck equation. SIAM Journal on Mathematical Analysis, 29(1):1-17, 1998.
Hamid Kamkari, Brendan Ross, Rasa Hosseinzadeh, Jesse Cresswell, and Gabriel Loaiza-Ganeni A geometric view of data complexity: Efficient local intrinsic dimension estimation with diffusion models. In Advances in Neural Information Processing Systems (NeurIPS), 2024.
Dongjun Kim, Seungjae Shin, Kyungwoo Song. Wanmo Kang, and Il-Chul Moon. Soft truncation: A universal training technique of score-based diffusion model for high precision score estimation. In International Conference Machine Learning (ICML), 2021.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint 1412,6600 2014.
Andrew Kirjner et al. Improving protein optimization with smoothed fitness landscapes. In Interna- tional Conference on Learning Representations (ICLR), 2024.
Alex Krizhevsky and Geoffrey E. Hinton. Learning multiple layers of features from tiny images. Technical report, University of Toronto, Department of Computer Science, Toronto, Ontario, Canada, 2009.
Nicolas Lanzetti, Antonio Terpin, and Florian Dörfler. Variational analysis in the wasserstein space. arXiv preprint arXiv:2406. 10676, 2024.
Nicolas Lanzetti, Saverio Bolognani, and Florian Dörfler. First-order conditions for optimization in the wasserstein space. SIAM Journal Mathematics of Data Science, 7(1):274-300, 2025.
12

---

Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and F Huang. A tutorial on energy based learning, Predicting Structured Data, 1(0), 2006.
Hankook Lee, Jongheon Jeong. Sejun Park, and Jinwoo Shin. Guiding energy- -baed models via contrastive latent variables. arXiv preprint arXiv:2303.03004 2023.
Minji Lee, Luiz Felipe Vecchietti, Hyunkyu Jung, Hyun Joo Ro, Meeyoung Cha, and Ho Min Kim. Robust optimization in protein fitness landscapes using reinforcement learning latent space, In International Conference on Machine Learning (ICML), 2024.
Yaron Lipman, Ricky TQ Chen, Heli Ben- Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In International Conference on Learning Representations (ICI.R), 2023.
Xingchao Liu, Chengyue Gong, al. Flow straight and fast: Learning to generate and transfer data with rectified flow. In International Conference on Learning Representations (ICLR), 2023.
Ziwei Liu, Ping Luo,, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In International Conference on Computer Vision (ICCV), 2015.
Nanye Ma, Shangyuan Tong, Haolin Jía, Hexiang Hu, Yu-Chuan Su, Mingda Zhang, Xuan Yang. Yandong Li, Tommi Jaakkola, Xuhui Jia, and Saining Xie. Inference-time scaling for diffusion models beyond scaling denoising steps. arXiv preprint, #rXiv:2501.09733 2025.
Morteza Mardani, Jiaming Song, Jan Kautz, and Arash Vahdat. variational perspective on solving inverse problems with diffusion models. In International Conference on Learning Representations (ICLR), 2024.
Roberto Molinaro, Samuel Lanthaler, Bogdan Raonié, Tobias Rohner, Victor Armegioiu. Stephan Simonis, Dana Grund, Yannick Ramic, Zhong Yi Wan, Fei Sha, et al. Generative ai for fast and accurate statistical computation of fluids. arXiv preprint urXiv:2409, 20:::28999,, 2024.
Kirill Neklyudov, Rob Brekelmans, Daniel Severo, and Alireza Makhzani. Action matching: Learning stochastic dynamics from samples. In International Conference on Machine Learning (ICML), 2023.
Aliaksandra Shysheya, Cristiana Diaconu, Federico Bergamin, Paris Perdikaris, José Miguel Hemández-LLobato, Richard Turner, and Emile Mathieu. On conditional diffusion models for pde simulations. In Advances in Neural Information Processing Systems (NeurIPS), 2024.
Sam Sinai, Richard Wang, Alexander Whatley, Stewart Sloeum, Elina Locane, and Eric D Kelsic. Adalead: A simple and robust adaptive greedy search algorithm for sequence design. arXiv preprint 201:000000.22411 2020.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Advances in neural information processing systems (NeurIPS), 2019.
Yang Song and Diederik P Kingma. How to train your masey models. arXiv preprint wrX.v:2101.03288 2021.
Yang Song, Jascha Sohl- Dickstiic,,,,, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations (ICLR), 2021.
Jan Pawel Stanezuk, Georgios Batzolis, Teo Deveney, and Carola-Bibiane Schönlieb. Diffusion models encode the intrinsic dimension of data manifolds. In International Conference on Machine Learning (ICML), 2024.
Qiao Sun, Zhicheng Jiang, Hanhong Zhao, and Kaiming He. noise conditioning necessary for denoising generative models? In International Conference Machine Learning (ICML), 2025.
Antonio Terpin, Nicolas Lanzetti, Martín Gadea, and Florian Dorfler. Learning diffusion at lightspeed.
In Advances in Neural Information Processing Systems (NeurIPS), 2024.
Tijmen Tieleman. Training restricted boltzmann machines using approximations to the likelihood gradient. In International Conference Machine Learning (ICML), 2008.
13

---

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. Transactions on Machine Learning Research, 2023.
Jonas Weidner, Ivan Ezhov, Michal Balcerak, Marie-Christin Metz, Sergey Litvinov, Sebastian Kaltenbach, Leonhard Feiner, Laurin Lux, Florian Kofter, Jana Lipkova, et al. A learnable prior improves inverse tumor growth modeling. IEEE Transactions Medical Imaging, 2024.
Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In International conference on machine learning (ICML), 2011.
Lemeng Wu, Chengyue Gong, Xingchao Liu, Mao Ye, and Qiang Liu. Diffusion-based molecule generation with informative prior bridges. Advances in Neural Information Processing Systems (NeurlPS), 2022.
Chen Xu, Xiuyuan Cheng, and Yao Xie. Normalizing flow neural networks by jko scheme. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
Sangwoong Yoon, Himchan Hwang, Dohyun Kwon, Yung-Kyun Noh, and Frank Park. Maximum entropy inverse reinforcement learning of diffusion models with energy-based models. In Advances in Neural Information Processing Systems (NeurIPS), 2024.
Bingliang Zhang, Wenda Chu, Julius Berner, Chenlin Meng, Anima Anandkumar, and Yang Song. Improving diffusion inverse problem solving with decoupled noise annealing. In IEEE/CVF Computer Vision and Puttern Recognition (CVPR), 2025.
Yasi Zhang, Peiyu Yu, Yaxuan Zhu, Yingshan Chang, Feng Gao, Ying Nian Wu, and Oscar Leong. Flow priors for linear inverse problems via iterative corrupted trajectory matching. arXiv preprint arXiv:2405. 18816, 2024.
Yaxuan Zhu, Jianwen Xie, Ying Nian Wu, and Ruiqi Gao. Learning energy based models by cooperative diffusion recovery likelihood. In International Conference on Learning Representations (ICLR), 2024.
14

---

a) Phase (warm-up):
b) Phase (main training):
Reference EBM:
Lor
Lor + LCD
CCD
Figure 5: Visualization of the energy (i( landscapes driving the samples from eight Gaussians to two moons. See Figure for the 2D perspective. (a) The OT flow loss enforces zero curvature in Vo(x)) along the trajectories to the target. (b) Around the 2 Moons, the curvature of Vo(zz) is adjusted to approximate log Prosons (I) x Vz(z) while remaining close to the pretrained landscape elsewhere. Combining these objectives yields potential energy landscape that is both efficient for sampling and representative of the underlying target data distribution. (c) An EBM is shown for comparison, trained using contrastive divergence loss. Visible mode collapse that slows down the equilibration. Less regular landscape away from the data as needs many simulations to explore it.
A Additional details on Energy Matching
In this section, we provide additional studies and visualizations method.
A.1 Energy landscape during training
In Figure 5 visualize how the potential Vz(x) transitions from flow-like regime, where the OT loss enforces nearly zero curvature away from the data manifold (a), to an EBM-like regime, where the curvature around the new data geometry (here, two moons) is adaptively increased to approximate log Pama(aa (x) (b). This two-stage design yields a well-shaped landscape that is both efficient to sample (thanks to mostly flat potential between clusters) and accurate for density estimation near the data modes. For comparison, (c) shows an EBM trained solely with contrastive divergence, exhibiting sharper but less globally consistent basins.
A.2 Ablation the sampling time
Here, provide ablation studies on CIFAR-10 unconditional generation. Specifically, first pretrain using Cor, and then fine-tune with (Cor + LCD). producing a stable Boltzmann distribution from which one sample. Figure F illustrates the FID function of sampling time т, for models trained under these different regimes. the case of pure Lor, the quality measure drops (FID increases) sharply when sampling T > this occurs because, data manifold, there is no Boltzmann-like potential well to keep them from drifting away. Because the fidelity slope the data manifold steep with respect to sampling time, methods lacking explicit time-conditioning can easily overshoot or undershoot, significantly impacting fidelity. This behavior might explain why some models degrade in performance when made iimeenndeeendett, as recently reported by Sun et al. (2025).
the samples move close the
In Figure 6 we also report results for different values of the temperature-swicching parameter 7",
which influences the sampling along the paths towards the data manifold (see Equation (3)).
A.3 Ablations on the OT Solver
We evaluate the computational overhead and sensitivity to solver choice of the OT solver employed in our experiments.
15

---

FID vs sampling time
10
8
6
FID
4
= 0.9, (Cor + LCD)
2
1.0, (Cor + LCD)
Lor
1
2
Sampling time Ts
Figure 6: CIFAR-10 unconditional generation FID vs. sampling time T, when sampling from models trained under different scenarios: pure Lot and combined (CoT + LCD), with temperature regime switching parameter r* {0.9, 1.0] during sampling. Lower FID indicates better generative quality, All results for Euler-Heun with At 0.01.
Computational overhead. the CIFAR-10 10 experiment, the OT solver accounts for roughly 1.5% of the training iteration time during Phase 1 Algorithm [ This overhead decreases to negligible level (approximately 0.01 %) Phase Algorithm where computational costs are predominantly
dominated by the generation of negative samples.
Impact of solver accuracy and complexity. Let us (over-idealise) and model a standardised CIFAR-10 image vector E Rd. drawn from (0, Ia), with dimensionality d = 32 x 32 3 3072, and paired with an independent Gaussian-noise vector 2 (O0,, Ia). Each coordinate of the difference (z x)) thus follows (0, 2), and the squared Euclidean distance distribution is
>z 1 ~ 2x3, which has mean 2d, standard deviation VSd, and relative spread
0.025. This
2d
demonstrates the "thin-shell" phenomenon, implying that all entries of the cost matrix Ct5 122-2)
concentrate around nearly identical values, consistent with the distance-comceentatiin effect observed by (Aggarwal et al. 2001]. Consequently, the choice among exact linear programming (LP), entropic regularisation (Sinkhorn), or even random matching should yield nearly identical cumulative optimal- transport costs, despite their different complexities: (n log n) for LP, O n²) for Sinkhorn (with regularisation strength x), and (n) for random matching.
Empirical evidence supporting this is summarised below:
Our CIFAR-10 runs: FID degrades from 3.34 (LPP to 3.37 (random).
Tong al.
2023
Table 5: FID scores of 4.44 (LPP) vs 4.46 (random) 100 steps.
Tong al.
2023
Fig. D.2: solver accuracy saturates beyond batch size 16 in the 2D task.
App. C.2: LP and Sinkhorn methods produce indistinguishable results unless regularisation extreme.
Terpin et al.
2024
We employ the LP approach for robustness, negligible cost, and no additional hyperparameters.
A.4 Sampling Time Analysis
Computing the gradient V~ V_Vo(z) via automatic differentiation (autograd Ansel et al. 2024 introduces additional computational overhead compared to directly evaluating Vo(x). Specifically, on the CIFAR-10 10 network architecture (see Figure 7), gradient evaluation is approximately 25
16

---

slower, as it requires both forward and backward passes, In contrast, flow matching and diffusion models directly parameterize the velocity field. thus only requiring forward computations during sampling.
Nevertheless, despite this per-step computational cost, method achieves competitive overall sampling efficiency due to reduced number of integration steps needed for high-quality generation. As demonstrated in Table Energy Matching achieves lower FID (3.34) in 173 seconds per batch, outperforming OT-FM (FID 3.74 in 136 seconds per batch) and DDPM++ (FID 3,45 in 183 seconds per batch). Our results thus indicate a favorable balance between computational overhead per step and total sampling runtime.
Table 4: Comparison of sampling efficiency and quality on CIFAR-IO (batch size 128, NVIDIA R6000 48GB GPU). Despite gradient computation overhead (V, (z(x) via backward pass), Energy Matching achieves superior FID scores with competitive wall-clock sampling time.
Method
Params
Steps
Sampling Time [s]l
FIDJ
Flow/DDiffusion-base Models
OT-FM (Tong al. 2023
37M
1000
136
3.74
DDPM+ Kim et al. 2021
62M
1000
183
3,45
Energy-based Models
Energy Matching (Ours)
50M
325
173
3.34
B Details on AAV inverse design protein generation
We optimize protein fitness for adeno-associated virus (AAV) sequences in the medium and hard data
regimes proposed by [Kijjnrr et al. 2024]. using latent encodings and backbone architectures from [Bogensperger al.
2025 Conditional sampling employs classifier guidance via learned predictor
networks Io or 9s to steer samples toward high-fitness regions. The CNN-based fitness predictors
from [Kirjner et al.
2024 trained only the limited training data for each regime.
Training follows Algorithm I and Algorithm] 2
We sample 128 sequences using Algorithm 3 keeping the same batch size across all baselines; detailed hyperparameters are given in Section D Generated sequences are evaluated for fitness using the learned oracle from [Krjner et al.] [222], ]. and further assessed for both intra-set diversity and novelty relative to the training sequences (Jin et al. 2022)]
Our approach achieves state-of-Ne th-thhe-rt fitness while improving diversity (see Table 5). Incorporating
interaction energy in Algorithm 3 further enhances diversity with manageable impact fitness.
Table 5: AAV optimization results. For VLGPO (flow-based) and Energy Matching, medium difficulty uses Iø. hard difficulty uses ga. Metries (Fitnesst, Diversityf. Noveltyt). Reported uncertainty of Fitness is expressed as standard deviation.
AAV medium
AAV hard
Method
Fitness
Diversity?
Novelty
Fitness
Diversity?
Novelty?
Learning Unnormalized Likelibood
Energy-based Models
Energy Matching (Ours)
(0.0)
5.86
(0.0)
4.77
Energy Matching (+repulsion) (Owrr)
0.58 (0.0)
6.22
5.0
0.60 (0.0)
5.22
6.6
Learning Transport/Score Along Trajecturies
Flow-based Models
VLGPO [Bogensperger et
2025
(0.0)
5.58
5.0
(0.0)
4.29
Diffusion Models
RE-JWJS [Tram et
2024
0.48 (0.0)
9.48
4.2
0.33 (0.0)
14.3
5.3
Methods
LatProrRI [Lee et
2024
0.57 (0.0)
3.00
5.0
0.57 (0.0)
3.00
5.0
GGS [Kirjerr at
2024
0,51 (0.0)
4.00
5.4
0.60 (0.0)
4.50
7.0
Adal.am
Stmmn
2020
0.46 (0.0)
8.50
2.8
(0.0)
8.53
3.4
CbAS
Brookes et 2019
0.43 (0.0)
12,70
7.2
0,36 (0.0)
14.4
8.6
GWG
Grallwwoli
(0.1)
6.60
7.7
(0.D)
GFN AL.
Jani cl dl 21497
0.20 (0.1)
9.60
19.4
0.10
11.6
19.6
17

---

C
Details on LID estimation
Definition.
To start, need to introduce the concept of local mass, defined as
M(r) S B((
p(a)) dx,
where p(x) is the local density and (Tamaa: r) is a ball of radius r around Tdaaas i.e. B( 7) {a Rd || Idatsll r}. The LID then given by:
LID(zaata) = d lim log (gg(M()))
log(r)
Intuitively, (r) measures how much probability mass is concentrated in ball of radius around Edata- As we shrink this ball, the growth rate of (r) in terms reveals the local dimensional structure of the data.
Assumptions. In the context of contrastive divergence, we assume that data points Zdata lie in well-like regions [Hyvärinen] 2006], i.e.:
VV (Idata) 2 0 and V is positive semidefinite (or nearly so).
Conceptually, (zx) can be thought of as an energy function; points where (V(aaata) = are near local minima of this energy, and the Hessian VVVaaaatal provides information about local curvature (see Figure for a qualitative illustration).
Energy-based density.
We define an energy- density
p(x) exp v(())
where e is temperature parameter. Near a data point ddata satisfying V (xdata) O, we can approximate (zx) by its second-order Taylor expansion:
(ax) (Ztata) + N ( rdata) V (Idata) Idasa).
Consequently, in view of the assumptions above,
p(z) x exp
1- (x Idata) T (Fdata))())) x'data) 2s
Local mass derivation and the rank of the energy Hessian. Substituting the local quadratic form of p(x) near Tdata into the definition of the local mass M (r), we obtain:
p(x)dxx x f
(()
( exp (- 2e (x - xdaaaa (arataa) - Ddaata)) dx.
B(zan.r)
For small r. the dominant contribution depends on the rank of the Hessian VZVaua))) Let Then, - one can show that M (r) Cr where C does
rank V depend on We take the logarithm on both sides and divide by log(r) to get
log(M(r)) log(C) + & log(r) = k + log(r) log(r) and the second term vanishes as r + 0. Hence,
log(C)
log(r)
(IData))) = d k.
Practical estimation. In practice, the LID at data point Vdata be estimated through the following procedure:
1. Train V(x) with Energy Matching.
Compute the Hessian = (Iraaa))..
Perform an eigenvalue decomposition H.
Then the estimated local data-manifold dimension corresponds the number of directions with negligible curvature (smaller magnitude than some 7).
18

---

D Training details
Below, we detail the training configurations for CIFAR-I0, ImageNet 32x32, CelebA, MNIST, and AAV. Additionally, we provide intuitions for practical hyperparameter choices to facilitate effective training across additional datasets. We recommend using SiLU activation functions wherever possible, as they smooth out the energy landscape and improve the numerical stability of the V. (x) computation. The gradient of the potential, Vz V(r), is computed using automatic differentiation via PyTorch's 's autograd [Ansel et al., 2024). We optimize all models using the Adam optimizer [Kingma and Ba 2014] and maintain an exponential moving average (EMA) of the model weights.
While we specifically adopt (i) a one-sided trimmed mean of negative sample energies and (ii) clamping of the contrastive loss for stability, any commonly used EBM technique (e.g., persistent contrastive divergence [Tieleman] 2008), replay buffers, multi-scale negative sampling) could be readily employed.
In our approach, we introduce two hyperparameters, a and B, to control these stabilizing techniques:
a fraction of negative energies discarded to remove outliers that skew the mean (e.g., top 10%),
B clamp threshold for LCD (i.e., we clamp LCD be 2 -3).
The architecture is shown Figure 7 We use the same UNet from Tong et al. 2023 (with fixed 1. = 0.0, making it effectively time-independent) followed by small vision transformer
CIFAR-10:
(ViT) Dosovviskky et al.) 2020] to obtain a scalar output. Hyperparameters are: T. = 3.25, T" 1.0, At 0.01, Langoviin 200. We train for 145k iterations using Algorithm I with EMA 0.9999 and then 2k more with Algorithm 2 and EMA 0.99 on 4xA The batch size is 128, learning rate is
1.2 10 Emax 0.01, ACD
0.1, and B 0.02. Negatives initialized on the data manifold follow the same temperature schedule as those initialized from the noise.
10
ImageNet 32x32: The architecture is shown in Figure 7 (same as for CIFAR-10). Hyperparameters are: т, = 2.5, T* = =100, At = 0.01, Langovin = 200. We train for 640k iterations using Algorithm T with EMA 0.9999 and then 1k more with Algorithm 2 and EMA 0.99 on 7xA100. The batch size 128, learning rate is 6 x 10 Emax = 0.01, ACD x 10 3 a = 0.1, and 3 = 0.02. Negatives
initialized on the data manifold follow the same temperature schedule those initialized from the noise.
CelebA: We scale the CIFAR-10 model by 2x; see Figure We set 2.0, 7" 1.0, At = 0.01, Mangevin = 200, and train for 250k iterations using Algorithm I with EMA 0.9999 then 4k with Algorithm 2 and EMA 0.99 on 4xA100. The batch size 32, learning rate x 10-1,
max = 0.05, ACD 1 10-4
We downscale the CIFAR-10 model (Figure 7 to 2M parameters by reducing the UNet base width to 32 channels, using channel multipliers [1, 2, 21, setting the number of attention heads in the UNet to 2, simplifying the Transformer head to an embedding dimension of 128, 2 Transformer
MNIST:
layers, 2 attention heads, and adjusting the output scale to 100.0. We set T, = 2.0, T" 1.0, Al 0.025, Maangevin 75, and train for 50k iterations using Algorithm with EMA 0.999 then 3.3k with Algorithm 2 and EMA 0.99 on single A100. The batch size is 128, learning rate is 1 x 10-4, maax = 0. ACD = x 10-3 30-3,_ a = 0.0, and = 0.05. Negatives initialized on the data manifold follow the same temperature schedule as those initialized from the noise.
AAV: We adopt the dieennnnnio CNN architecture as used [oogensperger et al.. 2025 sum-
ming the final-layer activations to obtain the potential. We train for 10k iterations using Algorithm T and for 1k iterations using Algorithm| 2 on single A 100. The batch size is 128, learning rate is
1 x 10-1 6 max = 0.1, M Langovin 200, A/- = 0.01, and ACD = 1 x 10-1, For Algorithm 3 we use T, = 1.7 for AAV medium and T2 = 1.3 for AAV hard, 7" = 0.9, < = 0.01 for AAV medium and
< 0.009 for AAV hard. We set the target fitness to y = 1 to aim for the maximum fitness in the Sequeeneuuee sequences.
generated
Intuition for other datasets: is essential for negative samples to reach the equilibrium distribution induced by the model or least the proximity of the data manifold. The condition M angeevin X Al > is critical to achieving this, with At small enough to ensure that negative samples remain of sufficient
19

---

meality-typically the same At as used in flow matching for the given generation task. In practice, we set Miangevin x At across most experiments. We recommend starting with r" 1.0 to ensure optimal transport regularization near the data manifold, thereby enhancing training stability. If additional conditions are required during sampling, exploring lower values (7" < 1.0) may be beneficial, as this parameter does not need to remain consistent between training and sampling (as shown in Figure (6). Training with r* < 1.0 is advised only in special cases, such as extremely low-dimensional problems like that shown in Figure L where it is possible to simultaneously be fär from the data manifold (from the perspective of the target mode) and close to it (from the perspective of another mode). The parameter Emax controls how extensively negative samples explore the space. For unconditional generation, we use mmax = 0.01, but for inverse problems or design tasks, higher values (e.g., mmxx 0.05) can improve robustness. The parameter To significantly depends on the task (unconditional or conditional) and thus must be tuned accordingly post-training. Without explicit tuning, selecting T, M Mangevin At is a reasonable default. Finally, parameters ÄCD, a, and /3 influence the stability of contrastive training and must be empirically determined alongside appropriate early stopping.
Input (3x32x32)
Hyperparameter
Value
Image size
3x32x32
Base channels (UNet)
128
ResBlocks
2
UNet (37M params)
channel mult
[1.
2, 2]
Attention resolution
16
Attention heads (UNet)
Head channels (UNet)
64
Dropout
0.1
Transformer (PatchEmbed+ VIT, 12M params)
Transformer (VIT) Head
Patch size
4
Fimbedding dim
384
Transformer layers
8
Scalar Potential
Transformer heads
4
Output scale
1000.0
Figure 7: Diagram of our UNet+ Transformer EBM for CIFAR-I0 10 and ImageNet 32x32. A UNet (37M params)
processes xx32x 32 image; its output is fed into Transformer head (PatchEmbed 8-layer ViT, 12M params) that produces scalar potential. Here we employ the identical UNet architecture as in [Tong et 2023). but with the time parameter fixed at = 0 to render the model -ndepnndeen.dnt.
Input
Hyperparameter
Value
(3x64x64)
Image size
3x64x64
Base channels (UNet)
128
ResBlocks
2
UNet (83M params)
channel_ mult
[1. 3.
Attention resolution
16
Attention heads (UNet)
4
Head channels (UNet)
64
Dropout
0.1
Transformer (PatchEmbed++ VIT, 25M params)
Trunsformer ( (VI)) Head
Patch size
4
Embedding dim
512
Transformer layers
8
Scalar
Transformer heads
8
Output scale
1000.0
Potential
Figure 8: Diagram of our UrassTrantrrmmer EBM for CelebA. A UNet (83M params) processes 3 64 image: its output fed into Transformer head (PatchEmbed+8-lay llyerr VIT. 25M params) that produces scalar potential.
20

---

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]
Every substantive claim in the abstract and introduction is substantiated by the theory, experiments, and ablation studies reported in the main text; limitations are stated explicitly. The "Outlook" (Section 4 contains clearly flagged aspirational goals that presented as future directions rather than accomplished results, ensuring the scope is not overstated.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in 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 or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: A dedicated "Limitations" paragraph (Section ! clearly details two concrete constraints of the method. Hence the paper meets the requirement to openly discuss its limitations.
Guidelines:
The answer NA means that the paper has no 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 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 (e.g., independence assumptions, noiseless settings, model vell-specifiiattin, 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 on the scope of the claims made, e.g., if the approach was only tested few datasets 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 is low or images are taken in low lighting. Or speech-to-text 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 faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged 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.
21

---

3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: Each theorem lists its full assumptions and is accompanied by complete proof-intuitiv intuitive sketches the main text and formal details in the appendix. When we rely known results, cite the originals and restate the necessary arguments for self- containment.
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 the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs 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 information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: The manuscript specifies architectural details, ayper-parametes and evaluation methods needed for an independent researcher to replicate the reported results.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a 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.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to 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 hosted model (e.g., in the of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS 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) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution a new model (e.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 instructions for how to construct the dataset).
22

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to 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, as described in supplemental material?
Answer: [Yes]
Justification: The experiments rely only on publicly available datasets, and an anonymized code bundle illustrating the full training evaluation pipeline is attached the submission; the repository and all ancillary assets needed for reproduction will be made openly available upon acceptance.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not 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 for more details.
/ wttppsscc/ppbllcc/..cpcccpcpppee
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 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 much information possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: All key choices data splits, hyper-parameters, optimizers, and selection criteria documented in the "Training Details" (Section D), while remaining lowelevel implementation specifics are available in the accompanying code. We follow standard procedures much possible.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core 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, 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?
23

---

Answer: [No]
Justification: Due to high computational costs involved in training and evaluating generative models, metrics like FID typically computed only few times, and thus standard deviations are usually not reported. Nevertheless, we do provide standard deviations for the Fitness results obtained the AAV dataset, where feasible.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars 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).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report I-sigma error bars, but 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 show in tables figures symmetric error bars that would yield results that are out of range (c.g. negative error rates).
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: The hardware specifications and number of training iterations (compute resources) are detailed in the "Training Details" (Section
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as 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 't make it into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics http:://wnourrs...ccppppppcctact
Answer: [Yes]
The study uses only publicly licensed datasets, releases anonymized code for full reproducibility, involves no human subjects or sensitive personal data, and discloses compute resources and limitations. These practices align with the NeurIPS Code of Ethics on transparency, privacy, attribution, and social responsibility; no deviations are required.
Justification:
Guidelines:
The answer NA means that the authors have not 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 sure to preserve anonymity (e.g., there a special consid-
eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [NA]
Justification: The method augments an existing class of generative models with a tunable flexibility parameter aimed inverse problems for science (c.g. protein design) rather than at highest-fidelity large-scale content generation. (i) trains solely on low-resolution widely used public benchmarks, (ii) releases profile prior open-source generators. Hence poses no incremental societal risks beyond those already introduced by e.g. flow matching, warranting [NA].
data, and (iii) mirrors the well-studied risk
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact 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), fairness considerations (e.g.. deploymeett 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 not tied to particular applications, let alone deployments. However, if there is a 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 as 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 or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (c.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 ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (c.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The work trains and evaluates its model exclusively on widely used public datasets (MNIST, CIFAR-10, 10, ImageNet32x32,, CelebA) and does not release any new data. Because many comparable generative models for these benchmarks are already available, the paper introduces incremental misuse risk and thus requires additional safeguards.
Guidelines:
The answer NA means that the paper poses no 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 to 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 relcasing unsafe images.
25

---

We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
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?
Question:
Answer: [Yes]
Justification: All external assets used in this work are standard, mublicly-availaan bench- marks. For each dataset we cite the original creators in the manuscript. Below we summarise the licence associated with every dataset use:
MNIST Creative Commons Attribution Shiiionn-SaaaeAkkee 3.0 (CC BY-SA 3.0) CIFAR-10 Custom licence from the authors: free use for research and educational purposes provided the dataset cited
ImageNet32x32
-Images inherit the original ImageNet non-commercial research licence Y
-Down-sampling scripts released under the MIT License
-
CelebA - Custom non commerccllaal research agreement (no redistribution or commercial
use; dataset webpage) 6
AAV Public-domain molecular data (NCBI places no restrictions provided the dataset is cited]
Guidelines:
The answer NA means that the paper does not use 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 is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should 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.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, pwpppsswi whocodeethcoee..ooo com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should provided.
this information is not available online, the authors are encouraged to reach out to the asset'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: The newly introduced assets are fully documented: detailed algorithm de- scriptions and the complete implementation are bundled an anonymized ZIP archive that accompanies this submission.
Guidelines:
The answer NA means that the paper does not release new assets.
http::/aarsiv.or///mm11111...11
httpe://imgonnet.orrg
26

---

Researchers should communicate the details of the dataset/code/mmodel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember 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 as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does 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 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 an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on 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.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is 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.
Answer: [NA]
Justification: LLMs only used for editing.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: ptttps://neuripr http://eurrpps.cc/ooooteeee http:///eeeeeeecee//0//////1 for what should or should not be described.
27


Reviews:
[{"review_id": "CJk7GQXW9I", "reviewer": "Reviewer_4N5t", "summary": "This paper aims to combine the benefits of flow-matching and energy-based models, i.e., they aim to learn a time-independent energy function that determines the data distribution. When noisy samples are far away from the data manifold, the energy descent is more like flow-matching/optimal-transport, displaying straight trajectories; when the samples get close to the data manifold, the sampling process transitions to the Langevin dynamics, which converges to a Boltzmann distribution closer to the data distribution. The method is shown to display advantages in generating images, proteins, and estimating local intrinsic dimensionality.", "questions": "* Why is \"time-independence\" preferred? While it is true that the energy function does not explicitly depend on time, the $\\epsilon$ schedule implicitly plays the role of time. \n* If the first phase (OT) is accurate, why is the second phase (contrastive objective) even necessary? What is the overhead of the OT solver, as a function of data points? How does it compare to the overhead of distillation (train a first model that builds noise-data pairs, instead of solving the OT problem explicitly)?\n* In inference time, the Langevin-type sampling can be slow. How does the performance scale with the number of Langevin steps?\n* The Hessian of energy, although expensive to compute, contains valuable information. Using Hessian to estimate the LID is a good example. What other things become possible with energy matching that are impossible/challenging with flow matching?", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n* The paper is well-written, pleasant to read\n* The idea is novel, interesting, and seems to work\n\nWeaknesses:\n* Energy matching needs extra overhead compared to flow matching\n* The motivation can be made clearer (see questions)", "was_revised": true, "final_justification": "I would like to thank authors for their clarification. The motivation appears clearer to me. I maintain my positive rating.", "rebuttal": "Thank you for the thoughtful review and valuable feedback. Below, we provide point-by-point responses addressing the raised weaknesses (W) and questions (Q). We appreciate your comments on the motivation of certain aspects; we agree that this could be clarified further and commit to improving the paper accordingly in the revised version.\n\n**W1:  Energy matching needs extra overhead compared to flow matching**\n\nWe agree with the reviewer that Energy Matching introduces additional overhead compared to flow matching, primarily due to the contrastive-divergence training phase (details in **Q2/Q3**, including complexity and timings). However, this overhead is justified by enabling functionalities beyond unconditional generative tasks. Specifically, the contrastive phase refines the energy landscape to explicitly represent data likelihood, which provides essential advantages including conditional posterior sampling and data manifold curvature information (see detailed motivation in **Q2** and **Q4**). In contrast, flow-based models lose explicit likelihood representation, limiting their applicability in broader probabilistic inference scenarios.\n\n**W2: The motivation can be made clearer**\n\nWe appreciate this feedback. To clarify our motivation, we will highlight explicitly in the introduction that while flow- and diffusion-based models effectively generate high-quality samples, they inherently lack explicit modeling of the unconditional data score. This absence complicates controlling the generation, likelihood estimation, and data manifold curvature analysis (more details in **Q4**) —limitations that Energy Matching directly addresses.\nWe will further emphasize this motivation clearly in the revision.\n\n**Q1: Why is \"time-independence\" preferred? While it is true that the energy function does not explicitly depend on time, the schedule implicitly plays the role of time.**\n\nThe data density we ultimately care about is time-independent; thus, constraining the energy to be time-independent eliminates unnecessary degrees of freedom. A single set of weights covers all noise scales, leading to fewer parameters and improved regularization—our model uses 50M parameters compared to 150M in [C1], yet achieves better (lower) FID scores: 3.34 (our latest CIFAR-10 result, to appear in the camera-ready, more in **Q3**) versus their 3.68. At inference, this approach also removes the need to precisely estimate the noise level, a challenging issue in inverse problems [C2], since estimating time for the schedule $\\varepsilon(t)$ is not critical: it is mostly a piecewise constant schedule and inaccuracies in estimating $\\varepsilon(t)$ should not significantly impact sample behavior for $t \\gg 1$. Thus, parameterizing the potential energy as time-independent provides a minimal representation from a theoretical perspective and robustness in practice.\n\n- [C1] Y. Zhu, et al. “Learning Energy-Based Models by Cooperative Diffusion Recovery Likelihood” (NeurIPS 2023)\n- [C2] H. Chung, et al.  “Diffusion Posterior Sampling for General Noisy Inverse Problems” (ICLR 2023)\n\n**Q2: If the first phase (OT) is accurate, why is the second phase (contrastive objective) even necessary? What is the overhead of the OT solver, as a function of data points? How does it compare to the overhead of distillation?**\n\nWe thank the reviewer for this important question. Indeed, as shown in Figure 5a, the OT phase alone effectively aligns transport paths, efficiently guiding samples from noise toward the data distribution. However, the resulting potential from this first phase is not directly meaningful in terms of data likelihood—it merely represents a low-curvature landscape that facilitates efficient transport. Crucially, this learned potential after the OT phase does not encode likelihood of the data in any meaningful way.\nThe second phase (CD objective) is therefore necessary to explicitly encode likelihood information into the energy landscape. This is achieved by exploring the data manifold from multiple directions via negative sampling, creating distinct \"energy wells\" around regions of high data density (Figure 5b). A meaningful potential, proportional to negative log-likelihoods, is the central advantage of energy-based models over flow-based models. This explicit likelihood encoding enables:\n- Flexible conditional generation (e.g. inverse problems, interactions), through posterior sampling. \n- Explicit (in a single forward pass) log-likelihood estimation.\n- Analysis of the data manifold structure (e.g., intrinsic dimensionality estimation).\n\nWe agree this motivation could be emphasized more clearly and will explicitly add this discussion to the revised manuscript.\n\nRegarding the overhead of the OT solver as a function of data points and its comparison to distillation:\nthe OT solver we use (POT linear program) has complexity $O(B^2 + d)$, where $B$ is the minibatch size and $d$ is the data dimensionality. In practice, the scaling is mainly driven by the $B$, as $d$ only contributes linearly through the computation of the cost matrix, which is relatively inexpensive. For batches well beyond the size of 128 used here, regularized Sinkhorn methods can reduce complexity to $O(B\\log B + d)$. Ablation studies in existing literature [C3, Appendix C.2, Figure 10] indicate minimal differences between these OT methods, provided the regularization does not severely degrade the transport map.\nThe OT solver accounts for **1.5%** of the training iteration time in phase 1 (Algorithm 1). This drops to **0.01%** during the main training phase (Algorithm 2), after the contrastive-divergence term, which is significantly more expensive due to the simulation-based generation of negative samples, is activated.\nTraining a dedicated distilled model to approximate the global OT map presents a promising direction for achieving sampling with fewer steps. Indeed, precomputing the OT map upfront could eliminate sampling errors that arise when the transport map is supervised only on minibatches. We agree this warrants a deeper discussion and promise to add an additional remark on this topic in the revised version.\n\n**Q3: In inference time, the Langevin-type sampling can be slow. How does the performance scale with the number of Langevin steps?**\n\nPlease refer to Figure 6 in our manuscript, which shows the relationship between FID and the number of Langevin steps (fixed step size $dt = 0.01$). We observe rapid improvement initially, with performance saturating after about 300–325 steps. While Langevin dynamics: $x_m^{(n+1)} \\leftarrow x_m^{(n)} -  \\nabla_x V_\\theta(x_m^{(n)}) \\Delta t + \\eta\\sqrt{2\\varepsilon^{(n)}\\Delta t},\\quad\\eta\\sim\\mathcal{N}(0,I)$ can require many steps to reach equilibrium, our training ensures rapid convergence near the data manifold (via the OT flow objective), after which Brownian motion efficiently explores data modes.\n\n**Concrete timing results:** Unconditional CIFAR-10 sampling (batch size 128, NVIDIA R6000 48GB GPU), using Euler-Heun (325 steps) and Euler-Maruyama (1000 steps):\n\n| Method           | Params | Steps | Sampling time [s] ↓ | FID ↓               |\n|------------------|--------|-------|---------------------|---------------------|\n| OT-FM            | 37M    | 325   | **88**              | 4.04                |\n| DDPM++           | 62M    | 1000  | 183                 | 3.74                |\n| OT-FM            | 37M    | 1000  | 136                 | 3.45                |\n| Energy Matching  | 50M    | 325   | 173                 | **3.34**$^\\dagger$  |\n\n$^\\dagger$ FID 3.34 (lower is better) is our latest CIFAR-10 result using energy matching. The main improvements from the originally submitted result (FID 3.97) are due to adjusted hyperparameters: $\\tau$* increased to 1.0, $\\lambda_{cd} = 1×10^\\{-3}$, and the Langevin steps increased from 300 to 325. \n\nWe thank the reviewer for the important suggestion. We'll highlight our method's competitive sampling efficiency compared to SOTA flow- and score-based baselines in the revision.\n\n**Q4: The Hessian of energy, although expensive to compute, contains valuable information. Using Hessian to estimate the LID is a good example. What other things become possible with energy matching that are impossible/challenging with flow matching?**\n\nWe appreciate this insightful question highlighting a key distinction: energy-based methods explicitly model a scalar field representing data likelihood, unlike flow-based methods. This explicit potential enables probabilistic reasoning tasks (e.g., intrinsic dimensionality estimation, inverse problems, out-of-distribution detection, model selection), which flow-matching inherently lacks. Energy Matching addresses this by learning a potential proportional to negative log-likelihood, supporting likelihood-based computations.\nExplicit energy modeling enables various Hessian-based probabilistic tasks, e.g.:\n- Active-subspace dimension inference [C3]: accelerating inference by identifying data-relevant directions.\n- Laplace approximations/Bayes factors [C4]: efficient model selection via marginal likelihood estimation.\n- Stochastic-Newton preconditioned MCMC [C5]: accelerated sampling, especially useful in high-dimensional inverse problems.\n\nWe thank the reviewer for highlighting this important consideration. While we acknowledge Energy Matching’s higher training cost due to the CD phase (as further elaborated in our response to **W3** of reviewer **FNAJ**), we agree that clearly outlining this trade-off is important. We will emphasize the probabilistic advantages, particularly in likelihood-based reasoning tasks, more explicitly in the revised manuscript.\n\n- [C3] P. Constantine, et  al.  “Accelerating Markov Chain Monte Carlo with Active Subspaces” (SIAM 2015)\n- [C4] E. Daxberger, et al. \"Laplace redux-effortless Bayesian deep learning.\" (NeurIPS 2021)\n- [C5] T. Bui-Thanh, et al. “A computational framework for infinite-dimensional Bayesian inverse problems” (SIAM 2015)"}, {"review_id": "H9s7giug64", "reviewer": "Reviewer_XTSh", "summary": "This paper introduces \"Energy Matching,\" a generative modeling framework that combines aspects of EBMs and flow-matching approaches. The method proposes learning a single, time-independent scalar potential field that moves samples from a noise distribution to the data distribution. This occurs via a dual-regime process: far from the data manifold, samples follow curl-free, optimal transport-like paths, while near the manifold, an entropic energy term transitions them into a Boltzmann equilibrium, intended to capture the data's likelihood structure. Contributions include reported performance improvements over existing EBMs in generation fidelity on benchmarks like CIFAR-10 and ImageNet. The framework is presented as a prior for solving inverse problems, demonstrated with applications such as controlled inpainting and protein design, including an interaction energy term for diversity. It also enables local intrinsic dimension (LID) estimation by analyzing the learned energy landscape, reportedly using fewer approximations than diffusion models. The authors state these results are achieved without time-conditioning, auxiliary generators, or additional networks.", "questions": "Figure 6 in Appendix A.2 illustrates that the choice of the temperature-switching parameter $\\tau^* $\n  (from Eq. 3) can impact FID scores. Could you elaborate on the sensitivity of the model's final performance and training stability to the precise value of $\\tau^* $ and the overall shape/schedule of $\\varepsilon(t) $? Is the selection of $\\tau^*$ primarily empirical, or are there more principled guidelines for its determination based on dataset characteristics or model behavior?\n\nWhat is the inference process of this method? Does it need both flow matching part and Langevin MCMC part? If so, would it cost too much time? A standalone, explicit sampling algorithm in the main text would improve the clarity and reproducibility of this paper.", "limitations": "Yes.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths\n1. The core concept of \"Energy Matching\" to unify flow matching and Energy-Based Models (EBMs) by learning a single, time-independent scalar potential is a novel contribution.\n\n2. The theoretical development, rooted in the Jordan-Kinderlehrer-Otto (JKO) scheme and first-order optimality conditions, provides a principled foundation for the proposed method. The empirical evaluation is quite comprehensive, covering unconditional image generation (CIFAR-10, ImageNet 32x32), inverse problems (inpainting, protein inverse design), and Local Intrinsic Dimension (LID) estimation.\n\nWeaknesses\n\n1. The paper fails to state how much the modeling relies on the OT solver. The paper does not provide a clear definition of the OT solver. Is the method employed in the paper the unique OT solver? If not, can authors provide ablations on the OT solver?\n\n2. The paper does not provide an analysis of the computational overhead introduced by the OT solver (e.g., what percentage of the training time per iteration is spent on this step). The OT solver itself can be computationally expensive, especially as the batch size increases. This cost is a critical factor for the scalability of the method, especially when training involves the MCMC sampling.\n\n3. The paper forthrightly mentions in its limitations that the method requires extra gradient computation with respect to the input, potentially increasing GPU memory usage (20-40%). This is an important practical consideration for scalability.", "was_revised": true, "final_justification": "The paper proposed a novel algorithm for EBM. The results show obvious improvement. The theoretical argument about distance concentration in high dimensions is compelling. Yet, it's relatively inefficient in training. So I suggest a borderline accept.", "rebuttal": "Thank you for your valuable review and constructive comments, recognizing the strengths of our approach, including its novelty and comprehensive evaluation, and for pointing out areas for improvement regarding the OT solver discussion. Below we address the weaknesses (W) and questions (Q) raised, and we commit to clarifying these points further in the revised version.\n\n**W1: The paper fails to state how much the modeling relies on the OT solver. The paper does not provide a clear definition of the OT solver. Is the method employed in the paper the unique OT solver? If not, can authors provide ablations on the OT solver?**\n\nRemark 2.1 (\"OT solver\") on page 4 of our manuscript mentions that we use the \"POT\" linear-programming (LP) backend (Flamary & Courty, JMLR 2021), which computes an exact 1-to-1 OT plan per minibatch (e.g., size 128). We thank the reviewer for pointing out that this explanation might not be sufficiently clear. This choice is not unique; approximations such as Sinkhorn (with thresholding) or random coupling can also be effective.\n\nLet us (over‑idealise) and model a standardised CIFAR‑10 image as a vector $x \\in \\mathbb{R}^d$ drawn from $\\mathcal{N}(0,I_d)$ with  $d = 32 \\times 32 \\times 3 = 3072$, and pair it with an independent Gaussian‑noise vector $z \\sim \\mathcal{N}(0,I_d)$. \nEach coordinate of the difference $(z - x)$ is distributed as $\\mathcal{N}(0,2)$, thus the squared Euclidean distance satisfies $\\|z - x\\|\\_2^2 \\sim 2\\chi\\_d^2$, with mean $2d$, standard deviation $\\sqrt{8d}$, and relative spread $\\frac{\\sqrt{8d}}{2d} \\approx 0.025$. This thin shell means all entries of the cost matrix $C_{ij} = \\|z_i - x_j\\|_2^2$ are nearly identical, exemplifying the distance‑concentration effect highlighted by [D0]. Consequently, whether we solve the transport exactly (LP), entropically (Sinkhorn), or even use a random matching, the accumulated optimal‑transport cost should be similar. This has been examined empirically:\n- [D1] Table. 5: FID 4.44 (LP) vs 4.46 (random) at 100 steps; \n- our new CIFAR-10 runs: 3.97 → 4.00 at 300 steps.\n- [D1] Fig. D.2: for batch > 16, solver accuracy saturates on the 2‑moons task.\n- [D2] App. C.2: LP and Sinkhorn are indistinguishable unless regularisation is extreme.\n\nWe kept POT LP solver for zero hyperparameters, robustness, and negligible cost (see **W2**). We thank the reviewer and will add a dedicated appendix with the discussion.\n\n- [D0] Aggarwal, et al. “On the Surprising Behavior of Distance Metrics in High Dimensional Space” (ICDT 2001)\n- [D1] Tong, et al. “Improving and generalizing flow-based generative models with minibatch optimal transport” (TMLR 2024)\n- [D2] Terpin, et al. “Learning diffusion at lightspeed” (NeurIPS 2024)\n\n\n**W2: The paper does not provide an analysis of the computational overhead introduced by the OT solver (e.g., what percentage of the training time per iteration is spent on this step). The OT solver itself can be computationally expensive, especially as the batch size increases. This cost is a critical factor for the scalability of the method, especially when training involves the MCMC sampling.**\n\nWe acknowledge that the original manuscript did not quantify the computational overhead from the OT solver. Additional experiments show that the OT solver accounts for **1.5%** of the training iteration time in Phase 1 (Algorithm 1), dropping to a negligible (**0.01%**) level in Phase 2 (Algorithm 2), where the computational cost is dominated by negative samples generation. \n\nWe thank the reviewer and will include this analysis in the revised manuscript.\n\n**W3: The paper forthrightly mentions in its limitations that the method requires extra gradient computation with respect to the input, potentially increasing GPU memory usage. This is an important practical consideration for scalability.**\n\nThe extra costs buy a key EBM benefit: $V(x)$ supplies explicit curvature and likelihood information, unlocking inverse‑problem solvers and manifold analysis. Flow/score models lack this and instead resort to heavier inference tricks: full‑trajectory differentiation in [D3] or repeated noise/data alternations in  [D4] which come with their own limitations. \n- [D3] H. Ben-Hamu, et al., “D-Flow: Differentiating through flows for controlled generation,” (ICLR 2024)\n- [D4] M. Mardani, et al., “A variational perspective on solving inverse problems with diffusion models,” (ICLR 2024)\n\n\n**Q1: Figure 6 in Appendix A.2 illustrates that the choice of the temperature-switching parameter  (from Eq. 3) can impact FID scores. Could you elaborate on the sensitivity of the model's final performance and training stability to the precise value of  and the overall shape/schedule of ? Is the selection of primarily empirical, or are there more principled guidelines for its determination based on dataset characteristics or model behavior?**\n\nThank you for raising this insightful question regarding sensitivity to the temperature-switching parameter ($\\tau^* $).\nThe parameter $\\tau^* $ marks the activation time ($t>\\tau^* $) of the Brownian-motion component in our sampling dynamics; prior to this ($t\\le \\tau^* $), sampling follows deterministic optimal-transport paths toward high-likelihood regions. As described in Algorithm 3 (Appendix B), conditional sampling terms depend explicitly on the temperature $\\varepsilon$, vanishing as $\\varepsilon \\to 0$.\nLowering $\\tau^* $ (e.g., 0.8) allows earlier conditioning, potentially aiding conditional tasks, but typically reduces performance in unconditional generation. Conversely, setting $\\tau^* $ closer to 1.0 reinforces OT regularization near the data manifold, stabilizing training and enhancing unconditional generation quality.\nInitially, we presented only a brief ablation of $\\tau^* $ (Figure 6, Appendix A.2). Encouraged by your feedback, we've expanded these experiments, individually optimizing hyperparameters (e.g., $\\lambda_{cd}$) for different $\\tau^* $ values. We now report two hyperparameter sets: Set #1 optimized at $\\tau^* =0.9$, and Set #2 optimized specifically at $\\tau^* =1.0$.\n\nOur updated CIFAR-10 results are:\n| $\\tau^*$ | FID ↓  | Notes                                |\n|----------|--------|--------------------------------------|\n| 0.8      | 4.01   | Plateau at ~275 steps (HP set #1)    |\n| 0.9      | 3.97   | Plateau at ~300 steps (HP set #1)    |\n| 1.0      | 3.98   | Plateau at ~300 steps (HP set #1)    |\n| 1.0      | 3.34   | Plateau at ~325 steps (HP set #2)    |\n\nBased on these expanded experiments, we improved upon our originally submitted results (FID 3.97 $\\to$ 3.34); the updated result (FID 3.34) will be included in the camera-ready version.\n\nOur recommended heuristic is to always start with $\\tau^* =1.0$ for unconditional generation tasks, ensuring the best balance of training stability and final performance. If additional conditions are present, one may explore lower values ($\\tau^* < 1.0$) later, carefully examining if they provide benefits. We will explicitly incorporate this heuristic rationale into the camera-ready version.\n\n\n**Q2: What is the inference process of this method? Does it need both flow matching part and Langevin MCMC part? If so, would it cost too much time? A standalone, explicit sampling algorithm in the main text would improve the clarity and reproducibility of this paper.**\n\nWe thank the reviewer for the valuable suggestion; our inference procedure smoothly transitions from deterministic sampling (akin to flow-matching, OT paths and no randomness) to stochastic Langevin MCMC sampling, following a discretized temperature schedule $\\epsilon^{(n)}$. Specifically, when the temperature is zero ($\\epsilon^{(n)}=0$ for $n \\Delta t<\\tau^* $), the updates are deterministic (flow-matching), and when the temperature becomes positive ($\\epsilon^{(n)}>0$ for $n \\Delta t \\geq \\tau^* $), the updates introduce stochastic noise for modes exploration, becoming Langevin MCMC. This unified scheme is explicitly described in Algorithm 3 (currently in Appendix B.1, which we will move to the main text for clarity).\nSpecifically, at each inference step, the updates are:\n- **Unconditional update:**\n\n  $x_m^{(n+1)} \\leftarrow x_m^{(n)} - \\nabla_x V_\\theta(x_m^{(n)})\\Delta t + \\eta\\sqrt{2\\epsilon^{(n)}\\Delta t}$\n\n- **Conditional update:**  \n  The conditional update is identical in form to the unconditional one but uses a composite potential:\n\n  $ V_\\theta^{(composit)}(x_m^{(n)}) = V_\\theta(x_m^{(n)}) +\\frac{\\epsilon^{(n)}}{\\zeta^2}\\|y - A(x_m^{(n)})\\|^2 + \\frac{\\epsilon^{(n)}}{\\sigma^2}\\sum_{k \\neq m} W(x_m^{(n)}, x_k^{(n)})$\n\nHere, briefly:\n- $V_\\theta(x)$ is the learned scalar potential ($\\sim$ negative log-likelihood).\n- $A(\\cdot)$ denotes a forward measurement operator used in inverse problems.\n- $W(\\cdot,\\cdot)$ represents an optional interaction energy between different samples, e.g. encouraging diversity among generated solutions.\n- $\\varepsilon^{(n)}$ controls the stochasticity of sampling ($\\varepsilon^{(n)}>0$ corresponds to Langevin MCMC).\n- $\\eta\\sim\\mathcal{N}(0,I)$ introduces Gaussian noise at each Langevin step.\n\nCIFAR-10 unconditional (batch size of 128, NVIDIA R6000 GPU), using Euler-Heun (325 steps) and Euler-Maruyama (1000 steps):\n\n| Method           | Params | Steps | Sampling time [s] ↓ | FID ↓               |\n|------------------|--------|-------|---------------------|---------------------|\n| OT-FM            | 37M    | 325   | **88**              | 4.04                |\n| DDPM++           | 62M    | 1000  | 183                 | 3.74                |\n| OT-FM            | 37M    | 1000  | 136                 | 3.45                |\n| Energy Matching  | 50M    | 325   | 173                 | **3.34**$^\\dagger$  |\n\n$^\\dagger$ as in **Q1**: latest improved result for our Energy Matching model (to appear in camera-ready, same architecture).\n\nWe thank the reviewer for suggesting this improvement and will explicitly include Algorithm 3 in the main text to enhance clarity and reproducibility."}, {"review_id": "lFFHRlasf1", "reviewer": "Reviewer_Rsu1", "summary": "The paper \"Energy Matching: ...\" introduces a method that combines energy based models with flow type models using optimal transport theory. By doing that they can learn a single time invariant scalar field that can be used for unconditional sampling, inverse problems or local intrinsic dimension estimation. Results demonstrate strong performance, sometimes only within energy based models.", "questions": "- Please improve the general accessibility of the material to readers outside the OT field. \n- The exposition of the results could be clearer. You write that you surpass state-of-the-art EBMs. This is technically true, but somehow omits the fact that other methods in the comparison achieve even better results. I am not in favor of only accepting results that beat state of the art, but you need to argue why your results are still relevant and what advantage/disadvantage other methods have compared to yours.", "limitations": "Are discussed in section 4.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strength:\n- The motivation of learning an energy based model is not time dependent, allowing for an efficient sampling scheme as well as inverse problems is important. The authors seem to make a relevant contribution to this question. \n\nWeaknesses\n- The exposition is very dense and barely understandable for someone outside the OT field. \n- The paper title seems overstated. I does not read like a unifying theory of flow matching and energy based models, but rather a combination of the two. \n- The experiments are limited to smaller scale data (imagenet 32x32) and its not clear how the methods scales. The authors do acknowledge that in the limitations.", "was_revised": true, "final_justification": "In light of the response of the authors to my and other questions, I believe it's a valuable and interesting contribution to NeurIPS. Regarding my review: The authors clarified the unifying aspect and put the results to non-EBM into context. I thus raised my score.", "rebuttal": "We thank the reviewer for their constructive suggestions. We provide point-by-point responses to the raised weaknesses (W) and questions (Q).\n\n**W1/Q1: The exposition is very dense and barely understandable for someone outside the OT field. / Please improve the general accessibility of the material to readers outside the OT field.**\n \nThank you for raising this concern. We fully agree that providing intuitive explanations is important, especially for readers less familiar with optimal transport.\nIn particular, we will add a concise paragraph clearly emphasizing the following key insight: a central OT goal in our method is to produce a potential landscape with minimal curvature—effectively a clean, gently sloped \"path\"—between the noise and data distributions. This smooth slope makes it easier and faster to generate high-quality samples using only a few ODE or SDE integration steps. This capability is essential, as generating accurate, high-quality samples allows us to produce reliable negative examples. These negative examples are then used by the contrastive objective to better capture the geometric structure, or \"curvature,\" of the data manifold close to the actual data points. Thus, away from the data manifold, the OT component naturally regularizes the contrastive objective, enhancing both training stability and final generation quality.\n\nWe will incorporate this intuitive perspective clearly in the method section before formally introducing the solution via the JKO scheme. Additionally, we will include a brief summary of essential OT concepts in the appendix for quick reference. We welcome any further suggestions to improve clarity.\n\n\n\n**W2: The paper title seems overstated. I does not read like a unifying theory of flow matching and energy based models, but rather a combination of the two.**\n \nThank you for this valuable remark. We would like to clarify why we chose this specific title and agree that this deserves a clearer explanation. We will incorporate the following rationale explicitly into the camera-ready version to better communicate why we believe Energy Matching unifies flow matching and energy-based models.\nOur work is framed entirely in the JKO variational picture: one free‑energy functional is minimised at every step, and the temperature parameter $\\varepsilon$ decides which regime we are in. For any positive $\\varepsilon$ the updates coincide with the dynamics used to train energy‑based models; as $\\varepsilon$ approaches zero the very same updates degenerate to the displacement‑only objective that flow matching employs.  We are, to the best of our knowledge, the first to show explicitly that flow matching is the zero‑temperature member of this JKO family, and the title is meant to foreground that link.\n\nExamining Algorithm 3 (p. 15) (also discussed in our response to **Q2** from reviewer **XTSh**) reveals a key limitation of flow matching which operates in the $\\varepsilon=0$ limit. Every term that would steer the sampler toward a specific observation is multiplied by $\\varepsilon$; setting $\\varepsilon$ to zero (precisely the flow‑matching limit) makes those terms disappear. The model is then locked into a single transport path dictated by the prior and cannot produce samples from a conditioned posterior without outside tricks and approximations. Highlighting this drawback, and the way our framework overcomes it, is a novel contribution of the paper.\n\nThe optimal‑transport portion of the loss provides positive guidance on where mass should move, while the contrastive-divergence loss contributes a negative signal that discourages excursions into low‑density regions. Rather than merely stitching two concepts together we show that they are part of the same variational picture with complementary losses for different regions of the trajectory.\n\nWe agree that explicitly clarifying this rationale will strengthen our contribution and help prevent potential confusion regarding the title. To address your valuable feedback, we will prominently emphasize this unified variational interpretation in the abstract and introduction.\n\n**W3: The experiments are limited to smaller scale data (imagenet 32x32) and it's not clear how the methods scales. The authors do acknowledge that in the limitations.**\n\nThank you for highlighting this important point regarding scalability. We agree that our current experiments focus primarily on image datasets of moderate resolution. This scale, however, is a widely studied benchmark, making it well-suited for demonstrating clearly the strengths and relative performance of our approach. Indeed, our method significantly surpasses the previous state-of-the-art FID scores for EBMs by achieving more than a two-fold improvement. We believe this demonstrates the potential to extend our method to higher-dimensional images in future work.\nImportantly, our approach is also well-suited to problems outside image generation, such as protein design, where typical data dimensionality is substantially lower than even moderate-resolution images. This is a significant and practically relevant domain that benefits directly from the explicit scalar-potential structure inherent in our approach, as demonstrated in our protein-generation experiments.\nRegarding computational complexity of some of the components:\n\n- Optimal‑transport couplings add only linear $O(d)$ overhead. For very high‑dimensional data there is both theoretical and empirical evidence suggesting dropping the solver and using random couplings that do not cost. (please see response **W1** for the reviewer **XTSh**).\n \n- The full Hessian spectrum with $O(d^3)$ is used exclusively for the Local Intrinsic Dimensionality diagnostic and is neither required for training nor for unconditional/controlled generation; practical deployments can drop this step or approximate it as suggested in the limitation section.\n\nWe appreciate the reviewer’s comment and will incorporate a deeper exploration of scalability to larger-scale datasets in our outlook and future research directions.\n\n**Q2: The exposition of the results could be clearer. You write that you surpass state-of-the-art EBMs. This is technically true, but somehow omits the fact that other methods in the comparison achieve even better results. I am not in favor of only accepting results that beat state of the art, but you need to argue why your results are still relevant and what advantage/disadvantage other methods have compared to yours.**\n\nWe thank the reviewer for pointing out that our comparison to non-EBM methods could indeed be clarified further. The inclusion of non-EBM results primarily serves as a contextual benchmark to illustrate that our method significantly advances energy-based modeling toward competitive generation quality, which was previously challenging for EBMs. Within the category of EBMs, our approach clearly achieves state-of-the-art performance.\nThe relevance of EBMs lies in their unique advantages: they inherently provide explicit likelihood estimates (in a single forward pass) for any given sample, enabling straightforward and flexible conditional sampling. Conditions or new interactions between samples can be incorporated by simply adding energies, resulting in a composite energy landscape (explicitly detailed in Algorithm 3 of Appendix B: Posterior Sampling,  Section 3.2 on Inverse Problems, and our response to reviewer **XTSh**'s question **Q2**). Additionally, EBMs naturally provide Hessian-based curvature information about the data manifold, enabling applications (such as LID estimation or more in the response to reviewer **4N5t**(**Q4**)) that flow- or score-based methods cannot directly offer without additional approximations. \n\nHowever, these advantages come with additional computational overhead during training, which we explicitly acknowledge (see our detailed discussion in **W3** of reviewer **FNAJ**). Whether this overhead is justified depends on the application. For use-cases requiring precise control over generation—such as many scientific applications already outlined in the outlook of our manuscript—this additional cost is typically justified. Similarly, scenarios where conditions, interactions, or labels for conditional generation may not be available at training time, necessitating a model encoding a versatile prior that can accommodate new constraints on demand, further highlight the practical value of our approach.\n\nWe sincerely thank the reviewer for emphasizing this important point; we will use the additional page in the camera-ready version to explicitly clarify how to interpret comparisons between EBM and non-EBM methods at the beginning of the results section."}, {"review_id": "1ha5PUgyyx", "reviewer": "Reviewer_FNAJ", "summary": "The manuscript proposes a hybrid generative model formulation that combines flow matching and energy-based modeling. The core concept behind such formulation is the Jordan–Kinderlehrer–Otto (JKO) scheme that selects a probability distribution in the next timestep as minimiser of linear combination of: Wasserstein distance between the target distribution and the distribution in the current timestep, potential energy with parameterized scalar potential, and internal energy expressed as negative entropy tuned with a temperature-like parameter $\\epsilon(t)$. Following previous work [Terpin et al., 2024], the JKO scheme is analysed under its first-order optimality conditions at each timestep $t$ (defined with Eq. 4). Different from previous works, the manuscript studies the resulting formulation depending on the time parameter $t$ and reveals two distinct regimes. In the case of early timesteps ($\\epsilon(t)$ \\approx 0), Eq. 4 approximately behaves like a flow-matching OT with time-independent potential. Such straight paths are useful for fast space traversal far from the data manifold. In the case of later timesteps nearer to the data manifold ($\\epsilon(t)$ \\approx 1), the scalar potential can be reformulated into Boltzmann's energy. Thus, the sampling trajectories are allowed to curve, which adds to the model's expressivity and matching of the data manifold. The parameterised scalar potential is trained in two stages. Warmup stage optimises the OT objective characteristic for flow matching, while the main training stage combines the OT objective with a contrastive objective typical for EBMs. Experimental evaluation indicates strong performance in unconditional image generation, inverse problems, and  estimation of local intrinsic dimension.", "questions": "See weaknesses.", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths:**\n\nS1. The proposed hybrid generative formulation elegantly combines OT-based flow matching and EBMs through time-dependent analysis.\n\nS2. The manuscript is well written and easy to follow.\n\nS3. Experimental evaluation indicates competitive performance for three different applications.\n\n**Weaknesses:**\n\nW1. The assumption $\\rho_t \\approx \\rho_{eq}$ (Line 84) may not hold in practice, especially early in the training. The manuscript should elaborate more on the feasibility of such an assumption. Perhaps visualise contributions of the tree components in Eq. 4 depending on $t$. My intuition is that this is why the warmup stage is needed in the first place.\n\nW2. The manuscript reports only the FID score for unconditional image generation. Reporting precision and recall [a] would further strengthen the results.\n\nW3. The proposed training procedure (Algorithm 2) increases training time compared to the standard flow matching due to the sampling required by the contrastive objective. This should be quantified and discussed in ablations (and possibly added to limitations).\n\nW4. The manuscript parameterises the scalar potential $V_{\\theta}$ as a deep model. This implies that both forward and backward passes are needed for sampling (as in EBMs). This is significantly slower than in the case of direct parameterisation of vector fields $\\nabla_x V_{\\theta}$  (typicall for flow matching/score matching). The manuscript should discuss this. For example present wallclock time vs FID graph for different baselines.\n\n[a] Tuomas Kynkäänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, Timo Aila: Improved Precision and Recall Metric for Assessing Generative Models. NeurIPS 2019", "was_revised": true, "final_justification": "The authors presented compelling answers to all raised concerns. Thus, I recommend acceptance.", "rebuttal": "We thank the reviewer for their thoughtful comments and helpful suggestions, which allow us to clarify important aspects of our manuscript. Below, we provide detailed responses addressing the raised weaknesses (W) and questions (Q):\n\n**W1: The assumption  $p_t \\approx p_\\{eq}$ (Line 84) may not hold in practice, especially early in the training. The manuscript should elaborate more on the feasibility of such an assumption. Perhaps visualise contributions of the tree components in Eq. 4 depending on $t$. My intuition is that this is why the warmup stage is needed in the first place.**\n\nWe thank the reviewer for highlighting this important point, as it reveals a missed opportunity in our manuscript to clearly explain the motivation behind the warm-up phase.\nThe statement from line 84 that $p_t$ is close to $p_{eq}$ (in the context of $p_t$ being ‘near data manifold’) indeed may not hold in practice, particularly in the early stages of training. However, line 84 does not imply any guarantees for untrained $V(x)$. Instead, the subsection \"First-order optimality conditions\" (pp. 3–4) describes properties of the desired solution (such as $p_{data} = p_{eq}$ and Eq. 4) rather than assumptions that must hold throughout training.\nDuring training, we explicitly measure how far the current distribution $p_t$ deviates from these optimal conditions and use this imbalance as our learning signal.\n\nWe welcome the reviewer's suggestion to visualize how the three forces evolve during training. We will include this visualization in supplementary material.\n\nEarly in training, when the potential $V$ is still untrained, contrastive terms provide limited guidance because the negative sample distribution $p_t$ may either be too distant from the data manifold (if initialized in noise) or too close due to poor mixing (if initialized directly on the manifold). This perspective clearly highlights the necessity of the warm-up phase identified by the reviewer. As briefly mentioned around lines 143–144, it is crucial to have negative samples properly distributed around the data manifold to probe low-density regions, thereby constructing energy walls that clearly separate these regions from the high-density data manifold. Once the warm-up phase sufficiently aligns $p_t$ with $p_{data}$, the contrastive terms can more effectively enforce the condition $p_t \\approx p_{data} \\Leftrightarrow p_t \\approx p_{eq}$, actively pulling $p_{eq}$ closer to $p_{data}$.\n\nWe thank the reviewer and will explicitly add a sentence at the start of the \"First-order optimality conditions\" subsection, clarifying that these conditions represent the optimization goal (features of the solution) and highlighting the critical role of the warm-up phase.\n\n\n**W2: The manuscript reports only the FID score for unconditional image generation. Reporting precision and recall [a] would further strengthen the results.**\n\nWe thank the reviewer for this valuable suggestion. While our results measured by FID significantly surpass existing EBMs (halving the best reported FID scores), we acknowledge that FID alone provides a limited perspective. For proteins, we indeed report additional metrics such as diversity and novelty alongside fidelity measures (see Table 4 in Appendix B, \"Posterior sampling\"). In the revised manuscript, we will explicitly highlight that, given our FID scores are now on par with state-of-the-art methods, future works on image generation with EBMs should incorporate complementary metrics such as precision and recall [a] for a more complete evaluation.\n\n**W3: The proposed training procedure (Algorithm 2) increases training time compared to the standard flow matching due to the sampling required by the contrastive objective. This should be quantified and discussed in ablations (and possibly added to limitations).**\n\nWhile the proposed training procedure indeed adds computational overhead compared to standard flow matching due to the contrastive sampling, we view this overhead as justified given the additional capabilities provided by our approach. As illustrated clearly in Figure 5a, the OT phase alone (without contrastive sampling) achieves results similar to flow matching (same objective, different parameterization), and does not require sampling negative examples. The contrastive objective refines the learned landscape, explicitly enabling functionalities that standard flow-based models cannot achieve directly, such as evaluating log-likelihoods in a single forward pass and performing posterior sampling under arbitrary conditions. Thus, if these functionalities are desired, the additional computation cannot be seen purely as overhead but as necessary to enable these key features. \n\nWe fully agree that this computational cost should be clearly quantified and discussed. To address this, we provide the following timing breakdown (per batch on CIFAR-10 using a single NVIDIA R6000):\n\n| Batch | Phase | Total [s] | OT coupling [s] (%) | Flow loss [s] (%) | CD loss [s] (%) |\n|-------|-------|-----------|---------------------|-------------------|-----------------|\n| 128   | 1     | 0.260    | 0.004 (1.5%)      | 0.256 (98.5%)   | 0 (0%)          |\n| 128   | 2     | 33.030   | 0.004 (0.01%)     | 0.270 (0.82%)   | 32.755 (99.17%)|\n\n**Additional Timings:** EM potential evaluation: **0.148 s**; EM velocity generation: **0.317 s**\n\n\nIn short, the flow component is roughly 2x slower compared to directly training a velocity field via flow matching due to the additional backward pass. While the CD loss computation itself is approximately 100x slower than the flow loss, CD iterations constitute only about 1–5% of total flow iterations (see Training Details). Using our current negative-sample generation approach (no replay buffers, small fixed step integration), the total training time can increase by up to 12 times in the extreme cases, due to both the curl-free parametrization and negative-sample simulations. \n\nAlthough this overhead remains negligible for shorter tasks (e.g., protein generation, small images), it could become significant in larger settings. Future improvements could incorporate techniques such as persistent contrastive divergence [E1, E2], replay buffers, or adaptive Langevin steps to significantly reduce generation costs of the negative samples. We omitted these common engineering tricks here in favor of clarity of presentation. \n\nWe thank the reviewer for highlighting this point. The overhead described here is characteristic of most traditional EBMs, which typically rely on simulating negative samples during training. Although simulation-based training lost popularity in recent years due to its computational demands and the rise of simulation-free approaches, improvements in hardware capabilities are now bringing simulated negatives back into practical relevance, as they provide complementary signals to simulation-free supervision. Importantly, this overhead is incurred mainly during training and the method remains competitive in terms of inference speed (addressed further in **W4**). Additionally, Energy Matching can naturally extend to discrete and structured domains, such as molecule generation, where recent work [E3] explores replacing Langevin-based MCMC negative sampling with analytical solutions derived from heat equations on structured spaces. \n\nWe will clearly emphasize this context in both the manuscript (Introduction) and appendix (Training Details), making it accessible especially to readers less familiar with EBMs. \n- [E1] T. Tieleman and G. Hinton, “Using Fast Weights to Improve Persistent Contrastive Divergence” (ICML 2009)\n- [E2] C. Jarzynski, “Nonequilibrium equality for free energy differences”\n- [E3] T. Schröder et al., “Energy-Based Modelling for Discrete and Mixed Data via Heat Equations on Structured Spaces” (NeurIPS 2024)\n\n**W4: The manuscript parameterises the scalar potential  as a deep model. This implies that both forward and backward passes are needed for sampling (as in EBMs). This is significantly slower than in the case of direct parameterisation of vector fields  (typicall for flow matching/score matching). The manuscript should discuss this. For example present wallclock time vs FID graph for different baselines.**\n\nWe fully agree that directly reporting parameter counts and sampling steps serves only as a proxy, and that explicitly comparing wall-clock times against achieved FID scores would offer clearer insights. Indeed, our scalar-potential parameterization requires both forward and backward passes for sampling, similar to traditional EBMs, which roughly doubles (as seen in **W3**)  sampling steps time compared to directly parameterizing vector fields.\n\nWe sincerely thank the reviewer for emphasizing this important point and will include a clear wall-clock time vs. FID comparison plot in the revised manuscript. This visualization will clearly compare our method with baseline methods we successfully reproduced. Below, we present the measurements obtained on CIFAR-10 (batch size of 128, NVIDIA R6000), using Euler-Heun (325 steps) and Euler-Maruyama (1000 steps):\n\n| Method           | Params | Steps | Sampling time [s] ↓ | FID ↓               |\n|------------------|--------|-------|---------------------|---------------------|\n| OT-FM            | 37M    | 325   | **88**              | 4.04                |\n| DDPM++           | 62M    | 1000  | 183                 | 3.74                |\n| OT-FM            | 37M    | 1000  | 136                 | 3.45                |\n| Energy Matching  | 50M    | 325   | 173                 | **3.34**$^\\dagger$  |\n\n- $\\dagger$ *FID 3.34 achieved by Energy Matching differs from the original submission (FID 3.97) due to adjusted hyperparameters: $\\tau^\\*=1.0$ (vs $0.9$), $\\lambda_{cd}=1\\times10^{-3}$ (vs $8\\times10^{-3}$), and 325 Langevin steps (vs 300). The architecture is unchanged; updated result in camera-ready.*"}]

Author rebuttal:

Meta-review: Tl;dr: Based on the reviews, rebuttal and ensuing discussion I recommend accept.

### Paper Summary

A novel generative modeling framework, Energy Matching, is introduced, that unifies flow matching and energy based models. Proposed method learns a single time-independent scalar potential energy function that governs the generative process. Paper claims that this approach allows for efficient, transport-like sampling far from the data manifold and explicit likelihood modeling near the data manifold. The empirical results show that Energy Matching significantly outperforms previous EBMs on standard image generation benchmarks (CIFAR-10, ImageNet 32x32) and is competitive with state-of-the-art diffusion and flow-matching models, while offering the additional benefits of a learned likelihood.

### Key strengths and weaknesses

Strenghts: 1) Novelty. Proposed method is novel and principled. The formulation is grounded in the JKO scheme, providing a solid theoretical foundation, 2) Strong empirical results: substantial improvements over existing EBMs are shown with competitive performance against other SOTA methods. Additional results on controlled inpainting and protein design further demonstrate the strengths of the approach.

Weaknesses: 1) Writing:  The main initial weakness (also pointed out by the reviewers) was the clarity of the presentation, particularly for people not familiar with optimal transport theory, 2) Overhead: The computational overhead of the proposed method, especially during training, compared to pure flow-matching models, was another concern, 3) Scaling to higher resolution: The scaling of the method to higher-resolution images was another concern.

### Decision justification

The decision to recommend acceptance is based on the novelty, theoretical elegance, and strong empirical results of the proposed method. The reviewers had concerns regarding the clarity of the method, the computational overhead, the role of the OT solver, and the scalability of the approach. The authors provided detailed and convincing responses to all the concerns. Authors have committed to significant revisions to the paper to improve its clarity, including adding a more intuitive explanation of the method, a more detailed discussion of the computational overhead, and moving the sampling algorithm to the main text. The reviewers were satisfied with the author responses and the planned revisions, and were unanimously in favor of acceptance.
Decision: Accept (poster)