Multigrid Training for Molecular Generation using Graph Neural Networks
This work addresses the computational cost and instability of high-resolution molecular modeling for practitioners in drug discovery and computational chemistry.
The authors propose a multigrid training strategy for molecular generation using graph neural networks and 3D grids, which transfers parameters from coarse to fine resolutions via biased random walk upsampling or convolutional parameter transfer. On receptor-conditioned 3D ligand generation, this method accelerates convergence and improves generalization compared to training from scratch.
Deep learning has demonstrated significant success for modeling biochemical molecular systems, where inputs are commonly represented as graphs or 3D grids. A major challenge is that computational cost scales with resolution, making full graph/grid computation of molecular densities expensive and often unstable. We introduce a multigrid training strategy that leverages low-resolution optimization to accelerate learning at higher resolution through parameter transfer across discretizations. For graph molecular representations, we progressively transfer parameters learned from a coarse graph to a sequence of increasingly finer graphs via biased random walk upsampling. For 3D molecular generation, we voxelize the molecular structures at multiple resolutions, pretrain a coarse-resolution conditional Variational Autoencoder (CVAE), and initialize a fine-resolution CVAE by transferring shape compatible convolutional parameters from the coarse model. Numerical experiments on receptor-conditioned 3D Ligand generation show that multigrid training accelerates convergence and improves generalization compared to training from scratch.