Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Learning 3D Anisotropic Noise Distributions Improves Molecular Force Field Modeling

Xixian Liu, Rui Jiao, Zhiyuan Liu, Yurou Liu, Yang Liu, Ziheng Lu, Wenbing Huang, Yang Zhang, Yixin Cao

TL;DR

A novel denoising framework AniDS: Anisotropic Variational Autoencoder for 3D Molecular Denoising 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.

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.

Learning 3D Anisotropic Noise Distributions Improves Molecular Force Field Modeling

TL;DR

A novel denoising framework AniDS: Anisotropic Variational Autoencoder for 3D Molecular Denoising 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.

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.
Paper Structure (34 sections, 30 equations, 4 figures, 17 tables, 2 algorithms)

This paper contains 34 sections, 30 equations, 4 figures, 17 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Preliminary: Coordinate Denoising and DenoiseVAE
  3. Methodology
  4. Overview of the AniDS Framework
  5. Structure-Aware Anistropic Noise Generation
  6. AniDS Learns Molecular Force Fields
  7. Adapting AniDS to Different Training Schemes
  8. Experiments
  9. Results of pre-training and Fine-tuning
  10. Results of Supervised Learning with Partial Corruption and Auxiliary Denoising
  11. Ablation Study and Hyperparameter Analysis
  12. Case Study: Validating Anisotropic Noise Behavior on Crystal and Molecular Structures
  13. Related Works
  14. Conclusion and Future Work
  15. Acknowledgement
  16. ...and 19 more sections

Figures (4)

  • Figure 1: Comparison between different denoising approaches. The opaque spheres represent noise distributions. Our approach captures noise distribution that is both anisotropic and heteroscedastic.
  • Figure 2: Overview of the AniDS framework.
  • Figure 3: Visualization of the H$_3$In$_{12}$O$_{48}$Pd$_{12}$ crystal. We select oxygen atoms at indices $\{32,36\}$, and a palladium atom at index $21$. Figure (b) presents the eigenvectors of the oxygens. Figure (c) shows the relationship between the structural energy and applied noise.
  • Figure 4: Analysis on the SNPH$_4$ molecule. (a) Local geometry showing atom indices. (b) Per-atom directional alignment between predicted eigenvalues and ground-truth energy sensitivities (1/sMAPE). (c) Atom-wise comparison of averaged eigenvalues and energy sensitivities.

Theorems & Definitions (2)

  • proof
  • proof