Full Equivariance in Score-Based Generative Models Towards Molecules and Materials Inverse Design
Please login to view abstract download link
In this work, we present a score-based generative model for 3D atomic structures that enables high-quality, symmetry-aware generation of both molecules and periodic solids (i.e., materials). Our model overcomes key limitations of existing methods by using a fully E(3)-equivariant architecture that can operate on arbitrary irreducible representations, allowing it to capture rich geometric and physical relationships in data. Evidence from supervised learning indicates that increasing the body-order correlation of the features leads to improved data efficiency and lower error. Furthermore, unlike prior models that treat graph structure, node features, and positions as separate generation targets, our method jointly denoises all components, leading to coherent and physically realistic outputs. We implement our framework in JAX for performance and scalability, and evaluate it on benchmark datasets using standard metrics such as validity and uniqueness. Our approach builds upon the framework of score-based diffusion models, where training data are progressively corrupted through a noise process defined by a stochastic differential equation. A neural network is trained to estimate the gradient of the log-density of the data distribution – referred to as the score – which is then used to reverse the diffusion process and reconstruct samples from noise. Originally developed for image generation, this methodology has recently gained traction in domains such as molecular and materials design. When combined with E(3)-equivariant graph neural networks, score-based diffusion enables models to respect the symmetries of 3D space, which is essential for accurate and data-efficient modeling of atomic systems. Our methodology integrates geometric deep learning principles and paves the way for controlled, symmetry-aware generation of complex atomistic structures across chemistry and materials science.
