---
license: mit
tags:
  - molecular-generation
  - 3d-molecules
  - latent-diffusion
  - variational-autoencoder
  - chemistry
datasets:
  - QM9
  - GEOM-Drugs
library_name: pytorch
---

# TopVAE — checkpoints and preprocessed datasets

Artifacts for **[Smoothing Dark Areas in Molecular Latent Diffusion](https://github.com/ComDec/TopVAE)**.

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

```bash
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.

```bibtex
@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](https://github.com/ComDec/TopVAE/blob/main/LICENSE).