| """Is the QM9 result real, or an artifact of our own representation? |
| |
| We encode a molecule as 9 fixed node slots plus 36 fixed edge slots. That forces |
| the model to learn slot-specific facts - slot 8 is never occupied, slot 7 is |
| occupied 80% of the time - which is precisely the position-dependence a relative |
| encoding cannot represent at the fully-masked state. Our RoPE baseline caps below |
| 80% validity, and we attributed that to the encoding. |
| |
| But DiGress and the graph-diffusion line do NOT generate this way. They sample |
| the number of atoms n from the training distribution FIRST, then generate an n-node |
| graph with a permutation-equivariant model. "Slot 8 is empty" never arises for |
| them. If that is the whole of our baseline's disadvantage, then our comparison is |
| against a strawman of our own construction and the QM9 claim does not stand. |
| |
| This script removes exactly that disadvantage, without retraining. For each |
| sample it draws n from the empirical training distribution, fixes node slots |
| n..8 and every edge touching them to "none" as OBSERVED context, and decodes only |
| the remaining positions. That is the DiGress setup expressed in our |
| representation, and it is strictly generous to the baseline. |
| |
| RoPE recovers -> the finding was representation-specific and the QM9 claim |
| must be withdrawn |
| RoPE still fails -> the disadvantage is not the padding slots, and the finding |
| survives a fair baseline |
| |
| Either way the answer goes in the paper. |
| """ |
| import argparse, glob, json, os, sys |
| import numpy as np |
| import torch |
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from model import Transformer |
| import diffusion as dfn |
| from qm9 import EDGES, MAX_ATOMS, SEQ_LEN, MolTokenizer, load_qm9, validity |
|
|
| p = argparse.ArgumentParser() |
| p.add_argument("--runs", default="runs") |
| p.add_argument("--root", default="data/qm9") |
| p.add_argument("--n_gen", type=int, default=2048) |
| p.add_argument("--gen_bs", type=int, default=512) |
| p.add_argument("--out", default="figures/mol_fair.json") |
| a = p.parse_args() |
|
|
| dev = "cuda" if torch.cuda.is_available() else "cpu" |
| X, tok, train_smiles = load_qm9(a.root, 40000) |
| ALLOWED = torch.from_numpy(tok.allowed_mask()).to(dev) |
|
|
| |
| counts = np.array([(row[:MAX_ATOMS] != tok.none).sum() for row in X]) |
| sizes, freqs = np.unique(counts, return_counts=True) |
| probs = freqs / freqs.sum() |
| print(f"atom-count distribution from training data: " |
| f"{dict(zip(sizes.tolist(), np.round(probs, 3).tolist()))}", flush=True) |
|
|
|
|
| def scaffold(bs, rng): |
| """A canvas with n sampled first: slots beyond n are OBSERVED as 'none'. |
| |
| This is the advantage DiGress has and our flat encoding withholds: the model |
| never has to discover which slots exist, because the unused ones are given. |
| """ |
| x = torch.full((bs, SEQ_LEN), tok.mask, dtype=torch.long) |
| fill = torch.ones((bs, SEQ_LEN), dtype=torch.bool) |
| ns = rng.choice(sizes, size=bs, p=probs) |
| for b, n in enumerate(ns): |
| for j in range(int(n), MAX_ATOMS): |
| x[b, j] = tok.none |
| fill[b, j] = False |
| for k, (i, j) in enumerate(EDGES): |
| if i >= n or j >= n: |
| x[b, MAX_ATOMS + k] = tok.none |
| fill[b, MAX_ATOMS + k] = False |
| return x, fill |
|
|
|
|
| def load(run_dir): |
| cfg = json.load(open(os.path.join(run_dir, "result.json")))["args"] |
| m = Transformer(len(tok), cfg["d"], cfg["layers"], cfg["heads"], cfg["pe"], |
| causal=False, max_len=SEQ_LEN + 8).to(dev) |
| m.load_state_dict(torch.load(os.path.join(run_dir, "model.pt"), map_location=dev)) |
| m.eval() |
| return m, cfg |
|
|
|
|
| @torch.no_grad() |
| def run(model, steps, seed): |
| rng = np.random.default_rng(seed) |
| outs = [] |
| for _ in range(max(1, a.n_gen // a.gen_bs)): |
| x, fill = scaffold(a.gen_bs, rng) |
| y, _ = dfn.fixed_k_decode(model, x.to(dev), tok, dev, steps, |
| fillable=fill.to(dev), allowed=ALLOWED) |
| outs.append(y.cpu()) |
| y = torch.cat(outs).numpy() |
| v, smis = validity(y, tok) |
| uniq = len(set(smis)) / max(1, len(smis)) |
| return v, uniq |
|
|
|
|
| res = {} |
| print(f"\n{'arm':<28}{'steps':>7}{'validity':>11}{'unique':>9}", flush=True) |
| for d in sorted(glob.glob(os.path.join(a.runs, "mol1-*"))): |
| if not os.path.exists(os.path.join(d, "model.pt")): |
| continue |
| m, cfg = load(d) |
| name = os.path.basename(d) |
| key = f"{cfg['mode']}-{cfg['pe']}" |
| for steps in (1, 4, 8, 16, 32, 45): |
| v, u = run(m, steps, cfg["seed"]) |
| res.setdefault(key, {}).setdefault(str(steps), []).append(v) |
| print(f"{name:<28}{steps:>7}{v*100:>10.2f}%{u*100:>8.1f}%", flush=True) |
| del m |
| torch.cuda.empty_cache() |
|
|
| print(f"\n{'='*70}") |
| print("FAIR BASELINE: atom count given, unused slots observed (the DiGress setup)") |
| print(f"{'='*70}") |
| print(f"{'arm':<22}" + "".join(f"{s:>9}" for s in (1, 4, 8, 16, 32, 45))) |
| for key in sorted(res): |
| print(f"{key:<22}" + "".join(f"{np.mean(res[key][str(s)])*100:>8.1f}%" |
| for s in (1, 4, 8, 16, 32, 45))) |
| os.makedirs(os.path.dirname(a.out) or ".", exist_ok=True) |
| json.dump(res, open(a.out, "w"), indent=2) |
| print(f"\nwrote {a.out}") |
| print("\nRead: if the relative-encoding arms now track the absolute ones, the QM9") |
| print("claim was an artifact of the flat fixed-slot encoding and must be withdrawn.") |
|
|
