1import json
  2import os
  3import sys
  4from pathlib import Path
  5
  6import numpy as np
  7import torch
  8import torch.nn as nn
  9
 10# Import the read-only shared benchmark package.
 11sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
 12from bench import get_dataset, train_model, evaluate, sweep_baseline, make_report
 13
 14SEEDS = [0, 1, 2, 3, 4, 5, 6, 7]
 15# The union of baseline and idea learning rates is identical.
 16GRID = [{"lr": 1e-3}, {"lr": 3e-3}, {"lr": 6e-3}]
 17EPOCHS = 8
 18BATCH = 128
 19NTRAIN = 400
 20NTEST = 400
 21
 22
 23def make_phi(device):
 24    z = torch.linspace(-1.0, 1.0, 8, device=device)
 25    z1, z2 = torch.meshgrid(z, z, indexing="ij")
 26    return torch.stack((torch.ones_like(z1), z1, z2, 0.5 * (z1 * z1 + z2 * z2)))
 27
 28
 29def conserve_hidden(h, rank=2):
 30    """Rank compress each 8x8 hidden state and restore four weighted moments."""
 31    b = h.shape[0]
 32    x = h.reshape(b, 8, 8)
 33    # Batched truncated SVD is the low-rank compression intervention.
 34    u, s, vh = torch.linalg.svd(x, full_matrices=False)
 35    xt = (u[:, :, :rank] * s[:, None, :rank]) @ vh[:, :rank, :]
 36    phi = make_phi(h.device)
 37    flat_phi = phi.reshape(4, -1)
 38    flat_x = x.reshape(b, -1)
 39    flat_t = xt.reshape(b, -1)
 40    gram = flat_phi @ flat_phi.T
 41    delta = (flat_phi @ (flat_x - flat_t).T).T
 42    coeff = torch.linalg.solve(gram + 1e-8 * torch.eye(4, device=h.device), delta.T).T
 43    return (flat_t + coeff @ flat_phi).reshape(b, 64)
 44
 45
 46class IdeaRNN(nn.Module):
 47    """Same rnn_small architecture, with conservative hidden compression."""
 48    def __init__(self, hidden=64, rank=2):
 49        super().__init__()
 50        self.rnn = nn.GRU(3, hidden, batch_first=True)
 51        self.head = nn.Linear(hidden, 1)
 52        self.rank = rank
 53        self.last_signature = {}
 54
 55    def forward(self, x):
 56        seq = x.view(x.shape[0], -1, 3)
 57        try:
 58            _, h = self.rnn(seq)
 59        except RuntimeError:
 60            old = torch.backends.cudnn.enabled
 61            torch.backends.cudnn.enabled = False
 62            try:
 63                _, h = self.rnn(seq)
 64            finally:
 65                torch.backends.cudnn.enabled = old
 66        raw = h[-1]
 67        out = conserve_hidden(raw, self.rank)
 68        if not self.training:
 69            with torch.no_grad():
 70                phi = make_phi(raw.device).reshape(4, -1)
 71                self.last_signature = {
 72                    "raw_norm": float(raw.norm().cpu()),
 73                    "compressed_relative_error": float((out - raw).norm().cpu() / (raw.norm().cpu() + 1e-12)),
 74                    "moment_residual": float(((out - raw) @ phi.T).abs().max().cpu()),
 75                }
 76        return self.head(out)
 77
 78
 79def make_baseline(cfg):
 80    def fn(seed):
 81        torch.manual_seed(seed)
 82        np.random.seed(seed)
 83        from bench import make_model
 84        ds = get_dataset("dynamics", seed, NTRAIN, NTEST)
 85        model = make_model("rnn_small", ds["input_shape"], ds["out_dim"])
 86        _, metric, _ = train_model(model, ds, epochs=EPOCHS, lr=cfg["lr"], batch=BATCH, log=lambda *_: None)
 87        return metric
 88    return fn
 89
 90
 91def make_idea(cfg):
 92    def fn(seed):
 93        torch.manual_seed(seed)
 94        np.random.seed(seed)
 95        ds = get_dataset("dynamics", seed, NTRAIN, NTEST)
 96        model = IdeaRNN(rank=2)
 97        _, metric, _ = train_model(model, ds, epochs=EPOCHS, lr=cfg["lr"], batch=BATCH, log=lambda *_: None)
 98        return metric
 99    return fn
100
101
102def inspect_signature(seed, cfg):
103    """Measure the proposed preservation on a trained benchmark model."""
104    torch.manual_seed(seed)
105    ds = get_dataset("dynamics", seed, NTRAIN, NTEST)
106    model = IdeaRNN(rank=2)
107    trained, _, _ = train_model(model, ds, epochs=EPOCHS, lr=cfg["lr"], batch=BATCH, log=lambda *_: None)
108    trained.eval()
109    device = next(trained.parameters()).device
110    with torch.no_grad():
111        _ = trained(ds["xte"].to(device))
112    sig = dict(trained.last_signature)
113    sig.update({"target": "hidden moments", "predicted_residual": 0.0, "observed_residual": sig.get("moment_residual", float("nan"))})
114    sig["confirmed"] = bool(sig.get("observed_residual", 1.0) < 1e-5)
115    return sig
116
117
118def main():
119    # Cheap numerical verification before any benchmark training.
120    torch.manual_seed(329)
121    h = torch.randn(5, 64)
122    before = h.reshape(5, 8, 8)
123    phi = make_phi(h.device).reshape(4, -1)
124    after = conserve_hidden(h).reshape(5, -1)
125    math_resid = float(((after - before.reshape(5, -1)) @ phi.T).abs().max())
126    assert math_resid < 1e-5
127
128    base = sweep_baseline(make_baseline, GRID, seeds=SEEDS)
129    idea_candidates = []
130    for cfg in GRID:
131        r = evaluate(make_idea(cfg), seeds=SEEDS)
132        idea_candidates.append((r["mean"], cfg, r))
133    _, best_cfg, idea = min(idea_candidates, key=lambda q: q[0])
134    trained_sig = inspect_signature(SEEDS[0], best_cfg)
135    report = make_report(
136        "dynamics", "rnn_small", base, idea,
137        {"mechanism_signature": {
138            "math_check_max_moment_residual": math_resid,
139            "trained_model": trained_sig,
140            "predicted": "four hidden-state moments preserved after rank-2 compression",
141            "observed": "test-time hidden-state moment residual measured from trained model",
142            "confirmed": bool(math_resid < 1e-5 and trained_sig.get("confirmed", False)),
143        }, "protocol": {"epochs": EPOCHS, "n_train": NTRAIN, "n_test": NTEST, "grid": GRID, "seeds": SEEDS}})
144    report["math_verification"] = {"max_moment_residual": math_resid, "passed": math_resid < 1e-5}
145    Path("bench_report.json").write_text(json.dumps(report, indent=2))
146    print(json.dumps(report, indent=2))
147
148
149if __name__ == "__main__":
150    main()