File size: 7,562 Bytes
ae02605 | 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 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 | """Does exact verifier feedback fix the absorbing-commit failure?
Diagnosis from the trained models: 94.9% of cells are correct but only 87.9% of
boards. Failed boards therefore carry roughly three wrong cells out of
fifty-seven - they are near misses, not collapses. Masked diffusion decoding is
absorbing, so those three cells can never be revised and each one destroys a
whole board.
Sudoku constraints are exactly checkable, so the repair signal is not a
heuristic: we know precisely which cells conflict. Remask exactly those (widened
to their row, column and box, since re-deciding one cell in an unchanged context
reproduces the same value), re-decode, repeat.
Every forward pass is counted, including the repair rounds, so the comparison
against fixed-K budgets is honest rather than flattering: a method that needs
sixty passes to repair is reported at sixty.
The control that matters is RANDOM remasking of the same number of cells. If
random repair helps as much as verifier-guided repair, then the gain is just
extra compute and the verifier is doing nothing.
"""
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 sudoku import (CELLS, SudokuTokenizer, expand_to_units, violation_mask)
p = argparse.ArgumentParser()
p.add_argument("--hard", default="data/sudoku_hard")
p.add_argument("--runs", default="runs")
p.add_argument("--pattern", default="sudh-*")
p.add_argument("--eval_bs", type=int, default=250)
p.add_argument("--out", default="figures/sudoku_remask.json")
a = p.parse_args()
dev = "cuda" if torch.cuda.is_available() else "cpu"
tok = SudokuTokenizer()
P = np.load(os.path.join(a.hard, "test_puz.npy"))
S = np.load(os.path.join(a.hard, "test_sol.npy"))
clues = (P > 0).sum(1)
print(f"benchmark: {len(P)} puzzles, clues {clues.min()}-{clues.max()} "
f"(mean {clues.mean():.1f})", flush=True)
Pt, St = torch.from_numpy(P), torch.from_numpy(S - 1)
blank = Pt == 0
def load(d):
cfg = json.load(open(os.path.join(d, "result.json")))["args"]
m = Transformer(len(tok), cfg["d"], cfg["layers"], cfg["heads"], cfg["pe"],
causal=False, max_len=CELLS + 8).to(dev)
m.load_state_dict(torch.load(os.path.join(d, "model.pt"), map_location=dev))
m.eval()
return m, cfg
def random_remask_like(bad_counts):
"""Control: remask the same NUMBER of cells, chosen at random."""
def fn(x):
B, L = x.shape
out = torch.zeros_like(x, dtype=torch.bool)
for b in range(B):
k = int(bad_counts[b])
if k:
idx = torch.randperm(L, device=x.device)[:k]
out[b, idx] = True
return out
return fn
@torch.no_grad()
def evaluate(m, fn):
preds, nfes = [], []
for i in range(0, len(Pt), a.eval_bs):
gold = St[i:i + a.eval_bs].to(dev)
bl = blank[i:i + a.eval_bs].to(dev)
x = torch.where(bl, torch.full_like(gold, tok.mask), gold)
y, nfe = fn(x, bl)
preds.append(y.cpu().numpy() + 1)
nfes.append(nfe)
pred = np.concatenate(preds)
board = float((pred == S).all(axis=1).mean())
cell = float((pred[P == 0] == S[P == 0]).mean())
return board, cell, float(np.mean(nfes)), pred
res = {}
dirs = sorted(glob.glob(os.path.join(a.runs, a.pattern)))
dirs = [d for d in dirs if os.path.exists(os.path.join(d, "model.pt"))]
print(f"evaluating {len(dirs)} models\n", flush=True)
for d in dirs:
m, cfg = load(d)
name = os.path.basename(d)
print(f"--- {name}", flush=True)
print(f" {'decoder':<34}{'NFE':>8}{'BOARD':>10}{'cell':>9}", flush=True)
row = {}
for K in (16, 61):
b, c, n, _ = evaluate(m, lambda x, bl, K=K: dfn.fixed_k_decode(
m, x, tok, dev, K, fillable=bl))
row[f"fixedK_{K}"] = {"board": b, "cell": c, "nfe": n}
print(f" {'fixed-K K=' + str(K):<34}{n:>8.1f}{b*100:>9.2f}%{c*100:>8.2f}%",
flush=True)
b, c, n, _ = evaluate(m, lambda x, bl: dfn.entropy_budget_decode(
m, x, tok, dev, 0.01, max_steps=CELLS + 2, fillable=bl))
row["entbudget"] = {"board": b, "cell": c, "nfe": n}
print(f" {'entropy-budget B=0.01':<34}{n:>8.1f}{b*100:>9.2f}%{c*100:>8.2f}%",
flush=True)
# Repair must start from the STRONGEST available decode, not a weak one.
# An earlier version used base_steps=16, which decodes to ~83% before repair,
# so the repaired result never beat plain 61-pass decoding at ~88% - the
# method was handicapped by its starting point, not by the idea.
for base in (32, 61):
for rounds in (4, 8):
b, c, n, _ = evaluate(
m, lambda x, bl, r=rounds, bs=base: dfn.verifier_remask_decode(
m, x, tok, dev, violation_mask, base_steps=bs, rounds=r,
max_revisits=3, fillable=bl, expand_peers=expand_to_units))
row[f"verifier_b{base}_r{rounds}"] = {"board": b, "cell": c, "nfe": n}
print(f" {'VERIFIER base=' + str(base) + ' rounds=' + str(rounds):<34}"
f"{n:>8.1f}{b*100:>9.2f}%{c*100:>8.2f}%", flush=True)
# narrow repair: remask ONLY the conflicting cells, not their whole units
b, c, n, _ = evaluate(m, lambda x, bl: dfn.verifier_remask_decode(
m, x, tok, dev, violation_mask, base_steps=61, rounds=8,
max_revisits=3, fillable=bl, expand_peers=None))
row["verifier_narrow"] = {"board": b, "cell": c, "nfe": n}
print(f" {'VERIFIER base=61 narrow (no units)':<34}{n:>8.1f}{b*100:>9.2f}%"
f"{c*100:>8.2f}%", flush=True)
# control: same budget, but remask at random instead of where it is wrong
b, c, n, _ = evaluate(m, lambda x, bl: dfn.verifier_remask_decode(
m, x, tok, dev,
lambda z: (torch.rand_like(z, dtype=torch.float) < 0.12),
base_steps=61, rounds=8, max_revisits=3, fillable=bl,
expand_peers=None))
row["random_remask_control"] = {"board": b, "cell": c, "nfe": n}
print(f" {'CONTROL: random remask, 8 rounds':<34}{n:>8.1f}{b*100:>9.2f}%"
f"{c*100:>8.2f}%", flush=True)
res[name] = row
del m
torch.cuda.empty_cache()
print(flush=True)
if res:
print("=" * 76)
print("SUMMARY (mean over models)")
print("=" * 76)
keys = list(next(iter(res.values())).keys())
print(f"{'decoder':<34}{'NFE':>8}{'BOARD':>10}")
base = None
for k in keys:
bs = [r[k]["board"] for r in res.values()]
ns = [r[k]["nfe"] for r in res.values()]
if k == "fixedK_61":
base = float(np.mean(bs))
print(f"{k:<34}{np.mean(ns):>8.1f}{np.mean(bs)*100:>9.2f}%")
if base is not None:
best_v = max((np.mean([r[k]["board"] for r in res.values()])
for k in keys if k.startswith("verifier_")), default=0)
ctrl = np.mean([r["random_remask_control"]["board"] for r in res.values()])
print(f"\n best verifier-guided : {best_v*100:.2f}%")
print(f" fixed-K at 61 passes : {base*100:.2f}% "
f"(verifier gains {(best_v-base)*100:+.2f} pts)")
print(f" random-remask control: {ctrl*100:.2f}% "
f"({(ctrl-base)*100:+.2f} pts)")
print("\n If the control gains as much, the verifier is doing nothing and")
print(" the improvement is only the extra compute.")
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}")
|