"""Masked diffusion on text8, and the experiment that tests the unified claim. Addition isolates failure mode (i): every conditional is a point mass, so all of the loss from parallel decoding comes from marginals that were never trained at high mask ratios, and no inference-time rule recovers any of it. The free-choice probes isolate failure mode (ii): the conditionals are maximum-entropy, nothing is trainable, and committing less is the only remedy. Natural text contains both in the same sequence. Finishing a word, a suffix or a function word is nearly deterministic; choosing the next content word is not. So text is where the two fixes have to compose, and this script measures whether they do, as a 2x2: fixed-K commit entropy-budgeted commit uniform-t training (a) (b) matched training (c) (d) The prediction is (d) < (b), (c) < (a) on generative perplexity at equal compute, and that the two gains are largely independent - one fixes what the other cannot touch. Metrics, all standard for this benchmark: * validation bits-per-character (the MDLM bound) - likelihood, comparable to published text8 numbers, and the check that the schedule change does not buy sample quality by wrecking the model. * generative perplexity under a held-out autoregressive evaluator, against the number of function evaluations actually used. * entropy and distinct-n-gram rate of the samples, reported alongside, because generative perplexity alone is trivially gamed by degenerate repetition. """ import argparse, json, math, os, time import numpy as np import torch import torch.nn.functional as F from model import Transformer from text8 import load_text8, batches import diffusion as dfn def get_args(): p = argparse.ArgumentParser() p.add_argument("--mode", choices=["mdlm", "matched", "full", "anybudget", "anybudget_nc"], default="mdlm") p.add_argument("--budget_bins", type=int, default=10) p.add_argument("--K", type=int, default=8, help="parallelism the schedule is matched to") p.add_argument("--lam", type=float, default=1.0) p.add_argument("--root", default="data/text8") p.add_argument("--seq_len", type=int, default=256) p.add_argument("--d", type=int, default=512) p.add_argument("--layers", type=int, default=12) p.add_argument("--heads", type=int, default=8) p.add_argument("--bs", type=int, default=64) p.add_argument("--steps", type=int, default=80000) p.add_argument("--lr", type=float, default=1e-4) p.add_argument("--warmup", type=int, default=4000) p.add_argument("--seed", type=int, default=0) p.add_argument("--n_gen", type=int, default=256, help="samples per decoding setting") p.add_argument("--gen_bs", type=int, default=64) p.add_argument("--evallm", default="runs/evallm") p.add_argument("--out", default="runs/t8dev") 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) tr, va, te, tok = load_text8(a.root, a.seq_len) print(f"text8: train {tr.shape} valid {va.shape} test {te.shape} vocab {len(tok)}", flush=True) COND = a.mode == "anybudget" model = Transformer(len(tok), a.d, a.layers, a.heads, "rope", causal=False, max_len=a.seq_len + 8, budget_bins=a.budget_bins if COND else 0).to(dev) print(f"denoiser params: {sum(q.numel() for q in model.parameters())/1e6:.1f}M", flush=True) opt = torch.optim.AdamW(model.parameters(), lr=a.lr, weight_decay=0.01, betas=(0.9, 0.95)) sched = torch.optim.lr_scheduler.LambdaLR( opt, lambda s: min((s + 1) / a.warmup, max(0.0, (a.steps - s) / max(1, a.steps - a.warmup)))) it = batches(tr, a.bs, seed=a.seed) t0 = time.time() for step in range(a.steps): x = torch.from_numpy(next(it)).to(dev) with torch.autocast("cuda", dtype=torch.bfloat16, enabled=(dev == "cuda")): if a.mode in ("anybudget", "anybudget_nc"): loss = dfn.uncond_any_budget_loss(model, x, tok, condition=COND) else: loss = dfn.text_loss(model, x, tok, a.mode, a.K, a.lam) opt.zero_grad(set_to_none=True) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) opt.step(); sched.step() if step % 2000 == 0: print(f"step {step:6d} loss {loss.item():.4f} {time.time()-t0:.0f}s", flush=True) model.eval() res = {"args": vars(a)} # ---- likelihood: the standard text8 number --------------------------------- vb = [torch.from_numpy(va[i:i + a.bs]) for i in range(0, min(64 * a.bs, len(va)), a.bs)] res["valid_bpc"] = dfn.nll_bpc(model, vb, tok, dev, mc=8) print(f"\nvalidation BPC (MDLM bound): {res['valid_bpc']:.4f}", flush=True) # ---- the evaluator --------------------------------------------------------- ev = None ev_path = os.path.join(a.evallm, "model.pt") if os.path.exists(ev_path): ec = json.load(open(os.path.join(a.evallm, "result.json")))["args"] ev = Transformer(len(tok), ec["d"], ec["layers"], ec["heads"], "rope", causal=True, max_len=ec["seq_len"] + 8).to(dev) ev.load_state_dict(torch.load(ev_path, map_location=dev)) ev.eval() print(f"loaded evaluator from {a.evallm}", flush=True) else: print(f"WARNING: no evaluator at {ev_path}; generative perplexity skipped", flush=True) @torch.no_grad() def gen_ppl(x): """Per-character perplexity the held-out autoregressive model assigns.""" tot, cnt = 0.0, 0 for i in range(0, len(x), a.gen_bs): b = x[i:i + a.gen_bs].to(dev) ce = F.cross_entropy(ev(b)[:, :-1].reshape(-1, len(tok)), b[:, 1:].reshape(-1), reduction="sum") tot += float(ce); cnt += b[:, 1:].numel() return math.exp(tot / cnt) def sample_stats(x): """Entropy and distinct-n-gram rate: the degeneracy guard. A model that emits the same low-perplexity string every time wins on generative perplexity while being useless, so the unigram entropy and the fraction of distinct 4-grams are reported next to it. Reference values from real text8 windows are computed below for comparison. """ f = np.bincount(x.reshape(-1).numpy(), minlength=len(tok)).astype(float) f = f[f > 0] / f.sum() ent = float(-(f * np.log(f)).sum()) grams = set() tot = 0 for row in x.numpy(): for i in range(len(row) - 3): grams.add(tuple(row[i:i + 4])); tot += 1 return ent, len(grams) / max(tot, 1) ref = torch.from_numpy(te[:a.n_gen]) res["reference"] = {"gen_ppl": gen_ppl(ref) if ev else None, "entropy": sample_stats(ref)[0], "distinct4": sample_stats(ref)[1]} print(f"real text8 reference: ppl {res['reference']['gen_ppl']}, " f"entropy {res['reference']['entropy']:.3f}, " f"distinct-4 {res['reference']['distinct4']:.3f}", flush=True) def blank(): return torch.full((a.gen_bs, a.seq_len), tok.mask, dtype=torch.long) def evaluate(tag, fn): """Run one decoding rule over n_gen samples and record quality vs compute.""" outs, nfes = [], [] for _ in range(max(1, a.n_gen // a.gen_bs)): x, nfe = fn(blank()) outs.append(x.cpu()); nfes.append(nfe) x = torch.cat(outs) ent, d4 = sample_stats(x) row = {"nfe": float(np.mean(nfes)), "gen_ppl": gen_ppl(x) if ev else None, "entropy": ent, "distinct4": d4, "sample": tok.decode(x[0][:120].tolist())} print(f" {tag:<26} nfe {row['nfe']:6.1f} ppl {row['gen_ppl']:9.2f} " f"H {ent:.3f} d4 {d4:.3f}", flush=True) return row print("\n=== decoding: quality against compute ===", flush=True) res["decode"] = {} def bud(K): """The budget to declare to the model, or None if it is not conditioned.""" return (torch.full((a.gen_bs,), float(K), device=dev)) if COND else None print("fixed-K (non-adaptive reference):", flush=True) for K in (1, 2, 4, 8, 16, 32, 64, 128, 256): res["decode"][f"fixedK_{K}"] = evaluate( f"fixed-K K={K}", lambda b, K=K: dfn.fixed_k_decode(model, b, tok, dev, K, budget_cond=bud(K))) print("confidence-threshold (published inference-time method):", flush=True) for tau in (0.5, 0.7, 0.9, 0.95, 0.99): res["decode"][f"conf_{tau}"] = evaluate( f"confidence tau={tau}", lambda b, tau=tau: dfn.confidence_threshold_decode( model, b, tok, dev, tau, budget_cond=bud(8))) print("entropy-budget (ours, fix (ii)):", flush=True) for B in (0.05, 0.1, 0.25, 0.5, 1.0, 2.0, 4.0, 8.0): res["decode"][f"entbudget_{B}"] = evaluate( f"entropy-budget B={B}", lambda b, B=B: dfn.entropy_budget_decode( model, b, tok, dev, B, budget_cond=bud(8))) 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)