High-order Equivariant Flow Matching for Density Functional Theory Hamiltonian Prediction

Dear Reviewer UyjR,

Thank you for your thoughtful and encouraging review. We greatly appreciate your recognition of the clarity of our writing, the originality of applying SE(3)-equivariant flow matching to Hamiltonian prediction, and the significance of our results in both Hamiltonian accuracy and SCF step reduction. Below, we address your comments in detail.

W1, Q1: It is unclear if the performance gains are due to the generative task or improvements in model architecture.

We appreciate the reviewer’s question regarding the source of performance improvement. To address this concern, we conducted an ablation study comparing our flow-based framework to a regression-based baseline under identical architectural conditions:

These results demonstrate that both architectural modifications and the adoption of flow matching contribute meaningfully to our performance gains. Notably, switching from regression to flow-based training yields consistent improvements across all metrics, even under the same model architecture. Furthermore, the choice of prior distribution (TE vs. GOE) and the use of a fine-tuning strategy further highlight the trade-off between Hamiltonian accuracy and energy prediction.

We will include the complete ablation results in the final version and publicly release our codebase to support further research in this area.

Q2-1: (Figure 3) Given QHFlow is evaluated with "3-step ODE-based sampling," wouldn't we expect the inference time ratio for QHFlow to be roughly 3 times that of QHNet?

This is due to additional latency that is constant across QHFlow and QHNet. The latency is caused by system-level overheads such as data movement and Python runtime costs. Once we make the (unrealistic) assumption that the latency is not part of the computational cost, QHFlow would cost roughly three times that of QHNet.

More importantly, we emphasize that the primary computational bottleneck in DFT workflows is the SCF iteration process. Since our method substantially reduces the number of SCF iterations, as shown in Figure 3, the added inference cost becomes negligible in the context of end-to-end simulation time.

Q2-2: While the bar colors of "Inf T Ratio" can be inferred by order, it would be helpful to modify this figure so that this is more obvious.

Thank you for the helpful suggestion. In the final version, we will revise the "Inf T ratio" part to be more visible and clear. Specifically, we will add a small inset above each “Inf T Ratio” bar group to zoom in on it and better highlight the differences.

As external link sharing is not allowed, we are unable to upload the revised figure. We will ensure to update it in the final manuscript.

Q3: What datapoints are the "71%" M17 and "53%" QH9 Hamiltonian derived from?

Thank you for the question. We report the mean relative improvement in Hamiltonian MAE compared to the previous state-of-the-art model, as reducing Hamiltonian error is one of the primary objectives. We measure the improvement by $1-\frac{A}{B}$ where $A$ is the Hamiltonian MAE of our model and $B$ is that of the previous best model.

For example, we observed the following improvements on MD17:

• Water: 58% = $1 - \frac{4.93}{11.7}$
• Ethanol: 73% = $1 - \frac{5.33}{20.09}$
• Malondialdehyde: 82% = $1 - \frac{3.80}{20.67}$
• Uracil: 80% = $1 - \frac{3.68}{18.65}$
• Average: 71%

And we observed the following improvements on QH9 (without WA-FT):

• Stable-id: 50% = $1 - \frac{22.95}{45.48}$
• Stable-ood: 54% = $1 - \frac{20.01}{43.33}$
• Dynamic-300k-geo: 50% = $1 - \frac{25.94}{52.18}$
• Dynamic-300k-mol: 58% = $1 - \frac{45.91}{108.19}$
• Average: 53%

We will revise the manuscript accordingly to reflect clarify how the improvements were calculated, as follows:

(line 15-16) QHFlow achieves state-of-the-art performance, reducing Hamiltonian error by 71% on MD17 and 53% on QH9 compared to the previous best model.

(line 289-290) With up to a 71% reduction in Hamiltonian prediction error on average across the entire MD17 dataset compared to the previous state-of-the-art model, QHFlow demonstrates strong predictive performance.

(line 297-298) As shown in Table 2, QHFLOW consistently outperforms all prior methods in all metrics up to a 53% reduction in Hamiltonian prediction error on average over all QH9 splits.

