Energy-Based Flow Matching for Generating 3D Molecular Structure
Authors: Wenyin Zhou
Christopher Iliffe Sprague
Vsevolod Viliuga
Matteo Tadiello
Arne Elofsson
Hossein Azizpour
Venue: ICML 2025 poster (2025)
Keywords: flow matching
molecular docking
protein backbone generation
energy-based models
Primary area: applications->chemistry_physics_and_earth_sciences
TL;DR: An energy-based formulation of flow matching with application to generating 3D molecular structure

Abstract: Molecular structure generation is a fundamental problem that involves determining the 3D positions of molecules' constituents. It has crucial biological applications, such as molecular docking, protein folding, and molecular design.
Recent advances in generative modeling, such as diffusion models and flow matching, have made great progress on these tasks by modeling molecular conformations as a distribution.
In this work, we focus on flow matching and adopt an energy-based perspective to improve training and inference of structure generation models. Our view results in a mapping function, represented by a deep network, that is directly learned to \textit{iteratively} map random configurations, i.e. samples from the source distribution, to target structures, i.e. points in the data manifold. This yields a conceptually simple and empirically effective flow matching setup that is theoretically justified and has interesting connections to fundamental properties such as idempotency and stability, as well as the empirically useful techniques such as structure refinement in AlphaFold. 
Experiments on protein docking as well as protein backbone generation consistently demonstrate the method's effectiveness, where it outperforms recent baselines of task-associated flow matching and diffusion models, using a similar computational budget.

Energy-Based Flow Matching for Generating 3D Molecular Structure
Wenyin Zhou' Christopher Iliffe Sprague 123 Vsevolod Villuga" 245 Matteo Tadiello 24 Hossein Azizpour 12
Arne Elofsson
Abstract
& Ben-Tal, 2018). It finds various applications, such as protein docking, which generates the bound structure of protein-ligame complex, and de novo protein design, which creates novel proteins fulfilling certain desirable biologi- cal functions. Regression- els (Jumper et al., 2021; Baek et al., 2021; Lin al., 2023) have demonstrated remarkable performance by using a deep neural network to make point estimate of the structure. However, these methods do not capture aleatoric uncertainty due to the possibility of multiple molecular conformations.
Molecular structure generation is a fundamental problem that involves determining the 3D posi- tions of molecules' constituents. It has crucial biological applications, such molecular dock- ing, protein folding, and molecular design. Re- cent advances in generative modeling, such as diffusion models and flow matching, have made great progress on these tasks by modeling molec- ular conformations as a distribution. In this work, we focus on flow matching and adopt an energy- based perspective improve training and infer- ence of structure generation models. Our view results in mapping function, represented by a deep network, that is directly learned to itera- rively map random configurations, iee samples from the source distribution, to target structures, i.e. points in the data manifold. This yields a con- ceptually simple and empirically effective flow matching setup that is theoretically justified and has interesting connections to fundamental prop- erties such as idempotency and stability, as well as the empirically useful techniques such as struc- ture refinement AlphaFold. Experiments on protein docking as well as protein backbone gen- eration consistently demonstrate the method's ef- fectiveness, where it outperforms recent baselines of task-associated flow matching and diffusion models, using similar computational budget.
structure prediction mod-
In contrast to regression-based models, generative meth- ods, such as diffusion (Song & Ermon, 2019; Ho et al., 2020; Song et al., 2021) and flow matching models (Lip- man et al., 2023; Albergo & Vanden-Eijnden 2023; Liu et al., 2022), can generate several possible configurations modeling the positional distribution of molecules. Owing to this advantage, sampling-base based approach has recently attracted huge interest in molecular structure generation, in- cluding molecular conformer generation (Hoogeboom et al., 2022), molecular docking (Corso et al., 2023), and protein design (Watson al., 2023; Campbell et al., 2024).
Here, importantly, first observe that these families of generative modeling are generally reminiscent of the computational models for molecular structure determina- tion. Particularly, in computational models, a stable equi- librium structure is sought as the local minimum of an energy function governed by the Boltzmann distribution p(x) x exp [==))) which resembles modeling (and gen- erating from) the data distribution by estimating the score a corresponding energy of its density, In this work, we make this connection explicit and propose simple modifica- tion of the objective in flow matching that directly minimizes an energy function.
1. Introduction
Structure prediction, the task to determine a molecule's 3D structure, is a fundamental problem in structural biology for understanding diverse biological mechanisms (Kessel
Specifically, we take the reconstruction error || -|| as the energy function and locally construct the energy landscape using contrastive predictions sampled from the flow model. Leveraging data parameterization (Stärk et al., 2023; Yim et al., 2023a) of flow matching, minimizing our devised energy, quite intriguingly. translates into training a neural network that iteratively predicts and refines the generated sample as shown in Fig. rendering the neural network, an approximated idempotent function.
'KTH Royal Institute of Technology. Stockholm, Sweden "Science for Life Laboratory, Sweden 'The Alan Turing Institute, London, United Kingdom 'DBB Stockholm University, Sweden ' MMax Planck Institute Polymer Rescarch, Mainz, Germany. Correspondence to: Wenyin Zhou Ewenyinz@kthh. @kth.se>.
Proceedings of the ndnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 20255 by the author(s).

---

Flow Matching: Idempotent Flow Map
fo(%))
(o(
Low
X
X
Palx)
Po(X)
Pi()
Pi(X)
(a) Flow matching data parameterization.
(b) Flow Matching: idempotent flow map
Figure 1. Training paradigm of standard flow matching and IDFlow. zo and æ1 are samples from the source and target distribution (z) and pi (z), Te is the linear interpolant between the source and target sample, â1 is the prediction of the network, fo($1) is the refined prediction, LerM is the conditional flow matching loss and La is the refinement loss. (a) Flow matching: directly predict the data &1 given the interpolant Te- (b) IDFlow: mapping function represented by deep neural network fo, iteratively predicts and refines the prediction. The energy is estimated to be the input output structure difference after refinement.
Notably, trained idempotent function (IDFlow) simulta- neously serves both neural sampler and a refiner, leading to the characterization of a predictor-rehe refinerr sampler, analo- gous to the predictor-corrector sampler diffusion models (Song et al., 2021). This approach iteratively refines the sampling trajectory toward points where the energy function approximates zero, effectively enhancing flow matching as an energy-baan mhergy-based model but crucially without significant in- crease in training cost. The key contributions are follows.
point Learning within this paradigm requires devising an energy function and loss function to shape the energy landscape. For instance, the L2 regression loss could be seen as directly using the energy function for the loss and the exact format the energy function (referred to as the energy architecture hereafter) the L2 norm between the network output fe(x)) and the label y:
Energy:
Eg(x) = llfe(z) - y/1 Loss Lo(x) = Fo(x),
(1)
We adopt an energy-based model perspective for flow matching sampling based molecular structure gen- eration tasks which respects the energy minimization.
(2)
We provide a simple instance of such framework with negligible computational burden, IDFlow, which itera- tively predicts and refines the sample and, importantly, draws interesting connections related approaches, such as structure refinement.
where (æ., y) is a training pair of input and annotation. Train- ing with this strategy primarily reduces the energy of the training examples. general, there are two ways of train- ing the EBMs: 1) contrastive approaches which maximize the energy margin between positive and negative training examples, and 2) regularized methods that minimize the vol- ume of low-energy space while avoiding the collapse of the overall energy landscape, c.g., regularizing the latent space to be close to the standard Gaussian as low energy space in variational autoencoders (Kingma Welling, 2013). Par- ticularly, the collapse occurs when the network produces identical outputs regardless of the input. One contrastive method of training EBMs is learning function that maps the points off the data manifold to the points on the data manifold (LeCun, 2022). It is this latter approach that we adopt in this work to improve Flow Matching for structure prediction. For a detailed review EBMs, see (LeCun et al., 2006).
We conduct extensive experiments based two recent flow matching models tackling three structure predic- tion tasks of (single and multi-ligand) docking and pro- tein design using four public datasets. Our results show that IDFlow improves the structure generation perfor- mance over the recent task fask-assoeiated flow matching and diffusion models.
2. Background
2.1. Energy-Basee Models
Energy-Basee Models (EBMs) consider a positive energy function Eo (x) E RR assigning a sealar value to each data
2

---

