| """Masked diffusion on structured graphs, where the ceiling is known exactly. |
| |
| This is the experiment the project could not run before. On every earlier task |
| the best achievable parallel-decoding accuracy was unknown, so a gap could |
| always be blamed on undertraining. Here it is computable in closed form: |
| |
| V* = sum over valid graphs of prod_i p_i(g_i) and V* >= 2^(-TC), |
| |
| with equality when the family has a fixed edge count. V* is the probability |
| that independent draws from the TRUE marginals land on a valid graph, so it is |
| the one-pass ceiling for any model whatsoever - it already assumes the marginals |
| are perfect. |
| |
| That makes the logic airtight in a way it was not on addition or text: |
| |
| * if a trained model reaches V* at one pass, the remaining gap to 100% is |
| PROVABLY information-theoretic - mode (ii) - and no training objective can |
| recover it; |
| * if it falls short of V*, the shortfall is exactly mode (i), the trainable |
| part, and budget conditioning should close it. |
| |
| The two failure modes are therefore separated numerically rather than argued |
| about, on the same task, at the same time. |
| |
| A graph on n nodes is a binary sequence of n(n-1)/2 edge indicators, so the |
| masked-diffusion machinery is used unchanged - only the data and the validity |
| check are new. |
| """ |
| import argparse, json, math, os, time |
| import numpy as np |
| import torch |
|
|
| from model import Transformer |
| import diffusion as dfn |
| from graphs import (EdgeTokenizer, exact_stats, is_valid, n_edges, sample) |
|
|
|
|
| def get_args(): |
| p = argparse.ArgumentParser() |
| p.add_argument("--family", default="matching", |
| choices=["matching", "2regular", "tree", "trianglefree", |
| "bipartite"]) |
| p.add_argument("--n", type=int, default=6) |
| p.add_argument("--mode", choices=["mdlm", "anybudget", "anybudget_nc"], |
| default="mdlm") |
| p.add_argument("--budget_bins", type=int, default=8) |
| p.add_argument("--pe", default="rope", choices=["rope", "ape", "sin", "alibi", "nope"], |
| help="position encoding; decisive at the fully-masked state") |
| p.add_argument("--d", type=int, default=256) |
| p.add_argument("--layers", type=int, default=6) |
| p.add_argument("--heads", type=int, default=8) |
| p.add_argument("--bs", type=int, default=256) |
| p.add_argument("--steps", type=int, default=30000) |
| p.add_argument("--lr", type=float, default=1e-4) |
| p.add_argument("--warmup", type=int, default=1000) |
| 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/gdev") |
| 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) |
|
|
| E = n_edges(a.n) |
| tok = EdgeTokenizer() |
| stats = exact_stats(a.family, a.n) |
| print(f"family={a.family} n={a.n} edges={E} |family|={stats['size']}", flush=True) |
| print(f" sum_i H(e_i) = {stats['sum_marginal_entropy_bits']:.3f} bits (the bound)", flush=True) |
| print(f" H(S) = {stats['joint_entropy_bits']:.3f} bits", flush=True) |
| print(f" TC = {stats['total_correlation_bits']:.3f} bits", flush=True) |
| print(f" V* ceiling = {stats['one_pass_ceiling']*100:.4f}% " |
| f"(2^-TC = {2**-stats['total_correlation_bits']*100:.4f}%)", flush=True) |
|
|
| X = torch.from_numpy(sample(a.family, a.n, 200_000, seed=a.seed)) |
|
|
| COND = a.mode == "anybudget" |
| model = Transformer(len(tok), a.d, a.layers, a.heads, a.pe, causal=False, |
| max_len=E + 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)))) |
|
|
| t0 = time.time() |
| for step in range(a.steps): |
| idx = torch.randint(0, len(X), (a.bs,)) |
| xb = X[idx].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=tuple(2 ** i for i in range(int(math.log2(E)) + 1)), |
| condition=COND) |
| 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), "exact": {k: v for k, v in stats.items() if k != "marginals"}} |
|
|
|
|
| @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, E), tok.mask, dtype=torch.long) |
| x, nfe = fn(blank) |
| outs.append(x.cpu()); nfes.append(nfe) |
| x = torch.cat(outs).numpy() |
| ok = is_valid(x, a.family, a.n) |
| uniq = len({tuple(r) for r in x[ok]}) / max(1, int(ok.sum())) |
| return float(ok.mean()), float(np.mean(nfes)), float(uniq) |
|
|
|
|
| def bud(K): |
| return torch.full((a.gen_bs,), float(K), device=dev) if COND else None |
|
|
|
|
| print("\n=== validity against decoding budget ===", flush=True) |
| print(f"{'rule':<22}{'NFE':>7}{'validity':>11}{'vs V*':>10}{'uniq':>8}", flush=True) |
| res["decode"] = {} |
| V = stats["one_pass_ceiling"] |
|
|
| for K in sorted({1, 2, 4, 8, E}): |
| v, nfe, u = gen(lambda b, K=K: dfn.fixed_k_decode(model, b, tok, dev, K, |
| budget_cond=bud(K))) |
| res["decode"][f"fixedK_{K}"] = {"validity": v, "nfe": nfe, "unique": u} |
| rel = f"{v/V:8.2f}x" if K == 1 else f"{'—':>9}" |
| print(f"{'fixed-K K=' + str(K):<22}{nfe:>7.1f}{v*100:>10.2f}%{rel}{u*100:>7.1f}%", flush=True) |
|
|
| for B in (0.05, 0.25, 1.0, 4.0): |
| v, nfe, u = gen(lambda b, B=B: dfn.entropy_budget_decode(model, b, tok, dev, B, |
| max_steps=E + 2, |
| budget_cond=bud(8))) |
| res["decode"][f"entbudget_{B}"] = {"validity": v, "nfe": nfe, "unique": u} |
| print(f"{'entropy-budget B=' + str(B):<22}{nfe:>7.1f}{v*100:>10.2f}%{'—':>9}" |
| f"{u*100:>7.1f}%", flush=True) |
|
|
| one = res["decode"]["fixedK_1"]["validity"] |
| res["reached_ceiling"] = one >= 0.9 * V |
| print(f"\nV* (exact one-pass ceiling) : {V*100:.4f}%", flush=True) |
| print(f"achieved at one pass : {one*100:.4f}% " |
| f"({one/V*100:.1f}% of the ceiling)", flush=True) |
| print("reading: " + ("the model is at the information-theoretic limit, so the " |
| "residual gap is mode (ii) and NOT trainable" |
| if res["reached_ceiling"] else |
| "the model falls short of the ceiling, so part of the gap " |
| "is mode (i) and IS trainable"), flush=True) |
|
|
| 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) |
|
|