We mark the text in bold to highlight its relevance to your concerns. They will be removed from the final manuscript.

Limitations: Lack of discussion on limitations; weak correlation between Hamiltonian prediction error and SCF acceleration.

Thank you for pointing this out! We originally included the discussion in Appendix F.4, but now realize that its visibility is low despite its importance. To alleviate your concern, we will move the limitation to the main text and revise the contents.

Furthermore, we agree with your points on the lack of correlation between Hamiltonian prediction error and SCF acceleration. We think this is very insightful and a very interesting venue for future research!

In the final manuscript, we will include the following text:

We did not fully verify the generalizability of our approach, e.g., extension to new architectures (WANet, SPHNet) or datasets (PubChemQH), due to the lack of resources, public codes, and public datasets. Our flow-based formulation introduces computational overhead for solving ODEs. In the future, one could consider removing the overhead with the one-step generative models [1]. Our experiments are limited to gas-phase molecules with up to 29 atoms and two commonly used XC functionals (PBE and B3LYP). The algorithms are left to be verified for more general settings, such as periodic solids or higher-accuracy functionals (e.g., hybrid or double-hybrid). We also acknowledge that our datasets consist of relatively small molecules. One could further test the scalability of our approach once new datasets on larger systems are available. We also note that our work does not give a theoretical explanation of how Hamiltonian prediction error translates to SCF acceleration. Developing a theoretically grounded algorithm specifically for SCF acceleration would be an interesting future research.

We mark the text in bold to highlight its relevance to your concerns. They will be removed from the final manuscript.

We hope our responses have addressed your questions and concerns. Please let us know if there is anything we may have missed or if any points require further clarification.

[1] Kornilov, Nikita, et al. "Optimal flow matching: Learning straight trajectories in just one step." 2024 NIPS.

Dear Reviewer upG2,

Thank you for your thoughtful and encouraging review. We sincerely appreciate your positive feedback on the clarity of the manuscript, the novelty and motivation of our SE(3)-invariant priors, and the rigor of our experimental evaluations. Below, we address your comments in further detail.

W1, Q1: Lack of evaluation on larger molecules such as MD22.

We fully resonate with the reviewer's concern. However, this is infeasible at the moment since there are no public datasets with large molecule sizes to evaluate our method. Note that the MD22 dataset is not a quantum Hamiltonian dataset. We also note the existence of PubChemQH, which consists of molecules with 40–100 atoms, but the dataset is not public.

Nevertheless, we strongly believe that our improvements will generalize to larger systems. The improvements may be larger since flow matching (and continuous-time generative models) has proven to be quite scalable for large systems. We note that our computational cost is similar to that of predictive models for each training step. For inference, the cost of multiple ODE integration steps is negligible compared to SCF iterations.

W2: The novelty of the approach is limited as it largely builds on existing methodologies and does not significantly advance methodology.

We resonate with this concern. Nevertheless, we would like to re-emphasize that our core contribution is not to methodological novelty. Instead, we propose to approach the Hamiltonian prediction problem with generative models to incorporate the underlying structure, which demonstrates strong empirical performance. We consider this idea to be novel, since the existing works mainly focused on developing new predictive architectures instead of generative models.

Q2: Is there any intuition why the TE prior generally performs better than the GOE prior?

Regrettably, this is more of an empirical finding rather than an insightful design. We hypothesize that the observed performance difference arises from the structural alignment between the TE prior and the Hamiltonian matrices. Specifically, the TE prior utilizes tensor expansion, which mirrors the structure and operations used to construct the Hamiltonian matrix in our model. Rather, GOE is constructed by the element-wise Gaussian distribution, which lacks this structural inductive bias.

Q3: There is generally a one-to-one mapping between structure and the ground-state Hamiltonian. Clarification on “distributional structure and uncertainty.

Thank you for pointing this out. We agree that the Hamiltonian is deterministic for a given structure. Our goal is not to model the uncertainty in data, but to better capture the predictive uncertainty arising from model approximation. In what follows, we provide a similar argument to Han et al. [1].

