Smoothing Dark Areas in Molecular Latent Diffusion

2026-06-11

AI Summary

This work addresses the challenge that molecular latent diffusion models often fall into “dark regions” of the latent space during sampling, producing 3D molecules with broken structures or chemically invalid geometries. To mitigate this issue, the authors propose Topology-Optimized Variational Autoencoder (TopVAE), which, for the first time, incorporates topological constraints directly into the training of molecular VAEs. This enables the decoder to internalize chemical and structural priors, effectively smoothing the latent space and eliminating dark regions. Integrated with a DiT architecture, the method generates high-quality, chemically valid 3D molecules without requiring post-processing. Experiments demonstrate significant improvements: FCD-3D scores are reduced by 77% on QM9 and 52% on GEOM-Drugs, the number of stable, connected molecules in zero-shot scaffold completion increases by 1.29×, and the model achieves state-of-the-art validity and connectivity (V&C) scores.

Abstract

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.

Problem

• latent diffusion
• molecular generation
• dark areas
• VAE
• chemical validity

Innovation

• latent diffusion
• molecular generation
• TopVAE
• dark areas
• chemical validity