May the Force be with You: Unified Force-Centric Pre-Training for 3D Molecular Conformations

Recent works have shown the promise of learning pre-trained models for 3D molecular representation. However, existing pre-training models focus predominantly on equilibrium data and largely overlook off-equilibrium conformations. It is challenging to extend these methods to off-equilibrium data because their training objective relies on assumptions of conformations being the local energy minima. We address this gap by proposing a force-centric pretraining model for 3D molecular conformations covering both equilibrium and off-equilibrium data. For off-equilibrium data, our model learns directly from their atomic forces. For equilibrium data, we introduce zero-force regularization and forced-based denoising techniques to approximate near-equilibrium forces. We obtain a unified pre-trained model for 3D molecular representation with over 15 million diverse conformations. Experiments show that, with our pre-training objective, we increase forces accuracy by around 3 times compared to the un-pre-trained Equivariant Transformer model. By incorporating regularizations on equilibrium data, we solved the problem of unstable MD simulations in vanilla Equivariant Transformers, achieving state-of-the-art simulation performance with 2.45 times faster inference time than NequIP. As a powerful molecular encoder, our pre-trained model achieves on-par performance with state-of-the-art property prediction tasks.

Paper

References (32)

Scroll for more · 20 remaining

Similar papers

Peer review

Reviewer 1y4N7/10 · confidence 4/52023-07-06

Summary

1. proposing a pretraining framework for conformation space including equilibrium and off-equilibrium. 2. this method also shows highly accurate in MD or relax, including molecules and polymers. 3. providing the community with a diverse set of DFT dataset of polymers.

Strengths

1. This paper is well-written and well-organized, effectively presenting the significance of force-centric pretraining for molecular systems. The proposed framework demonstrates comprehensive and impressive results in molecular dynamics simulations of small molecules and polymers, as well as in quantum chemistry property predictions. 2. The authors also introduce an innovative concept in their work: the de-noising equilibrium balance and off-equilibrium co-training method.

Weaknesses

1. As previously mentioned, I am particularly intrigued by the efficiency of the de-noising equilibrium tasks. It would be better to conduct an detail ablation study of de-noise equillibrium and force regularization.

Questions

1. Would it be possible to incorporate the PCQM4Mv2 homo-lumo prediction results in the downstream analysis? This benchmark is well-defined and provides a competitive comparison for evaluating the performance of various models. 2. About code and data access, I believe that the Poly24 dataset is of significant value to the research community for future investigations.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

In the discussion section, it is mentioned that both the backbone and time efficiency play crucial roles in the performance of the proposed AI model. Besides I recommend incorporating partial charge prediction during the pretraining or downstream stages of your framework, if feasible. This addition could potentially enhance the overall effectiveness and applicability of your model.

Reviewer 8Ywg4/10 · confidence 4/52023-07-07

Summary

This paper incorporates off-equilibrium molecular conformations into consideration, in which the energy is not minimized. This paper defines pre-training objectives for force field learning on sampled off-equilibrium molecular conformations. Experimental results validate the effectiveness of pre-training on additional off-equilibrium data.

Strengths

(1) The presentation of the paper is quite clear and easy to follow. The writing is well-organized; (2) The proposed idea is quite interesting and necessary. How to utilize off-equilibrium data is a useful and important research question for pre-training 3D molecular property prediction. This reviewer agrees with the significance of the proposed idea.

Weaknesses

(1) The novelty of the proposed method is not very high. Considering there are already published papers researching denoised pretraining strategies on 3d molecules, which lowers the novelty of the proposed method a bit. Then the major contribution of the proposed method is mainly about introducing additional data. This reviewer suggests that this paper might be more suitable for a domain journal; (2) The contribution of this paper might not be enough. Currently, the major contribution includes introducing additional data, and additional pre-training objectives. There are some important works like SCN (Spherical channel networks for modeling atomic interactions) utilizing separate force predictions. So separately predicting forces cannot be regarded as an important contribution; (3) The experimental results do not demonstrate the effectiveness of the proposed method. It seems that the proposed method seldom outperforms previous baseline methods, which makes this reviewer feel a little bit confused.

Questions

It seems the reported results on MD17 are not that superior. Probably this reviewer misunderstands the experimental tables but it seems the sota method is usually not the proposed method. Does the proposed method outperform the previous baseline methods?

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

The major limitation of the proposed method is that the force prediction might only play a role as a regularizer since learning force field is also a very challenging task since force is a vector feature instead of a scalar feature like energy.

Authorsrebuttal2023-08-17

We would like to kindly remind the reviewer to read our response that we have made available for all reviewers in response to questions on the presentation of the paper. Thank you!

Authorsrebuttal2023-08-20

