File size: 6,408 Bytes
1c9c214 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 | """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)
|