Learning 3D Anisotropic Noise Distributions Improves Molecular Force Fields

Thank you for your feedback. Your suggestions have enhanced the reliability of our work and the readability of the paper. Below are responses to your questions one by one.

Q1: 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.

Response: Thank you for pointing this out. We agree that Table 1 could be clearer in indicating the backbone models used. In our experiments, EquiformerV2 was used as the base model for all variants, including AniDS and the compared DeNS results. We will update the table caption and column headings in Table 1 to explicitly reflect this and avoid ambiguity.

Q2: Although energy and force prediction task is included, but dynamic simulation tasks (e.g., using predicted energies and forces to relax structures) are missing.

Response: We thank the reviewer for highlighting the importance of assessing model robustness in realistic atomistic simulation settings. We conducted a geometry optimization test using our model on a subset of the Aspirin molecule from MD17.

Specifically, we randomly selected 15 equilibrium structures and used our force predictor (EquiformerV2 + AniDS) to optimize the geometries. We then compared the optimized structures to those obtained using DFT geometry optimization and evaluated the RMSD between atomic positions.

Due to licensing restrictions with FHI-aims (used in the original dataset with the PBE-vdW-TS functional and tight settings), we performed DFT-based geometry optimization using PySCF at the PBE-D3BJ/def2-SVP level, which provides comparable accuracy in molecular geometries.

The results are summarized below:

The significant reduction in RMSD after optimization indicates that the learned force field can guide the system toward physically meaningful minima, showing promise for its application in downstream atomistic simulations.

We appreciate the reviewer’s suggestion and agree this direction is valuable for future extensions.

Q3: Analysis about additional computational cost of using AniDS beyond previous noise distributions is currently missing.

Response:

Thank you for raising this important point. To assess the additional computational cost introduced by AniDS, we conducted a comparison of training efficiency and model complexity across different noise modeling approaches using the rMD17 dataset (Aspirin molecule, 50 train data). The results are summarized below:

As shown, AniDS introduces a moderate increase in training time and model size, primarily due to the more expressive structure-aware noise generator that outputs full covariance matrices. However, this cost is justified by the significant improvement in performance (lowest MAE among all baselines).

Thank you for your feedback. Your suggestions have enhanced the reliability of our work and the readability of the paper. Below are responses to your questions one by one.

Q1: 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?

Response: Thank you for the insightful question. While the model does not explicitly encode prior knowledge about anisotropy, it learns reasonable anisotropic correction terms through supervised signals from the denoising objective, which is theoretically grounded in force field learning. This objective guides the model to assign higher noise along directions that are energetically less sensitive, and lower noise in directions where small perturbations lead to large energy changes. Over time, this enables the model to discover direction-specific stiffness patterns purely from data.

The key intuition is that our anisotropic noise generator is conditioned on molecular structure via a graph neural network encoder. As the model observes repeated patterns of atomic environments and their response to perturbations (e.g., rigid bonds vs. flexible chains), it learns to modulate the noise distribution accordingly. The learned anisotropic correction terms are not handcrafted but emerge from optimizing the denoising loss, which penalizes inaccurate reconstructions—particularly in directions where force gradients are steep.

Regarding the KL term: although it encourages the learned noise distribution to remain close to an isotropic Gaussian prior, its role is primarily regularization, not constraint. We carefully balance the KL loss with the denoising loss using a tunable coefficient (λ_KL), ensuring that the model retains flexibility to learn meaningful anisotropic components while avoiding degenerate or trivial solutions (e.g., collapsing to zero noise). Our ablation study (Table 3c) shows that setting λ_KL appropriately (e.g., λ_KL = 1.0) allows the model to learn informative, structure-aware anisotropic noise while maintaining stability during training.

Q2: 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?

Response: Thank you for this valuable suggestion. We agree that directly comparing the learned force field with DFT-computed reference forces provides a stronger validation of the model’s physical fidelity.