We first clarify what we meant by distributional structure and uncertainty (in the prediction). Regression models typically minimize element-wise losses (e.g., MSE), implicitly modeling conditional independence between matrix elements: $p(\mathbf{H}|\mathcal{M}) = \prod_{i,j} p(H_{ij}|\mathcal{M})$, where $\mathbf{H}$ is the Hamiltonian, $H_{ij}$ is its $(i,j)$-th element, and $\mathcal{M}$ denotes the input molecule. In contrast, conditional generative models, such as diffusion or flow-based models, introduce a prior $z$ that globally conditions the output (i.e. $p(\mathbf{H}|\mathcal{M}) = \int_z p(z)\prod_{i,j} p(H_{ij}|\mathcal{M},z)dz$). This allows the model to better capture structural correlations among Hamiltonian elements. Even in deterministic tasks, this structure-aware generation leads to more coherent predictions, as evidenced in recent advances like AlphaFold3 [2], DiffCSP [3], image segmentation [4], and an object detection model [5].

We further demonstrate how better capturing the uncertainty leads to improved performance. When the model is uncertain between two plausible Hamiltonians $\mathbf{H}_1$ and $\mathbf{H}_2$, a regression model often predicts their average $(\mathbf{H}_1 + \mathbf{H}_2)/2$, which may not correspond to a physically meaningful solution and results in high uncertainty. A generative model, however, can represent both modes explicitly, yielding sharper, more reliable outputs.

Empirically, our flow-based model achieves consistently lower MAE compared to regression baselines as shown in Tables 1 and 2, supporting its advantage in modeling both the structural dependencies and predictive uncertainty of Hamiltonian prediction.

Nevertheless, we fully agree that our earlier phrasing was ambiguous. We will revise the manuscript to clarify this point:

(lines 38-39) These approaches rely on pointwise regression, which may lead to high predictive uncertainty and often struggle to capture the structural correlations present in the Hamiltonian matrix.

Suggestions and Typos

Thank you for your helpful suggestions. We will revise Section 5.4 to clearly indicate that the results shown in Table 3 correspond to QHFlow with the TE prior, as reported in Table 2.

Specifically, we will update section 5.4 as follows:

(line 333-334) Table 3 compares the performance of GOE and TE priors, where the TE prior results are identical to those of QHFlow reported in Table 2. We found that the TE prior consistently yields lower errors than the GOE prior across all QH9 splits. This highlights the importance of designing task-aligned and symmetry-aware priors for accurate Hamiltonian prediction.

We mark the text in bold to highlight its relevance to your concerns. They will be removed from the final manuscript.

We also appreciate the correction regarding the molecular name in Table 1 and will correct it in the final version.

We hope our responses have addressed your questions and concerns. Please let us know if there is anything we may have missed or if any points require further clarification.

References

[1] Han, Xizewen, Huangjie Zheng, and Mingyuan Zhou. "Card: Classification and regression diffusion models." NIPS 2022

[2] Abramson, J., Adler, J., Dunger, J. et al. Accurate structure prediction of biomolecular interactions with AlphaFold 3. Nature 2024

[3] Jiao, Rui, et al. "Crystal structure prediction by joint equivariant diffusion." NIPS 2023.

[4] Chen, Shoufa, et al. "Diffusiondet: Diffusion model for object detection." ICCV 2023.

[5] Baranchuk, Dmitry, et al. "Label-efficient semantic segmentation with diffusion models." ICLR 2022.

Thank you for your detailed and thoughtful response to my comments, particularly the clarification regarding “distributional structure and uncertainty.”

Regarding the large-scale evaluation, I was referring to the experimental setting in [1] in section 3.3, where the authors evaluated models on 500 randomly sampled conformations and applied the same SCF setup used in MD17 for the ground truth. However, since I did not cite this work in my original review, I do not expect you to incorporate this result within the discussion period though this could be a nice addition for the camera-ready as it would strengthen your argument of generalizing to larger systems.

