Title: Learning 3D Anisotropic Noise Distributions Improves Molecular Force Fields
Venue: NeurIPS 2025
Accepted: yes
Keywords: 3D molecular pre-training

molecular force field

coordinate denoising

equivariant neural networks
Abstract: Coordinate denoising has emerged as a promising method for 3D molecular pretraining due to its theoretical connection to learning molecular force field. However, existing denoising methods rely on oversimplied molecular dynamics that assume atomic motions to be isotropic and homoscedastic. To address these limitations, we propose a novel denoising framework AniDS: Anisotropic Variational Autoencoder for 3D Molecular Denoising. AniDS introduces a structure-aware anisotropic noise generator that can produce atom-specific, full covariance matrices for Gaussian noise distributions to better reflect directional and structural variability in molecular systems. These covariances are derived from pairwise atomic interactions as anisotropic corrections to an isotropic base. Our design ensures that the resulting covariance matrices are symmetric, positive semi-definite, and SO(3)-equivariant, while providing greater capacity to model complex molecular dynamics. Extensive experiments show that AniDS outperforms prior isotropic and homoscedastic denoising models and other leading methods on the MD17 and OC22 benchmarks, achieving average relative improvements of 8.9% and 6.2% in force prediction accuracy. Our case study on a crystal and molecule structure shows that AniDS adaptively suppresses noise along the bonding direction, consistent with physicochemical principles. Our code is available at https://github.com/ZeroKnighting/AniDS.


Peer reviews:
comments: SUMMARY:
This paper proposed a new denoising framework called AniDS for machine learning based molecular force fields. The proposed method addressed the limitations of previous noise injection design in terms of isotropic and homoscedastic. It is more physics informed than random noise. AniDS is claimed to be SO(3) equivariant, and the experimental results on MD17 and OC22 S2EF show that it is better than previous noise distributions.

STRENGTHS AND WEAKNESSES:
## Strengths:
1. This paper proposed a new, better, and physics informed noise distribution for MLFF pretraining.
2. The proposed noise distribution satisfies key physical restrictions including SO(3) equivariance. 
3. The experimental results show that the proposed AniDS is better than DeNS, and can be used to boost EquiformerV2's performance. 

## Weaknesses:
1. Clarity suggestions: AniDS is a general framework that can apply to many models, but the base model name is missing in table 1. It is not very clear what base models have been used by DeNS by looking at the table. 
2. Although energy and force prediction task is included, but dynamic simulation tasks (e.g., using predicted energies and forces to relax structures) are missing. 
3. Analysis about additional computational cost of using AniDS beyond previous noise distributions is currently missing.

QUESTIONS:
As listed above in Weaknesses.

LIMITATIONS:
Limitations are not provided in the main paper but in the appendix. This might not be a good practice.
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-06-25

comments: SUMMARY:
This paper proposes AniDS, a novel denoising-based training framework to assist molecular force field prediction. Building upon prior denoising frameworks, the method introduces atom-specific full covariance matrices and designs the neural network to ensure that the learned covariance matrices are symmetric, positive definite, and equivariant. The approach is evaluated on multiple benchmarks including MD17, OC22, and MPtrj.

STRENGTHS AND WEAKNESSES:
Strengths:

1. Denoising pre-training to enhance force field prediction is a hot direction, and refining noise design to make it consistent with physics is promising.

2. The method is sound, and the experimental results are comprehensive, clearly demonstrating the effectiveness of the proposed approach.

3. The paper is well written, and the figures are clear and informative.

Weaknesses:

More evidence is needed to justify that the learned anisotropic noise is physically meaningful (see Questions 1–2 below).

QUESTIONS:
1. Without of prior knowledge, why is the model able to learn reasonable anisotropic correction terms? Could the authors provide an intuitive explanation? Also, since the KL term aligns the learned noise with an isotropic prior, does it suppress the learning of anisotropic components?

2. While a case study is provided to illustrate the quality of the learned force field, I believe this is not sufficient. Could the authors compare the learned force field (i.e., the regression target in Eq. (9)) with DFT-computed force fields, for instance by computing correlation coefficients as in the SliDe paper?

3. Could you provide a comparison of training efficiency with existing baselines?