Following the approach in the SliDe paper, we conducted a same analysis between the predicted forces and DFT-calculated forces. Due to computational constraints, we randomly selected 1,000 structures from PCQM4Mv2 with fewer than 32 atoms, and computed their DFT ground-truth forces. We then pre-trained models using supervised force regression on these structures and fine-tuned them on MD17 using only 50 training samples per molecule.

The table below reports the Force MAE(kcal/mol/Å), comparing AniDS with DeNS and models trained from scratch:

Interestingly, we observe that pretraining with DFT-supervised labels does not lead to significant improvement over the unsupervised methods. This may be due to the fact that PCQM4Mv2 primarily consists of equilibrium structures, which provide limited force variation, or inconsistencies in the level of DFT theory used during force computation.

Nevertheless, even under such conditions and with a limited number of training samples, AniDS still achieves the best performance, demonstrating its strong capability to learn accurate force fields from small datasets.

Q3: Could you provide a comparison of training efficiency with existing baselines?

Response: Thank you for raising this important point. To assess the additional computational cost introduced by AniDS, we conducted a comparison of training efficiency and model complexity across different noise modeling approaches using the rMD17 dataset (Aspirin molecule, 50 train data). The results are summarized below:

As shown, AniDS introduces a moderate increase in training time and model size, primarily due to the more expressive structure-aware noise generator that outputs full covariance matrices. However, this cost is justified by the significant improvement in performance (lowest MAE among all baselines).

Thank you for the extensive efforts on reviewing our submission and the valuable suggestion. Here we present a detailed response to address your concern.

Due to space limitations, we do not include the full reviewer question here. Thank you.

Q1:

Response: Thanks for the insightful comment. While modeling the energy response to perturbations is a promising strategy, our work follows the research line of learning denoising force fields, which offers several advantages. First, denoising objectives are theoretically equivalent to learning the score function, which further corresponds to the negative force field under the Boltzmann assumption. This allows the model to explicitly approximate a pseudo harmonic force field. Second, according to this theoretical connection, we can specifically design the noise distribution to reflect properties of the underlying force field. In AniDS, we introduce structure-aware noise covariances to allow the model capture the anisotropic flexibility of the local structures, which is nontrivial to explicitly design for energy-response-based methods. Last, the denoising objective can be applied in purely unsupervised pretraining on equilibrium structures like PCQM4Mv2, without requiring energy or force labels, while energy-response-based methods typically rely on zeroth-order labels (i.e. energies) for training. We agree that energy-response-based methods offer an important perspective and we will extend a discussion of such approaches in the revised related work section.

Q2:

Response: We thank the reviewer for this insightful suggestion. To better evaluate the effectiveness of AniDS under more challenging conditions, we followed the protocol introduced in recent works such as NequIP, MACE, and ICTP, and conducted experiments on four subsets of the revised MD17 (rMD17) dataset. For each molecule, we used 50 training samples, following the ICTP setup. The results are summarized in the table below.

(Unit: meV/Å; values are force MAE. Full experimental details and hyperparameter settings will be included in the updated manuscript.“–” indicates results not completed due to time constraints.)

As shown above, AniDS achieves competitive performance on most datasets, and demonstrates notable improvements over the EquiformerV2 baseline trained from scratch. This confirms that the structure-aware anisotropic noise modeling introduced by AniDS leads to improved learning efficiency and accuracy.

Regarding the second point, we appreciate the interest in assessing the performance of AniDS when training from scratch. Our method is designed to learn structure-dependent, atom-wise noise distributions through pretraining on diverse molecular configurations. As such, training AniDS on a single molecule (or limited structural diversity) may restrict its ability to learn meaningful noise patterns, potentially diminishing its performance—a limitation that has also been observed in other VAE-based methods. Nevertheless, for completeness, we include a from-scratch variant of our model trained without pretraining in the ablation study (Table 3(b)), where it still provides notable improvements over non-pretrained baselines. These results further suggest that AniDS is flexible and effective, even without large-scale pretraining.

We hope these additions address the reviewer’s concerns and strengthen the empirical claims of our work.

Q3:

