DEQuify your force field: Towards efficient simulations using deep equilibrium models

Authors: Andreas Burger, Luca Thiede, Alan Aspuru-Guzik, Nandita Vijaykumar
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: Machine Learning Force Fields, Deep Equilibrium Models

Abstract:
Machine learning force fields show great promise in enabling more accurate force fields than manually derived ones for molecular dynamics simulations. 
State-of-the-art approaches for ML force fields stack many equivariant graph neural network layers, resulting in long inference times and high memory costs. This work aims to improve these two aspects while simultaneously reaching higher accuracy.
Our key observation is that successive states in molecular dynamics simulations are extremely similar, but typical architectures treat each step independently, disregarding this information.
We show how deep equilibrium models (DEQs) can exploit this temporal correlation by recycling neural network features from previous time steps. 
Specifically, we turn a state-of-the-art force field architecture into a DEQ, enabling us to improve both accuracy and speed by $10\%-20\%$ on the MD17, MD22, and OC20 200k datasets. 
Compared to conventional approaches, DEQs are also naturally more memory efficient, facilitating the training of more expressive models on larger systems given limited GPU memory resources.

Decision: reject
Reviewer scores: 5, 6, 6, 3

Reviews and discussion:
<thread>
creation date:
1730099634392

thread_type:
official_review

<thread_object>
creator: reviewer_edd9

<content>
summary:
This work utilizes deep equilibrium models to enhance the accuracy, speed, and memory efficiency of EquiformerV2. Furthermore, DEQuiformer could accelerate molecular dynamics simulations by leveraging the similarity between successive states.

strengths:
* This paper proposes incorporating DEQ into MLFF.
* Compared to EquiformerV2, DEQuiformer shows advantages.
* The paper demonstrates experiments on the rapid convergence of DEQuiformer's fixed points.

weaknesses:
The related issues will be elaborated in the "questions" section.

questions:
* Lack of Innovation: Although this paper is, to my knowledge, the first to combine DEQ with MLFF, the concept of DEQ is based on the fixed-point property of neural network hidden states. This work seems to merely change the application scenario to MLFF without actual innovation.
* In the introduction, the paper misses a lot of related works in the field, such as [1]-[6].
* The paper highlights that in molecular dynamics simulations, because consecutive frames have similar configurations, the fixed point from the previous frame can be used as the initial trial for the next frame, thus accelerating the process. My concerns are as follows:
1. The experiments provided to support this point are insufficient. For example, it would be useful to compare the number of iterations with and without reuse. Additionally, molecular dynamic simulations should be performed to claim your points.
2. The paper mentions using a validation set of 50 samples from MD17, but Fig. 4b states it uses the validation set. Were these 50 samples selected in consecutive order? Additionally, the paper mentions selecting an extra 1000 consecutive samples. If these were used, they should not be considered part of the validation set. The authors should clarify to avoid reader confusion.
3. Fig. 4b is confusing. What is the percentage relative to? The figure needs a clearer explanation.
4. Is the acceleration effect of DEQuiformer primarily due to reuse or the reduced number of model layers? I suspect the latter is the main factor.
* In conventional MLFF experiments, the MD17 test set includes all remaining data, while rmd17 provides a specific test set. However, this paper uses only 1000 data points from the MD17 test set, which is not a fair comparison.
* The paper claims that the data split for MD22 is the same as for MD17, which is inconsistent with the original MD22 paper [7] and conventional MLFF splits[5,6,8,9].

* The paper lacks results from other benchmark models for MD17 and MD22. Benchmarking with additional methods could provide a better evaluation of the results.

* Why aren’t the energy and force tables for MD17 and MD22 combined? This makes it inconvenient for readers to compare the results. Additionally, the energy results for EquiformerV2 are exceedingly poor. Was the training fully converged? This could imply a bias in lowering the benchmark, making the comparison unfair.

* I greatly appreciate that the authors included a limitations section, but I have some questions regarding certain points mentioned. The limitations section mentions uncertainty about whether DEQ is applicable to other MLFFs because EquiformerV2 uses a separate force output head. I don’t quite understand the logic here. Is DEQuiformer based on the assumption that forces have a fixed-point property? It seems to me that DEQuiformer assumes the hidden layers of molecular representations have a fixed-point property, so DEQ should theoretically be applicable to any model predicting any property, even if forces are obtained via autograd.

* The presentation of the paper could be improved. For readers not very familiar with the field, more introductory content on concepts like equivariance and Equiformer should be added in the preliminary section.

[1] Schütt, Kristof T., et al. "Schnet–a deep learning architecture for molecules and materials." The Journal of Chemical Physics 148.24 (2018).

