fedeotto/kronos-models

KRONOS — checkpoints

Trained weights for KRONOS, a latent autoregressive diffusion model for 3D molecule generation, together with the Unified AutoEncoder (UAE) checkpoints its latent space is built on.

Contents

uae_qm9/{last.ckpt, latent_stats.pth}          # autoencoder + latent statistics (QM9)
uae_geom/{last.ckpt, latent_stats.pth}         # autoencoder + latent statistics (GEOM-Drugs)
kronos-{qm9,geom}-layers{8,10,12,16}-heads{8,10,12,16}-d{512,640,768,1024}/manual_last.ckpt

The layers16-heads16-d1024 checkpoints (283 M parameters) are the main models reported in the paper; the three smaller ones are the scaling ablation.

Usage

Download into the checkpoints/ directory of the code repository — the checkpoints store vae_path = checkpoints/uae_{qm9,geom}/last.ckpt internally, so the layout must be preserved.

python -m kronos.scripts.test \
    --model_path checkpoints/kronos-qm9-layers16-heads16-d1024 \
    --ckpt_name manual_last.ckpt --num_samples 10000 --sampler ddim

Citation

@article{ottomano2026autoregressive,
  title={Autoregressive latent diffusion for 3D molecule generation},
  author={Ottomano, Federico and Ren, Gaopeng and Li, Yingzhen and Jelfs, Kim E and Ganose, Alex M},
  journal={arXiv preprint arXiv:2607.09277},
  year={2026}
}