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

Wenhan Gao, Jingxiang Qu, Yi Liu
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:33460-33482, 2026.

Abstract

Diffusion models typically operate in fixed-dimensional metric spaces, whereas 3D geometric molecular data vary in dimensionality because molecules differ in size (number of atoms). A common adaptation in diffusion models for 3D molecular generation is to use models that handle variable-sized inputs, such as graph neural networks and transformers. However, these approaches ignore that molecular size also sets the spatial scale of atomic coordinates, causing inconsistent generative trajectories. In 3D molecular diffusion, generation can be seen as forming a coarse structure first and then refining atomic positions. Larger molecules form coarse structures earlier than smaller ones because their spatial scales are larger relative to the noise. This makes the generative process inconsistent across sizes, with trajectories driven by molecular size rather than by a unified generative pattern. We are the first to identify and analyze this size-induced inconsistency by decomposing denoising dynamics, showing how spatial scale shapes formation of both 3D structure and atom types. Based on this, we propose Scaling the Prior (StP), which rescales the prior distribution by molecular size to normalize learning across sizes, harmonize denoising trajectories, and generate high-quality molecules. The code is available at https://github.com/wenhangao21/ICML26-StP.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-gao26v, title = {Scaling the Prior: Size-Consistent Geometric Diffusion for 3{D} Molecular Generation}, author = {Gao, Wenhan and Qu, Jingxiang and Liu, Yi}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {33460--33482}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/gao26v/gao26v.pdf}, url = {https://proceedings.mlr.press/v306/gao26v.html}, abstract = {Diffusion models typically operate in fixed-dimensional metric spaces, whereas 3D geometric molecular data vary in dimensionality because molecules differ in size (number of atoms). A common adaptation in diffusion models for 3D molecular generation is to use models that handle variable-sized inputs, such as graph neural networks and transformers. However, these approaches ignore that molecular size also sets the spatial scale of atomic coordinates, causing inconsistent generative trajectories. In 3D molecular diffusion, generation can be seen as forming a coarse structure first and then refining atomic positions. Larger molecules form coarse structures earlier than smaller ones because their spatial scales are larger relative to the noise. This makes the generative process inconsistent across sizes, with trajectories driven by molecular size rather than by a unified generative pattern. We are the first to identify and analyze this size-induced inconsistency by decomposing denoising dynamics, showing how spatial scale shapes formation of both 3D structure and atom types. Based on this, we propose Scaling the Prior (StP), which rescales the prior distribution by molecular size to normalize learning across sizes, harmonize denoising trajectories, and generate high-quality molecules. The code is available at https://github.com/wenhangao21/ICML26-StP.} }
Endnote
%0 Conference Paper %T Scaling the Prior: Size-Consistent Geometric Diffusion for 3D Molecular Generation %A Wenhan Gao %A Jingxiang Qu %A Yi Liu %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-gao26v %I PMLR %P 33460--33482 %U https://proceedings.mlr.press/v306/gao26v.html %V 306 %X Diffusion models typically operate in fixed-dimensional metric spaces, whereas 3D geometric molecular data vary in dimensionality because molecules differ in size (number of atoms). A common adaptation in diffusion models for 3D molecular generation is to use models that handle variable-sized inputs, such as graph neural networks and transformers. However, these approaches ignore that molecular size also sets the spatial scale of atomic coordinates, causing inconsistent generative trajectories. In 3D molecular diffusion, generation can be seen as forming a coarse structure first and then refining atomic positions. Larger molecules form coarse structures earlier than smaller ones because their spatial scales are larger relative to the noise. This makes the generative process inconsistent across sizes, with trajectories driven by molecular size rather than by a unified generative pattern. We are the first to identify and analyze this size-induced inconsistency by decomposing denoising dynamics, showing how spatial scale shapes formation of both 3D structure and atom types. Based on this, we propose Scaling the Prior (StP), which rescales the prior distribution by molecular size to normalize learning across sizes, harmonize denoising trajectories, and generate high-quality molecules. The code is available at https://github.com/wenhangao21/ICML26-StP.
APA
Gao, W., Qu, J. & Liu, Y.. (2026). Scaling the Prior: Size-Consistent Geometric Diffusion for 3D Molecular Generation. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:33460-33482 Available from https://proceedings.mlr.press/v306/gao26v.html.

Related Material