[2] Gasteiger, Johannes, Janek Groß, and Stephan Günnemann. "Directional Message Passing for Molecular Graphs." International Conference on Learning Representations.

[3] Schütt, Kristof, Oliver Unke, and Michael Gastegger. "Equivariant message passing for the prediction of tensorial properties and molecular spectra." International Conference on Machine Learning. PMLR, 2021.

[4] Coors, Benjamin, Alexandru Paul Condurache, and Andreas Geiger. "Spherenet: Learning spherical representations for detection and classification in omnidirectional images." Proceedings of the European conference on computer vision (ECCV). 2018.

[5] Wang, Yusong, et al. "Enhancing geometric representations for molecules with equivariant vector-scalar interactive message passing." Nature Communications 15.1 (2024): 313.

[6] Wang, Zun, et al. "Efficiently incorporating quintuple interactions into geometric deep learning force fields." Advances in Neural Information Processing Systems 36 (2024).

[7] Chmiela, Stefan, et al. "Accurate global machine learning force fields for molecules with hundreds of atoms." Science Advances 9.2 (2023): eadf0873.

[8] Kovács, Dávid Péter, et al. "Evaluation of the MACE force field architecture: From medicinal chemistry to materials science." The Journal of Chemical Physics 159.4 (2023).

[9] Li, Yunyang, et al. "Long-Short-Range Message-Passing: A Physics-Informed Framework to Capture Non-Local Interaction for Scalable Molecular Dynamics Simulation." The Twelfth International Conference on Learning Representations.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730640792702

thread_type:
official_review

<thread_object>
creator: reviewer_emzx

<content>
summary:
This work develops and demonstrates a strategy for converting a deep network to deep equilibrium models for the task of machine learning inter-atomic potentials. The strategy is interesting and novel compared to other strategies in that area, and leads to improvements in both the accuracy and inference speed. By re-using solutions in the fixed point solver in subsequent time steps in a simulation, the inference speed is further improved. The results are demonstrated on a few diverse datasets including MD17, MD22, and OC20.

strengths:
* The strategy is unlike other current approaches to improve ML for molecular simulations using typical message passing systems.
* The strategy works for both molecules and catalysts
* The performance (accuracy and inference speed) can be improved over a current leading model (EquiformerV2), though the EqV2 baselines were not thoroughly tuned for every experiment
* The authors considered possible edge cases and questions about the impact of the fixed point convergence settings on the final performance. 

The authors focus on the benefits of DEQ for accelerating MD simulations by re-using fixed point guesses, but I think there are even more important and influential directions that this paper will unlock. 
1. The underlying simulations for the training data are themselves ground state calculations that come from a self-consistent solver for the electronic structure, and there is a strong analogy between the methods in this work and the methods in the underlying solver. 
2. A subset of the community has been focusing on computing additional equilibrium properties in a system such as charge distribution, and those methods typically work by using standard message passing GNNs to compute an electronegativity and then solving for an equilibrium distribution of charge (e.g. Wai Ko, Behler et al. Nat Com 2021, or Deng, Cedar et al. Nat. Mat. Int. 2023 among many others). I think there is huge potential (pun intended) for this work to lead to large increases in performance in predicting these other ground state problems that require equilibrium across extended systems. Further, there’s probably connections between this work and other methods to accelerate the convergence of the underlying simulations themselves, perhaps by predicting the electron density itself.

weaknesses:
* As the authors point out, computing gradients of the resulting properties w.r.t. the atomic positions is difficult (but not impossible) and by not addressing this only direct-force models can be improved currently. 
* The baseline comparisons use model architectures tuned for larger systems. It would be preferable to compare results for the model here on published baseline models for the original (larger) datasets like OC20-2M so that it is clear the results are not simply due to better tuning of the DEQ-architecture compared to the baseline EqV2 architectures
* The authors state that fixed-point re-use cannot be tested for OC20, but it is not clear to me why. Specifically, OC20 has an MD subset that could be used. Further, the authors could simply run a short-timescale MD simulation using the potential (say in ASE) to compare the inference time savings. 
* Large variations in Table 1 are a bit suspicious (see question below) and decrease confidence in the conclusions for MD22.

questions:
1. In Table 1 MD22, the force MAEs seem highly variable and perhaps highlight convergence/training stochasticity rather than intrinsic model differences. For example, in MD22/Stachyose, all EqV2 results are ~11, as well as DEQ (2-layer), but DEQ (1-layer is 0.31, 30X smaller). Similarly Ac-Ala3-NHME has a 10X larger force MAE for EqV2(1-layer). Any AT-AT/DHA have ~10-20X smaller force MAEs for DEQ than EqV2. Can the authors confirm the results in this table and/or improve the training consistency? I understand one of the points of DEQ is better training dynamics, but the Stachyose results suggest even with DEQ there’s significant stochasticity similar to the baseline.

