architecture validation suite

validate.py

@@ -0,0 +1,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)