We would like to remind the reviewer to please take a look at our previous reply and comments, especially regarding the contribution of the paper and further experiments and explanations of the improved performance of ET-OREO.

Reviewer EQHq5/10 · confidence 5/52023-07-08

Summary

The paper propose a set of additional objectives for training machine learning interatomic potentials, in particular for equilibrum conformations as well as a denoising term. They show that this improves stability, when pretrained on a large dataset. They also introduce a new dataset of polymers.

Strengths

The improvement from ET-ORE to ET-OREO is interesting and surprising. It doesn't really make sense that equilibrum forces should help The polymer data set is great!

Weaknesses

Probably the biggest weakness is that the authors chose a weak baseline model with TorchMDNet. The entire paper hinges on work from Fu et al. that showed that in their experiments all methods except the NequIP potential were unstable in MD simulation. With many additional tricks and an enormous amount of data they now match NequIP. But will this entire improvement just go away when you apply the ideas to a stronger baseline model like NequIP? I would strongly suggest the authors to repeat the experiments with a better baseline and see if there are still improvements. Otherwise why wouldn't anyone just use NequIP to begin with? The speedup is small and as explained below mostly an artefact of hyperparameters. Finally, a central weakness it that ET-OREO is pretrained on a massive dataset while the baseline methods that are stable can do this from a much much much smaller dataset. Another big weakness is that this unexplained improvement from ET-ORE to ET-OREO simultaneously adds the equilibrum and the denoising objective. Why no ablation study to understand which of them is the important one? The inference speed comparison to NequIP is not thorough. NequIP has shallow pareto curves on accuracy vs efficiency which is well known, i.e. a L=2 models gets almost exactly the same accuracy as L=3, but is much faster (see original NequIP paper). The paper compares to a NequIP model that is solely optimized for accuracy, not for efficienc. This would be helpful to compare to. In addition, NequIP modern versions of NequIP now exists that are faster, e.g. Allegro or MACE. The presentation in addition has several obvious, outlined item-by-item below: - "they are locally well-defined for atoms and generalizable across various molecules" --> this statement is clearly not true in this generality. In what way is a force generalizable across molecules? Needs clarification what is meant here? - following immediately: "they reflect quantum mechanical interactions among atoms and expose the underlying quantum structures of molecules;" --> extremely generic, if even true, atoms should be electrons here... - "By harnessing these factors, we posit that learning from atomic forces can 57 guide the model to learn a comprehensive and unified landscape of molecular conformations across 58 different molecules, datasets, and equilibrium states" --> this has been explored for decades in every single paper training interatomic potentials, leveraging atomic forces is in now way a novel idea, this needs a strong rephrasing. - "We introduce a novel force-centric molecular conformation pretraining paradigm that trains 87 a unified conformation representation encompassing both equilibrium and off-equilibrium 88 molecules" this is a vast overstatement, see above - 'The above-mentioned papers followed the same learning paradigm, where different models are 116 trained for different molecules. Consequently, these methods cannot be easily generalized to unseen 117 molecules" --> there have been many efforts at building general-purpose force-fields, see e.g. the ANI series. - ". Despite their impressive accuracy in static prediction, their ability 119 of running simulations were not reported until [26] showed that the increased static accuracy does not 120 translate to better simulation performance" again, horrendously false statement, many of the MLIP papers you cite ran various simulations, across diverse and difficult phenoma from ionic diffusion to heterogeneous catalysis to small molecules. - "In molecular dynamics (MD) simulations, the molecular 149 conformations are moved according to forces and the thermostat of the simulation." what about NVE simulations that don't have a thermostat? - "For equilibrium molecular conformations x ∼ E , we assume their forces to be zero and impose a zero-force regularization:" --> what is the reasoning behind this, if you *have* the forces anyway, they will be zero at equilibrium, so you're just again training on forces, again what this community has been doing since the 60s. - the denoising objective seems like it has obvious failure cases: perturbing a relaxed structure and argueing the force must go into the direction of the GS completely ignores the many-bodyness of the PES, discussing these would be helpful - a lot of the work seems to be based on improving one or two specific previously introduce pretraining objectives for equilibrium data. In the opening of the manuscript, these are cited several times, but never explained. This makes the manuscript difficult to read.

Questions

N/A

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

2 fair

Presentation

2 fair

Contribution

3 good

Limitations

N/A

Authorsrebuttal2023-08-19

