MULTIGRID TRAINING FOR MOLECULAR GENERATION USING GRAPH NEURAL NETWORKS

Deep learning models for molecular design, such as Graph Neural Networks (GNNs) and 3D voxel-based generative models, face significant computational challenges as the resolution of the molecular representation increases. The authors investigate whether a multigrid training strategy—traditionally used in numerical analysis to solve multiscale systems—can be adapted to deep learning to improve training efficiency and model performance.

The paper proposes a multigrid framework that constructs a hierarchy of representations, moving from coarse to fine resolutions. For graph-based models, the authors utilize a biased random walk upsampling operator to progressively expand coarse subgraphs into finer ones. For 3D molecular generation, they employ a conditional Variational Autoencoder (CVAE) that is pretrained on coarse voxel grids and then initialized at finer resolutions using parameter transfer. To maintain consistency across these levels, the framework incorporates a guidance loss that penalizes discrepancies between predictions at different resolutions, ensuring that the finer models benefit from the structural knowledge captured by the coarse models.

Numerical experiments on receptor-conditioned 3D ligand generation demonstrate that the multigrid training approach significantly accelerates convergence compared to training models from scratch. Furthermore, the authors show that this hierarchical initialization improves the generalization capabilities of the models. Theoretical analysis of the scheme suggests that, under conditions of good alignment and sufficient coverage, the transferred weights provide a well-conditioned starting point for the fine-level optimization, effectively reducing the loss and gradient norms at initialization.