1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6import torch.nn.functional as F
  7
  8SEED = 1729
  9
 10def seed_all(seed=SEED):
 11    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 12    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 13
 14
 15def toy_check():
 16    # A zero-conditional-mean increment and a history-dependent biased increment.
 17    rng = np.random.default_rng(SEED)
 18    n = 200000
 19    h = rng.normal(size=n)
 20    noise = rng.normal(size=n)
 21    # bounded history test functions, including a nonlinear one
 22    hs = np.stack([np.ones(n), np.tanh(h), h/(1+np.abs(h))], axis=1)
 23    unbiased = 0.08 * noise
 24    biased = 0.08 * noise + 0.08 * np.tanh(h)
 25    norms = []
 26    for d in (unbiased, biased):
 27        moments = (d[:, None] * hs).mean(axis=0)
 28        norms.append(float(np.sqrt(np.sum(moments**2))))
 29    return {"unbiased_moment_norm": norms[0], "biased_moment_norm": norms[1],
 30            "ratio_biased_to_unbiased": norms[1] / max(norms[0], 1e-12),
 31            "passed": bool(norms[1] > 5 * norms[0])}
 32
 33
 34def make_data(n, steps=6, obs_dim=8, seed=SEED):
 35    rng = np.random.default_rng(seed)
 36    y = rng.integers(0, 3, size=n)
 37    # Each revealed observation is a noisy class-dependent vector; later steps add evidence.
 38    means = np.array([[1.0, 0.0, 0.0, 0.0, .3, -.2, 0, 0],
 39                      [0.0, 1.0, 0.0, 0.0, -.2, .3, 0, 0],
 40                      [0.0, 0.0, 1.0, 0.0, 0, 0, .3, -.2]], dtype=np.float32)
 41    x = means[y, None, :] + rng.normal(0, 1.15, size=(n, steps, obs_dim)).astype(np.float32)
 42    return torch.tensor(x), torch.tensor(y, dtype=torch.long)
 43
 44
 45class SeqClassifier(nn.Module):
 46    def __init__(self, obs_dim=8, hidden=32, classes=3):
 47        super().__init__()
 48        self.gru = nn.GRU(obs_dim, hidden, batch_first=True)
 49        self.head = nn.Linear(hidden, classes)
 50    def forward(self, x):
 51        z, _ = self.gru(x)
 52        return self.head(z)
 53
 54
 55def mart_loss(logits):
 56    p = logits.softmax(-1)
 57    d = p[:, 1:] - p[:, :-1]                 # B,T-1,K
 58    hist = p[:, :-1]
 59    # bounded finite test family: constant, each current coordinate, and tanh coordinates
 60    feats = [torch.ones_like(hist[..., :1]), hist, torch.tanh(3.0 * (hist - 1/3))]
 61    feat = torch.cat(feats, dim=-1)           # B,T-1,R
 62    moments = torch.einsum('btk,btr->tkr', d, feat) / p.shape[0]
 63    return (moments.square().mean(), moments.detach())
 64
 65
 66def ece(probs, y, bins=10):
 67    conf, pred = probs.max(-1)
 68    out = 0.0
 69    for lo, hi in zip(torch.linspace(0, 1, bins+1)[:-1], torch.linspace(0, 1, bins+1)[1:]):
 70        mask = (conf >= lo) & ((conf < hi) if hi < 1 else (conf <= hi))
 71        if mask.any(): out += mask.float().mean().item() * (conf[mask].mean().item() - (pred[mask] == y[mask]).float().mean().item())**2
 72    return math.sqrt(out)
 73
 74
 75def evaluate(model, x, y):
 76    model.eval()
 77    with torch.no_grad():
 78        logits = model(x); p = logits.softmax(-1)
 79        nll = F.cross_entropy(logits[:, -1], y).item()
 80        acc = (logits[:, -1].argmax(-1) == y).float().mean().item()
 81        step_nll = F.cross_entropy(logits.reshape(-1, 3), y[:, None].expand(-1, logits.shape[1]).reshape(-1)).item()
 82        step_ece = [ece(p[:, i], y) for i in range(p.shape[1])]
 83        ml, moments = mart_loss(logits)
 84        # Per-step held-out moment norm, with the same finite feature family.
 85        norms = moments.square().sum(dim=(1,2)).sqrt().cpu().tolist()
 86    return {"final_nll": nll, "final_accuracy": acc, "all_step_nll": step_nll,
 87            "final_ece": step_ece[-1], "step_ece": step_ece,
 88            "heldout_moment_norm_mean": float(np.mean(norms)),
 89            "heldout_moment_norm_last_transition": float(norms[-1])}
 90
 91
 92def train(lam, xtr, ytr, xte, yte, seed=SEED):
 93    seed_all(seed)
 94    model = SeqClassifier()
 95    opt = torch.optim.Adam(model.parameters(), lr=3e-3)
 96    # Full-batch is deliberately small and makes the empirical moment definition transparent.
 97    for epoch in range(90):
 98        model.train(); opt.zero_grad()
 99        logits = model(xtr)
100        loss = F.cross_entropy(logits.reshape(-1, 3), ytr[:, None].expand(-1, 6).reshape(-1))
101        ml, _ = mart_loss(logits)
102        (loss + lam * ml).backward(); opt.step()
103    return evaluate(model, xte, yte)
104
105
106def main():
107    seed_all()
108    toy = toy_check()
109    xtr, ytr = make_data(768, seed=SEED)
110    xte, yte = make_data(768, seed=SEED+1)
111    results = {"toy_check": toy, "runs": {}}
112    for lam in (0.0, 0.1, 1.0):
113        results["runs"][str(lam)] = train(lam, xtr, ytr, xte, yte)
114    results["comparison"] = {
115        "baseline": results["runs"]["0.0"],
116        "idea_lambda_0.1": results["runs"]["0.1"],
117        "idea_lambda_1.0": results["runs"]["1.0"]}
118    Path("results.json").write_text(json.dumps(results, indent=2))
119    print(json.dumps(results, indent=2))
120
121if __name__ == "__main__": main()