Smoothing Dark Areas in Molecular Latent Diffusion
For researchers in molecular generation, this work addresses a critical bottleneck in latent diffusion by ensuring latent space smoothness, leading to more reliable and chemically valid molecule generation.
Latent diffusion for 3D molecular generation suffers from 'dark areas'—latent regions that decode to invalid molecules. The proposed TopVAE reduces these dark areas by embedding structural constraints during training, achieving up to 77% lower FCD-3D on QM9 and 52% lower on GEOM-Drugs, with improved stability and connectivity.
Latent diffusion is a promising framework for scalable 3D molecular generation, but it requires a latent space that remains smooth, valid, and navigable beyond posterior samples. Existing molecular VAEs, however, are typically learned through reconstruction-based objectives, which do not guarantee such a latent space. We show that this leads to dark areas: regions of latent space that are reachable during diffusion sampling but decode to disconnected or chemically invalid molecules. Unlike in image generation, molecular decoding requires strict structural and chemical precision, so even small latent perturbations can produce catastrophic failures. We therefore propose TopVAE, a topology-optimized VAE that reduces dark areas by making the decoder internalize structural and chemical constraints during training, eliminating the need for test-time chemical correction. TopVAE greatly improves off-posterior robustness, and when paired with a standard DiT, achieves $77\%$ lower FCD-3D on QM9, the highest V&C, $52\%$ lower FCD-3D on GEOM-Drugs, and $1.29{\times}$ more stable and connected molecules on zero-shot scaffold inpainting.