I have no further questions or concerns.

[1] Zhang et al. “Self-Consistency Training for Density-Functional-Theory Hamiltonian Prediction.” ICML 2024

Thank you for your thoughtful comments and questions throughout the review process. Your feedback has helped us clarify and strengthen our work. We have reviewed the paper you mentioned and now have a clearer understanding of the context you referred to. We will incorporate the suggested evaluation in the camera-ready version, as we believe it will also be valuable for future follow-up studies.

Dear Reviewer NZgB,

Thank you for your thoughtful review. We appreciate your comments on the clarity of the writing with necessary background for understanding our work, and the proposed application of flow matching. Below, we address your points in detail.

W1: While the empirical performance motivates the work very well, the qualitative reasoning provided is weak. There is neither a distribution of Hamiltonian nor an uncertainty for a given structure.

Thank you for pointing this out. We agree that the Hamiltonian is deterministic for a given structure. Our goal is not to model the uncertainty in data, but to better capture the predictive uncertainty arising from model approximation. In what follows, we provide a similar argument to Han et al. [1].

We first clarify what we meant by distributional structure and uncertainty (in the prediction). Regression models typically minimize element-wise losses (e.g., MSE), implicitly modeling conditional independence between matrix elements: $p(\mathbf{H}|\mathcal{M}) = \prod_{i,j} p(H_{ij}|\mathcal{M})$, where $\mathbf{H}$ is the Hamiltonian, $H_{ij}$ is its $(i,j)$-th element, and $\mathcal{M}$ denotes the input molecule. In contrast, conditional generative models, such as diffusion or flow-based models, introduce a prior $z$ that globally conditions the output (i.e. $p(\mathbf{H}|\mathcal{M}) = \int_z p(z)\prod_{i,j} p(H_{ij}|\mathcal{M},z)dz$). This allows the model to better capture structural correlations among Hamiltonian elements. Even in deterministic tasks, this structure-aware generation leads to more coherent predictions, as evidenced in recent advances like AlphaFold3 [2], DiffCSP [3], image segmentation [4], and an object detection model [5].

We further demonstrate how better capturing the uncertainty leads to improved performance. When the model is uncertain between two plausible Hamiltonians $\mathbf{H}_1$ and $\mathbf{H}_2$, a regression model often predicts their average $(\mathbf{H}_1 + \mathbf{H}_2)/2$, which may not correspond to a physically meaningful solution and results in high uncertainty. A generative model, however, can represent both modes explicitly, yielding sharper, more reliable outputs.

Empirically, our flow-based model achieves consistently lower MAE compared to regression baselines as shown in Tables 1 and 2, supporting its advantage in modeling both the structural dependencies and predictive uncertainty of Hamiltonian prediction.

Nevertheless, we fully agree that our earlier phrasing was ambiguous. We will revise the manuscript to clarify this point:

(lines 38-39) These approaches rely on pointwise regression, which may lead to high predictive uncertainty and often struggle to capture the structural correlations present in the Hamiltonian matrix.

W2: The technical novelty of this work is constrained to the two new base densities.

We resonate with this concern. Nevertheless, we would like to re-emphasize that our core contribution is not to methodological novelty. Instead, we propose to approach the Hamiltonian prediction problem with generative models to incorporate the underlying structure, which demonstrates strong empirical performance. We consider this idea to be novel, since the existing works mainly focused on developing new predictive architectures instead of generative models.

W3-1, Q2: Hyperparameters vary quite a bit between datasets. How sensitive is your approach to the choice of hyperparameters? Is there a good way to estimate these a priori?

We appreciate the reviewer’s concern and would like to clarify that our framework is generally robust to hyperparameter choices. Across seven out of eight datasets, we use nearly identical training hyperparameters, except for the minimum learning rate and the use of S-block embedding. The only exception was MD17–water, which contains only 500 training samples; For this case, we adopted a smaller learning rate to ensure stable optimization while keeping the other model hyperparameters fixed.

