File size: 7,282 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 151 152 153 154 155 156 157 158 159 160 161 162 163 164 | """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)
|