| """QM9 molecule generation: validity against denoising budget. |
| |
| Phase 2. Where the synthetic graph families give an exactly computable ceiling |
| but an artificial task, QM9 gives a real benchmark whose metric is unambiguous: |
| RDKit sanitisation either succeeds or it does not, and it enforces valency at |
| every atom simultaneously. That makes it a joint constraint of exactly the kind |
| parallel decoding breaks - mode (ii) - on data nobody designed for this project. |
| |
| The axis of comparison is the one the graph-diffusion literature competes on: |
| how much validity survives as the number of denoising steps falls. Discrete |
| graph diffusion conventionally uses hundreds of steps; the claim under test is |
| that budget-conditioned training plus an entropy-budgeted commit rule holds |
| validity at far fewer. |
| |
| All arms share the model, the data and the training budget, so the comparison |
| isolates the schedule and the decoding rule rather than scale. |
| """ |
| import argparse, json, math, os, time |
| import numpy as np |
| import torch |
|
|
| from model import Transformer |
| import diffusion as dfn |
| from qm9 import MolTokenizer, SEQ_LEN, load_qm9, validity |
|
|
|
|
| def get_args(): |
| p = argparse.ArgumentParser() |
| p.add_argument("--mode", |
| choices=["mdlm", "matched", "anybudget", "anybudget_nc"], |
| default="mdlm") |
| p.add_argument("--K", type=int, default=8, |
| help="budget the schedule is matched to (mode=matched)") |
| p.add_argument("--budget_bins", type=int, default=8) |
| p.add_argument("--pe", default="rope", |
| choices=["rope", "ape", "sin", "alibi", "nope"], |
| help="QM9 marginals are strongly position-dependent (slot 8 is " |
| "never occupied, some edge slots never bond), so this is " |
| "predicted to matter a great deal at low step counts") |
| p.add_argument("--root", default="data/qm9") |
| p.add_argument("--limit", type=int, default=0, help="0 = all molecules") |
| p.add_argument("--d", type=int, default=384) |
| p.add_argument("--layers", type=int, default=8) |
| p.add_argument("--heads", type=int, default=8) |
| p.add_argument("--bs", type=int, default=256) |
| p.add_argument("--steps", type=int, default=60000) |
| p.add_argument("--lr", type=float, default=1e-4) |
| p.add_argument("--warmup", type=int, default=2000) |
| p.add_argument("--seed", type=int, default=0) |
| p.add_argument("--n_gen", type=int, default=4096) |
| p.add_argument("--gen_bs", type=int, default=512) |
| p.add_argument("--out", default="runs/moldev") |
| return p.parse_args() |
|
|
|
|
| a = get_args() |
| torch.manual_seed(a.seed); np.random.seed(a.seed) |
| dev = "cuda" if torch.cuda.is_available() else "cpu" |
| os.makedirs(a.out, exist_ok=True) |
|
|
| X, tok, train_smiles = load_qm9(a.root, a.limit or None) |
| print(f"qm9 sequences {X.shape} vocab {len(tok)} train SMILES {len(train_smiles)}", |
| flush=True) |
| Xt = torch.from_numpy(X) |
| ALLOWED = torch.from_numpy(tok.allowed_mask()).to(dev) |
|
|
| COND = a.mode == "anybudget" |
| model = Transformer(len(tok), a.d, a.layers, a.heads, a.pe, causal=False, |
| max_len=SEQ_LEN + 8, |
| budget_bins=a.budget_bins if COND else 0).to(dev) |
| print(f"params {model.n_params()/1e6:.2f}M mode={a.mode} pe={a.pe}", flush=True) |
|
|
| opt = torch.optim.AdamW(model.parameters(), lr=a.lr, weight_decay=0.01) |
| sched = torch.optim.lr_scheduler.LambdaLR( |
| opt, lambda s: min(1.0, (s + 1) / a.warmup) * |
| 0.5 * (1 + math.cos(math.pi * min(1.0, s / a.steps)))) |
| BUDGETS = tuple(2 ** i for i in range(int(math.log2(SEQ_LEN)) + 2)) |
|
|
| t0 = time.time() |
| for step in range(a.steps): |
| xb = Xt[torch.randint(0, len(Xt), (a.bs,))].to(dev) |
| opt.zero_grad(set_to_none=True) |
| with torch.autocast("cuda", dtype=torch.bfloat16, enabled=(dev == "cuda")): |
| if a.mode in ("anybudget", "anybudget_nc"): |
| loss = dfn.uncond_any_budget_loss(model, xb, tok, choices=BUDGETS, |
| condition=COND) |
| elif a.mode == "matched": |
| loss = dfn.uncond_mdlm_loss(model, xb, tok, "matched", a.K) |
| else: |
| loss = dfn.uncond_mdlm_loss(model, xb, tok, "uniform") |
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) |
| opt.step(); sched.step() |
| if step % 5000 == 0: |
| print(f"step {step:6d} loss {loss.item():.4f} ({time.time()-t0:.0f}s)", flush=True) |
|
|
| model.eval() |
| res = {"args": vars(a), "decode": {}} |
|
|
|
|
| @torch.no_grad() |
| def gen(fn): |
| outs, nfes = [], [] |
| for _ in range(max(1, a.n_gen // a.gen_bs)): |
| blank = torch.full((a.gen_bs, SEQ_LEN), tok.mask, dtype=torch.long) |
| x, nfe = fn(blank) |
| outs.append(x.cpu()); nfes.append(nfe) |
| x = torch.cat(outs).numpy() |
| v, smis = validity(x, tok) |
| uniq = len(set(smis)) / max(1, len(smis)) |
| novel = sum(s not in train_smiles for s in set(smis)) / max(1, len(set(smis))) |
| return {"validity": v, "nfe": float(np.mean(nfes)), "unique": uniq, |
| "novel": novel} |
|
|
|
|
| def bud(K): |
| return torch.full((a.gen_bs,), float(K), device=dev) if COND else None |
|
|
|
|
| print("\n=== validity / uniqueness / novelty against denoising budget ===", flush=True) |
| print(f"{'rule':<24}{'NFE':>7}{'valid':>9}{'unique':>9}{'novel':>9}", flush=True) |
|
|
|
|
| def report(tag, key, r): |
| res["decode"][key] = r |
| print(f"{tag:<24}{r['nfe']:>7.1f}{r['validity']*100:>8.2f}%" |
| f"{r['unique']*100:>8.1f}%{r['novel']*100:>8.1f}%", flush=True) |
|
|
|
|
| for K in (1, 2, 4, 8, 16, 32, SEQ_LEN): |
| report(f"fixed-K K={K}", f"fixedK_{K}", |
| gen(lambda b, K=K: dfn.fixed_k_decode(model, b, tok, dev, K, |
| budget_cond=bud(K), allowed=ALLOWED))) |
|
|
| for B in (0.05, 0.25, 1.0, 4.0): |
| report(f"entropy-budget B={B}", f"entbudget_{B}", |
| gen(lambda b, B=B: dfn.entropy_budget_decode( |
| model, b, tok, dev, B, max_steps=SEQ_LEN + 2, |
| budget_cond=bud(8), allowed=ALLOWED))) |
|
|
| for tau in (0.7, 0.9, 0.99): |
| report(f"confidence tau={tau}", f"conf_{tau}", |
| gen(lambda b, tau=tau: dfn.confidence_threshold_decode( |
| model, b, tok, dev, tau, max_steps=SEQ_LEN + 2, |
| budget_cond=bud(8), allowed=ALLOWED))) |
|
|
| torch.save(model.state_dict(), os.path.join(a.out, "model.pt")) |
| json.dump(res, open(os.path.join(a.out, "result.json"), "w"), indent=2) |
| print(f"\nsaved to {a.out}", flush=True) |
|
|