As a general guideline, we recommend selecting the largest batch size based on available GPU memory, since increasing batch size or training steps consistently leads to improved performance. For other hyperparameters, we generally recommend following our setting. However, hyperparameter tuning is inevitable for achieving the best performance in machine learning. One can start by tuning the learning rates to further improve performance.

W3-2: WA Loss provides inconsistent results.

While we do not fully understand the reviewer’s concern, we believe it refers to the inconsistent improvements of QHFlow+WA-FT on Hamiltonian MAE in Table 2. For this case, we would like to clarify that this is, in fact, consistent with the motivation of WA-FT loss.

We first note that the WA-FT loss was mainly designed to enhance energy prediction performance, which is consistently improved across datasets. When minimizing the weighted combination of the original QHFlow loss (for Hamiltonian prediction) and the WA-FT loss (for energy prediction), there exists a natural trade-off between the metrics. Our experiments show that one can find a "sweet spot" that yields high accuracies for both metrics, but the trade-off is inevitable.

Q1: Unclear explanation of Figure 3 (SCF accelerating performance)

Thank you for pointing this out. We will revise the explanation of Figure 3 (both the caption and the main text) to improve clarity.

Specifically, we will update the caption as follows:

(Figure 3 caption) The SCF Iter Ratio reflects the relative improvement in the number of SCF iterations compared to conventional DFT. All time-related values (Inf T Ratio, SCF T Ratio, Total T Ratio) are expressed as percentages relative to the total runtime of conventional DFT initialized with minao, which is set to 100%. Lower values indicate faster convergence. For example, an “Total T Ratio” of 46% means that QHFlow converges in 46% of the time required by standard DFT, including the model inference time.

In the main text, we will revise Section 5.3 as follows:

(line 321-323) As shown in Figure 3, QHFlow significantly accelerates SCF convergence. For instance, on QH9-stable (id), QHFlow achieves convergence in just 46% of the runtime compared to conventional DFT, corresponding to a 54% relative speedup. Compared to prior ML-based initializations such as QHNet (53%), QHFlow yields an additional 7%p reduction. These results demonstrate the practical utility of QHFlow for accelerating real-world DFT simulations.

While we originally reported relative ratios (i.e., 13% $\approx 1 - \frac{0.46}{0.53}$), we agree that we should present the improvement more clearly, such as expressing it as a percentage-point difference (e.g., a 7%p gain from 53% to 46%).

Limitations: Lack of discussion on limitations in terms of scalability

Thank you for pointing this out! We originally included the discussion in Appendix F.4, but now realize its low visibility despite the importance. To alleviate your concern, we will move the limitation to the main text and revise the contents.

To be specific, we will include the following text:

We did not fully verify the generalizability of our approach, e.g., extension to new architectures (WANet, SPHNet) or datasets (PubChemQH), due to the lack of resources, public codes, and public datasets. Our flow-based formulation introduces computational overhead for solving ODEs. In the future, one could consider removing the overhead with the one-step generative models [6]. Our experiments are limited to gas-phase molecules with up to 29 atoms and two commonly used XC functionals (PBE and B3LYP). The algorithms are left to be verified for more general settings, such as periodic solids or higher-accuracy functionals (e.g., hybrid or double-hybrid). We also acknowledge that our datasets consist of relatively small molecules. One could further test the scalability of our approach once new datasets on larger systems are available. We also note that our work does not give a theoretical explanation of how Hamiltonian prediction error translates to SCF acceleration. Developing a theoretically grounded algorithm specifically for SCF acceleration would be an interesting future research.

We mark the text in bold to highlight its relevance to your concerns. They will be removed from the final manuscript.

We hope our responses have addressed your questions and concerns. Please let us know if there is anything we may have missed or if any points require further clarification.

References

[1] Han, Xizewen, Huangjie Zheng, and Mingyuan Zhou. "Card: Classification and regression diffusion models." NIPS 2022

[2] Abramson, J., Adler, J., Dunger, J. et al. Accurate structure prediction of biomolecular interactions with AlphaFold 3. Nature 2024