Dear Reviewer, We appreciate the time you've taken to review our work. We kindly direct your attention to the points we made in our reply. - Our contribution is to the pre-training framework, not the model architecture itself. ET-OREO has significantly improved the performance of ET, and as we have shown in the tables we submitted in the reply, outperforms NequIP by as much as 40%, with significantly reduced training time and inference time and greater versatility to finetune on different molecules. This should be enough for one to say that the pre-training framework of ET-OREO is effective. - Again, our contribution is the pre-trained framework that hinges on forces. This includes the joint training of forces on off-equilibrium data and the de-noising objective on equilibrium data as an approximation. We will make changes to our claims regarding training with forces, see the end of this reply. - It is also well-known that pre-trained models work better in larger datasets with deeper models and more parameters. Of course, NequIP is shown in the existing paper to be a strong baseline when for a small model data are limited, but neither is the case in the context of pre-training. We re-iterate that our main contribution is the pre-training framework, which in principle is independent of the model architecture of choice, and the choice of ET is based on the consideration of scalability, inference speed, and the size of model parameter space. The inference time of NequIP makes it more difficult for a pre-training task. - We are still training a NequIP model with our pre-training framework. It takes a significantly longer time, and we will share the results as soon as possible. In addition to the results we reported previously, upon further tuning ET-OREO to 2000 epochs, we recorded even better performance metrics on ET-OREO than previously shared: - Aspirin: 1.0 meV/A (52% improvement over NequIP's 2.1 meV/A) - Ethanol: 1.0 meV/A (28% improvement over NequIP's 1.4 meV/A) - Naphthalene: 0.9 meV/A (18% improvement over NequIP's 1.1 meV/A) - Salicylic Acid: 0.9 meV/A (10% improvement over NequIP's 1.0 meV/A) Thank you for your constructive feedback, and we are looking forward to your reply. ---------------------- Here's the revised version of the third paragraph in the Introduction where the reviewer expressed several concerns about the claim of the paper and language. We have clarified the concept of forces being generalizable and removed the use of language that suggests that we were the first to leverage forces. In this paper, we incorporate both off-equilibrium and equilibrium conformations into a unified representation learner. We propose a new pre-training model, "ET-OREO'" for 3D molecular conformations. This model integrates both equilibrium and off-equilibrium data from multiple large-scale datasets. Its training objective hinges on \emph{atomic forces}, defined as the negative gradient of a molecule's potential energy with respect to atomic coordinates. Atomic forces exhibit several notable properties: (1) they are \emph{physically well-defined} observable, i.e., the force acting on an atom is determined solely and uniquely from its local environment, defined as the real-space distribution of its neighboring atoms; (2) they are generalizable across various molecules in the sense that atoms from different molecules that have the same local environment should experience the same atomic forces; and (3) they can unify equilibrium and off-equilibrium data, as equilibrium data can be conceptualized as local minima in the latent (configuration) space with zero forces, while off-equilibrium data aid the model in more accurately characterizing the high-energy chemical space beyond equilibrium. Among these points, (2), which is the direct consequence of (1), makes atomic forces fundamentally different from potential energy which is defined for the whole system under consideration (molecules) and can only be determined up to an additional constant. Therefore, a predictive model of atomic forces is transferable, i.e., in principle, it can be used directly for molecules of any size (number of atoms). This advantage makes the "learning atomic forces"' fundamentally different from the traditional approach in which a fictitious concept of ``atomic energy'' must be defined, predicted, and combined to obtain the total potential energy of the whole system.

Reviewer EQHq2023-08-19

Thank you for the comments. I have adjusted my score. I believe there remain issues on how useful this is in practice and the weak baseline chosen, as outlined extensively above. This is a lot of pretraining for an improvement of 2.1 meV/A which likely wouldn't even be noticeable under DFT (the errors reported here are already way beyond DFT accuracy, no practitioner would ever care about this difference, let alone be willing to do the pretraining). As such I believe this method will see limited adoption, but may still be a valuable contribution to improve pretraining strategies and objectives.

Authorsrebuttal2023-08-20

Thank you for the comment and score adjustment. We agree that continuing to improve accuracy with respect to DFT has little incremental value as the latter is after all itself a computational approximation, which was why we put much emphasis on the relatively short time required to develop an accurate forcefield and versatility of ET-OREO being a pre-trained model. Further, we are trying to explore more applications of our method and better evaluation metrics. On Ethanol and Naphthalene, further finetuned ET-OREO can achieve lower h(r), signifying more similar bond distribution to that of a DFT simulation. We are also evaluating if our model can improve other more interesting applications other than force prediction accuracy. We have some initial positive results, and while we can't share the results in this paper in time, we are excited to include them in future work. Thank you.

Reviewer EQHq2023-08-18

I acknowledge that I have read the author's rebuttal. The authors have not provided new data or arguments that address the raised concerns, in particular about the weak baseline and the overstatement re using forces. In particular, "should ideally work for arbitrary deep learning models with sufficient expressivity, including NequIP." is not enough, this needs actual data, otherwise we're all just guessing. You're still comparing weak-baseline (ET) + your idea vs. strong baseline (NequIP) without your idea as far as I can tell? The score therefore remains the same.

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC