Autoregressive latent diffusion for 3D molecule generation
2026-07-10 • Machine Learning
Machine Learning
AI summaryⓘ
The authors present KRONOS, a new model that generates 3D molecules by working in a simplified latent space, allowing it to model both the shape and connections of molecules. Unlike previous methods that need a fixed molecule size or separate models for different tasks, KRONOS can generate molecules of varying sizes and can fill in missing parts based on a given fragment. They use a special training technique that helps the model handle both starting from scratch and completing partial molecules with one single approach. Tests show KRONOS performs very well compared to existing methods, especially among autoregressive models, without losing quality when mixing different generation tasks.
3D molecule generationdiffusion modelsautoregressive modelslatent spaceautoencodermolecular graph topologyFill-in-the-Middle (FIM)QM9 datasetGEOM-Drugs dataset
Authors
Federico Ottomano, Gaopeng Ren, Yingzhen Li, Kim E. Jelfs, Alex M. Ganose
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.