[3] Jiao, Rui, et al. "Crystal structure prediction by joint equivariant diffusion." NIPS 2023.

[4] Chen, Shoufa, et al. "Diffusiondet: Diffusion model for object detection." ICCV 2023.

[5] Baranchuk, Dmitry, et al. "Label-efficient semantic segmentation with diffusion models." ICLR 2022.

[6] Kornilov, Nikita, et al. "Optimal flow matching: Learning straight trajectories in just one step." 2024 NIPS.

Dear Reviewer YiLn,

Thank you for your thoughtful and encouraging review. We appreciate your recognition of the clarity of our presentation, the strength of our experiments, and the relevance of our contribution within the existing literature. Below, we address your comments in detail.

W1, Q1: Are the improvements insightful or chemically meaningless?

Our improvements are chemically meaningful since they (1) reliably accelerate DFT calculations and (2) maintain chemical accuracy even for out-of-distributions.

First, since QHFlow accelerates DFT (as shown in Fig. 3), it reduces the overall computational cost in cases where DFT was previously required. Even for out-of-distribution data, the additional SCF iterations ensure that the solver reliably achieves DFT-level accuracy. This is particularly impactful for chemical simulations, such as ab initio molecular dynamics (AIMD), which heavily rely on DFT calculations and often encounter out-of-distribution inputs.

Furthermore, even without additional SCF iterations, QHFlow achieves the level of chemical accuracy (error < 1 kcal/mol ≈ 1,600 μEₕ [1,2]) for energy-related properties such as HOMO, LUMO, and the HOMO–LUMO gap, where baseline models fall short. Notably, as shown in Table 2, this level of accuracy is maintained even on unseen molecular configurations from molecular dynamics trajectories, specifically the QH9-dynamic (300k-mol) set.

W2, Q2: Lack of discussion on limitations in terms of scalability, generalizability, computational cost, or domain-specific applicability.

Thank you for the suggestion! We originally included the discussion in Appendix F.4, but now realize that its visibility is low despite its importance. To alleviate your concern, we will move the limitation to the main text and add more content to incorporate your comments.

To be specific, we will include the following text:

(Generalizability) We did not fully verify the generalizability of our approach, e.g., extension to new architectures (WANet, SPHNet) or datasets (PubChemQH), due to the lack of resources, public codes, and public datasets. (Computational cost) Our flow-based formulation introduces computational overhead for solving ODEs. In the future, one could consider removing the overhead with the one-step generative models [3]. (Domain-specific applicability) Our experiments are limited to gas-phase molecules with up to 29 atoms and two commonly used XC functionals (PBE and B3LYP). The algorithms are left to be verified for more general settings, such as periodic solids or higher-accuracy functionals (e.g., hybrid or double-hybrid). (Scalability) We also acknowledge that our datasets consist of relatively small molecules. One could further test the scalability of our approach once new datasets on larger systems are available. We also note that our work does not give a theoretical explanation of how Hamiltonian prediction error translates to SCF acceleration. Developing a theoretically grounded algorithm specifically for SCF acceleration would be an interesting future research.

We mark the keywords, e.g., generalizability and computational cost, as bold to highlight their relevance to your concerns. They will be removed from the final manuscript.

We hope our responses have addressed your questions and concerns. Please let us know if there is anything we may have missed or if any points require further clarification.

Reference

[1] Pople, John A. "Nobel lecture: Quantum chemical models." Reviews of Modern Physics 71.5 (1999): 1267.

[2] Kuang, Yang, Yedan Shen, and Guanghui Hu. "Towards chemical accuracy using a multi-mesh adaptive finite element method in all-electron density functional theory." Journal of Computational Physics 518 (2024): 113312.

[3] Kornilov, Nikita, et al. "Optimal flow matching: Learning straight trajectories in just one step." 2024 NIPS.

Thank you for your kind support of our work. We truly appreciate your constructive comments throughout the review process, which have helped us improve the overall quality of our work.