Learning 3D Anisotropic Noise Distributions Improves Molecular Force Field Modeling

TL;DR

AniDS introduces an anisotropic noise modeling framework for 3D molecular denoising, significantly improving force field predictions by capturing directional atomic interactions and structural variability.

Contribution

This work presents AniDS, a novel anisotropic variational autoencoder that models atom-specific, full covariance Gaussian noise for better molecular dynamics representation.

Findings

AniDS outperforms prior models on MD17 and OC22 benchmarks.

Achieves 8.9% and 6.2% improvements in force prediction accuracy.

Effectively suppresses noise along bonding directions, aligning with physical principles.

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.