TopVAE — checkpoints and preprocessed datasets

Artifacts for Smoothing Dark Areas in Molecular Latent Diffusion.

Code, configs and reproduction instructions live in the GitHub repository: https://github.com/ComDec/TopVAE. This Hub repository holds only the bytes that are too large for git — the trained weights and the preprocessed molecule lists.

Values quoted below and in the code repository are what these artifacts produce; where they differ from the paper's printed table, the value here is the measured one.

Download

pip install -U "huggingface_hub[cli]"

# from the root of a TopVAE clone — the layout below is what the configs expect
hf download EscheWang/TopVAE --local-dir .

# then verify -- the manifest's paths are relative to checkpoints/, so run it from there
(cd checkpoints && sha256sum -c SHA256SUMS)      # macOS: shasum -a 256 -c SHA256SUMS

A mismatch means you do not have the file the provenance table describes. Do not proceed.

What is here

checkpoints/
  QM9/
    topvae_dit_cf_ep9999/          Table 2 — the QM9 generation model (TopVAE + DiT)
      .hydra/config.yaml             frozen config; its model.net.vae.run_dir points at ↓
      checkpoint/epoch_9999.ckpt
    topvae_cf_v4b09/               Table 1 — QM9 TopVAE, and the frozen VAE under the model above
      .hydra/config.yaml
      checkpoint/last.ckpt
      checkpoint/latent_distribution.pth    latent whitening statistics — part of the model
    uae_qm9.ckpt                   Table 1 — QM9 UAE baseline
    uae_udm_qm9.ckpt               Table 2 — UAE + UDM baseline
  GEOM_Drugs/
    topvae_ditB_iso_ep5149/        Table 3 — the GEOM-Drugs generation model (TopVAE + DiT-B)
      .hydra/config.yaml
      checkpoint/epoch_5149.ckpt
    topvae_v18alt_ep0749/          Table 3 — the frozen VAE backbone underneath it
      .hydra/config.yaml
      checkpoint/archive/epoch_0749.ckpt
      checkpoint/archive/latent_distribution.pth
    uae_geom.ckpt                  Table 1 — GEOM-Drugs UAE baseline
  SHA256SUMS

data/
  QM9/mol_list/{train,valid,test}.pt              heavy-atom only, kekulized
  GEOMDrugs/mol_list/{train,valid,test}.pt
  */mol_list/processed/                           derived caches, shipped so nothing is recomputed

A diffusion checkpoint cannot be loaded on its own. It needs the frozen Hydra configuration it was trained with, and the paired frozen VAE, which that config resolves by a repository-relative path. So the weights must sit inside your clone at exactly these paths, and commands must be run from the repository root. Give a UDM its run-directory path (.../topvae_dit_cf_ep9999/checkpoint/epoch_9999.ckpt) — the loader walks up from the checkpoint looking for .hydra/config.yaml.

latent_distribution.pth is part of the model, not a cache. It carries the latent whitening statistics the DiT prior was trained against. The code regenerates it from the training set if it is absent, which is silent and slow, and for anisotropic variants it is the difference between 0.197 and 0.979 validity.

Two GEOM weights are present under both a flat name (topvae_ditB_iso_ep5149.ckpt, topvae_v18alt_ep0749.ckpt) and their run-directory path. They are the same bytes, shipped twice on purpose so that sha256sum -c reports OK on all 15 manifest entries and every command in the code repository resolves. Only the run-directory copy can load a UDM; the flat name is a convenience for src/eval.py, which needs nothing but ckpt_path=. Replace one with a symlink or hardlink locally if you would rather not store 1 GB twice.

What these artifacts reproduce

Single documented commands, from these weights, on one A100. Full transcripts are in the code repository's docs/reproduction.md.

Verified end to end on 2026-07-28. Both commands below were re-run from a fresh checkout of the released code, against these weights after sha256sum -c, and reproduced every metric in the table: the largest absolute deviation was 1e-6 on QM9 (FCD_3D) and 5e-6 on GEOM-Drugs (FCD_3D); every other cell matched to all six printed decimals. The QM9 command was then repeated a third time from a pristine git clone of the public repository, on a machine holding nothing but that clone and these weights, and agreed to 2.6e-6 — again only on FCD_3D, whose coordinate→graph bond inference classifies one or two molecules in 10,000 differently between environments. Every other metric was identical to all printed digits.

	QM9 (Table 2)	GEOM-Drugs (Table 3)
FCD (deduplicated — the reported variant)	0.184929	3.876948
FCD_3D	0.197121	8.178909
AtomStab 2D	1.000000	0.999486
MolStab 2D	1.000000	0.977300
Validity & Connectivity	0.997800	0.945400
Uniqueness	0.967200	1.000000
AtomStab 3D	0.990289	0.831988
MolStab 3D	0.925100	0.021600

Sampling protocol: seed 42, n = 10,000, greedy valence repair. Temperature is a per-dataset setting — GEOM-Drugs uses 0.95 and 100 steps, QM9 uses none. Temperature, not repair, is what moves FCD: the same QM9 checkpoint at temperature 0.95 gives FCD 0.450886, a factor of ~2.4.

FCD_3D carries a coverage denominator of its own — the coordinate→graph rule returns a SMILES for 3,398 of the 10,000 generated GEOM molecules and for 264 unique reference molecules — so the two FCD rows are computed against different reference sets by construction and are not comparable to each other.

What is not here

Stated rather than left to be discovered.

• Table 4 (scaffold inpainting) is not part of this release. Its only entry point resolved models from a registry of the authors' own run directories.
• Table 5 (component ablations) is not part of this release. Its checkpoints were not located.
• The GEOM-Drugs TopVAE row of Table 1 has no shipped checkpoint; the other three of that table's four model × dataset cells do.

Citation

The paper is not yet published; cite it as unpublished.

@unpublished{wang2026topvae,
  title  = {Smoothing Dark Areas in Molecular Latent Diffusion},
  author = {Wang, Xi and Li, Jiahan and Xia, Yuxuan and Wu, Yingcheng and
            Zheng, Shaoyi and Wang, Shengjie},
  year   = {2026},
  note   = {Preprint.}
}

License

MIT. See LICENSE.