Flow Matching: Idempotent Flow Map
2.2. Flow Matching
output as either the vector field uy o or the data i1 (Stärk et al., 2023; Jing et al., 2024). From the EBMs perspective, the latter results in objective function that directly pushes down the energy of the trajectory samples, which is desirable:
Flow matching (Lipman et al., 2023; 2024; Albergo & Vanden-Eijnden, 2023; Liu et al., 2022) is simulation-free training method for continuous normalizing flows (Chen et al., 2018) 8) that connects two arbitrary distributions with flexible design choices. This allows for building straight paths between any source and target sample, which has great potential for accelerated sampling. Specifically, flow matching aims learn a deeedeeendnnt vector field ey (2) evolving the sample from the prior po (2r) the target distri- bution p1 (1)) via an ordinary differential equation (ODE);
LCEM (f) = ()), - x 11.-
The ø1 parameterization yields an Euler-like step sampling algorithm:
dx
[as(w)
(8)
dt
dæ dt
where fe is the mapping between point from the sampling trajectory to the corresponding target data point, and is, here- after, referred to as the flow map. The general format of the vector field can be found in the App. B.3. A more intuitive understanding of fo is the "denoiser" that approximates the expectation of p P [z)]:
(x))
(3)
where x could either be defined over Euclidean space Idd or Riemannian Manifold M (Chen & Lipman, 2024). To achieve this, the flow matching objective (Lipman et al., 2023) regresses the predicted vector field to the true vector field (I)::
E |Xx, = R2 fe.(x))
(9)
(PM(0)= == 12 (1o. - ue (2) (x))
(4)
In this work, we focus the Iy parameterization for its connections to the EBMs loss and generative adversarial nets (Goodfellow et al., 2014).
However, this objective function cannot be leveraged for training as the true vector field is intractable in practice. Pre- vious work (Lipman et al., 2023; Tong et al., 2024) shows that with the construction of the conditional probability paths, the conditional flow matching (CFM) objective shares the same parameter gradient the flow matching objective V LFM = LCLC such that learning the conditional vec- tor field would also optimize for the marginal vector field:
Riemannian flow matching. Flow matching is extended to the Riemannian manifold by constructing a premetric that measures the proximity of a point æ to the target in for defining the vector field (Chen & Lipman, 2024). The core idea is to define the premetric as the geodesic distance, such that the geodesic interpolant 24 can be efficiently com- puted with the exponential and logarithmic map on simple manifolds (e.g. the N-D sphere) in closed-form:
(O) =
(-1(0,1) [[00,((x)- - u1(x uo, x1)
Euclidian:
Ie = t + (1 - t)x0,
(10)
so)pp(z))
z1ng((11)
Riemannian: I2 = exPa, (tlog (I1)).
(11)
(5)
where ue (2( 21 is the conditional vector field that gen- erates the conditional probability path (r mo, 21). and To and ii are sampled from the prior (or source) distribu- tion and the training (or target) set respectively. Notably, flow matching does not depend on any prior distribution (standard Gaussian in diffusion models) and allows flexible construction of probability paths. To reduce the transport cost, conditional optimal transport builds the probability path that linearly transfers the source sample to the target:
Specifically, the exponential and logarithmic map of the manifold SO(3) can be calculated using the well-known Rodrigues' formula. This has been applied to generating protein backbones (Yim et al., 2023a; Bose et al., 2024; Huguet et al., 2024; Wagner et al., 2024) and robot motion learning (Braun et al., 2024).
3. Method
Pe (1 mo) (x)))(()) of)
px(x)= = txı + (1 t)
3.1. Task Formulation and Notation
(6) where o; is predefined standard deviation which could be either constant or a time-dependen noise scheduler (e.g. o, /211 - a)),,, and (x) is the mean of the condi- tional distribution.
We consider two widely adopted tasks, molecular docking and protein backbone generation for the investigation of the method on generating 3D molecular structure.
Protein docking. For the molecular docking task, we rep- resent the protein structure using Cartesian coordinates as y E 3x and the ligand structure as x E 33xm, where np and ny are the number of protein residues and ligand atoms,
Ti parameterization. Similar to diffusion models (Wat- son et al., 2023), one can opt to parameterize the network
3

---

Flow Matching: Idempotent Flow Map
respectively. The aim is to learn the conditional distribution p(zll)) i.e. the docking distribution of the ligand æ given the protein structure y.
tive (high-accuracy) and negative (low-accuracy) examples (Corso et al., 2023). Then the logits of such a network can be used as the confidence score (energy function) for an unseen example. In contrast to these two approaches, in this work, we propose to directly construct and learn an energy function within the Flow Matching model. However, from objective Eq. 7, is evident that only the energy associated with the trajectory sample I; is pushed down, while the contrastive samples @1 generated by the flow model, unlike 31, remain unaffected by this process. Therefore, the idea of the energy-based flow matching is precisely to better shape the energy landscape to improve training with contrastive samples 21 and to encourage reaching the minimum of the energy function as shown in Fig. 1.
Single ligand docking assumes a unique protein-ligany pair (2r, y), such that the generative model learns the bind- ing modes of the protein y with a single sample from the conditional distribution p(xl))
Multi-ligand docking assumes multiple pairs of protein y and ligands zy; (Ii, y), (22, ..., (I.,, v). where is the number of ligands for protein y. The training signal of multi-ligand docking can be enriched with several docking conformations 2, provided for learning the conditional dis- tribution p(aly), but also increases the learning difficulty. This task finds application in multi-ligand binding pocket design, such as the enzyme design for multiple reactants.
Shaping the loss landscape with ②t. We first need to define energy function E(+1) IR" + R4, such that high-probability data lies around its minima. Following (Zhao et al., 2017), a simple energy architecture could be the reconstruction error involving a certain function G:
Protein backhone generation represents each protein residue by a frame representation (r, E SE(3), rotation r E (3) and a translation s E T(3). Hence, a protein y is modeled by the N compositions of SE(3) group, where N is the number of protein residues. As such, the genera- tive model is defined over the manifold y M = SEE(3) for learning the protein distribution ply). The translation places each protein residue, particularly the Ca carbon, relative to the reference frame, and the rotation captures the local orientation of the residue. The detailed definition of the reference frame can be found in App. C.2.
Eo((1) = [|G(#t) - 11
(13)
Since the conversion from energy to probability could be achieved through the Boltzmann distribution, this energy function also aligns with the likelihood function in the con- text of conditional flow matching with the Gaussian assump- tion over the training example æ1, where the exponent is quadratic function with a constant variance (App. B.4). Since the energy function shares similar L2 form as the CFM loss, the loss landscape in the conformation space, which was originally only influenced by the observed in (3te, as t + 1). now is further shaped by the estimated Φ1. As- suming G is a function that perfectly maps any क to 2 on the data manifold, Eq. 13 will assign high energy (larger reconstruction error) to 'bad' o1 and low energy (smaller reconstruction error) to 'good" 21, However, in principle, G could be any neural network that keeps the dimensionality of the input data. Crucially, following this, during training, the flow map (x) receives the gradient not only from the CFM loss but also from the energy loss of Eq. 13 for the contrastive samples.
3.2. Energy-Band Flow Matching
Energy relaxation, confidence model and the EBMs. Data-driven generated structures often tend to be subop- timal when evaluated against physiochemical energy func- tions. A common solution, known as energy relaxation, is to post-optimize the molecular poses with the energy func- tion, so that the structure after postprocessing corresponds to the local energy minimum (Buttenschoen et al., 2024). In sampling- based molecular docking, a confidence model ranks the samples by outputing their corresponding confi- dence scores. In both cases, a scalar is produced for the molecular structure and can be collectively understood as the energy. To draw connection, we can write down the energy of the generated sample dt as:
Sample from the energy function. With the energy func- tion devised, we can define an energy-based density with the Boltzmann distribution:
Eg(ā1)= Do())).
= fo,(x,,
(12)
p(x) exp( E(x(-
where Dg is the energy architecture and ét is the sample produced by the flow map fe. For structure relaxation, the energy architecture Dg is a specific biophysics-informe en- ergy function, for example, the Amber energy (Duan al., 2003), which involves the bond, angular, Lennard Jones, and Coulomb energy. Under the scenario of molecular docking, Da could be a separate neural network trained by generating samples for every training example and using the training annotations to produce binary labels to distinguish the posi-
(14)
7
where Z is some unknown normalizing constant. Taking the derivative of the log likelihood with respect to x yields:
v log(p(z)) = [().....
(15)
Now, assume that we define gradient flow vector field:
(xx) = - F...).
(16)
4

---

Flow Matching: Idempotent Flow Map
Algorithm 1 Idempotent Flow Map Training
Algorithm 2 Predictor Refiner Sampler
Require: prior distribution Po(z), data distribution P1 (a) while Training do
Require: prior distribution Po. number of integration steps T. and trained function fe
Zg Po(za), T1 P1 (21).
steps t 1 AL + 7
~ U (0, 1), m 10,
4 x1 +
-0
f 0
x ~ (M,, r
20 ~ Po(Fo)
if m S 0.5 then h ~ randint(1, m
2t t 2u
while steps S T 1 do
With corch.no-gree ():
21 x f0,1
1 fo.((z)) 21 lisl = []
21 - f0,((01)
34 x 3 + Al Ba
for = 0.... do
t イー t + At
11 + -- (&) ( detach() Đι sillis.appeedd
steps + steps +
end while
end for
21 +- 24
Jn lest) Eny |,:1 x1|12 0 + Update(O, VOTRR
LR
i.., iterating the flow map infinitely many times yields the same result. Under the assumption that our dataset of equi- librium configurations is contained in M, we would always want to query f((x) at / xc. Ideally, after training a model to learn the idempotent flow map. we would have (x) = 21 for any x. However, imperfections may re- main, in which case (xr = & would land somewhere close to 21, but not exactly there. In this case, an iterative refinement procedure (Stärk al., 2023; Jing et al., 2024), can also be applied during inference. The vector from to c = fxx) can be used as a step direction in an integration scheme, such as Euler's 's method. Heuristically, if the mag- nitude of the vector is large, the integrator will take a large step in that direction, and vice versa, eventually leading to stabilization around the final prediction.
else
LCRM + (x) - x: 112
0 e Updateec VoLCFM)
end if
end while
Interestingly, as the energy in Eq. 13 is lower bounded by O, the convergence of the flow map is further guaranteed as:
(x) E (a E TR V (x) ( = 0} = M, (17)
where M a subset of the equilibrium configurations. This implies that, starting from any x E Rdd the gradient flow will converge to some local minima of the energy function E, where the likelihood of the data is locally maximized.
Furthermore, enforcing the flow map fo be idempotent draws informative connection to the structure refinement regression model (Jumper et al., 2021) which recycles the output for iterative refinement. Hence, in tandem with the CFM loss, we propose our idempotent objective as:
3.3. Idempotent Flow Map
In general, the energy function defined in Eq. 13 renders G to be "neural refiner". To learn the refiner we can simply define the loss for G as:
L (0) E
[ .deacchc))h()) xr 11
(-11
I (=
Lg [[G(fz) - -111
(18)
2 AL=M(BBllloeet))001
(21) where é1 is dynamically sampled from the flow model dur- ing training.
Intuitively, the refiner G is easier to learn than the flow map f. and thus the flow map could have the capacity to also refine the mapped sample. This, interestingly, results in learning idempotent flow map to its prediction:
3.4. Training and Inference
fx) = (f)
(19)
The idempotent objective function enables the network re- fine samples iteratively. Theoretically, the sample 9 could be refined for an infinite number of times. However, exces- sive refinements increase inference time, a key limitation of sampling- based methods. To mitigate this, we perform only one refinement per step. Specifically, the predictor-refiner sampler makes a prediction é and refines it, resulting in two
Hence, the energy landscape is shaped by mapping the off- the- manifold points á to on-manifold points ø1 (LeCun, 2022). The idempotence of the flow map is clear from the Eq. 17:
Soe (oc(So(()) = (x),
(20)
5

---

Flow Matching: Idempotent Flow Map
Table 1. SINGLE F. LIGAND DOCKING. Structure generation (ten samples average for each example) comparison of methods the
PDBBind for pocket-level docking.
SEQUENCE SIMILARITY SPLIT
TIME SPLIT
METHOD
DISTANCE-POCKET
RADIUS-POCKET
SETTANCE-POCKER
RADIUS-POCKET
<<2
MED.
%<2
MED.
%<2
MED.
%<2
MED.
PRODUCT SPACE DIFFUSION
27.2
3.2
16.1
4.0
20.8
3.8
15.2
4.3
HARMONICFL.OW
30.1
3.1
20.5
3.4
42.8
2.5
28.3
3.2
IDFLOW
35.6
2.9
21.0
3.7
44.3
2.4
34.7
3.1
attention layer with the attention weights to be invariant to the structure input. The architecture is kept the same as the baseline to maintain a fair comparison.
Table 2. MULTI-LIGAND DOCKING. Structure generation (ten samples average each test example) comparison of methods on the Binding MOAD.
METHOD
% <
% <
5
MED.
4. Experiments
EIGENFOLD DIFFUSION
39.7
73.5
2.4
HARMONICFLOw
44.4
75.0
2.2
In this section, we investigate several research questions about the proposed framework. Q1: How does IDFlow perform compared with the standard flow matching setup? Q2: Does the training setup also apply to Riemannian flow matching? Q3: Is IDFlow transferable for generating molecules of various sizes (both ligand and protein)? Q4: How does IDFlow perform for sampling the modes of dis- tribution? Q5: How do different setups affect the IDFlow? To answer these questions, we design experiments for two widely applicable tasks, molecular docking (Q1, Q4, Q5) and protein backbone generation (Q2, Q3, Q5), which tar- get predicting the binding pose of the ligand in a protein- ligand complex and generating physically plausible protein structures. Specifically, we build the proposed IDFlow from two recent flow matching- matching-based structure generation models (Stärk et al., 2023; Yim et al., 2023a).
IDFL.W
43.8
83.1
2.1
Number of Function Evaluations (NFEs) per step. To better align training with sampling, we adopt a strategy similar to self-conditioning (Chen et al., 2023), where the training of the sampler and refiner is separated. For 50% of the training time, the network undergoes flow matching training. In the remaining 50%, the network first predicts & with the gradient detached, and then trains for the idempotent objective Eq. 21. Empirically, tuning the maximum number of iterations max helps achieve gradual refinement. In- creasing Kmax beyond provides diminishing returns, with smaller variations observed at Kmax = 2. For the training of 'refiner', ideally, the timestep t should be set to as &1 represents data, but in practice we find it doesn't impact the method's performance. This approach also reduces memory overhead, only K outputs need be stored when looping the network K times. The entire approach only in- troduces one extra hyperparameter K Kmaz and remains to be highly orthogonal to the existing methods. Further details are provided in Algorithms 1 and 2.
4.1. Molecular Docking
Dataset and baselines. To evaluate the structure genera- tion capability of IDFlow for pocket-level docking, we train the model and its ablations on the PDBBind v2020 dataset (Liu et al., 2017) for both the time and the 30% sequence similarity split and the Binding MOAD (Hu et al., 2005) with 30% sequence similarity split, as proposed in (Stärk et al., 2023). We follow the dataset preprocessing steps of the HarmonicFlow for both PDBBind and BindingMOAD (details in E. For both single and multi-ligand docking, we consider HarmonicFlow (Stärk et al., 2023) and the prod- uet space diffusion model over ligand geometry (rotation, translation and torsion angles) (Corso et al., 2023) as the baselines. HarmonicFlow is the standard flow matching setup which samples the harmonic prior for ligand atoms and leams the flow map for predicting the bound structure of ligand. The reported results are averaged over three runs.
3.5. Architecture
For the docking task, the network is parameterized as the refinement tensor field network (Stärk et al., 2023; Thomas et al., 2018), predicting the ligand atom 3D positions 3x The building block for equivariance is the tensor product convolution layer, which constructs the message the ten- product between the node and edge embedding. Besides, each layer is designed to predict the update of the previ- ous layer through the auxiliary loss. For protein backbone generation, we leverage the invariant point attention (IPA)- based structure transformer (Jumper et al., 2021) to predict the protein frames. The IPA differs from the standard self-
Sampling time. While the number of sampling steps sig-
6

---

Flow Matching: Idempotent Flow Map
Table PROTEIN BACKBONE GENERATION RESULTS PDB. Comparison of methods for 200 generated proteins each length [100, 150, 200, 250, 300]. The mean and standard deviation are reported over three runs for GAFL and FrameFlow. The results of FoldFlow2 are from the public weight, and other results are extracted from (Wagner et al., 2024). The number in the parentheses FrameFlow and IDFlow is the NFEs for generating one sample. The time required to generate backbone of length 100 is also reported. Pretrained weights from folding model trained on dataset larger than PDB.
METHOD
DESIGNABILITY (t)
DIVERSITY (1)
NOVELTY (L)
HELIX CONTENT
STRAND CONTENT
TIME [S]
PDB DATASET (300)
0.39
0.23
FRAMEDIFF
0.54
0.45
0.71
0.53
0.20
24.3
FOLDFLOW-SFM
0.69
0.44
0.77
0.91
0.01
24.3
FOLDFLOW-OT
0.82
0.44
0.79
0.88
0.00
24.3
FOLDFLOW2
0.94
0.37
0.69
0.87
0.01
2.8
RFDIFFUSION
0.89
0.37
0.74
0.58
0.24
21.0
GAFL
0.866 + 0.021
0.35 +: 0.01
0.71 + 0.01
0.53 = 0.01
0.24 +: 0.01
8.8
FRAMEFLOW (200)
0.824 0.037
0.35 0.01
0.70 + 0.00
0.57 + 0.01
0.19 0.00
6.6
IDFLOW (100)
0.871 + 0.013
0.35 1 0.00
0.70 1 0.02
0.60 = 0.07
0.16 0.06
3.3
IDFLOW (200)
0.927 + 0.020
0.35 1: 0.01
0.72 + 0.01
0.61 = 0.08
0.16 +: 0.07
6.6
nificantly impacts performance, we maintain consistent computational budget with HarmonicFlow, using 10 steps (20 NFEs), compared to HarmonicFlow's 's 20 steps.
following (Yim et al., 2023b). This results in total 19327 proteins of various lengths for training. The main base- line is FrameFlow (Yim et al., 2023a) which we extend to build IDFlow. FrameFlow the standard Riemnanian flow matching setup which samples from the source distribution to be the composition of uniform distribution on SO(3) and standard Gaussian and learns the flow map to predict the protein structure. The results on SCOPe can be found in Table 7 (App. F). Again, to keep the sampling time the same as the baseline, we run inference with IDFlow using 50 steps (100 NFEs) on SCOPe and 100 steps (200 NFEs) on PDB.
Evaluation metrics. Following (Stärk et al., 2023; Corso al., 2023), we use the fraction the test samples that have root mean squared deviation (RMSD) below 2 or 5 (% 2 and % < 5) and the RMSD median (Med.) for evaluating the docking performance.
Results. We first investigate the method's performance on the single ligand docking and report metrics in Table IDFlow consistently outperforms or on par with both HarmonicFlow and product space diffusion models at the same inference time. Notably, IDFlow achieves 5.5% per- formance increase on RMSD < 2 for sequence similarity split distance pocket docking and 6.4% on RMSD < 2Á for the time split radius pocket docking, demonstrating the improvement over different problem definition and datasets. Table 6 (App. F) shows great improvements in top-40, top- 10, and top-5 accuracy for both RMSD < 1Â and RMSD < 2A, highlighting the effectiveness of the approach for enhanced mode coverage. Table 2 shows the results multi-ligand docking. Also in this setting, IDFlow main- tains the same level of performance with HarmonicFlow on RMSD < 2 and demonstrates 8.1% performance increase for RMSD < 5Å.
Evaluation metries. We adopt the commonly used metrics designability, diversity and novelty as in (Wu et al., 2024; Yim et al., 2023b; Bose et al., 2024; Huguet et al., 2024; Wagner et al., 2024). Designablility can be seen as proxy for biophysical consistency and reports whether an amino acid sequence can be found that folds into the generated protein structure. Diversity and novelty report how struc- turally different the generated structures and to the existing protein structures, respectively. For more detailed definition of metrics we refer App. E.2.
each other,
Results on PDB. Table 3 summarizes the experiments' re- sults with models trained on the PDB dataset. With 100 NFEs (50 steps), IDFlow achieves 5.5% performance in- crease in designability over FrameFlow while being com- petitive at diversity and novelty. With 200 NFEs (100 steps), IDFlow obtains designability with 10.3% increase and even approximates the performance of FoldFlow2, which relies on more advanced architecture. data curation, and is condi- tioned on sequence information, on most metrics. Remark- ably, contrast to FoldFlow2, IDFlow does not tend to collapse to generating only helical protein structures. Fig. visualizes a set of randomly selected designable protein structures generated by IDFlow. Proteins arranged in non-redundant topologies and are comprised by both a-
4.2. Protein Backbone Generation
Datasets and baselines. In this section we evaluate ID- Flow on the task of unconditional protein structure genera- tion. The experiments are first conducted on a small curated dataset SCOPe (Fox al., 2014; Chandonia al., 2022) comprised of 3928 protein structures filtered by lengths between 60 and 128 residues. Next, we evaluate the ID- Flow on the subset of PDB, with maximum protein length 512 and maximum coil content of 50 % filtering process

---

Flow Matching: Idempotent Flow Map
(a) N=300
(b) N=250
(c) N=200
(d) N=150
(e) N=100
Figure 2. Designable protein backbones generated by IDFlow at various length N {100, 150, 200, 250, 300)-
Table Ablation on timesplit radius pocket docking.
IDFlcw with Different Number of Refinements During sampling
% <
MED.
35
Kmaz- (HARMONICFLOW))
28.3
3.2
Kmaz
2.42
3,0 = 0.05
Kmaz
3.17
3.1 0.09
32
helices and 8-strands.
30
4.3. Ablations
28
Number of iterations. We first ablate the number of itera- tions Kmaz training for idempotency the timesplit radius pocket docking in Table 4 for IDFlow. A single iteration improves over baseline, with more iterations resulting in marginal performance gains but less variation. Fig. shows the ablation of number of refinements of sampling at test time. Since more refinements increase the total number of NFEs, we keep the overall NFEs around 20. The perfor- mance peaks at refinement. We attribute the perfor- mance drop with more refinement steps to the discretization errors of the ODE.
Figure 3. Ablation the number of refinements & at test time inference.
6. Conclusion
We present an energy-based formulation for flow matching, enhanced training framework for flow matching bined with the EBMs for 3D molecular structure generation. We provide specific instance of the proposed framework, IDFlow, which considers the reconstruction error as the en- ergy function, shaping the loss landscape with contrastive samples produced by the flow model. builds the flow map predict and refine the sampling trajectory iteratively toward the target structure. We validate the improvements of IDFlow over existing flow matching methods in both Riemannian and Euclidean spaces as an alternative to the standard flow matching training and sampling paradigm, demonstrating its effectiveness over task-associated flow matching and diffusion model baselines in molecular dock- ing and protein backbone generation.
5. Related Work
Generative modeling for molecular structure generation. Generative modeling, mostly diffusion and flow matching models have been applied to generate molecular structures of ligands (Stärk et al., 2023) and proteins (Yim et al., 2023b;a). Molecular conformer generation (Xu et al., 2022a: Jing et al., 2022; Wang et al., 2024; Hassan et al., 2024) maps the molecular graph to the 3D position of atoms by learning the posing distribution. Sampling-base
docking method (Corso et al., 2023; 2024; Qiao et al., 2024; Lu et al., 2024) models the bound structure of ligand given the recep- tor structure. Several generative models for protein structure (Watson et al., 2023; Wu et al., 2024; Campbell et al., 2024; Bose et al., 2024; Huguet et al., 2024; Yim et al., 2023b;a; Wagner et al., 2024) have been shown to sample physically plausible and functional protein structures de novo.
One limitation of the work is the increased training cost, multiple forward passes are required. Besides, extra refine- ment incurs a larger discretization error for the ODE sam- pling. Future work can combine the EBMs with biophysics- informed energy to generate chemically plausible structures and use EBMs as scoring functions for realistic docking.
8

---

Flow Matching: Idempotent Flow Map
Impact Statement
Enabling multimodal flows with applications to protein co-design. arXiv preprint arXiv:2402.04499, 2024.
This paper presents work whose goal to advance the field of Machine Learning. There are many potential societal consequences of our work, none which feel must be specifically highlighted here.
Chandonia, J.-M., Guan, L., Lin, S., Yu, C., Fox, N. K., and Brenner, S. E. Scope: improvements to the structural classification of proteins extennded database to facilitate variant interpretation and machine learning. Nucleic acids research, 50(D1): DDSDD:DD553--59 2022.
Acknowledgements
This work was partially supported by the Knut and Alice Wallenherg Foundation. The computations were enabled by resources provided by the National Academic Infrastructure for Supercomputing in Sweden (NATSS), partially funded by the Swedish Research Council through grant agreement no. 2022-06725. Experiments were performed using the su- percomputing resource Berzelius provided by the National Supercomputer Centre at Linköping University and the Knut and Alice Wallenberg foundation.
Chen, R. T., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Chen, R. Q. and Lipman, Y. Flow matching on general geometries. In The Twelfth International Conference Learning Representations, 2024. URL https:// openreview net wwrrmmormmemmemm i dd=g7ohDITITL
Chen, T., ZHANG, R., and Hinton, G. Analog bits: Gen- erating discrete data using diffusion models with self- conditioning. In The Eleventh International Confer- ence on Learning Representations, 2023. URL https: / /pperreview net forum? d=33 jRRjxFFF..
References
Albergo. M. S. and Vanden-Eijndesn E. Building normaliz- ing flows with stochastic interpolants. In The Eleventh International Conference on Learning Representations, net/ forum?
Cheng, A., Lo, A., Lee, K. L. K., Miret, S., and Aspuru- Guzik, A. Stiefel flow matching for noment-constraing structure elucidation. arXiv preprint arXiv:2412. 2024.
2023. URL btctps://opereett id=1i7geBbCRIt.
Back, M. DiMaio, F., Anishchenko, I., Dauparas, J., Ovchinnikov, S., Lee, G. R., Wang, J., Cong, Q., Kinch, L. N., Schaeffer, R. D., al. Accurate prediction of pro- tein structures and interactions using a three-track neural network. Science, 373(6557):871-887, 2021.
Corso, G., Stärk, H., Jing, B., Barzilay, R., and Jaakkola, T. S. Diffdock: Diffusion steps, twists, and turns for molecular docking. In The Eleventh International Confer- ence on Learning Representations, 2023. URL https: / //openreveew. net/ www.moidddddir
Ben-Hamu, H., Puny, O., Gat, I., Karrer, B., Singer, U., and Lipman, Y. D-flow: Differentiating through flows for controlled generation. arXiv preprint 4rX1v.2402.14017, 2024.
mBbs.
Corso, G., Deng, A., Polizzi, N., Barzilay, R., and Jaakkola, T. S. Deep confident steps to new pockets: Strate- gies for docking generalization. In The Twelfth In- ternational Conference Learning Representations, 2024. URL htt.tt/t:/// / net/ forrm?? id=UfBIXpTK10.
Bose, J., Akhound-Sadegh. T., Huguet, G., FATRAS, K., Dector-Brooks, J., Liu, C.-H., Nica, A. C., Korablyov, M., Bronstein, M. M., and Tong, A. SE(3)-stochastic flow matching for protein backbone generation. In The Twelfth International Conference on Learning Representations, 2024. URL https: //tp:://oppaaee http://wppeeeeeeeeeeeeeee. n forum? Id=kJFIH23hXb.
Corso, G., Somnath, V. R., Getz, N., Barzilay, R., Jaakkola, T., and Krause, A. Composing unbalanced flows for flexible docking and relaxation. In The Thirteenth In- lernational Conference on Learning Representations, 2025. URL https: wwwp//eeeeeeeeeeeoooeeoeee net nwt//orunrt forum? i dd
Braun, M., Jaquier, N., Rozo, L., and Asfour, Riemannian flow matching policy for robot motion learning. arXiv preprint 0oo:24003.10622 2024.
Dauparas, J., Anishchenko, I., Bennett, N., Bai, H., Ragotte, R. J., Milles, L. F., Wicky, B. I. M., Courbet, A., de Haas, R. J., Bethel, N., Leung, P. Y., Huddy, T. F., Pellock, S., Tischer, D., Chan, F., Koepnick, B., Nguyen, H., Kang, A., Sankaran, B., Bera, A. K., King, N. P., and Baker,
Buttenschoen, M., Morris, G. M., and Deane, C. M. Pose- busters: Ai-based docking methods fail to generate physi- cally valid poses or generalise to novel sequences. Chem- ical Science, (5(9)3330-31133 2024.
Campbell, A., Yim, J., Barzilay, R., Rainforth. T., and Jaakkola, T. Generative flows on discrete state-spaces:
D. Robust deep learning- based protein sequence design using ProteinMPNN, Science, 378(6615):49-66 2022.

---

Flow Matching: Idempotent Flow Map
Duan, Y., Wu, C., Chowdhury, S., Lee, M. C., Xiong, G., Zhang, W., Yang, R., Cieplak, P., Luo, R., Lee, T., al. A point-charge force field for molecular mechanics sim- ulations of proteins based on condensed sondensed-phase quantum mechanical calculations. Journal of computational chem- istry, (191))19999200000 2003.
Jing, B., Berger, B., and Jaakkola, T. Alphafold meets flow matching for generating protein ensembles. arXiv preprint (...i::2402.04844 2024.
Jumper, J., Evans, R., Pritzel, A., Green, T., Figurnov, M., Ronneberger, O., Tunyasuvunakool, K., Bates, R., Zídek, A., Potapenko, A., et al. Highly accurate protein structure prediction with alphafold. nature, 596(7873):583-58) 2021.
Engh, R. and Huber, R. Structure quality and target parame- ters. 2012.
Fox, N. K., Brenner, S. E.,, and Chandonia, J.-M. Scope: Structural classification of -xtriins-eeteeeddd integrat- ing scop and astral data and classification of new struc- tures. Nucleic acids research, 42(D1):D304-D33 2014.
Kessel, A. and Ben-Tal, N. Introduction to proteins: struc- ture, function, and motion. Chapman and Hall/CRC, 2018.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.698)) 2014.
Geiger, M. and Smidt, T. e3nn: Euclidean neural networks. arXiv preprint wrXiv:2207.09453, 2022.
Kingma, D. P. and Welling, M. Auto-encoding variational bayes, 2013.
Goodfellow, I. J., Pouget-Abadie, J., Mirza, M., Xu, B., Warde- Farley, D., Ozair, S., Courville, A., and Bengio, Y. Generative adversarial networks, 2014. URL https: //arxiv.org org/ abs 1406.266222 2661
LeCun, Y. From machine learning to autonomous
intelligence: Lecture 2, 2022. URL https:
// //Leshouches2022221 http:ees2222....hh... githu...t.i... iot..ibe..////.sss......
Hassan, M., Shenoy, N., Lee, J., Stark, H., Thaler, S., and Beaini, D. Et-flow: Equivariant flow-matching for molecular conformer generation. arXiv preprint arXiv:2410.22388. 2024.
lecun (elun-2020 cessed: 2024-10-02.
02, pdf. Ac-
LeCun, Y., Chopra, S., Hadsell, R., Ranzato, M., Huang, F., et al. tutorial on energy-based learning. Predicting structured data, 1(0), 2006.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion proba- bilistic models. Advances in neural information process- ing systems, 33:684006851, 2020.
Lin, Z., Akin, H., Rao, R., Hie, B., Zhu, Z., Lu, W., Smetanin, N., Verkuil, R., Kabeli, O., Shmueli, Y., et al. Evolutionary-scale prediction of atomic-level protein structure with a language model. Science, 379(6637): 1123-1130, 2023.
Hoogeboom, E., Satorras, V. G., Vignac, C., and Welling, M. Equivariant diffusion for molecule generation in 3d. In International conference on machine learning, pp. 8867- 8887. PMLR, 2022.
Lipman, Y., Chen, R. T. Q., Ben-Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling. In The Eleventh International Conference on Learning Represen- tations, 2023, URL https: / openreview net/ forum? -d=-@@MRDDP.........
Hu, I.., Benson, M. L.,, Smith, R. D., Lerner, M. G., and Carlson, H. A. Binding moad (mother of all databases). Proteins: Structure, Function, and Bioinformaties, 60(3): 333-340. 2005.
Huang, Y., Zhang, O., Wu, L., Tan, C., Lin, H., Gao, Z., Li, S., Li, S., et al. Re-dock: Towards flexible and realistic molecular docking with diffusion bridge. arXiv preprint arXiv:2402.11459 2024.
Lipman, Y., Havasi, M., Holderrieth, P., Shaul, N., Le, M., Karrer, B., Chen, R. T., Lopez-Paz, D.. Ben-Hamme H., and Gat, I. Flow matching guide and code. arXiv preprint 2024.
Huguet, G.. Vuckovic, J., Fatras, K., Thibodcau-Laufer E., Lemos, P., Islam, R., Liu, C.-H., Rector-Brooks, J., Akhound-Sadegh, T., Bronstein, M., et al. Sequence- augmented se (3)-flow matching for conditional protein backhone generation. Advances neural information processing systems, 2024.
Liu, X., Gong, C., and Liu, Q. Flow straight and fast: Learning generate and transfer data with rectified flow. arXiv preprint 2.X0002009033003, 2022.
Liu, 7.,, Su, M., Han, I.., Liu, J., Yang. Q., Li, Y., and Wang, Forging the basis for developing protein- inter- action scoring functions. Accounts of chemical research, 50(2):302-309, 2017.
Jing, B., Corso, G., Chang, J., Barzilay, R., and Jaakkola, T. Torsional diffusion for molecular conformer generation. Advances in Neural Information Processing Systems, 35: 2022.
Lu, W., Wu, Q., Zhang, J., Rao, J., Li, C., and Zheng, S. TANKBind: Trigonometry-awan neural networks
10

---

Flow Matching: Idempotent Flow Map
for drug-protein binding structure prediction. In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K. (eds.), Advances in Neural Information Processing Systems, 2022. URL pettpa://open.eeeee wd=MSHDFwGYwwt.
Song, Y., Dhariwal, P., Chen, M., and Sutskever, I. Consis- tency models. In International Conference on Machine Learning, pp. 3221 PMLR, 2023.
net/ forum?
Sprague, C. I., Elofsson, A., and Azizpour, H. Incorporat- ing stability into flow matching. In ICML 2024 Work- shop on Structured Probabilistic Inference Generative Modeling, 2024. URL https: https:t//ppeeneeeeeeeee openrevien .nt// forum?dd=FRRiee jev.
Lu, W., Zhang, J., Huang, W., Zhang, Z., Jia, X.,, Wang, Z., Shi, I.., Li, C., Wolynes, P. G., and Zheng, S. Dynam- icbind: predicting ligand-specific protein-ligand complex structure with a deep equivariant generative model. Na- ture Communications, 15(1):1071, 2024.
Stärk, H., Ganea, O., Pattanaik, L., Barzilay, R., and Jaakkola, T. Equibind: Geometric deep learning for drug binding structure prediction. In International conference on machine learning, pp. 20503-20521. PMLR, 2022.
Méndez-Lucio, O., Ahmad, M., del Sio-Chanona, E. A., and Wegner, J. K. geometric deep learning approach to predict binding conformations of bioactive molecules. Nature Machine Intelligence, 3(12):1033 1009, 2021.
Stärk, H., Jing, B., Barzilay, R., and Jaakkola, T. Harmonic celf-conditioned flow matching for multi-ligand docking and binding site design. arXiv preprint arXiv:2310.0576,, 2023.
Pei, Q. Gao, K., Wu, L., Zhu, J., Xia, Y., Xie, S., Qin, T., He, K., Liu, T.-Y., and Yan, R. Fabind: Fast and accurate protein-ligand binding. Advances in Neural Information Processing Systems, 36, 2024.
Thomas, N., Smidt, T., Kearnes, S., Yang, L., Li, L., Kohlhoff, K., and Riley, P. Tensor field networks: Rotation-and and translation-equuivrian neural networks for 3d point clouds. arXiv preprint arXiv: UrXiv:/802.082111 2018.
Plainer, M., Toth, M., Dobers, S., Stärk, H., Corso, G., Mar- quet, C., and Barzilay, R. BiftDock-Pocker Diffusion for pocket-level docking with sidechain flexibility. 2023.
Tong, A., FATRAS, K., Malkin, N., Huguet, G., Zhang, Y., Rector- Beotor-Brooks, J., Wolf, G., and Bengio. Y. Improving and generalizing flow-based generative models with mini- batch optimal transport. Transactions on Machine Learn- ing Research, 2024. ISSN 2835-8856. URL https: // openreview net / forummorum//mmmm Soorumidd-CD9Snc73300 Ex- pert Certification.
Pooladian, A.-A., Ben-Hamu, H., Domingo-Ennich, C., Amos, B., Lipman, Y., and Chen, R. T. Multisample flow matching: Straightening flows with minibatch cou- plings. arXiv preprint arXiv:2304. 14772, 2023.
Qiao, Z., Nie, W., Vahdat, A., Miller III T. F., and Anandku- mar, A. State-specific protein-ligand complex structure prediction with a multiscale deep generative model. Na- fure Machine Intelligence, 6(2):195-208, 2024.
Trippe, B. L., Yim, J., Tischer, D., Baker, D., Broderick, T., Barzilay, R., and Jaakkola, T. Diffusion probabilistic mod- eling protein backbones 3d for the motif-scaffolding problem. arXiv preprint arXiv:2206.044 041 19,19, 2022.
Senior, A. W., Evans, R., Jumper. J., Kirkpatrick, J., Sifre, L., Green, T., Qin, C., Žídek, A., Nelson, A. W., Bridg- land, A., et al. Improved protein structure prediction using potentials from deep learning. Nature, 577(7792): 706-710, 2020.
Van Kempen, M., Kim, S. S., Tumescheit, C., Mirdita, M., Lee, J., Gilchrist, C. L., Söding, J., and Steinegger, M. Fast and accurate protein structure search with foldseek. Nature biotechnology, 42(2):243-246, 2024.
Shocher, A., Dravid, A. V., Gandelsman, Y., Mosseri, I., Rubinstein, M., and Efros, A. A. Idempotent gener- ative network. In The Twelfth International Confer- ence on Learning Representations, 2024. URL https: /
Wagner, S., Scute, L., Viliuga, V., Wolf, N., Gräter, F., and Stuehmer, J. Generating highly designable proteins with geometric algebra flow matching. In The Thirt--eighth Annual Conference on Neural Information Processing Systems, 2024. URL https: openreview net/ Eorum? Saauumdd--nAnnnxxyyy...
Song, Y. and Ermon, S. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
Wang, Y., Elhag, A. A., Jaitly, N., Susskind, J. M., and Bautista, M. A. Swallowing the bitter pill: Simplified scalable conformer generation. In Forty-first Interna- tional Conference on Machine Learning, 2024.
Song. Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. Score-based generative mod- cling through stochastic differential equations. In In- ternational Conference on Learning Representations, 2021. URL httph::////// httttp:://penrevvieee wtwpeeviwwienneeeeeeeeeeeeeeeeee net/forunt 1==
Watson, J. L., Juergens, D., Bennett, N. R., Trippe, B. L., Yim, J., Eisenach, H. E., Ahern, W. Borst, A. J., Ragotte,

---

Flow Matching: Idempotent Flow Map
R. J., Milles, L. F., al. De novo design of protein struc- ture and function with rfdiffusion. Nature, 620(7976): 1089-1100, 2023.
Wu, E., Yang, K. K., van den Berg, R., Alamdari, S., Zou, Y., Lu, A. X., and Amini, A. P. Protein structure generation via folding diffusion. Nature communications, 15(1):1059, 2024.
Xu, M., Yu, I.., Song, Y., Shi, C.,, Ermon, S., and Tang, Geodiff: geometric diffusion model for molecular conformation generation. In International Conference Learning Representattonss 2022a. URL https:// openreview net/ http://orrnnnddPe/E//EeEe 1d=22cxXMM2VVC..
Xu, Y., Liu, Z., Tegmark, M., and Jaakkola, T. Poisson flow generative models. Advances in Neural Information Processing Systems, 35:16782- 16795, 2022b.
Xu, Y., Liu, Z., Tian, Y., Tong, S., Tegmark, M., and Jaakkola, Pfgm++: Unlocking the potential of physics- inspired generative models. International Conference Machine Learning, pp.
PMLR, 2023.
Yim, J., Campbell, A., Foong, A. Y., Gastegger, M., Jiménez-Luma, J., Lewis, S., Satorras, V. G., Veeling, B. S., Barzilay, R., Jaakkola, T., et al. Fast protein back- bone generation with se (3) flow matching. arXiv preprint arXiv:2310.05297, 2023a.
Yim, J., Trippe, B. L., De Bortoli, V., Mathieu, E., Doucet, A., Barzilay. R., and Jaakkola, Se (3) diffusion model with application to protein backbone generation. arXiv 2023b.
preprint ar arXiv:2302.00221
Zhang, Y., Cai, H., Shi, C., and Tang. J. E3bind: An end- to-end equivariant network for protein-ligand docking. In The Eleventh International Conference Learning Representations, 2023, http::::: ppppreveen.eeeeeeeeeeeeee net/
Zhao, J., Mathieu, M., and LeCun, Y. Energy-based gen- erative adversarial networks. In International Confer- on Learning Representations, 2017. URL https: // //openreview..o net / forum? i /tt//ouummm/teee drr 9pmcee.
12

---

Flow Matching: Idempotent Flow Map
Appendix Overview
In the Appendix, we provide the details:
In section A, we further discuss the related work, encompassing the flow matching, sampling- based molecular docking and protein backhone generation.
In section B, More discussions including the idempotency and stability, additional energy function, the general format of vector field with data parameterization flow matching and the associated energy reconstruction error are provided.
In section C, we provide a short overview of Riemannian flow matching on modeling the protein backbone.
In section D, we summarize the architecture of IDFlow for docking and protein backbone generation.
In section F, we provide the additional results and analysis.
In section G, we visualize the generated samples.
A. More Related Work
Flow matching. Flow matching (Lipman et al., 2023; 2024; Albergo & Vanden-Eijndee 2023; Liu et al., 2022) is simulation-free training paradigm of continuous normalizing flows (Chen et al., 2018) generalizing the diffusion models into arbitrary priors and flexible construction of probability paths. It builds the transformation from source to target distribution via an ODE and supports the source target sample coupling through optimal transport within batch samples (Tong et al., 2024; Pooladian et al., 2023). This approach can also be extended to manifold where the trajectory simulation is avoided on simple geometry (Chen & Lipman, 2024).
Sampling-basse based molecular docking, The seminal work DIFFDOCK (Corso et al., 2023) defines the diffusion model on the ligand geometry, speeifically, the product space of global rotation, translation and torsion angles of ligand for efficiently modeling the docking conformation of ligand given the protein. (Corso al., 2024) further investigates the method's scalability and uses confidence bootstrapping for enhanced training. DYNAMICBIND (Lu et al., 2024) extends the existing approach to the flexibility of protein backbone and side chains, modeling the distribution shift from apo to holo structure. NEURALPLEXER (Qiao et al., 2024) is co-folding approach combining the diffusion models and structure prediction model for sampling the protein-ligand complex. Recently, to align with the realistic docking scenario, the diffusion- docking is constrained to within the pocket with the side chain angles flexibility (Plainer et al., 2023; Huang et al., 2024), HARMONICFLOw (Stärk et al., 2023) is the first flow matching method applied to the pocket-level docking for generating 3D structure by directly modeling the Cartesian coordinates of ligand. FLEXDOCK (Corso et al., 2025) proposes an unbalanced flow matching framework ease the learning of large conformational changes of protein and chain the flow for generating the energetically valid poses.
Protein backbone generation. Generating protein backbone has risen be promising with recent diffusion (Wu et al., 2024; Yim et al., 2023b; Watson et al., 2023) and flow matching-baat -base methods (Yim et al., 2023a; Bose et al., 2024; Huguet et al., 2024; Wagner et al., 2024). RFDIIFUSION (Watson et al., 2023) achieved significant success for biologically validated structures and functions in experiments, FOLDFLOW (Bose et al., 2024) and FRAMEFLOW (Yim et al., 2023a) adopted the protein frame represention for generation with SE(3) flow matching. FOLDFLOW2 (Huguet et al., 2024) extends the method with sequence input, architecture improvement and the Reinforced Finetuning on secondary structure. GAFL (Wagner et al., 2024) improves upon the FRAMEFLOW with Clifford frame attention rather than the IPA.
B. More Discussions
B.1. Idempotency and Stability
Idempotency and stability are two fundamental notions over many disciplines, for example, formation control, dynamical systems and molecular dynamics. has been recently explored for sampling from stable distribution (Sprague et al., 2024) or being used as the building block for generative models (Shocher et al., 2024). potentially has a broader impact in many scientific disciplines, not limited to generating molecules, but also structure elucidation (Cheng et al., 2024), etc.
13

---

Flow Matching: Idempotent Flow Map
Idempotency helps learn a robust function where additional function evaluations do not drastically alter the output. This relates stability in terms of the loss used during training such that the iteratively generated output stays the data manifold, desirable property for any generative model. Moreover, the concept of stability could be further expanded into any energy function. Combined with flow models, existing work (Xu et al., 2022b; 2023) proposed to learn a stable vector field to be a Poisson field implicitly governed by the potential energy induced by the training sample over the augmented space of the data. Consistency models (Song et al., 2023) achieves fast sampling with just one or two steps of denoising by learning a consistent function across multiple noise levels. The idempotency is only guaranteed at = 0 by design as the boundary conditions.
B.2. Other Energy Function
IDFlow proposes to use the same energy function as the CFM loss in the context of flow matching which resulting in a relatively smooth minimum in the loss landscape. In general, any type of energy function can be considered.
Distance potential. The relative distances among atoms deliver important information for structure prediction (Senior et al., 2020). Assuming certain distribution the distance can help derive a statistical approximation of the energy function (Méndez-Lucio et al., 2021). Specifically, the intramolecular potential enforces constraints on the ligand atoms pairs to be within certain ranges of distances, depending on the bond type connecting the atoms. For the protein-ligand complex, the intermolecular potential sets up the constraint of the distance between the ligand atoms and protein residues within the binding pocket. Hence, we can enforce the preservation of the protein-ligand distance and the bond distance of the ligand as an energy function:
E = Σ
||das dr +
Σ
||de, - dr.5ll
(22)
(boo)
E protein ligand atom pairs
Machine learning force fields. MLFF (e.g, openmm or fair-chem) can be applied to sample relaxation so that the optimized samples is energetically favorable. However, sample relaxation may lead to larger RMSD error as the projected gradient vx, (x1) (I1 refers the data) may leave the data manifold (Ben-Hamu al., 2024). Within the proposed framework, we apply the MLFF at each denoising step:
21 = fo,(())
2 Va,
è = FF( or = 81
E(t1)
+
2
E1
T+ = 34 +
1-t
Throughout the iterative procedure of sample prediction and relaxation at each step, intuitively, low energy valid sample can be generated. However, this is prohibitive in practice due to the huge increase in sampling time of running the MLFF, and thus one may only use the MLFF when the sample path is close to the data or use a neural refiner to approximate the low energy sample, such as the structure relaxation flow trained with the flat bottom potential (Corso et al., 2025). Another option is sample from the EBM through the Langevin MCMC, which intends to sample from the true distribution governed by the energy function Za (xxp(-E((t))), The c is the step size and is the standard Gaussian sample.
B.3. Vector Field of the Data Parameterization
With the data parameterization flow matching, the exact format of the vector field (x(x) could be diverse as long as the predicted data á fo. (x) could be reached following this vector field. As the trajectory sample 2, is often constructed as linear interpolation, the conditional vector field that generates a single sample at time / is linear in the following general format:
" (<(x/11, vO) A(f((z x1)
(23)
14

---

Flow Matching: Idempotent Flow Map
where A(L) is time-dependent rescaling coefficient. As the marginal vector field uy (3) is the convex combination of the conditional vector fields (Eq. from (Lipman al., 2023)), we can directly plug into the computation:
Uy (x) = /f (2 2o)
Pr Pr 1 2o (p(oo) dæidæo
Pa(x)
= /f A(t) x1)
p1(0o 2o [0)g( )p(()) dzydzo
(((2)
A(1) // (a--m))((0)
(p(p))((1)) 20 d
= A(t) / (x 1=))))
= (1)p)aa)))))) [x - 21|0 = x]]
= A(t) (z E (xn
x1)
where in the third equality we use the Bayes rule assuming ((z1) and p(zo) are independent. This yields general format of the marginal vector field supposing the conditional vector field is linear. In practice, the flow map fo,((z) is trained to predict the data, implicitly learning the expected value of x given x. Hence, the form of the marginal vector field with data parameterization is:
ux(x) = A(t)(x - Ber-mane = x]]
A(I)(z fo,((1)
With A(I) T- the marginal vector field trained by the conditional optimal transport is recovered as:
We (x) (0.
(24)
1-
Interestingly, with the learned flow map the expectation over (z)((z))z), the flow map's idempotency is clear as the expectation is an idempotent operator:
fo,x(fo,(()) = (n = x]] = Exy (en x].
(25)
B.4. Conditional Negative likelihood Energy
Conditional flow matching assumes a smooth delta function for the training example a ~ N(tle,,of
p(z|(1) ~ exp (- 20 (2 xn) (x (x-
(26)
where 21 is a single sample from the training set. However, the true x1 unknown during sampling, and Φ = fa,( is estimated to refine the trajectory the sampling path. We can approximately assume the d distribution follows a similar format:
p (m1) exp (~2d²
20 (ā1
2 21)
(27)
Ideally, each sampling step intends to find estimate of âi such that it lies in the high-density region and approximates be stable around certain area of the energy function. With the exact energy format the negative likeeiid of Eq. 27, the energy of ét could be computed as:
1 (01 20p(01 TT)T(a z1)
E(t1) llo (((0)/z))
(28)
20?
Furthermore, considering the energy minimum x' of Eq. 28 equivalent the gradient of the energy function equal zero, and we could obtain the z' by neural network G(ă)) which could be seen neural refiner that projects the i to its associate point on the data manifold. The gradient of the likelihood of Φ could be written as:
20 (x' x1) (at - 21) = Vlog x1) = of
G( z'
VE(1)) = V
(VG())))(((t)) =1) =0
(29)
15

---

Flow Matching: Idempotent Flow Map
Observing Eq. 29, the term could be optimized if the neural refiner G approximates the 21, which aligns with the proposed idempotent objective 21:
LR [G( =111².
(30)
C. An Overview of Riemannian Flow Matching
C.1. Riemannian Manifold
Considering a smooth manifold M that is locally Euclidean (i.e., each point has a neighborhood that can be smoothly mapped to an open subset of R)") and differentiable, each point of the manifold x .M has a neighborhood that resembles an open subset of TR.. The tangent plane attached each point x on the manifold is called the langent space TzMM. The collection of all the tangent spaces T..M for each point .M forms the tangent bundle of the manifold, denoted by T T.M.
The Riemannian metric g is a smoothly varying inner product on the tangent spaces of the manifold. Formally, at each point t E M, the tangent space T2M is equipped with a positive-definit inner product (s varying across the manifold. This is typically denoted as Oz for measuring the distance on the manifold. For instance, the metric on SO(3) typically requires a bilinear function ( IR$ x R + R to be symmetric positive definite, Hence, we could define positive definite quadratic form for the metric:
1 2
(ti 12)so(3))
- 12)
(31)
where t and t2 are elements of the Lie algebra of SO(3), which the tangent space at the identity element of SO(3) and consists of skew-symmetric matrices.
Geodesic generalizes the concept of straight line in Euclidean space Riemannian manifold, where it locally minimizes the distance between two points on the manifold. It preserves the Riemannian metric and allows for parallel transport of vectors along curves.
The exponential map (x) TTM + M takes a tangent vector e on T&M and maps it to another point y on the manifold along the geodesic. The inverse of the exponential map is called the logarithm map, denoted as log, (v) M - T,M. The logarithm map takes a point y E .M and returns the tangent vector at x that connects x and y through the geodesic distance.
C.2. Riemannian Flow Matching on Protein Geometry
As stated in 3.1, the Riemannian flow matching on protein geometry is defined on the SE(3) group. This backbone parameterization, proposed by AlphaFold2 (Jumper al., 2021), denotes a frame representation SE(3) for each protein residue. With N amino acids, the protein backbone lies on the manifold SE(3) V Each frame x = (r, s) = SE(3) consisting of a rotation r and translation s applies the rigid transformation to the reference frame. The reference frame has four backbone atoms with idealized coordinates centered at the C atom:
N = (-0.525, 1.363, 0.0)
C = (0.0, 0.0, 0.0)
C* = (1.526, 0.0)
O+ (0.627, 1.062,
The idealized coordinates are the experimental measurement of bond angles and lengths (Engh & Huber, 2012).
SF(3) Riemannian flow matching requires to devise a metric on SF(3). As the SE(3) is a semi-direct group, one suitable metric can decompose the computation into SO(3) and R²: (f1, t2)3)(3) = (ri, t2) + t2)s.. This results in the geodesic of SE(3) becomes the same as those on the product space of SO(3) and R Moreover, as the Lie algebra of S3(3) is skew-symmetric matrix, its matrix exponential has a closed-form solution, referred to as the Rodrigues formula. Therefore, its exponential and logarithms map could be computed efficiently without simulating the trajectory over the manifold. In particular, assuming is rotation vector, known the axis-angle representation, and ń is its uniquely identified element of the Lie algebra, its matrix exponential and logarithms can be expressed as (Bose et al., 2024):
Matrix Exponential:
R exp(u)) (1) = cos (w) I + sin(w)ew + (1 - cos(w))ene
(32)
Matrix Logarithms: à log( R) = 2sin(w) (R R
(33)
16

---

Flow Matching: Idempotent Flow Map
where w || Ew
T[ represents the angle and axis of the rotation, and R is the SO(3) group element on the manifold.
From Eq. 33, the angle w could be derived from the trace of the rotation matrix R: cos(w) Rr( For the translation
R.3 the computation trivial following the Eucledian formulation.
With the matrix exponential and logarithm hand, we can construct the linear interpolant on SE(3) the Riemannian conditional optimal transport:
Translation: s: = t.s + (1 t0
(34)
Rotation: r. exPro (tlogr, (r1))
(35)
where x (1, is the frame. Then, the flow matching objective extended to the SF(3)* is:
N
LSF(3) (0) = 1-
Σ
(Ivo.e(s) - &ellie + [Ivo,(()) - -
(36)
zoope (ar ma zom=p(2r)
2100(21)
where samples from the uniform distribution SO(3) and Gaussian on R and samples from the training set.
ti parameterization.
The objective can also be reparametrized as the flow map directly predicting æ1 = (s1,r)):
LSB(3) E t~1((0,1) prope (2 sompliz)) a)))(())
Σ (I.o,,(s) s (Ifo,((r) rı l1300s))
(37)
D. Architecture Details
The architecture used for the IDflow is kept to the same as the main baseline HARMONICFLOW and FRAMEFLOW, including node and edge features initialization. We summarize the architecture details here.
D.1. Refinement Tensor Field Networks
Tensor Field Networks (TFN) (Thomas et al., 2018) incorporates the equivariance by parameterizing the node feature update as the path weight of the tensor product & between scalar node features h2 and spherical harmonies of the edge vector. Denoting the invariant feature of the node as h; and coordinate as I;, the edge vector connecting the node and Pij Tz Tj. The node update built as:
Vi = V(eiss ,195 maj = (fiss) Bayy h4
h/t11 h BN
Σ mij
[N:
JEN,
We first compute the path weights Vhy the coefficients of the spherical harmonics, where eig is the edge embedding between node and node j, the superscript l in hy represents node feature of the 1tth layer, and T is the linear layer defined separately for four different types of edges (receptor to receptor, ligand to ligand, ligand to receptor, receptor to ligand). Then the message may is constructed by the tensor product, in which Y (fiz) is the evaluation of the spherical harmonics at Pig. All the aggregated messages are averaged over the number of neighbours and then pass to the equivariant batch normalization.
Each layer first updates the node feature, and then updates the coordinate via the 0(3) equivariant linear layer Φ:
++1 - (h*+1)
(38)
The learning paradigm involves feature and coordinate refinement, where each layer of the network predicts updates for the previous layer. Additionally, the coordinate output at each layer is "deeply supervised" by the training annotations through the auxiliary loss function. The implementation is based on the e3nn library (Geiger & Smidt, 2022).
17

---

Flow Matching:
Idempotent Flow Map
D.2. IPA-Based Structure Transformer
The invariant point attention mechanism for equivarance and backbone parameterization are initially proposed by (Jumper et al., 2021). The entire procedure involves the network prediction and the backbone update. More details can be found in (Yim et al., 2023b) of Appendix We list the algorithms here:
Node update. Each layer consists of an IPA embedding layer, a transformer layer and linear layer for feature transition.
hipa IPA 22, Tt)
h LayerNorm Ihipa + he)
hrans Transformer(hin)
hout Linear(huans) + he
he+1 Mh
Edge update. The edge update is computed by first projecting the updated node feature to half the dimension hdown- The projected features is then concatenated with the edge embedding feeding into the linear layer:
haown
Linear(he+a)
Z
cat(h
down: e
ZE+1
LayerNorm(Linear(zn
where hoown and bown features expanded into separate dimensions.
Backbone update. The network output is structured to be translation update Supdate and the quaternion representation (1, bi, C. d:) then being converted into a rotation matrix.
bi, C., d; st update Linner(hh)) (he)
a bi, C; (1, bi, Ci d:)/ by cî df
(a,)2 + (b,)² 3- (c,)² - (d.)²
2b,c; - 2a,d;
2b,d; + 2a,C;
nupdaate
26,4 + 24,d,
(a:)²
(b,)² + (c, 1 (d.)²
2c,d; 2a,
2b,d. - 2a,0,
2b,d, - 2a,c;
(a.) (b,)² - (c,) 2 + (d.)²
update
(Ruddate
update
2
S
update
CE+1 = 3y. xf
where is the index of protein residue.
E. Experimental Details
In this section, we provide additional details about the datasets, experimental setup and evaluation metries. The source code available at http:///gthubb.comm/aaaaaaette/e
E.1. Molecular Docking
E. DATASETS
The PDBBind v2020(Liu al., 2017) with total of 19k complexes timesplit a commonly used benchmark for molecular docking (Stärk et al., 2022: Lu al., 2022: Corso et al., 2023; Zhang et al., 2023; Pei et al., 2024; Corso al., 2024). The time split proposed by (Stärk et al., 2022) consists of 17k complexes before 2019 for training and validation and 363 complexes after 2019 for testing without the seen ligand in the training set. The 30% sequence similarity split is constructed from the same dataset but with the constraint of the chain-wise similarity less than 30%, which is considered as a more difficult split for the timesplit.
BindingMOAD (Hu et al., 2005) is another curated dataset from the PDB, with a different preprocessing pipeline from the PDBBind, ending up with 41k complexes. The dataset has been recently explored for more challenging benchmark
18

---

Flow Matching: Idempotent Flow Map
construction and multi-ligand docking. Similar the PDBBind, the maximum 30% sequence similarity split provides 56649, 1136 and 1288 for training, validation and testing. Only one biounit for each complex is used for training. The complex with only one contact (protein residue ligand atom distance less than 4) is further filtered out, retrieving 36203, 734 and 756 training, validation and test examples.
1.1. POCKET DEFINITION
Radius Pocket, The pocket center is the mean position of the protein residues of which the minimum distance to any ligand atoms is less than 8A. The radius is computed as the maximum between 5A and the ligand radius (half of the largest distance between the ligand atoms) plus the radius pocket buffer (set to 7À in all experiments). The pocket residues are selected based on the comparison between the residue pocket center distance and the radius. The residue pocket center distance is randomly flipped by the = 2.
Distance Pocket. The protein residue ligand atoms distances are extracted by the minimum distance between the residue and any ligand atom positions. Again, these protein-ligand distances are further randomly shifted with the a 2. The final pocket residues are the ones whose distances are below 14 Å.
E.1. EXPERIMENTAL SETUP
The number of vector features and scalar features for TFN is set to be 32 and 8, respectively. Hence, there is no higher order representation (> 1) being used in the experiment and we do not use batch normalization and residual connection for the aggregated messages, but only layernorm the input features for each layer. The batch is set to be 4 for each GPU. The model trained for 150 epochs using the Adam optimizer (Kingma & Ba, 2014) with the initial learning rate 1c-3 and a polynomial scheduler. The flow matching conditional standard deviation is set be constant o; = a 0.5. The training takes around 20 hours on 8 RTX A100 GPUs for single ligand docking and one day on multi Milti-iigand docking. The validation is conducted for every epoch, and the checkpoint with the largest RMSD <2Åis selected for inference. The number of function evaluations is set to be 20 consistent with the HarmonicFlow. More details can be found in the GitHub repository of HarmonicFlow (https://gtthub.com/) http://.......aaeees////o/oooo
E.2. Protein Backbone Generation
E.2.1. EVALUATION METRICS
Designability. In experiments follow established self-consistency evaluation pipeline introduced before (Trippe et al., 2022; Watson et al., 2023). For every sampled backhone, we perform inverse folding with ProteinMPNN (Dauparas et al., 2022) and generate 8 sequences, which are subsequently refolded with ESMfold (Lin et al., 2023). We then compute the consiseeny RMSD (scRMSD) by aligning each ESMfold-refolded candidate with the originally sampled backbone on C atoms and consider a backbone designable if scRSMD < 2.0 Å.
Diversity. This metric measures the similarity among the designable backbones by computing the pairwise Template Modeling (TM) score.
Diversity
TM(2;,
N
mi(m 1)
(
where N is the total number of protein lengths, ny is the number of designable backbones at each length 1, and j are the index of the protein at each length.
Novelty. The novelty is defined as the TM score between the designable backbone and its closest natural protein found in the PDB databases with FoldSeek (Van Kempen et al., 2024), averaging over all the designable proteins:
Novelly- 1 Σ
max; (axx;M(r;,
1=1
where n the number of designable proteins.
19

---

Flow Matching: Idempotent Flow Map
E.2.2. EXPERIMENTAL SETUP
We train the model on RTX A100 GPUs for 150 epochs (~ 22 hours) on SCOPe and 600 epochs 3 days) on PDB. After Nimin epochs of training, the checkpoints designability is selected for the final result. Respectively, Nuriin set to 100 and 300 for SCOPe and PDB, and Nsweep is set to 10 and 50. The number of iteration K maz is set to be 1, as any value beyond 1 fills out of the memory even on "fat" GPU. The other hyperparameters are kept to be the same as the FrameFlow report the key setting here:
swept every Neweep for inference. The checkpoint achieving the highest
We
Hyperparameter
Setup
learning rate node embedding dimension edge embedding dimension number of head for IPA number of query key points number of value channel number head for transformer layer number of layer
le-4
256
128
8
8
12
4
6
Table 5, Key hyperparameter setup for IDFlow on protein backbone generation.
F. Additional Results
Validation metrics on docking. Fig. 4 shows the validation metric curve of HarmonicFlow and IDFlow. IDFlow converges much faster and achieves higher RMSD < 2Á than HarmonicFlow.
val_rmsd=2
iDFlow HarmniiFlow
0.4
0.3
0.1
epoch
20
60
80
100
120
140
Figure 4. The validation metric curve of docking on PDBBind of radius pocket time split. The fraction of validation RMSD < 2A vs. epoch time for IDFlow and HarmonicFlow.
Top-k accuracy on docking. The top-k accuracy assumes the known ground truth and picks up the sample with the highest accuracy of "RMSD < 2", k is the number samples being generated. reflects the model's capability of mode sampling. Table 6 shows the results of single ligand docking on PDBBind. IDFlow consistently outperforms the HarmonicFlow on different metrics across the dataset, demonstrating the improved mode capturing of IDFlow.
20

---

Flow Matching: Idempotent Flow Map
Table Top-k (40, and accuracy comparison of methods the PDBBind splits for pocket level docking based different metrics. The number k is specified in parenthesis.
Sequence Similarity Split
Time Split
Distance- Distance-Pocket
Radius-Pocket
Distance- Pocket
Radius-Pocker
%<1
%<2
Med.
%<]
%<<
Med.
%<l
%<2
Med.
%<l
%<2
Med.
HarwonicFlow (40)
30.7
63.2
1.5
16.1
46.8
2.1
35.8
66.9
1.3
21.9
52.4
1.9
IDFlow (40)
34.3
64.0
1.4
18.6
48.0
2.1
37.6
67.1
1.2
27.8
53.4
1.8
HarmonicFlon Flow (10)
22.0
53.9
1.8
9.1
36.7
2.4
25.8
58.8
1.6
15.8
45.4
2.2
IDFlow (10)
25.5
56.2
1.7
14.2
40.9
2.4
29.6
60.9
1.5
22.3
47.4
2.2
HarmonicFlow (5)
15.4
47.5
2.1
7.7
31.2
2.8
19.3
55.0
1.7
11.8
42.0
2.5
IDFlow (5)
20.6
52.0
1.9
10.0
34.3
2.6
24.7
57.9
1.6
17.8
44.9
2.4
Energy reduction. Fig. presents 1.2-error energy minimization comparison during sampling for IDFlow and Harmon- icFlow averaged over the test of radius pocket time split. While the idempotency also not perfect for IDFlow with the continuous loss minimization objective can never achieve absolute idempotency, it's apparent that the energy converges much faster and better for IDFlow.
Estimated Energy during Sampling Averaging over the Testset
12
IDFlow
HarmonicFlow
10
7.5
Number of Step
Figure 5. L2 error at cach step during sampling averaging over the test set for HarmonicFlow and IDFlow. The energy of IDFlow is calculated by L.2 error after refinement each step, with total of 10 steps. For HarmonicFlow, we sample 20 steps with refinement for computing the L2 error, but use the output of the first iteration to simulate the ODE. The red dashed line marks the lowest energy of HarmonicFlow.
Results on SCOPe. Table 7, we find that IDFlow can generate more designable proteins compared to the baseline FrameFlow and appears to be equally diverse and novel. Besides, the results over three runs are also more stable than the FrameFlow in terms of standard deviation. This demonstrates the method's effectiveness in generating highly designable proteins (relatively small, length 128).
Table 7. Protein Backbone Generation SCOPe. 10 backbones for the protein length from 60 to 128, in total 690 proteins. denotes the retrained FrameFlow. The results of FRAMEFLOW and GAFL are extracted from (Wagner et al., 2024) recomputed for mean and standard deviation.
METHOD
DESIGNABILITY (t)
DIVERSITY (L)
NOVELTY (L)
FRAMEFLOW
81.2
0.37
FRAMEFLOW
90.1 + 1.7
0.39 + 0.01
0.80 + 0.01
GAFL
90.2 + 0.6
0.38 + 0.02
IDFLOW
90.9 + 1.4
0.38 + 0.02
0.78 + 0.00
Number of steps. Fig. analyzes the impact of NFEs on IDFlow and FrameFlow for the SCOPe dataset on both designability
21

---

Flow Matching: Idempotent Flow Map
and diversity. Both methods exhibit performance degradation when NFEs reduced below 50. As shown in Fig. 6a, under identical sampling budgets, IDFlow achieves higher designability metries and demonstrates greater robustness when sampling with fewer steps (30 NFEs). This suggests IDFlow's 's better stability in low-step regimes. Fig. 6b presents the diversity metrics versus the NFEs for IDFlow and FrameFlow. Again, similar to the designability, IDFlow appears to be more robust to the NFEs and generates relatively more diverse structures.
Designability Comparison SCOPe
Diversity (1) Comparison SCOPe
0.95
Method
0.955
-B- E3F ow FromoFFow
-
0.385
Designabiig
0.00
0.300
0.375
0.370
0.75
Method
0.70
0.800
1DFaw
FramFFlow
60
100
120
140
40
60
80
100
120
140
MFEa
NFEs
(a) Designability (1) comparison between IDFlow and Frame- Flow SCOPe dataset.
(b) Diversity (4) comparison between IDFlow and FrameFlow on SCOPe dataset.
Figure 6. Designability and diversity comparison vs. NFEs on SCOPc dataset.
G. Visualization
Figure exhibits some randomly selected generated molecules compared with the Ground Truth. The sample is generated by 20 NFEs.
Figure 8 shows the designable backbones generated by IDFlow of various protein lengths. The sample is generated by 200 NFEs.
22

---

Flow Matching: Idempotent Flow Map
(a) 3pwd
(b) 4dgm
(c) 5198
(d) 5mkx
(e) 6fap
(f) 3b19
Figure 7. Six randomly selected generated complexes the radius pocket docking 30% sequence similarity split. The one in blue is the ground truth, and the one in green is generated from IDFlow.
23

---

Flow Matching: Idempotent Flow Map
Figure 8. Designable Backbones at length [100, 150. 200, 250, 300] trained PDB. Protein length from top bottom [300, 200.
150, 100].
24


Reviews:
[{"review_id": "PErVb3615K", "reviewer": "Reviewer_Fdeb", "summary": "This paper proposes to enhance the flow-matching framework with an energy-based perspective to learn iterative mapping. Such an idempotent mapping, as demonstrated theoretically in the paper, has better stability during generation. Experiments on protein docking and generation demonstrate better generation quality.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the effort to review our work! Here is our answer for addressing your concerns.\n\n>For the comparison of time, the authors tried to unfairly compare models with different numbers of sampling steps.\n\nWe want to clarify that our comparison is fair. HarmonicFlow adopts 20 steps (equivalent to 20 NFEs) for sampling, specified in Appendix E-“Hyperparameter” [1], while our method uses 10 steps (also equivalent to 20 NFEs). We apologize for the confusion and will clarify in the next version.\n\n> Intuitively, the ratio should be important, as the zero mapping $f_{\\theta,t}(x)=0$ trivially satisfy the idempotency condition but does not give the correct clean data. The authors should also demonstrate the impact of the probability $m$ for balancing between these two losses.\n\nWe appreciate the point raised by the reviewer that idempotency training can lead to a trivial solution. To resolve this, we adopt a simple approach where during training network input is detached from the computation graph (Eq. 21) for $L_R$, so that there is no information leak from previous iterations affecting the training dynamics. With our 50% training setup, the two losses are equally balanced with m from uniform distribution. This strategy aligns with the sampling algorithms, which at each step we first predict the clean sample and then refine it.\n\n>It seems that in Figure 3, the proposed model still requires multiple NFEs to achieve descent. This seems to indicate that the learned model is not close to idempotent. The authors should verify this.\n\nWe acknowledge the fact that empirically idempotency can not be guaranteed, as the continuous loss minimization objective could not be perfectly optimized unless a rigid structure is imposed. We do not claim the idempotency is achieved but only encouraged through the loss function. In the link (https://anonymous.4open.science/r/ICML2025R-F85D/L2-Error-TestTime.png), we provide an L2 error reduction plot during sampling for HarmonicFlow and IDFlow averaged over the time split test set. This shows that even if the absolute idempotency is not achieved, the IDFlow yields better idempotency.\n\n>This work has a potentially broader impact on scientific domains including protein design, protein docking, and other generative tasks in AI4Science domains. I suggest the author also discuss such applications in downstream tasks.\n\nHere are some discussions. First, the proposed energy-based framework can potentially improve the chemical plausibility of generated molecules without too much training overhead. The refiner can be purposed to refine the sample to a distribution governed by a certain energy function that the practitioners are interested in. It can also be integrated with pre-trained forcefield. More details can be found in paragraph 4 of the rebuttal to the reviewer rcEp. These relate to some other AI4Science problems such as crystal structure or molecular structure elucidation in chemistry. Second, the idea of idempotency also has a potential impact as the idempotency encourages the network to traverse over the loss landscape to find a locally smooth solution, which can potentially lead to better generalization of generative models for being adversarial robust.\n\n> Specifically, on SO(3), the target prediction in Eq.36 differs from the Riemannian flow matching loss in Eq.35.\n\nWe thank the reviewer for raising a valid point regarding SO(3) parameterization. While Equations 35 (Euclidean) and 36 (manifold) parameterize flows differently, both leverage geodesics constructing the flow path---key to Riemannian flow matching’s theoretical strength. The rotation field can be computed as: $\\frac{\\log_{r_t}(\\hat{r}_1)}{1-t}$ ensures the predicted rotation $\\hat{r}_1$ aligned with SO(3)’s geometry to follow the geodesics. However, since rotation is parameterized through the quaternions, and quaternions’ double cover of $\\mathrm{SO}(3)$ introduces non-uniqueness, which may potentially increase the learning difficulty of the network.\n\n>Some notations can be improved to be more consistent. For example, the authors used different fonts in Algorithm 1 for the same variable. For multiple baselines with the same better performance score, only one was highlighted in bold, which may be misleading.\n\nThanks for highlighting the notation inconsistencies. We will revise Table 3 to separate baselines, clarify improvements over FrameFlow, boldface all state-of-the-art results, and correct all notation/writing inconsistencies in the next version.\n\n[1] Harmonic Self-Conditioned Flow Matching for joint Multi-Ligand Docking and Binding Site Design, ICML 2024.", "claims_and_evidence": "The authors proposed an alternative flow-matching objective inspired by energy-based models. In their experiments, the proposed model achieved superior performance over the baselines. A few more ablation studies (see \"Experimental Design\" section) would be more convincing.", "methods_and_evaluation_criteria": "The method in the paper is well-supported, with both good intuitions and theoretical results inspired by energy-based models. However, some evaluation metrics on unconditional protein generation were questionable.\n\n- It is unclear why the authors separated the models and the baselines into two categories in Table 3, and only compared models within the category.\n- For multiple baselines with the same better performance score, only one was highlighted in bold, which may be misleading.\n- For the comparison of time, the authors tried to unfairly compare models with different numbers of sampling steps.", "theoretical_claims": "The theoretical grounds of energy-based models are clearly stated and discussed in this work.", "experimental_designs_or_analyses": "While most of the presented results can demonstrate the superior performance of the proposed approach, the following aspects can be better verified (theoretically or empirically) for more convincing claims.\n\n- Algorithm 1 adopted the combination of the flow matching loss ($L_G$ in Eq.18) and the idempotent loss ($L_R$ in Eq.21). Intuitively, the ratio should be important, as the zero mapping $f_{\\theta,t}(x)\\equiv0$ trivially satisfy the idempotency condition but does not give the correct clean data. The authors should also demonstrate the impact of the probability $m$ for balancing between these two losses.\n- Intuitively, if the learned denoiser is perfectly idempotent, one-step (or few-step) generation can be performed as the model will generate a consistent prediction. However, it seems that in Figure 3, the proposed model still requires multiple NFEs to achieve descent. This seems to indicate that the learned model is not close to idempotent. The authors should verify this.", "relation_to_broader_scientific_literature": "This work has a potentially broader impact on scientific domains including protein design, protein docking, and other generative tasks in AI4Science domains. I suggest the author also discuss such applications in downstream tasks.", "essential_references_not_discussed": "I believe essential references have been discussed in this paper.", "other_strengths_and_weaknesses": "In the unconditional protein generation task, the authors implicitly applied the proposed framework to Riemannian manifolds, which could be substantially different. For example, the equivalence of the target prediction versus the vector field prediction only holds for the Euclidean manifold (or zero-curvature manifolds) but would fail for general manifolds. Specifically, on SO(3), **the target prediction in Eq.36 differs from the Riemannian flow matching loss in Eq.35**. The authors should explicitly note this difference, which does not lead to Riemannian flow matching or any of its theoretical benefits. In this way, it might be better to formulate the proposed approach as a standalone generative framework instead of a variant of flow matching.", "other_comments_or_suggestions": "Some notations can be improved to be more consistent. For example, the authors used different fonts in Algorithm 1 for the same variable.", "questions_for_authors": "See other sections for questions.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xruAaoHNwz", "reviewer": "Reviewer_Z5a6", "summary": "This paper introduces a new method to train flow matching models. They want to sample from an energy function and this comes with very interesting outcomes. By using an energy function based on a squared euclidean distance, they realize that it boils down to train an indepotent map. To make the training efficient, they follow the self-conditioning procedure where they train the indepotent map 50% of the time and use regular flow matching 50% of the time. The sampling also uses only two function evaluation to be compute efficient. They evaluated their method on protein backbone generation and molecular docking where it outperforms the methods it is built on.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the effort to review our work! Here is our answer for addressing your concerns.\n\n>The paper is not always well written. I got a little lost in the paragraph Energy relaxation, confidence model and the EBMs. Maybe the author can rewrite it.\n\nThanks for pointing out the writing about the paper, specifically the section 3.2. We would make the writing clearer in the next version. In section 3.1, we want to connect the structure relaxation techniques and the confidence models to the EBM concept that the generated samples $\\hat{x}_1$ should associate with the minimum of the energy function. This energy function could be the physical energy function, like the one used in structural relaxation, or the confidence model output used for ranking the sample.\n\n>Can we couple it with self-conditioning?\n\nWe thank the reviewer for bringing the self-conditioning into discussion. We want to stress that the idempotency is orthogonal to the self-conditioning. Self-conditioning is still doing input conditioning where the network is conditioned by the $\\hat{x}_1$ to predict the data $x_1$ (in our setting) $\\textbf{given the input $x_t$}$. The conditioning is usually achieved by concatenation for images or as edge features for molecules. In our case, since the self-conditional information is injected through the edge features, it can be nicely integrated with the IDFlow. The implementation requires a double “$50\\\\%$”: for each training step, we first decide if to activate the self-conditioning and then idempotency training. This results in $25\\\\%$ time self-conditioning FM training, $25\\\\%$ time non-self-conditioning FM training, $25\\\\%$ time self-conditioning idempotency training, and $25\\\\%$ time non-self-conditioning idempotency training.", "claims_and_evidence": "They claim to define a flow matching method that sample form a distribution and this claim is true. They claim to achieve better results than the existing flow matching method and this is also true.", "methods_and_evaluation_criteria": "The methods is very relevant and interesting. The evaluation follows the standard practice in the literature.", "theoretical_claims": "Not applicable", "experimental_designs_or_analyses": "THey follow the literature. The authors did a sensitivity analysis with respect to the number of function evaluation.", "relation_to_broader_scientific_literature": "The literature review is complete.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "The paper is not always well written. I got a little lost in the paragraph Energy relaxation, confidence model and the EBMs. Maybe the author can rewrite it.", "other_comments_or_suggestions": "See above.", "questions_for_authors": "Can we couple it with self-conditioning?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xSkCD5VQlb", "reviewer": "Reviewer_rcEp", "summary": "The authors propose an energy-based perspective of flow matching for the purpose of improving the quality of 3D structures predicted by generative models. This perspective leads to a modified training procedure called Idempotent Flow Map Training, which trains a network to produce predictions by iterative refinement, where the initial output of the network is again passed into the network. In particular, the network uses the “denoiser” or “x” parameterization. This simple modification enables consistently better prediction quality across experiments in docking and protein backbone generation, with no increase in inference cost, but at an increased training cost.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the effort to review our work! Here is our answer for addressing your concerns and a link to additional figures.\n\nhttps://anonymous.4open.science/r/ICML2025R-F85D/\n\n>Can you include Figure 4 but with wall-clock time as the x-axis?...\n\nThanks for pointing out the increased throughput of IDFlow. In the link we provide a figure of the validation metric  (ValidationMetrics.png) with the x-axis to be wall clock time. The HarmonicFlow-L model has 60 scalar and 10 vector features for the tensor field network, totaling 16.3M parameters, compared to 5.7M in HarmonicFlow and IDFlow. The figure shows the better validation performance with a larger model size, but still lower than the IDFlow. Following the reviewer’s advice, we compare the HarmonicFlow-L with IDFlow in the link (TimeSplitRadiusDocking.png), which demonstrates that idempotency does not simply come from the increased training throughput. It also shows the improvement over baseline with fewer model parameters.\n\n>Can you provide at least one experiment of how the framework of energy-based flow matching would work for an energy function that is not the L2 loss? For example, alanine dipeptide with either a classical force field or xTB, or use a neural force field such as https://fair-chem.github.io/core/quickstart.html (less priority)\n\nWe thank the reviewer for suggesting the use of pre-trained forcefields or DFT energies in our framework! We describe how to incorporate this approach: the pre-trained forcefield can be directly integrated into the sampling algorithm, with the Langevin diffusion process refining the sample as follows: $\\hat{x}_1 = \\hat{x}_1 - \\frac{\\epsilon^2}{2} \\nabla E(\\hat{x}_1) + \\epsilon z$ where $\\epsilon$ is the step size and $z$ is a sample from a standard Gaussian. The gradient can be computed using the pre-trained forcefield, and the energy can be derived from DFT. We will include this in the next version of the paper. However, introducing a pre-trained forcefield may unfairly advantage the baselines due to the added complexity and capacity. The core innovation of our approach is using a refiner $G$ parameterized by the same network to refine samples toward the minimum of the energy function with minimal training overhead. The energy function is defined by the loss on the refiner's output, and the idempotent flow map arises from using the same loss function as the CFM loss.\n\nIn the link, we also present additional experimental results (TimeSplitRadiusDocking.png) on time-split radius pocket docking, where we use a different energy function more aligned with the molecular setup. The new energy function (EB-FM) consists of three terms: 1) the reconstruction errors 2) the ligand bond distance as intramolecular potential 3) the protein ligand distance as intermolecular potential. Without further tuning the weights of different energy losses, the EB-FM achieves better performance on the RMSD median. We will include the full results in the next version.\n\n> Can you demonstrate a simple application to images, such as CIFAR-10?\n\nWe thank the reviewer for raising the potential application to image data. Idempotency is a general idea that can be applied to many different types of generative models and different modalities. Considering higher dimensionality and the weaker theoretical connection of images to the stability and energy assumption, we'd like to leave this to future work.\n\n>The concept of learning an idempotent map is very related to consistency models.\n\nWe mildly disagree with the opinion that IDFlow is an instance of consistency models (CM). Key distinctions are as follows. First, consistency training minimizes the discrepancy between noisy data at neighboring steps, while flow matching embeds consistency by outputting clean samples across noise levels. In this sense, CM (PF-ODE) is closer to standard flow matching data parameterization. Consistency training enables fast sampling, while flow matching still requires ODE simulation. Second, an idempotent flow should ideally have zero velocity at $𝑡=1$ whereas CM has a non-zero, unstable vector field at $t=1$. Idempotency can also be applied to CM, making the learned consistency function also idempotent.\n\n> Idempotent Flow Map training could just as easily be thought of as introducing a bit of simulation-during-training....\n\nFrom idempotent inference perspective, refining $\\hat{x}_1$ during training adds some simulation, whereas vanilla flow matching relies on ODE simulation without idempotent inference. Introducing idempotency into flow matching ties into key concepts like physical stability in molecular dynamics. Training with idempotent flow reduces network uncertainty by incorporating the domain knowledge that generated molecules should be stable.\n\nLastly, we thank the reviewer for highlighting the broader impact of 3D structure generation and for pointing out typos. We will address these in the next version by including more references and improving the writing.", "claims_and_evidence": "The claim of an energy-based perspective of flow matching is a bit strained since the paper only ever uses L2 reconstruction error as the energy function. At least one experiment using another energy function should be shown to support this claim. This would also help support the claim that the proposed method is designed for structure generation models. As far as I can tell, the proposed method is applicable to any diffusion/flow model which can use the “denoiser” parameterization, and so this method would be applicable for images and other modalities as well.\n\nThe proposed method does not seem to rely on the energy-based perspective. The ablation of $K_max$ in Table 4 indicates that setting $K_{max}$ greater than 0 is what is responsible for improved performance. However, Idempotent Flow Map training could just as easily be thought of as introducing a bit of simulation-during-training to the standard flow matching training procedure, and providing extra training supervision to each step of simulation-during-training.\n\nThe claim that IDFlow trains faster is uncertain to me. While the experimental comparison uses the same architecture and same number of epochs for training, IDFlow uses more compute per forward pass due to the requirement of simulation-during-training in the refinement loss. Figure 4 in Appendix F demonstrates that IDFlow trains with fewer epochs, but it would be valuable to see what these curves look like if the x-axis were switched to wall-clock time. This is important since it appears that both models are undertrained (validation metrics have not completely plateaued yet). It would also be an interesting baseline to compare to simply increasing the size of the model architecture.", "methods_and_evaluation_criteria": "In the area of molecular structure prediction, the proposed experiments are complete and make sense. However, the proposed method appears to be much more general than the application area of 3D structure prediction, and I am curious how the proposed training procedure would affect an image model.", "theoretical_claims": "I checked the derivations in the appendix and found no issues.", "experimental_designs_or_analyses": "I checked that the experiments on molecular docking and protein backbone generation closely follow experiments executed in previous work. The only concern is how much the training cost increases by.", "relation_to_broader_scientific_literature": "The paper provides a simple method for improving generation quality for flow/diffusion models, which is relevant to any work that applies these generative models.\n\nSpecifically, the key contribution of the paper is Idempotent Flow Map training, which is related to recycling/self-conditioning. This paper demonstrates that these tricks empirically improve performance, though I am still unsure of the full reason except that model expressivity is increased.\n\nThe concept of learning an idempotent map is very related to consistency models (see below).", "essential_references_not_discussed": "The focus on idempotency is very similar to Consistency Models [1], although Consistency Models focus on reducing the number of generation steps rather than increasing the quality of generated samples. Idempotency Flow Map Training can be seen as an instance of Consistency Training except that the training target sometimes comes from multiple steps of simulation during training.\n\nConsistency models are only idempotent at t=0, where the skip connection forces the network to be the identity function. In contrast, IDFlow is idempotent for every t.\n\nHowever, the sampling approach is different from Consistency Models.\n\n\n[1] Song, Y., Dhariwal, P., Chen, M., & Sutskever, I. (2023). Consistency models.", "other_strengths_and_weaknesses": "The paper provides thorough background knowledge on flow matching and experimental setup with molecular structure prediction.\n\nThe notation is sometimes confusing, with networks like $G_\\theta$ and $E_\\theta$ defined for the sake of abstracting out an energy function, only to always set the energy equal to the L2 loss.\n\nSimple and effective methods are valuable, but appear less novel when it is not as clear why these simple changes provide incremental improvements. If I were to critique the method, it appears that the proposed method simply adds simulation-during-training and provides extra training supervision by calculating loss", "other_comments_or_suggestions": "Typos:\nline 641: \"seminar\"\nline 240, left column: “L_CFM” should be “L_G”\nline 947: “We train the model 8 on” should be “We train the model on 8”\n\nuse of $K_{max}$ vs $k_{max}$ is not consistent\n\n3D structure prediction is relevant to many more areas of chemistry than just biomolecules. For example, it is relevant for moment-constrained structure elucidation [1] and crystal structure prediction [2] [3]. Crystal structure prediction in particular has traditionally focused on finding the lowest-energy structures, where energy is given by DFT.\n\n[1] Cheng, A., Lo, A., Lee, K. L. K., Miret, S., & Aspuru-Guzik, A. (2024). Stiefel Flow Matching for Moment-Constrained Structure Elucidation. arXiv preprint arXiv:2412.12540.\n[2] Jiao, R., Huang, W., Lin, P., Han, J., Chen, P., Lu, Y., & Liu, Y. (2023). Crystal structure prediction by joint equivariant diffusion. Advances in Neural Information Processing Systems, 36, 17464-17497.\n[3] Zeni, C., Pinsler, R., Zügner, D., Fowler, A., Horton, M., Fu, X., ... & Xie, T. (2023). Mattergen: a generative model for inorganic materials design. arXiv preprint arXiv:2312.03687.", "questions_for_authors": "Can you include Figure 4 but with wall-clock time as the x-axis? This would help clearly demonstrate the effectiveness of IDFlow.\nIf IDFlow has a higher wall-clock training time than original baselines, can you compare IDFlow to the performance of a larger model with similar memory and training throughput?\nCan you provide at least one experiment of how the framework of energy-based flow matching would work for an energy function that is not the L2 loss? For example, alanine dipeptide with either a classical force field or xTB, or use a neural force field such as https://fair-chem.github.io/core/quickstart.html\n(less priority) Can you demonstrate a simple application to images, such as CIFAR-10?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "N/A"}}, {"review_id": "pDpvcgOFwv", "reviewer": "Reviewer_AZnb", "summary": "This paper proposes an enhanced flow matching framework, improving the standard setup from the energy-based perspective.\nIt provides a specific instance called IDFlow, which employs the reconstruction error as the energy function during training and then iteratively predicts and refines the sample.\nExtensive experiments are conducted to validate the improvements of IDFlow over the standard setup, demonstrating its effectiveness.\n\n## update after rebuttal\nThe authors have addressed some of my concerns.\nHowever, I still think that comparing only one pair (one standard method with its corresponding IDFlow) for each task is insufficient to provide a comprehensive evaluation. After all, IDFlow is a modified flow matching that relies on the specific flow matching implementation, so it's necessary to validate its effectiveness and reliability in a broader context. Therefore, I decided to maintain my score but respect any final decision.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the effort to review our work! Here is our answer for addressing your concerns and a link to additional figures.\n\nhttps://anonymous.4open.science/r/ICML2025R-F85D/\n\n>Since the definition and goal of Eq.13 and Eq.18 are obviously different, why can the proposed method be interpreted from the energy-based perspective?\n\nFirst, we want to clarify that Eq. 18 is the proposed objective to learn the refiner G that maps the $\\hat{x}_1$ onto its corresponding point $x_1$ on the data manifold. One way of viewing Eq. 18 from the energy-based point of view is if the G and f shares the same network (our case), the landscape is not only shaped by the trajectory sample $x_t$ (Eq. 7), but also the generated sample $\\hat{x}_1$. This is associated with one way of training the EBMs, which maps the point off the manifold back to the data manifold [2]. Besides, the idempotency training also encourages the network to traverse over the L2 energy landscape to find a locally smooth solution, which can potentially lead to better generalization for being adversarial robust. \n\n>Based on Eq.18, it's not enough to claim the flow map is idempotent.\n\nWe acknowledge the fact that empirically idempotency can not be guaranteed, as the continuous loss minimization objective could not be perfectly optimized unless a rigid structure is imposed. We do not claim that idempotency is achieved but only encouraged through the loss function. In the anonymous link (L2-Error-TestTime.png), we add an L2 error reduction plot during sampling for HarmonicFlow and IDFlow averaged over the time split test set (363 examples). This shows that even if the absolute idempotency is not achieved (0 across the timestep), the IDFlow yields better idempotency that the L2 error converges faster and better than the baseline.\n\n> it's necessary to report the performance of IDFlow based on various implemented versions of FoldFlow.\n\nWe acknowledge that the additional experiments on FoldFlow could further strengthen the claim. However, FrameFlow is conceptually very similar to the core idea of FoldFlow-OT and FoldFlow-Base, which both apply the SE(3) flow matching to generate the protein backbones. The only deviation is the FoldFlow-SFM introducing the stochastness for sampling using the IGSO(3) distribution. Considering the high similarity of the methods, we would expect a similar performance improvement over FoldFlow.\n\n> According to lines 323-327, the sampling steps for HarmonicFlow and IDFlow are the same. If so, it would be unfair to HarmonicFlow, as IDFlow involves two NFEs in a single step.\n\nWe want to clarify that our comparison is fair. HarmonicFlow adopts 20 steps (equivalent to 20 NFEs) for sampling, specified in Appendix E-“Hyperparameter” [1], while IDFlow uses 10 steps (also equivalent to 20 NFEs). We apologize for the confusion and will clarify in the next version.\n\n>How to understand the relation between refiner and energy-based models?\n\nWe could think of the refiner as a means to map the generated sample to be minimal of a certain energy function. The energy function is imposed by defining a loss function on the refiner output. In the paper, we use L2-error because of its simplicity and alignment with performance metrics (RMSD) and the nice properties of idempotency. In this case, this refiner maps the point off the data manifold \\hat{x}_1 to the data manifold (similar to the denosing autoencoder), a contrastive approach for training the EBMs [2]. Generally speaking, we can impose different energy losses, such as the flat bottom potential in [3], on the refiner output to improve the chemical plausibility of the generated sample. Besides, it can also integrate with the pretrained force field that, instead of refining the sample to the minimum, one can also leverage the Langevin diffusion to sample from the distribution governed by the energy function. A detailed description can be found in the response to the reviewer rcEp in paragraph 3.\n\n>It's significant unmatched for the refiner operation between training and sampling.\n\nRefining the sample multiple times at each step increases the sampling budget. In the link (TestTime-K-Ablation.png), we ablate k while keeping the total NFEs constant (20 or 21), using more discretizations for smaller k. Ideally, an idempotent function wouldn’t require refinement at test time, but empirical results show that setting k=1 yields the best performance. We attribute this to large discretization errors for larger k, as more NFEs are needed for idempotency. Setting k=1strikes the best tradeoff between idempotency and discretization.\n\n[1] Harmonic Self-Conditioned Flow Matching for joint Multi-Ligand Docking and Binding Site Design, ICML 2024.\n\n[2] LeCun, Y. From machine learning to autonomous intelligence: Lecture 2, 2022. URL https://leshouches2022.github.io/SLIDES/lecun-20220720-leshouches-02.pdf.\n\n[3] Composing Unbalanced Flows for Flexible Docking and Relaxation, ICLR 2025.", "claims_and_evidence": "Some claims made in the submission are not supported by clear and convincing evidence: \n* The specific relationship between energy-based models and the proposed method. For example, since the definition and goal of Eq.13 and Eq.18 are obviously different, why can the proposed method be interpreted from the energy-based perspective?\n* The rationality of the idempotent flow map. Based on Eq.18, it's not enough to claim the flow map is idempotent.", "methods_and_evaluation_criteria": "The proposed method makes sense for the problem at hand.", "theoretical_claims": "This paper claims that the standard setup of flow matching can be improved from the energy-based perspective, but does not provide the corresponding complete proof.", "experimental_designs_or_analyses": "The experimental designs and analyses are suitable for validating the effectiveness of the proposed method.", "relation_to_broader_scientific_literature": "This paper proposes a modification to the flow matching objective by directly minimizing an energy function, making it relevant to various applications of flow matching.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "No.", "other_comments_or_suggestions": "1. Since IDFlow is a modified flow matching that relies on the specific flow matching implement, comparing only one pair (one standard method with its corresponding IDFlow) for each task is insufficient to provide a comprehensive evaluation. For example, it's necessary to report the performance of IDFlow based on various implemented versions of FoldFlow.\n2. According to lines 323-327, the sampling steps for HarmonicFlow and IDFlow are the same. If so, it would be unfair to HarmonicFlow, as IDFlow involves two NFEs in a single step.", "questions_for_authors": "1. The main contribution of this paper is just the incorporation of the refiner operation into flow matching, how to understand the relation between them and energy-based models?\n2. For the sampling algorithm (Algorithm 2), how does the performance change when the refiner operation is removed? Additionally, how does the performance vary when the iteration count of the refiner operation is increased?  Moreover, why is $k$ a random integer between 1 and $K_{max}$ during training, but fixed to 1 during sampling? It's significant unmatched for the refiner operation between training and sampling.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper presents an energy-based formulation of flow matching for 3D molecular structure generation. The authors instantiate their framework through IDFlow, which leverages reconstruction error as the energy function to better shape the loss landscape using contrastive samples generated by the flow model. IDFlow constructs a flow map that iteratively predicts and refines the sampling trajectory toward the target molecular structure. Experiments on protein binding and protein backbone generation are conducted to verify the effectiveness of the proposed model. 

The paper received four reviews. Following the rebuttal and discussion, a consensus was achieved among all four reviewers, culminating in the decision to accept the paper. This agreement is attributed to the paper's innovation, promising results, and potentially broader impacts. The AC agrees with the reviewers and recommends accepting the paper. To further improve the paper quality, the AC urges the authors to revise their paper by taking into account all the suggestions provided by the reviewers, including integrating extra experimental results conducted during the rebuttal phase.
Decision: Accept (poster)