File size: 5,847 Bytes
7115984 | 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 | """Architecture validation before any real training.
Checks:
1. Dataset integrity: answer marginals ~uniform per k (no majority-class
shortcut); LookupSolver scores 100% (pipeline exactness).
2. Gradient flow through recurrence: a 1-block model unrolled R=8 gets
non-zero gradients on every block parameter; loss at init is finite.
3. Overfitting a tiny fixed training set: both architectures reach ~100%
train accuracy on 256 fixed examples (gradients + capacity are real).
4. Interface checks: n_iters variation works, smaller state spaces evaluate
with the same weights (subset of vocab/positions), shapes correct.
Exit code 0 iff all checks pass. Run: python validate.py
"""
import sys
import time
import numpy as np
import torch
torch.set_num_threads(8)
import torch.nn.functional as F
from ptr.common import (LookupSolver, check_answer_marginals,
sample_eval_batch, sample_train_batch)
from ptr.models import PtrTransformer
DEV = "cpu"
S, D = 64, 128
def check_marginals_and_lookup():
rng = np.random.default_rng(0)
m = check_answer_marginals(rng, S=S, k_max=8, n=100_000)
print(f"[1a] answer-marginal TV per k: "
f"{ {k: round(v, 4) for k, v in m['tv_per_k'].items()} } -> "
f"{'PASS' if m['pass'] else 'FAIL'}")
assert m["pass"]
solver = LookupSolver()
b = sample_eval_batch(rng, 1000, S, k=8)
pred = solver.solve(b["f"], b["x"], 8)
acc = float((pred == b["answer"]).mean())
print(f"[1b] LookupSolver exact-match on random data: {acc:.4f} -> "
f"{'PASS' if acc == 1.0 else 'FAIL'}")
assert acc == 1.0
def check_gradients():
torch.manual_seed(0)
model = PtrTransformer(S, D, 4, 512, depth=1, zero_init=True).to(DEV)
rng = np.random.default_rng(1)
b = sample_train_batch(rng, 64, S)
f = torch.as_tensor(b["f"], device=DEV)
x = torch.as_tensor(b["x"], device=DEV)
k = torch.as_tensor(b["k"], device=DEV)
y = torch.as_tensor(b["answer"], device=DEV)
R = 8
# Zero-initialised output projections make every block the identity at
# step 0, so upstream parameters receive no gradient on the very first
# backward. Warm up a few AdamW steps first, then check that the full
# R=8 unroll propagates gradients to every block parameter.
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
for _ in range(3):
opt.zero_grad(set_to_none=True)
F.cross_entropy(model(f, x, k, n_iters=R), y).backward()
opt.step()
opt.zero_grad(set_to_none=True)
logits = model(f, x, k, n_iters=R)
loss = F.cross_entropy(logits, y)
loss.backward()
bad = [n for n, p in model.named_parameters()
if p.requires_grad and (p.grad is None or p.grad.abs().sum() == 0)]
print(f"[2] recurrent unroll R={R}: init loss {loss.item():.4f} "
f"(uniform would be {np.log(S):.4f}); "
f"params with zero/missing grad: {bad or 'none'} -> "
f"{'PASS' if not bad and torch.isfinite(loss) else 'FAIL'}")
assert not bad and torch.isfinite(loss)
def overfit_tiny(depth, n_iters, steps=2500, lr=1e-3, target=0.99):
torch.manual_seed(depth * 10 + n_iters)
model = PtrTransformer(S, D, 4, 512, depth=depth, zero_init=True).to(DEV)
rng = np.random.default_rng(7)
fixed = sample_train_batch(rng, 256, S, k_max=8)
f = torch.as_tensor(fixed["f"], device=DEV)
x = torch.as_tensor(fixed["x"], device=DEV)
k = torch.as_tensor(fixed["k"], device=DEV)
y = torch.as_tensor(fixed["answer"], device=DEV)
opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
t0 = time.time()
acc = 0.0
for i in range(steps):
opt.zero_grad(set_to_none=True)
logits = model(f, x, k, n_iters=n_iters)
loss = F.cross_entropy(logits, y)
loss.backward()
opt.step()
if i % 100 == 0 or i == steps - 1:
acc = (logits.argmax(-1) == y).float().mean().item()
print(f" depth={depth} R={n_iters} step {i}: "
f"loss {loss.item():.4f} acc {acc:.3f} "
f"({time.time()-t0:.0f}s)", flush=True)
if acc >= target:
break
acc = (logits.argmax(-1) == y).float().mean().item()
ok = acc >= target
print(f"[3] overfit tiny set (depth={depth}, R={n_iters}): "
f"final acc {acc:.3f} -> {'PASS' if ok else 'FAIL'}")
assert ok
def check_interface():
torch.manual_seed(3)
model = PtrTransformer(S, D, 4, 512, depth=4, zero_init=True).to(DEV)
rng = np.random.default_rng(11)
for n_iters in (1, 4, 12):
b = sample_eval_batch(rng, 8, S, k=16)
f = torch.as_tensor(b["f"], device=DEV)
x = torch.as_tensor(b["x"], device=DEV)
k = torch.as_tensor(b["k"], device=DEV)
logits = model(f, x, k, n_iters=n_iters)
assert logits.shape == (8, S)
# smaller state space with the same weights (subset of vocab + positions)
b = sample_eval_batch(rng, 8, 16, k=64)
logits = model(torch.as_tensor(b["f"], device=DEV),
torch.as_tensor(b["x"], device=DEV),
torch.as_tensor(b["k"], device=DEV), n_iters=4)
# head is sized to the trained S=64; a smaller eval state space S'=16
# yields 64 logits and must be MASKED to the first S' classes
assert logits.shape == (8, 64)
assert logits[:, :16].argmax(-1).shape == (8,)
print("[4] interface: n_iters 1/4/12 and S=16-subset eval (with head "
f"masking) all OK (logits {tuple(logits.shape)}) -> PASS")
if __name__ == "__main__":
check_marginals_and_lookup()
check_gradients()
overfit_tiny(depth=4, n_iters=1)
overfit_tiny(depth=1, n_iters=4)
overfit_tiny(depth=1, n_iters=8)
check_interface()
print("ALL VALIDATION CHECKS PASSED")
sys.exit(0) |