2. How stochastic is the fixed point solver inside of DEQ2? Specifically, what is the likelihood that you find different solutions for the same inputs? This is probably especially important for molecular dynamics, as multiple possible solutions for the local potential energy surface could lead to interesting dynamics, or possibly history-dependent artifacts. 

3. The authors might want to consider training a model to predict the partial charges or magnetic moments in the mptrj dataset (similar to CHGNet) to greatly improve the applicability of this work. Obviously this is a significant additional experiment, but if these results were competitive there too I think this work would be much more compelling.

4. Experiments for inference-time speedup for simulations in OC20 are also possible, can the authors shed insight on whether the same speedups are seen there too?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730731148185

thread_type:
official_review

<thread_object>
creator: reviewer_yjg1

<content>
summary:
This paper applies Deep Equilibrium Networks (DEQ) [1] to EquiformerV2 [2], leveraging DEQ’s advantages of significantly reduced memory overhead and faster inference. The authors further enhance DEQ’s fixed-point solution by integrating temporal correlation from MD datasets as an initialization strategy, which would otherwise be computationally prohibitive. To validate their approach, they compare the performance of various EquiformerV2 and DEQ variants on the MD22, MD17, and OC20 benchmarks.

References:

[1]https://arxiv.org/abs/1909.01377

[2]https://arxiv.org/abs/2306.12059

strengths:
1. Equivariant networks based on Transformer architectures are highly computationally intensive, requiring high-end GPUs and often taking multiple days to train. Any efforts to reduce this computational burden are valuable.

2. I commend the authors for applying the Deep Equilibrium Model (DEQ) approach to EquiformerV2. Integrating a fixed-point solver into the Transformer architecture—especially by adapting from TorchDEQ[1] and leveraging temporal correlation—is an impressive accomplishment.

References:

[1] https://github.com/locuslab/torchdeq

weaknesses:
1. Temporal correlation in molecular dynamics (MD) datasets has been used effectively to initialize fixed-point solvers. I like this idea, but it relies on an assumption of ordered data, which breaks down during phase transitions. For instance, in the simple case of water melting, the transformation from solid to liquid occurs abruptly, disrupting the expected order and causing this approach to fail at the transition point.


2. This approach also struggles with datasets that lack a clear temporal order, where entries are jumbled and any sense of sequence is lost (as noted by the authors for datasets like rMD17 and OC20). This limitation restricts the applicability of the method, impacting one of the core contributions of this work. Can you please discuss potential ways to address this limitation or expand on other potential applications where temporal ordering is preserved?

questions:
Please also see weakness.

1. To clarify what part was contributed in this paper and taken from[1] can the authors write a pseudo code and highlight the lines? For instance: side-by-side comparison of the original DEQ algorithm and the authors' modified version for equivariant networks.

2.  In Fig 2 why at 0th iteration(for test error) the models start at such different errors? Happens to be that the model starting with lower initial error saturates at lower overall error at end of epochs. I am unclear if this performance gain is due to less parameters (less overfit)in DEQ or due to the time correlation introduced in this paper? Could you please compare DEQ models with and without the temporal correlation initialization, while keeping the parameter count constant?

3. I suggest to add a table with number of parameters across all models.

references:

[1]https://arxiv.org/abs/1909.01377
</content>
</thread_object>

</thread>
<thread>
creation date:
1730941834145

thread_type:
official_review

<thread_object>
creator: reviewer_isqw

<content>
summary:
The authors follow techniques proposed in [1], [2], and [3] to convert EquiformerV2 into a Deep Equilibrium Model (DEQuiformer), leading to improvements in accuracy and inference speed on MD17, MD22, and OC20 200k datasets. The benefits of DEQuiformer are primarily derived from the ability to re-use previously inferenced fixed points as initializations for the next frame's inference.

strengths:
- Can trade off accuracy for simulation speed post-training very easily by modulating the fixed-point error threshold on the solver
- The model outperforms EquiformerV2 on training time
- The model outperforms EquiformerV2 on MD17, MD22, and OC20 200K in accuracy and inference time
- The authors empirically show that DEQuiformer successfully and stably converges to a fixed point during inference

weaknesses:
- Current work only applies to direct force prediction models
- No evidence that the technique could work on different model architectures
- No current comparisons with the inference times and accuracies of other model architectures