Response: Thank you for this insightful feedback. We agree that OC22, while still widely used, is beginning to show signs of performance saturation, and we appreciate the reviewer’s perspective on the limitations this may introduce in evaluating model improvements.

Regarding the unit of measurement: we reported energy errors in meV, consistent with previous works on OC22 such as the original benchmark paper[1] and recent baselines including E2Former[2] and DeNS[3], to ensure direct comparability.

Thank you again for pointing this out.

Reference: [1]Tran, Richard, et al. "The Open Catalyst 2022 (OC22) dataset and challenges for oxide electrocatalysts." ACS Catalysis 13.5 (2023): 3066-3084. [2]Li, Yunyang, et al. "E2Former: A Linear-time Efficient and Equivariant Transformer for Scalable Molecular Modeling." arXiv preprint arXiv:2501.19216 (2025). [3]Liao, Yi-Lun, et al. "Generalizing denoising to non-equilibrium structures improves equivariant force fields." arXiv preprint arXiv:2403.09549 (2024).

Q4:

Response: We thank the reviewer for the constructive suggestion. We agree that evaluating AniDS on smaller but challenging datasets with reduced model sizes would help demonstrate its practical utility and generalizability beyond OC22 and MD17.

Due to time and computational resource constraints, we conducted a preliminary evaluation on the AT–AT–CG–CG subset of the MD22 dataset, which is known for its chemical complexity and force field modeling difficulty. The results are summarized below:

(Unit: meV/Å; values are force MAE. Full experimental details and hyperparameter settings will be included in the updated manuscript.)

Despite the smaller scale of the dataset and model, AniDS continues to provide consistent improvements, confirming that the benefits of anisotropic, structure-aware noise modeling are not limited to large-scale settings.

We appreciate the reviewer’s helpful guidance, which has helped us broaden the scope of our evaluation.

Q5:

Response: We thank the reviewer for highlighting the importance of assessing model robustness in realistic atomistic simulation settings. We conducted a geometry optimization test using our model on a subset of the Aspirin molecule from MD17.

Specifically, we randomly selected 15 equilibrium structures and used our force predictor (EquiformerV2 + AniDS) to optimize the geometries. We then compared the optimized structures to those obtained using DFT geometry optimization and evaluated the RMSD between atomic positions.

Due to licensing restrictions with FHI-aims (used in the original dataset with the PBE-vdW-TS functional and tight settings), we performed DFT-based geometry optimization using PySCF at the PBE-D3BJ/def2-SVP level, which provides comparable accuracy in molecular geometries.

The results are summarized below:

The significant reduction in RMSD after optimization indicates that the learned force field can guide the system toward physically meaningful minima, showing promise for its application in downstream atomistic simulations.

Q6:

Response: We appreciate the reviewer’s thoughtful feedback. The case study in Section 4.4 focused on a few representative atoms in a periodic crystal to provide a clear and interpretable illustration of the relationship between learned covariance eigenvalues and energy sensitivity. However, we agree that presenting a more comprehensive, quantitative evaluation would strengthen our conclusions.

In the revised manuscript, we plan to include an additional analysis based on a single-molecule system, where the number of atoms is more manageable and not affected by periodic boundary conditions. Specifically, we will provide correlation plots between the eigenvalues of the learned covariance matrices and the measured energy sensitivity (e.g., sMAPE) along the corresponding eigenvectors, computed across all atoms in the molecule. This will allow us to more broadly validate the model’s ability to suppress noise in energy-sensitive directions.

Q7:

Response: We appreciate the reviewer’s question.We would like to emphasize that AniDS is designed to operate in both supervised and unsupervised settings. In particular, our model can be pre-trained purely on equilibrium structures without any force supervision, such as those found in PCQM4Mv2. In this setting, AniDS learns from coordinate perturbations alone using the denoising objective (Eq. 9), which is theoretically grounded in approximating the underlying force field via score matching.

This unsupervised pretraining can then be followed by fine-tuning on smaller force-labeled datasets (e.g., MD17), or directly applied to downstream tasks where only structural information is available. Moreover, AniDS can also serve as an auxiliary regularization objective when reference forces are partially available, helping improve model robustness and generalization.

