Energy-Based Flow Matching for Generating 3D Molecular Structure

Thank you for the effort to review our work! Here is our answer for addressing your concerns and a link to additional figures.

https://anonymous.4open.science/r/ICML2025R-F85D/

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?

First, 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.

Based on Eq.18, it's not enough to claim the flow map is idempotent.

We 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.

it's necessary to report the performance of IDFlow based on various implemented versions of FoldFlow.

We 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.

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.

We 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.

How to understand the relation between refiner and energy-based models?

We 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.

It's significant unmatched for the refiner operation between training and sampling.

Refining 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.

[1] Harmonic Self-Conditioned Flow Matching for joint Multi-Ligand Docking and Binding Site Design, ICML 2024.

[2] LeCun, Y. From machine learning to autonomous intelligence: Lecture 2, 2022. URL https://leshouches2022.github.io/SLIDES/lecun-20220720-leshouches-02.pdf.

[3] Composing Unbalanced Flows for Flexible Docking and Relaxation, ICLR 2025.

Thank you for the effort to review our work! Here is our answer for addressing your concerns and a link to additional figures.

https://anonymous.4open.science/r/ICML2025R-F85D/

Can you include Figure 4 but with wall-clock time as the x-axis?...

Thanks 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.

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)

We 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.

In 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.

Can you demonstrate a simple application to images, such as CIFAR-10?

We 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.

The concept of learning an idempotent map is very related to consistency models.

We 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.

Idempotent Flow Map training could just as easily be thought of as introducing a bit of simulation-during-training....

From 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.

Lastly, 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.

Thank you for the effort to review our work! Here is our answer for addressing your concerns.

For the comparison of time, the authors tried to unfairly compare models with different numbers of sampling steps.

We 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.

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.

We 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.

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.

We 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.

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.

Here 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.

Specifically, on SO(3), the target prediction in Eq.36 differs from the Riemannian flow matching loss in Eq.35.

We 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.

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.

Thanks 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.

[1] Harmonic Self-Conditioned Flow Matching for joint Multi-Ligand Docking and Binding Site Design, ICML 2024.

Thank you for the effort to review our work! Here is our answer for addressing your concerns.

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.

Thanks 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.

Can we couple it with self-conditioning?

We 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) $**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.