questions:
Questions
- On line 278: "Using consecutive samples would have the downside of a large variance in the results depending on the starting index, while 1000 uniformly spaced samples yield similar results to expensively testing on all > 100,000 data points." I'm not sure if I understand this. Depending on the starting index, the accuracy will have high variance compared to the accuracy on testing on all 100,000 data points? Given a fixed datapoint, does this mean the accuracy of the model is different depending on that datapoint's position in the overall sampled trajectory? I'm not sure if this should be the case.

- Do shock simulations where consecutive frames lead to vastly different atomic arrangements lead to different performance/timing results? Are there guarantees that fixed-point reuse leads to less solver steps? 


Limitations
- This is spoken in the limitations section, but I believe a very big drawback of the method is that it only applies to direct force prediction models. Real world molecular dynamics simulation necessitates forces to be formulated as gradients of energies to maintain important physical properties regarding the conservative nature of the force field. Although it's possible to run molecular dynamics simulation using a direct force prediction force field, it's still not clear whether this is trustworthy and safe to do so. Relaxations, on the other hand, don't as rigorously require conservative force fields but also don't require "millions to billions" of timesteps. Molecular dynamics is a fantastic beneficiary of the proposed speedup of inference, but the method isn't applicable at the moment.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732041712199

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We thank the reviewers for the constructive feedback and their questions. Some points were brought up by several reviewers, and we conducted additional experiments to address them: 
1. We ran MD simulations on OC20 to measure the model’s solver step and runtime improvement due to fixed point reuse. DEQuiformer with 1-layer goes from 29 iterations without fixed point reuse to 11 with fixed point reuse while significantly beating a 14-layer EquiformerV2 in accuracy, the biggest we were able to fit in memory.
2. We measured the difference between force prediction with and without fixed point reuse on the full test set of Aspirin and OC20 MD simulations to show that the predictions are still Markovian: The average relative difference is less than 1% for both Aspirin, and the MD simulations.
3. We also trained the OC20 models three times longer and MD17 Aspirin twice as long and with up to ten times more parameters to address convergence concerns: While all models got better, the performance difference in favor of DEQuiformer remained.

We will add these experiments with detailed explanations to the paper.
</content>
</thread_object>

</thread>


Metareview:
metareview:
In this work, authors propose DEQuiformer, which applies deep equilibrium models (DEQ) to machine learning force fields by modifying the EquiformerV2 architecture. The key innovation is leveraging temporal continuity in molecular dynamics (MD) simulations to accelerate convergence by reusing fixed points between timesteps. The work demonstrates that DEQuiformer achieves both higher accuracy and faster inference compared to EquiformerV2 on MD17, MD22, and OC20 datasets. Specifically, it shows 10x improvements while using fewer parameters. The authors validate that fixed point reuse reduces solver iterations from 29 to 11 steps on OC20 MD simulations while maintaining prediction accuracy within 1% relative difference.

The paper demonstrates several strengths as follows. The implementation achieves measurable improvements in both accuracy and computational efficiency for MD simulations, supported by comprehensive experiments added in response to reviewer feedback. The work integrates multiple existing techniques, including deep equilibrium models and optical flow estimation, while introducing necessary architectural modifications such as input injection normalization to optimize performance. A key innovation lies in leveraging temporal continuity in MD simulations, which Reviewer emZx particularly highlighted as having broader potential impact for modeling equilibrium properties in materials science. The authors' thorough response to reviewer concerns, including additional validation through MD simulations on OC20 and extensive convergence studies with scaled-up models and extended training periods, strengthens the paper's empirical foundation.

However, the work has several limitations that warrant consideration. As Reviewer edd9 points out, the current implementation is restricted to direct force prediction rather than energy gradient-based approaches, though the authors note this aligns with recent high-performing models in the field. However, it is well-known that this results in non-conservative forces. It raises questions about the stability of MD simulations and also the roughness of the underlying forcefield. More importantly, the consequence of the direct force prediction on basic physical laws such as energy and momentum conservation, an important aspect for MD simulations, remain questionable. While there are several direct force prediction models, most of them are applied for energy minimization or identifying stable structures rather than MD simulations. Additionally, the baseline comparisons rely on reduced datasets instead of full-scale training due to computational constraints, although scaling experiments indicate the relative performance improvements persist. The initial submission contained some inconsistencies in the MD22 evaluation results, though these were subsequently corrected during revision. These weaknesses, while notable, do not fundamentally undermine the paper's primary contributions and demonstrated practical utility.

Additional points are included below. Based on these points, the manuscript is not recommended for acceptance. For resubmission, the authors should focus on establishing theoretical foundations for DEQ in force fields, expanding to energy gradient-based approaches, and conducting thorough comparisons on standard benchmarks with complete training. This would help establish the approach's scientific merit beyond engineering optimization.