Scaling the Prior: Size-Consistent Geometric Diffusion for 3D Molecular Generation

Wenhan Gao, Jingxiang Qu, Yi Liu

Keywords: 3D Molecular Generation, Geometric Diffusion

Abstract: Diffusion models usually operate in fixed-dimensional metric spaces. In contrast, geometric molecular data naturally vary in dimensionality as molecules have different sizes (numbers of atoms). As a simple adaptation, existing diffusion models for geometric molecular generation employ network architectures that can handle variable-sized inputs, such as graph neural networks and transformers. **However, these approaches overlook the fact that the molecular size also determines the spatial scale of the atomic coordinates, which in turn induces inconsistent behaviors in the generative trajectories across different molecular sizes.** The generative process of geometric diffusion for 3D molecular generation can be viewed as first establishing a coarse structural target, followed by progressively refining the precise atomic positions. In particular, larger molecules tend to establish coarse structures earlier than smaller molecules due to their larger spatial scales relative to that of the noise. As a result, the reverse process becomes inconsistent across molecular sizes, with the denoising trajectories relying heavily on molecular sizes rather than on a unified generative pattern. In this work, we are the first to identify and analyze this size-induced inconsistency through a decomposition of the denoising dynamics, which reveals how spatial scale affects the progression of molecular formation, in both 3D structures and atom types. Building on this insight, we propose Scaling the Prior (StP), a simple yet effective approach that normalizes the learning and generative process across molecular sizes by rescaling the prior distribution based on molecular sizes. This adjustment harmonizes the denoising trajectories, enabling the model to learn a unified generative pattern and produce consistently high-quality molecules.

Primary Area: applications to physical sciences (physics, chemistry, biology, etc.)

Submission Number: 4182