TopVAE

Smoothing Dark Areas in Molecular Latent Diffusion — a topology-first molecular autoencoder whose latent space stays decodable in the regions a diffusion prior actually samples.

Xi Wang¹ · Jiahan Li¹ · Yuxuan Xia¹ · Yingcheng Wu² · Shaoyi Zheng¹ · Shengjie Wang^1
1New York University · ²Stanford University

In one line: on GEOM-Drugs, 10,000 prior samples decode at validity 1.0000 with 94.54 % of them a single connected molecule (V&C 0.945400) and 2D molecular stability 0.977300; on QM9 the same command gives validity 1.0000, V&C 0.997800 and FCD 0.184929. Those are measured, not quoted — see Results and Reproduction. This repository ships no baseline results, so it cannot rank itself against other methods; the paper does that.

Weights and preprocessed datasets: https://huggingface.co/EscheWang/TopVAE — anonymous download, no token.
Paper: no public URL yet — cite the entry below.

_{Paper Figure 1. Panel A diagnoses the problem, panel B is the architecture. The two scatter plots at the bottom of panel A are labelled "Previous VAE (32.6% Valid)" and "Ours (90.5% Valid)"; the 32.6% baseline is ADiT, not UAE — the labelled three-panel version with every baseline named is below.}

Important

What reproduces, before you spend an afternoon on it.

• Table 2 (QM9) and Table 3 (GEOM-Drugs) reproduce end to end, one command each — re-verified three separate ways on 2026-07-28, including from a bare machine holding nothing else, at a maximum deviation of 2.2e-6 across all twenty reported cells. → Reproduction
• The VAE posterior cells run but do not assemble into the paper's comparison.
• Table 4 (scaffold inpainting) and Table 5 (component ablations) are not in this release. → What this release does not cover

---

The problem

A dark area is a patch of the VAE's latent space that the diffusion model can land on but the decoder cannot handle. Decode a point there and you get a broken molecule — atoms in disconnected pieces, or bonds no real chemistry allows — even though that same decoder rebuilds training molecules essentially perfectly. Unlike images, where a slightly-off latent yields a slightly-off picture, a slightly-off molecular latent yields something that is not a molecule at all.

Formally, writing A_LDM for the latent region reachable during diffusion sampling and M_valid for the set of molecules that pass RDKit sanitization and form a single connected component, the dark areas are D_dark = { z in A_LDM : D(z) not in M_valid }. The paper names two sources: topological disconnection (the decoded graph fractures into isolated components) and chemical invalidity (predicted bonds violate valence or compatibility rules).

The consequence for evaluation is the sharpest sentence in the paper: per-atom stability metrics mask this failure by accepting fragmented atom clouds as stable. Which is why every table below reports connectivity beside stability, and never one under the other's name.

