Autoregressive latent diffusion for 3D molecule generation

2026-07-10

AI Summary

Existing 3D molecular generation methods struggle to simultaneously support unconditional and fragment-based conditional generation while often requiring predefined molecular sizes. This work proposes KRONOS, a unified framework that integrates autoregressive and diffusion mechanisms within the latent space of a pretrained molecular autoencoder to jointly model molecular graph topology and 3D geometry. Inspired by Fill-in-the-Middle, KRONOS employs a hybrid training strategy that enables efficient support for both generation modes within a single architecture without performance trade-offs. Experimental results demonstrate that KRONOS achieves superior unconditional generation performance on QM9 and GEOM-Drugs compared to existing autoregressive models and matches that of diffusion-based approaches, while preserving near-identical unconditional generation quality even when trained for fragment-conditioned synthesis.

Abstract

Three-dimensional (3D) molecule generation has been dominated by diffusion models, which achieve strong generation quality but typically require the molecular size to be specified a priori. Recent autoregressive approaches have substantially narrowed the performance gap while naturally supporting variable-length generation and conditioning on partial molecular context. However, balancing unconditional and context-conditioned generation remains challenging. We introduce KRONOS, a latent autoregressive diffusion framework that generates molecules in the latent space of a pre-trained autoencoder, jointly modeling molecular graph topology and geometry, while retaining the flexibility of autoregressive generation. We further introduce a mixed training strategy inspired by Fill-in-the Middle (FIM) paradigm, enabling both unconditional and fragment-conditioned molecular generation within a single left-to-right autoregressive model. Experiments on QM9 and GEOM-Drugs demonstrate that KRONOS achieves leading unconditional generation performance among autoregressive methods, while remaining competitive with diffusion models. Moreover, fragment-conditioned generation is achieved with negligible impact on unconditional generation performance, demonstrating that both generation paradigms can be supported within a single architecture.

Problem

• 3D molecule generation
• autoregressive modeling
• conditional generation
• unconditional generation
• molecular geometry

Innovation

• autoregressive latent diffusion
• 3D molecule generation
• fragment-conditioned generation
• mixed training strategy
• molecular graph and geometry modeling