Technical Summary: KRONOS – Autoregressive Latent Diffusion for 3D Molecule Generation

Problem Statement

Three-dimensional (3D) molecule generation is a critical task in drug discovery and molecular design. Current state-of-the-art approaches primarily rely on diffusion models and flow matching. While these methods achieve high generation quality, they typically require the molecular size (number of atoms) to be specified a priori, limiting their flexibility for variable-length generation.

Conversely, autoregressive (AR) models naturally support variable-length generation and conditioning on partial molecular structures (e.g., elaborating known scaffolds), which is highly valuable in medicinal chemistry. However, existing AR models face two main challenges:

1. Performance Gap: Historically, AR models have lagged behind diffusion models in generation quality, particularly regarding geometric fidelity and chemical validity.
2. Balancing Objectives: Achieving strong performance in both unconditional generation (de novo design) and fragment-conditioned generation (molecular completion) within a single architecture remains difficult. Many AR models struggle to generalize between these paradigms without architectural modifications or performance trade-offs.

Additionally, existing latent molecular generators have almost exclusively relied on diffusion, leaving the potential of autoregressive generation in latent space largely unexplored.

Methodology: KRONOS

The authors introduce KRONOS, a latent autoregressive diffusion framework that generates 3D molecules directly in the latent space of a pre-trained Unified Autoencoder (UAE). The method combines three core components:

1. Latent Autoregressive Generation

Instead of operating in the raw data space (predicting discrete atom types/bonds and continuous coordinates simultaneously), KRONOS operates in the continuous latent space of a pre-trained UAE (Luo et al., 2025).

• Representation: A 3D molecule is encoded into a sequence of latent tokens $Z = \{z_1, \dots, z_n\}$, where each token jointly encodes atomic identity, bonding environment, and 3D geometry.
• Factorization: The model learns a left-to-right autoregressive distribution:
  $$p_\theta(z_{1:n}) = \prod_{i=1}^n p_\theta(z_i | z_{<i})$$
• Variable Length: Since the latent space lacks a natural end-of-sequence token, KRONOS employs a dedicated stop classifier (a lightweight MLP) trained alongside the generator to predict sequence termination dynamically, eliminating the need for a predefined atom count.

2. Diffusion-Based Token Prediction

To model the conditional distribution of continuous latent tokens $p_\theta(z_i | z_{<i})$, KRONOS utilizes a Diffusion Loss formulation (Li et al., 2024).

• A causal transformer encodes the preceding tokens into a context representation $c_i$.
• A diffusion denoising network $\epsilon_\theta$ predicts the noise added to the current latent token $z_i$, conditioned on $c_i$ and the diffusion timestep.
• This approach replaces discrete token prediction with a denoising objective, enabling flexible modeling of continuous variables without discretization.

3. Mixed Training Strategy (Fragment-Conditioned Generation)

To support both unconditional and fragment-conditioned generation within a single model, the authors introduce a training strategy inspired by the Fill-in-the-Middle (FIM) paradigm.

• Fragment Decomposition: Molecules are decomposed into chemically meaningful fragments using a BRICS-based scheme.
• Sequence Construction: During training, each molecule is transformed into either:
  • A default-order sequence (for unconditional generation).
  • A fragment-conditioned sequence where a random subset of fragments acts as a "seed" (context), and the remaining fragments are the generation target.
• Mechanism: The model is trained to autoregressively generate the target fragments conditioned on the seed fragments. This is achieved by assigning a single deterministic ordering to the fragments and atoms within them, rather than marginalizing over all permutations (as in Order-Agnostic AR models). This allows the model to learn molecular completion capabilities without architectural changes.

Key Contributions

1. KRONOS Framework: A novel latent autoregressive diffusion framework that jointly models molecular graph topology and geometry in a unified latent space.
2. Unified Training Strategy: A mixed training approach enabling both unconditional and fragment-conditioned generation within a single left-to-right autoregressive model, inspired by the FIM paradigm.
3. Variable-Length Generation: A unified framework that naturally supports variable-length generation via a stop classifier, avoiding the need for a priori size specification.
4. Performance: Demonstration that latent autoregressive generation can achieve leading performance among autoregressive methods while remaining competitive with diffusion models, without sacrificing unconditional quality when enabling fragment conditioning.

Experimental Results

The authors evaluated KRONOS on the QM9 and GEOM-Drugs datasets.

Unconditional Generation

• QM9: KRONOS achieved the highest 3D Validity (97.3%) and strongest overall performance among autoregressive methods. It remained competitive with leading diffusion models (e.g., JODO, EDM) across metrics including 3D Atom Stability, Bond Length, and Bond Angle MMD.
• GEOM-Drugs: KRONOS achieved the strongest overall performance among all evaluated baselines (both autoregressive and diffusion).
• Physical Plausibility: Using semi-empirical geometry optimization (GFN2-xTB), KRONOS-generated conformations showed low Root Mean Square Distance (RMSD) and energy differences ($\Delta E$) relative to optimized structures, indicating they lie close to physically plausible regions of conformational space.

Fragment-Conditioned Generation

• On a dataset of ~15k molecule-fragment pairs, KRONOS substantially outperformed the autoregressive baseline QUETZAL in both Mean@3 and Best@3 metrics for validity and stability.
• Crucially, enabling fragment-conditioned generation had negligible impact on unconditional generation performance, demonstrating that both paradigms can be supported within a single architecture.

Significance and Claims

The paper claims that KRONOS provides a promising alternative to existing data-space autoregressive and diffusion-based approaches. Its significance lies in:

• Bridging the Gap: It substantially narrows the performance gap between autoregressive and diffusion models, achieving leading results among AR methods while matching diffusion baselines.
• Flexibility: It unifies the strengths of autoregressive models (variable length, partial conditioning) with the high-quality generation of diffusion models, all within a single latent space formulation.
• Practical Utility: By supporting fragment-conditioned generation without architectural modifications, it offers a practical tool for medicinal chemistry tasks where molecules are often designed by elaborating known scaffolds.

The authors note limitations, including reliance on the expressiveness of the pre-trained latent space and the use of a fixed BRICS decomposition scheme, suggesting that future work could explore learning fragment abstractions jointly with the generative model or extending the framework to property-guided design.