LIMITATIONS:
Yes
TITLE: Official Review
CLARITY: 4
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-06-28

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents AniDS, an anisotropic variational autoencoder for 3D molecular denoising, designed to overcome key limitations in previous denoising approaches, namely the assumptions of isotropic atomic motion and a uniform noise scale across all atoms. AniDS learns atom-specific, full covariance matrices for Gaussian noise distributions, incorporating geometric priors to capture directional dependencies in atomic perturbations. The authors establish a connection between their denoising approach and learning molecular force fields. The method is evaluated using multiple datasets, including MD17 (with PCQM4Mv2 pre-training), OC22, and MPTrj, showing improved accuracy in force prediction compared to baseline methods.\n\nSTRENGTHS AND WEAKNESSES:\nThis paper has several strengths:\n\n1. The proposed method introduces structure-aware noise generation and learns per-atom full covariance matrices, which is a non-trivial extension of previous work.\n2. The connection to force field learning is established and well-motivated.\n3. The analysis comparing eigenvalue magnitudes of learned covariance matrices with energy changes along corresponding perturbations provides valuable insight.\n\nHowever, before acceptance, several weaknesses should be addressed:\n\n1. The manuscript lacks a clear positioning of the proposed approach within the broader scope of related work. For example, a related strategy is to estimate energy response to perturbations using reference or predicted forces (e.g., https://www.nature.com/articles/s41524-020-0323-8, https://doi.org/10.48550/arXiv.2408.05215). The authors mayconsider discussing whether learning the perturbation itself is more advantageous than modeling the energy response to a random perturbation.\n2. The accuracy of the resulting models is primarily evaluated on the MD17 and OC22 datasets. While MD17 is widely used, state-of-the-art models typically achieve very low errors on this dataset, which makes any comparison less meaningful. Therefore, recent works such as NequIP, MACE, and ICTP have used a subset of 50 samples drawn from rMD17 to benchmark accuracy under more challenging conditions (i.e., https://www.nature.com/articles/s41467-022-29939-5, https://doi.org/10.48550/arXiv.2206.07697, and https://doi.org/10.48550/arXiv.2405.14253). To make stronger claims about the potential improvements offered by AniDS, the authors may consider evaluating it on MD17 (or rMD17) in a similar way. Furthermore, the authors may include results for the baseline model trained with AniDS from scratch, i.e., without pre-training.\n3. The OC22 dataset is an important benchmark; however, it appears to be approaching a saturation point similar to that of QM9, which may reduce its reliability for assessing model performance. Setting this aside, the reported improvements in force predictions are rather modest, reflecting limitations similar to those discussed regarding the MD17 dataset. Furthermore, the energy errors are reported in meV, complicating their assessment. The authors may consider presenting errors in meV/atom instead, as this aligns better with the established standards in computational materials science and enables a more straightforward assessment of machine-learned force field accuracy.\n4. The improvements over DeNS are generally modest. Given the increased model complexity (i.e., 129M vs. 127M parameters), it is unclear whether these gains are practically significant. The authors might consider applying their approach to smaller but challenging datasets (e.g., the HEA dataset in https://doi.org/10.48550/arXiv.2405.14253, but any other dataset can be used instead) using smaller models. For molecular systems, applying the approach to ANI datasets or MD22 (again, reducing the training set size, if necessary) could provide further insights.\n5. Since the primary goal of improving machine-learned force fields is their robust application in atomistic simulations, the authors should consider including results that demonstrate this, for example, by performing molecular dynamics simulations in an NVT ensemble and evaluating long-term energy conservation.\n6. The comparison of eigenvalue magnitudes of learned covariance matrices with energy changes is currently limited to a few selected atoms. Including more general correlation plots would better demonstrate the broader validity of these results.\n7. Finally, the manuscript does not address the applicability of the proposed approach to tasks where reference forces are unavailable. Discussion on this point would clarify the method’s scope and potential limitations.\n\nQUESTIONS:\nMy questions would relate to the issues raised above, so addressing them would suffice to change my evaluation of the paper.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work proposes a novel anisotropic denoising framework with a structure-aware generator for molecular force field learning. The framework contains a structure-aware noise generator that predicts a covariance matrix from the structure, a denoising autoencoder and a covariance-scaled reconstruction loss. A KL divergence regularizer is also implemented to avoid trivial solution. The noise generator is designed to control the anisotropic level based on the atomic context, and to ensure the covariance properties. The author proves the objective approximates the force field.\n\nThe framework is tested on several datasets and tasks, including MD17 force prediction and energy prediction in OC22, and demonstrated better performance than the benchmarks. Ablation study shows the advantage of pre-training on PCQM4Mv2 and the slight advantage of using subtractive correction than additive. The eigenvector of the learned covariance matrix is shown to correspond to the bond direction.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1. The work has a sound theoretical foundation. Relaxing the homoscedastic and isotropic constraints could potentially enhance the model capacity, and the design of the proposed noise generator captures the anisotropic property of atomic-specific and context-dependent dynamics. \n2. The authors provide comprehensive results on several datasets and tasks, ablation study, and case study and visualization, which offer strong evidences for the superiority of the proposed method. \n\nWeaknesses:\nI have only minor questions about the results. See Questions\n\nQUESTIONS:\n1. How to interpret the off-diagonal parts of the covariate matrix? How would they relate to atomic interactions?\n2. The relationship between the numbers in Fig 3c and the conclusions in 4.4 could be more clearly explained. Does small sMAPE (and large eigenvalue) mean the perturbation has less impact on the energy, and similar values on v1~3 mean more isotropic noise?\n3. Could the authors discuss briefly about how the method could be scaled to larger systems? Especially for pre-training and fine-tuning?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}