We will clarify these use cases in the revised manuscript to better highlight the flexibility and broad applicability of AniDS across different data regimes.

We sincerely thank you for your valuable comments and the positive feedback! They have helped us to enhance the related work of our paper to better summarize the current status of the research. To address your concerns, we present the following point-by-point responses.

Q1: How to interpret the off-diagonal parts of the covariate matrix? How would they relate to atomic interactions?

Response: Explanation. Thank you for the thoughtful question. The off-diagonal components of the covariance matrix capture the correlations between different spatial directions in the noise distribution for each atom. Unlike diagonal matrices—which assume that noise in each direction (x, y, z) is independent—the off-diagonal entries encode directional dependencies, effectively modeling how motion in one axis correlates with motion in another.

Eigenvector Analysis. To better interpret these terms, we perform eigenvalue decomposition on the covariance matrix. The resulting eigenvectors indicate principal directions of noise, while the corresponding eigenvalues determine the magnitude of the noise along each direction. These principal directions are not necessarily aligned with the coordinate axes, reflecting the fact that atomic interactions are inherently anisotropic—forces such as bond stretching or angle bending exhibit direction-specific stiffness.

By learning full covariance matrices that include off-diagonal terms, our model can capture these directional preferences induced by local chemical environments and interatomic forces. This allows the learned noise distributions to reflect the underlying anisotropic energy landscape, ultimately leading to more accurate force field modeling. We believe this is one of the key advantages of our structure-aware anisotropic noise generator.

Q2: 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?

Response: We appreciate the reviewer’s question and the opportunity to clarify the interpretation of Figure 3c and Section 4.4. In Figure 3c, the three directions v1, v2, and v3 correspond to the eigenvectors of the learned covariance matrix for a given atom. For each direction, we apply a small perturbation and measure its impact on the total energy using sMAPE (Symmetric Mean Absolute Percentage Error).

A larger sMAPE indicates that a perturbation in that direction causes a greater change in energy, i.e., the system is more sensitive along that direction. Conversely, a smaller eigenvalue (i.e., less noise injected by our model in that direction) reflects the model’s decision to suppress noise along directions where perturbations are energetically costly. This behavior is consistent with our goal of aligning the noise distribution with the underlying energy landscape.

Regarding isotropy: if the eigenvalues across all three directions are similar, it suggests that the model considers the local environment approximately isotropic, and thus applies similar noise levels in all directions. In contrast, discrepancies among eigenvalues indicate anisotropic noise, where the model adapts the magnitude of noise based on the directional stiffness of the local atomic environment.

This coupling between energy sensitivity (sMAPE) and learned noise scale (eigenvalue) highlights the structure-aware nature of our anisotropic noise generator, which aims to respect physical constraints and improve force field learning.

Q3: Could the authors discuss briefly about how the method could be scaled to larger systems? Especially for pre-training and fine-tuning?

Response: Thank you for raising this important point. Our method is designed to be scalable to larger systems, both during pre-training and fine-tuning. The key reason is that our anisotropic noise generator and denoising autoencoder are both implemented as neural networks with local and structure-aware message passing, which naturally generalize to systems with more atoms.

In particular, the noise generator predicts per-atom full covariance matrices based on local structural contexts, allowing the computation to scale linearly with the number of atoms. This design ensures that even for large molecules or materials, the model can make predictions in a parallel and efficient manner. Furthermore, since our framework is model-agnostic, it can be paired with scalable backbones such as EquiformerV2 or other graph-based architectures tailored for large-scale molecular systems.

During pre-training, we can leverage existing large-scale datasets (e.g., PCQM4Mv2, OC22, or MPTrj) that already contain millions of structures. Fine-tuning on specific tasks or systems involves freezing the noise generator and only training the denoising autoencoder, which further reduces computational cost and improves scalability.

In future work, we also plan to explore hierarchical representations and domain decomposition to further enhance scalability for very large systems such as proteins or crystalline solids.