_{Panel titles, verbatim: "(a) UAE (44.4% valid)", "(b) ADiT (32.6% valid)", "(c) TopVAE (90.5% valid)". Green = "Valid & Connected", red = "Invalid / Disconnected". These three percentages are read off the figure itself; none of them appears in the paper's text or tables, the figure it belongs to is captioned as a 2-D PCA of per-molecule mean latent vectors (n = 2,000, QM9 test set), and this release does not regenerate that experiment — see scope. Treat them as illustration, not as a result you can cite.}

---

The method

Four components. The paper's component names do not appear verbatim in the source, so this table is the mapping.

Component	What it does	Where it lives
AdjFirst decoder	Decodes topology first — adjacency, then atom types, then bond types, then coordinates — so geometry is conditioned on a graph that has already been committed and repaired. Bond prediction is hard-gated by adjacency.	src/models/vae/decoders/graphormer_adjfirst.py (GraphormerAdjFirstDecoder)
TopoBridge	Turns soft adjacency into a binary, guaranteed-connected graph: keep each atom's top-k neighbours, then repeatedly BFS from the largest component and insert the single most likely cross-component edge until every atom is reachable. Discrete forward, straight-through backward. The guarantee is on the adjacency TopoBridge
emits, not on the final molecule: ChemCO may later drop a bond to satisfy valence, which is why
9,454 of 10,000 GEOM generations are single-fragment rather than all 10,000.	src/models/vae/projectors.py (the *_bfs_* repair family)
ChemCO	Unrolled primal–dual solver over bond assignments — pair exclusivity, valence caps, minimum degree — including its greedy valence repair, which is an inner step of ChemCO's continuous→discrete→continuous loop and is not a separable post-processor. Bridge edges are protected, so ChemCO cannot undo TopoBridge.	src/models/vae/constraint_pp/
AGCL	The selective teacher. ChemCO's correction is distilled into the raw decoder only on molecules where it lowers the ground-truth bond NLL (positive advantage), so the decoder internalizes the constraints rather than depending on the projection to supply them. Note that both documented commands below still pass --repair, so the reported numbers are produced with ChemCO active at inference; the release does not measure it switched off. Config key loss_weight_pp_selective.	configs/experiment/QM9_noH_Keku/topvae.yaml L72, GEOM_noH_Keku/topvae.yaml L131

The multi-stage module wiring them together is src/models/vae/topvae_multistage.py.

Warning

The released QM9 checkpoint was not trained at the paper's stated AGCL weight. The paper's hyper-parameter table gives λ_AGCL = 1.0 for QM9 and 0.1 for GEOM-Drugs. Both shipped presets — and the frozen config inside the released QM9 checkpoint — set loss_weight_pp_selective: 0.1, so GEOM matches the paper and QM9 is 10× lower. The QM9 numbers in this README are what that checkpoint produces; they are not a reproduction of the paper's stated QM9 training recipe. A second, smaller divergence sits beside it: the paper gives ChemCO steps (train / eval) as 20 / 40 for QM9, while the config and the frozen checkpoint both use 20 / 50. Neither was reconciled before release, and both are flagged here rather than left for a reader to find.

_{Decoder robustness to latent perturbation. TopVAE is highest across both panels; the ADiT baseline collapses to the floor by σ ≈ 0.2–0.3; UAE decays but stays off the floor. Its source and plotting script are in the paper's own figures/ directory but the figure is never included in the compiled paper, so it carries no caption and no table cell, and no shipped command regenerates these curves. (The figure's legend spells the baseline "AADiT"; the paper spells it ADiT throughout, and so does this README.)}

---

Results

These are the values the released artifacts produce, each from one documented run of the command shown beside it. They are not transcriptions of the paper's printed cells. Some cells differ from the paper's printed table; where they do, the value here is the measured one and the paper's is the one to query with the authors. No baseline column appears here, because this release contains no baseline results — every baseline cell in the paper is copied from its own publication or re-trained by the authors, and none of that is reproducible from this repository.

QM9 — Table 2 (TopVAE + DiT)

--n_samples 10000 --seed 42, no temperature scaling, greedy repair.

metric	value
FCD (deduplicated — the reported variant)	0.184929
FCD (multiplicity-preserving, secondary)	0.190801
FCD_3D (FCD_coord→graph)	0.197121
AtomStab2D	1.000000
MolStab2D	1.000000
validity	1.000000
V&C — valid and connected	0.997800
SC — stable and connected	0.997800
V&U	0.967200
AtomStab3D	0.990289
MolStab3D	0.925100

GEOM-Drugs — Table 3 (TopVAE + DiT-B)

