Auf dem Weg zu einem einheitlichen latenten Raum für die latente Diffusionsmodellierung von 3D-Molekülen

Towards Unified Latent Space for 3D Molecular Latent Diffusion Modeling

March 19, 2025

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

cs.AI

Zusammenfassung

Die Generierung von 3D-Molekülen ist entscheidend für die Arzneimittelentdeckung und Materialwissenschaft und erfordert Modelle, die komplexe Multimodalitäten verarbeiten können, einschließlich Atomtypen, chemischer Bindungen und 3D-Koordinaten. Eine zentrale Herausforderung besteht darin, diese Modalitäten unterschiedlicher Formen zu integrieren, während die SE(3)-Äquivarianz für 3D-Koordinaten aufrechterhalten wird. Um dies zu erreichen, verwenden bestehende Ansätze typischerweise separate latente Räume für invariante und äquivariante Modalitäten, was die Effizienz sowohl beim Training als auch bei der Stichprobenentnahme verringert. In dieser Arbeit schlagen wir den Unified Variational Auto-Encoder for 3D Molecular Latent Diffusion Modeling (UAE-3D) vor, ein multimodales VAE, das 3D-Moleküle in latente Sequenzen aus einem einheitlichen latenten Raum komprimiert, während ein nahezu null Rekonstruktionsfehler aufrechterhalten wird. Dieser einheitliche latente Raum beseitigt die Komplexitäten der Handhabung von Multimodalität und Äquivarianz bei der Durchführung von latenten Diffusionsmodellierungen. Wir demonstrieren dies durch den Einsatz des Diffusion Transformers – ein allgemeiner Diffusionsmodell ohne molekulare Induktionsverzerrung – für die latente Generierung. Umfangreiche Experimente auf den Datensätzen GEOM-Drugs und QM9 zeigen, dass unsere Methode neue Maßstäbe sowohl in der de novo als auch in der konditionalen 3D-Molekülgenerierung setzt und führende Effizienz und Qualität erreicht.

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.