К унифицированному латентному пространству для моделирования латентной диффузии 3D молекул

Towards Unified Latent Space for 3D Molecular Latent Diffusion Modeling

March 19, 2025

Авторы: Yanchen Luo, Zhiyuan Liu, Yi Zhao, Sihang Li, Kenji Kawaguchi, Tat-Seng Chua, Xiang Wang

cs.AI

Аннотация

Генерация 3D молекул имеет ключевое значение для разработки лекарств и материаловедения, требуя от моделей обработки сложных мультимодальных данных, включая типы атомов, химические связи и 3D координаты. Основная задача заключается в интеграции этих модальностей различной формы с сохранением SE(3)-эквивариантности для 3D координат. Для достижения этого существующие подходы обычно поддерживают отдельные латентные пространства для инвариантных и эквивариантных модальностей, что снижает эффективность как обучения, так и генерации. В данной работе мы предлагаем Unified Variational Auto-Encoder for 3D Molecular Latent Diffusion Modeling (UAE-3D) — мультимодальный VAE, который сжимает 3D молекулы в латентные последовательности из единого латентного пространства, сохраняя при этом практически нулевую ошибку реконструкции. Это единое латентное пространство устраняет сложности обработки мультимодальности и эквивариантности при выполнении латентного диффузионного моделирования. Мы демонстрируем это, используя Diffusion Transformer — универсальную диффузионную модель без каких-либо молекулярных индуктивных предубеждений — для генерации латентных представлений. Эксперименты на наборах данных GEOM-Drugs и QM9 показывают, что наш метод значительно устанавливает новые стандарты как в de novo, так и в условной генерации 3D молекул, достигая лидирующих показателей по эффективности и качеству.

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.