Vers un espace latent unifié pour la modélisation de diffusion latente de molécules 3D

Towards Unified Latent Space for 3D Molecular Latent Diffusion Modeling

March 19, 2025

Auteurs: Yanchen Luo, Zhiyuan Liu, Yi Zhao, Sihang Li, Kenji Kawaguchi, Tat-Seng Chua, Xiang Wang

cs.AI

Résumé

La génération de molécules 3D est cruciale pour la découverte de médicaments et la science des matériaux, nécessitant des modèles capables de traiter des multi-modalités complexes, incluant les types d'atomes, les liaisons chimiques et les coordonnées 3D. Un défi majeur consiste à intégrer ces modalités de formes différentes tout en maintenant l'équivariance SE(3) pour les coordonnées 3D. Pour y parvenir, les approches existantes maintiennent généralement des espaces latents séparés pour les modalités invariantes et équivariantes, ce qui réduit l'efficacité à la fois lors de l'entraînement et de l'échantillonnage. Dans ce travail, nous proposons l'Unified Variational Auto-Encoder for 3D Molecular Latent Diffusion Modeling (UAE-3D), un VAE multi-modal qui compresse les molécules 3D en séquences latentes à partir d'un espace latent unifié, tout en maintenant une erreur de reconstruction proche de zéro. Cet espace latent unifié élimine les complexités liées à la gestion de la multi-modalité et de l'équivariance lors de la modélisation de diffusion latente. Nous démontrons cela en utilisant le Diffusion Transformer—un modèle de diffusion à usage général sans biais inductif moléculaire—pour la génération latente. Des expériences approfondies sur les ensembles de données GEOM-Drugs et QM9 montrent que notre méthode établit de nouveaux référentiels significatifs à la fois pour la génération de novo et conditionnelle de molécules 3D, atteignant une efficacité et une qualité de premier plan.

English

3D molecule generation is crucial for drug discovery and material science, requiring models to process complex multi-modalities, including atom types, chemical bonds, and 3D coordinates. A key challenge is integrating these modalities of different shapes while maintaining SE(3) equivariance for 3D coordinates. To achieve this, existing approaches typically maintain separate latent spaces for invariant and equivariant modalities, reducing efficiency in both training and sampling. In this work, we propose Unified Variational Auto-Encoder for 3D Molecular Latent Diffusion Modeling (UAE-3D), a multi-modal VAE that compresses 3D molecules into latent sequences from a unified latent space, while maintaining near-zero reconstruction error. This unified latent space eliminates the complexities of handling multi-modality and equivariance when performing latent diffusion modeling. We demonstrate this by employing the Diffusion Transformer--a general-purpose diffusion model without any molecular inductive bias--for latent generation. Extensive experiments on GEOM-Drugs and QM9 datasets demonstrate that our method significantly establishes new benchmarks in both de novo and conditional 3D molecule generation, achieving leading efficiency and quality.