--n_samples 10000 --seed 42 --temperature 0.95, 100 steps (the checkpoint's own default), greedy repair.

metric	value
FCD (deduplicated — the reported variant)	3.876948
FCD (multiplicity-preserving, secondary)	3.876945
FCD_3D (FCD_coord→graph)	8.178909
AtomStab2D	0.999486
MolStab2D	0.977300
validity	1.000000
V&C — valid and connected	0.945400
SC — stable and connected	0.923500
V&U	1.000000
AtomStab3D	0.831988
MolStab3D	0.021600

Denominators from that same run: 10000 sampled, 10000 decoded and evaluable, 9454 connected, validity 1.0000, 10000 unique generated SMILES against 1994 unique reference SMILES.

Warning

V&C and S^C are different quantities and only coincide on QM9. V&C is valid and connected (edm_2d.validity × edm_2d.connected_ratio); S^C is stable and connected (edm_2d.mol_stable_AND_connected). They agree on QM9 (both 0.997800) only because MolStab2D is exactly 1.000000 there. On GEOM-Drugs they differ: 0.945400 against 0.923500. Never print one under the other's name.

Three more things worth knowing before quoting any of this

• FCD and FCD_3D are not comparable to each other. FCD_3D pushes both sides through the coordinate→graph rule, which resolves a SMILES for only 3398 of the 10000 generated molecules and for 264 unique reference molecules on the GEOM run — a different reference set by construction. docs/metrics.md §2.
• Temperature is a per-dataset setting, and temperature is what moves FCD, not repair. The same QM9 checkpoint at --temperature 0.95 gives FCD 0.450886 instead of 0.184929 (~2.4×); dropping --repair moves it by 0.0001. GEOM-Drugs uses 0.95, QM9 uses none. This is the released configuration — the paper's appendix specifies temperature 1.0 for both.
• V&C and V&U are derived, not emitted. No JSON key is named v_and_c or v_and_u; they are edm_2d.validity × edm_2d.connected_ratio and edm_2d.validity × uniqueness.unique_ratio.

Full definitions, denominators and the output-JSON schema: docs/metrics.md §2.

---

Reproduction

This is the part of the release worth trusting, so here is exactly what was done.

On 2026-07-28 the two headline commands were executed the way a stranger would, three separate times:

run	setup	result
1	From the committed tree, with the released weights	Tables 2 and 3, max deviation 5e-6
2	From a pristine git clone, weights copied in	QM9 (Table 2), max deviation 2.6e-6
3	Pristine git clone on a machine holding nothing else; weights and datasets fetched by an anonymous, token-free hf download EscheWang/TopVAE --local-dir .; (cd checkpoints && sha256sum -c SHA256SUMS) → 15/15 OK and (cd data && sha256sum -c SHA256SUMS) → 10/10 OK	Tables 2 and 3: QM9 max deviation 2.1e-6, GEOM-Drugs 2.2e-6

In all three, the only metric that moved at all was FCD_3D, in the sixth decimal. Every other metric agreed to every digit printed. The cause is documented and is not a sampling difference: the coordinate→graph bond-inference rule classifies one or two molecules out of 10,000 differently between environments, which perturbs the ChemNet statistics. The decoded molecule set itself is byte-identical across GPUs.

Cost of run 3, the one that matters, measured on one A100, each figure including the one-time dataset-cache build that a first run pays:

	QM9	GEOM-Drugs
sampling	176 s	981 s
metrics + FCD	89 s	488 s
load + first-run cache build	239 s	1,185 s
end to end	9.6 min	47.8 min

Later runs reuse the cache and are correspondingly shorter. Add the ≈ 3.2 GB download and, for GEOM-Drugs, ≈ 12 GB of disk for the cache the datamodule builds.

The two commands

Important

These are shown here because this section is about evidence. Do Quickstart first — without the environment, the weights and the data they fail at import or at the first checkpoint path.

# QM9 -- Table 2
CUDA_VISIBLE_DEVICES=0 python scripts/eval/udm_prior_eval_with_fcd.py \
    --ckpt checkpoints/QM9/topvae_dit_cf_ep9999/checkpoint/epoch_9999.ckpt \
    --n_samples 10000 --batch_size 128 --dataset qm9 --seed 42 \
    --repair \
    --out_json logs/analysis/udm_eval/qm9_table2.json

# GEOM-Drugs -- Table 3
CUDA_VISIBLE_DEVICES=0 python scripts/eval/udm_prior_eval_with_fcd.py \
    --ckpt checkpoints/GEOM_Drugs/topvae_ditB_iso_ep5149/checkpoint/epoch_5149.ckpt \
    --n_samples 10000 --batch_size 128 --dataset geom --seed 42 \
    --temperature 0.95 --repair \
    --out_json logs/analysis/udm_eval/geom_table3_seed42.json

Each command does the whole table — sampling, decoding, EDM 2D/3D stability, uniqueness, FCD and FCD_3D — and writes every sampling knob into the output JSON under sampling_protocol. Transcripts and per-metric deviations: docs/reproduction.md §6.

Warning

Run from the repository root, and do not rename the checkpoint directories. A diffusion checkpoint's frozen config resolves its paired VAE by a repository-relative run_dir read literally out of .hydra/config.yaml, and there is no flag that relocates it. The checkpoint itself must either sit inside its own run directory (as in the commands above) or be pointed at one with --run_dir — the loader walks up looking for .hydra/config.yaml and raises FileNotFoundError when there is none. See scripts/eval/README.md.

---

Quickstart

From clone to a number, on one GPU.

1. Environment

git clone https://github.com/ComDec/TopVAE.git
cd TopVAE
conda env create -f environment.yaml -p ./envs/topvae   # canonical install path
conda activate ./envs/topvae
pip install -r requirements.txt                         # the complete dependency list
pip install -e . --no-deps                              # makes `import src` resolve
python -c "import torch, torch_geometric, lightning, rdkit; \
           import src.models.vae.topvae_multistage; \
           assert torch.cuda.is_available(), 'no CUDA device -- the eval scripts would silently run on CPU'; \
           print('env OK', torch.__version__, 'cuda', torch.version.cuda)"

The requirements.txt line is required, not insurance: environment.yaml is unmodified lightning-hydra-template boilerplate and installs none of rdkit, torch_geometric, fcd_torch or posebusters. Use the import line above rather than python -c "import src", which succeeds on an environment with no torch at all and therefore tells you nothing.

Resolved versions the transcripts were produced against, and the prefix-vs-named env choice

python 3.10.20      torch 2.4.1+cu121             lightning 2.4.0
rdkit 2026.03.4     torch_geometric 2.8.0.post1
fcd_torch OK        posebusters OK                cuda True

environment.yaml's own name: field is mol. Passing -p ./envs/topvae overrides it and gives a prefix environment activated by path — the form every command in this repository is written for. Drop -p and you get a named environment mol instead; then use conda activate mol throughout. The two are not interchangeable in a copy-pasted command. requirements.txt is also the pip-only fallback for anyone not using conda. Full rationale: docs/reproduction.md §0.

2. Weights and data

pip install -U huggingface_hub

# fetch INTO the clone: the layout on the Hub is exactly the layout the configs expect,
# and every directory name is load-bearing
hf download EscheWang/TopVAE --local-dir .

(cd checkpoints && sha256sum -c SHA256SUMS)      # expect 15/15 OK; macOS: shasum -a 256 -c
(cd data        && sha256sum -c SHA256SUMS)      # expect 10/10 OK

The cd is not optional — each SHA256SUMS lists its paths relative to its own directory. A mismatch means you do not have the file these tables describe; do not proceed.

hf download fetches the whole Hub repository, not just the checkpoints: ≈ 3.2 GB on disk — ≈ 2.6 GB of weights (nine .ckpt/.pth files, two of them the same bytes under a second name so that both documented path forms resolve) plus ≈ 533 MB of preprocessed mol_list datasets. Budget disk for that, not for the weights alone. Provenance for every file: checkpoints/README.md §1.

3. Verify the data cache before you trust a number

python scripts/data_preparation/verify_dataset.py                                # --level files
python scripts/data_preparation/verify_dataset.py --dataset geom --level cache   # after the cache exists

Do not skip either line, and note that the first does not catch the thing that has actually gone wrong on this project. --level files only hashes the distributed files; the cache-contamination signature is checked exclusively by --level cache, which builds the datamodule and is the only invocation that can hard-abort with exit 2. A contaminated GEOM cache once produced wrong published numbers here, and the failure mode was plausible-looking wrong numbers rather than a crash. Levels, per-dataset coverage and exit codes: docs/reproduction.md §0a.

4. Run one of the two commands above. That is the whole path.

Cost, hardware, and the one known operational issue

Reproducing the GEOM row needs the entire GEOM-Drugs training split (299,859 molecules), not just the test split, because the shipped configs set auto_split: false and infer position_std and the novelty SMILES set from train. Budget ~12 GB for the datamodule cache built on first use. All reference runs were on one A100-80GB; the minimum VRAM was never measured, and --batch_size is the knob on a smaller card.

All seven experiment presets default to the W&B logger — append logger=csv to any src/eval.py command to run without an account.

Known issue. On one run, src/eval.py on QM9 with data.num_workers: 4 blocked indefinitely after the test loop, on a shared GPU. It did not recur. If an evaluation appears to hang after Testing DataLoader 0: 100%, append data.num_workers=0. docs/reproduction.md §5.

Training (entry point only — not part of the verified reproduction)

src/train.py is the Hydra training entry point and composes the same seven presets:

python src/train.py experiment=QM9_noH_Keku/topvae logger=csv

No training run is verified by this release, no transcript is recorded, and nothing here claims this command reproduces the released checkpoints. Everything reproducible in this repository goes through the two evaluation commands above.

---

What this release does not cover

Stated plainly, because it decides whether this repository is any use to you.

Paper table	Status
Table 2 — QM9 generation	Reproduces end to end, one command
Table 3 — GEOM-Drugs generation	Reproduces end to end, one command
Table 1 — VAE dark-area diagnosis	Does not assemble. See below
Table 4 — scaffold inpainting	Not in this release. Its entry point resolved models from a registry keyed on the authors' local run directories, and that script is not in this tree
Table 5 — component ablations	Not in this release. None of the three checkpoints it cites were located

• Three of Table 1's five rows ship, and all three run — QM9 UAE, QM9 TopVAE, GEOM-Drugs UAE. The other two are the QM9 ADiT-VAE row, a third-party baseline this release does not contain, and the GEOM-Drugs TopVAE row, whose checkpoint was not located. But they do not assemble into the paper's comparison: TopVAE and UAE are different LightningModules and log different metric sets, so a cell-by-cell comparison needs a common scorer this release does not ship. The measured values are in docs/reproduction.md §6b, which says so plainly. The GEOM-Drugs TopVAE checkpoint that row used was not located; the shipped GEOM TopVAE is a different run and its posterior numbers must not be quoted as that cell. Separately, iFID and the noise-scale sweep are implemented under src/metrics/ but no shipped command emits them (docs/metrics.md §6.6).
• Valence repair is applied asymmetrically in that posterior comparison: ON for TopVAE, OFF for UAE, because UAE's decoder is already sanitize-clean. It is deliberate, it is an asymmetry in a head-to-head, and it is documented at docs/metrics.md §4.1.
• QM9/uae_udm_qm9.ckpt ships but cannot be loaded by the generation script — there is no frozen .hydra/config.yaml anywhere above it, so you get a FileNotFoundError. The weights are genuine and hash-verified; the config was not recovered.
• Only GEOM-Drugs can be rebuilt from raw data, with export_geom_drugs_to_mol_list.py. There is no QM9 exporter — its dataset module imported openbabel unconditionally, which neither requirements.txt nor environment.yaml installs, so both were dropped. QM9 mol_list comes from the download only.
• PoseBusters, MMFF clash and scaffold diversity are implemented under src/metrics/ but no shipped command emits them. They are library code, not results this release reproduces.

Limitations the paper states about the method itself

• The constraints are graph-level only. ChemCO enforces valence rules on the graph; 3D physical priors — van der Waals clashes, force-field energy bounds, steric strain — are named as future work. (Beyond what the paper states: ChemCO's constraint set is valence limits, bond exclusivity, atom-pair legality and degree rules — it carries no aromaticity, formal-charge, stereochemistry or ring-strain term.)
• Not uniformly state of the art, and the paper says so. On 2D FCD, TopVAE trails UDM-3D and JODO on both datasets. On 3D geometry MMD the picture is split: on QM9 the paper reports TopVAE best on bond length and bond angle and behind UDM-3D only on dihedral, which it attributes to the topology-first factorization, "which prioritizes graph correctness over torsional accuracy"; on GEOM-Drugs it trails UDM-3D on all three.
• 3D physical plausibility remains hard for everyone. the best GEOM-Drugs MolStab3D in the paper's table is TopVAE's, and the released artifacts give 0.021600 — still near zero in absolute terms. The paper's MMFF94 study concedes that both generative models sit farther from equilibrium than the reference conformers do.
• Compute overhead is real, and larger than the forward pass alone suggests. The paper reports TopVAE at 4.3× slower per forward+backward step and 5.1× more peak GPU memory than UAE (3.11× per step under 4-GPU DDP). The gentler 2.2–2.4× figure is forward-only and covers molecules of 15–35 heavy atoms, not the full GEOM-Drugs range, which reaches 90. TopoBridge is cheap; ChemCO is the bottleneck, which the paper defends on the grounds that ChemCO is primarily a training-time component. This release does not test that defence: both documented commands run with --repair.
• TopoBridge can in principle over-connect. The paper raises this itself, answers it with degree and ring-size distribution statistics, and concedes that the edges BFS does insert have atypically long bond lengths, reflecting their last-resort nature.
• Baseline comparability is limited. UDM-3D ships no checkpoints and was re-trained by the authors "for reference only"; every other baseline cell is taken from its own paper. None of it is in this release.

---

What it generates

_{24 molecules from the GEOM-Drugs model, titled in-figure "TopVAE+DiT — Generated Molecules (GEOM-Drugs)". Each is annotated with QED (quantitative estimate of drug-likeness, 0–1, higher is more drug-like), SA (synthetic accessibility, 1–10, lower is easier to make), MW (molecular weight, in daltons) and atom count. This is a qualitative panel selected for display, not a metric.}

---

Repository layout

Path	Contents
src/	The library plus the two Hydra entry points, src/train.py and src/eval.py
configs/	Hydra composition tree; configs/experiment/ holds 7 presets — QM9_noH_Keku/{topvae,uae,udm} and GEOM_noH_Keku/{topvae,uae,udm,udm_ditB}
scripts/	Five standalone scripts in three directories; three run as shipped, two do not — the table there says which and why
checkpoints/	SHA256SUMS, the frozen .hydra/config.yaml files and both latent_distribution.pth, committed. The .ckpt weights are downloaded into this tree
data/	SHA256SUMS. The preprocessed mol_list/*.pt datasets are downloaded here
docs/	metrics.md (protocol, denominators, JSON schema), reproduction.md (verified commands and transcripts), fcd_reference_provenance.md (which FCD number came from which reference set)
tests/	pytest modules covering metrics, decode modes, noise utilities and dataset filtering; run them with make test. There is deliberately no test CI

---

Citation

The paper is not yet published; cite it as unpublished. Machine-readable metadata is in CITATION.cff.

@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.

Contact

Questions, reproduction failures and corrections belong in GitHub Issues. If a number you obtain disagrees with one above, please include the sampling_protocol block from your output JSON and the (cd checkpoints && sha256sum -c SHA256SUMS) output — those two are enough to tell a real discrepancy from a different checkpoint or a different protocol.