1import json, math, os, sys
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6
  7sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
  8from bench import (get_dataset, make_model, train_model, evaluate,
  9                   sweep_baseline, make_report)
 10
 11# Sequence is the matched track: the task requires multi-token temporal
 12# correlations, which is the structure affected by progressive depth.
 13TRACK = "sequence"
 14MODEL = "transformer_tiny"
 15EPOCHS = 8
 16BATCH = 128
 17# The union of idea and baseline learning rates is explicitly shared.
 18LR_GRID = [0.0015, 0.003, 0.006]
 19WD_GRID = [0.0]
 20SEEDS = tuple(range(8))
 21N_TRAIN, N_TEST = 400, 200
 22
 23class DiagonalCoupling(nn.Module):
 24    """Exact additive coupling with channelwise (diagonal) f and g."""
 25    def __init__(self, dim, scale=0.10):
 26        super().__init__()
 27        assert dim % 2 == 0
 28        h = dim // 2
 29        self.f = nn.Parameter(torch.randn(h) * scale)
 30        self.g = nn.Parameter(torch.randn(h) * scale)
 31
 32    def forward(self, x):
 33        a, b = x.chunk(2, dim=-1)
 34        ap = a + b * self.f
 35        bp = b + ap * self.g
 36        return torch.cat((ap, bp), dim=-1)
 37
 38    def inverse(self, y):
 39        ya, yb = y.chunk(2, dim=-1)
 40        b = yb - ya * self.g
 41        a = ya - b * self.f
 42        return torch.cat((a, b), dim=-1)
 43
 44class IdentityPair(nn.Module):
 45    """C_theta followed by C_theta^{-1}, then untied during optimization."""
 46    def __init__(self, dim, scale=0.10):
 47        super().__init__()
 48        self.forward_block = DiagonalCoupling(dim, scale)
 49        self.inverse_block = DiagonalCoupling(dim, scale)
 50        self.inverse_block.load_state_dict(self.forward_block.state_dict())
 51
 52    def forward(self, x):
 53        return self.inverse_block.inverse(self.forward_block(x))
 54
 55class PairedTransformer(nn.Module):
 56    """Bench transformer structure with one cheap identity-paired block."""
 57    def __init__(self, input_dim, out_dim):
 58        super().__init__()
 59        d, depth = 64, 2
 60        self.win = input_dim
 61        self.inp = nn.Linear(1, d)
 62        self.pos = nn.Parameter(torch.zeros(1, input_dim, d))
 63        nn.init.normal_(self.pos, std=.02)
 64        layer = nn.TransformerEncoderLayer(d, nhead=2,
 65            dim_feedforward=128, batch_first=True, dropout=0.0)
 66        self.enc = nn.TransformerEncoder(layer, depth)
 67        self.pair = IdentityPair(input_dim * d)
 68        self.head = nn.Linear(input_dim * d, out_dim)
 69
 70    def forward(self, x):
 71        h = self.inp(x.unsqueeze(-1)) + self.pos[:, :x.shape[1]]
 72        z = self.enc(h).reshape(x.shape[0], -1)
 73        return self.head(self.pair(z))
 74
 75
 76def seed_all(seed):
 77    np.random.seed(seed)
 78    torch.manual_seed(seed)
 79    if torch.cuda.is_available():
 80        try:
 81            torch.cuda.manual_seed_all(seed)
 82        except Exception:
 83            pass
 84
 85
 86def ds(seed):
 87    return get_dataset(TRACK, seed, n_train=N_TRAIN, n_test=N_TEST)
 88
 89
 90def base_fn(cfg):
 91    def run(seed):
 92        seed_all(seed + 10000)
 93        net = make_model(MODEL, ds(seed)["input_shape"], 1)
 94        _, metric, _ = train_model(net, ds(seed), epochs=EPOCHS,
 95            lr=cfg["lr"], batch=BATCH, weight_decay=cfg["weight_decay"], log=lambda *_: None)
 96        return float(metric)
 97    return run
 98
 99SIGNATURES = []
100
101def idea_fn(cfg):
102    def run(seed):
103        seed_all(seed + 10000)
104        d = ds(seed)
105        net = PairedTransformer(d["input_shape"][0], 1)
106        # Quantify insertion behavior before any optimizer update.
107        net.eval()
108        with torch.no_grad():
109            x = d["xte"][:min(64, len(d["xte"]))]
110            h = net.inp(x.unsqueeze(-1)) + net.pos[:, :x.shape[1]]
111            z = net.enc(h).reshape(x.shape[0], -1)
112            insertion_rel = float((net.pair(z) - z).norm() / z.norm().clamp_min(1e-12))
113        trained, metric, _ = train_model(net, d, epochs=EPOCHS,
114            lr=cfg["lr"], batch=BATCH, weight_decay=cfg["weight_decay"], log=lambda *_: None)
115        if trained is None:
116            return float("nan")
117        trained.eval()
118        with torch.no_grad():
119            x = d["xte"][:min(64, len(d["xte"]))].to(next(trained.parameters()).device)
120            h = trained.inp(x.unsqueeze(-1)) + trained.pos[:, :x.shape[1]]
121            z = trained.enc(h).reshape(x.shape[0], -1)
122            residual = float((trained.pair(z) - z).norm() / z.norm().clamp_min(1e-12))
123        # Retest the predicted O(eta) untied response on the trained model.
124        trained.zero_grad(set_to_none=True)
125        out = trained.head(trained.pair(z))
126        loss = (out ** 2).mean()
127        loss.backward()
128        grads = [p.grad.detach().clone() for p in trained.pair.parameters()]
129        original = [p.detach().clone() for p in trained.pair.parameters()]
130        vals = []
131        for eta in (1e-5, 1e-4, 1e-3):
132            with torch.no_grad():
133                for p, q, g in zip(trained.pair.parameters(), original, grads):
134                    p.copy_(q - eta * g)
135                changed = (trained.head(trained.pair(z)) - out.detach()).norm() / out.detach().norm().clamp_min(1e-12)
136                vals.append((eta, float(changed)))
137        with torch.no_grad():
138            for p, q in zip(trained.pair.parameters(), original): p.copy_(q)
139        valid = [(a,b) for a,b in vals if b > 0]
140        slope = float(np.polyfit(np.log([a for a,b in valid]), np.log([b for a,b in valid]), 1)[0]) if len(valid) == 3 else float("nan")
141        SIGNATURES.append({"seed": seed, "insertion_relative_output_change": insertion_rel,
142                           "trained_pair_residual": residual, "untied_eta_response": vals,
143                           "trained_loglog_slope": slope})
144        return float(metric)
145    return run
146
147
148def main():
149    # Baseline sweep includes every idea learning rate, satisfying search parity.
150    grid = [{"lr": lr, "weight_decay": wd} for lr in LR_GRID for wd in WD_GRID]
151    base = sweep_baseline(base_fn, grid)
152    # Evaluate all three idea settings on all eight paired seeds, then report best.
153    idea_results = []
154    for cfg in grid:
155        r = evaluate(idea_fn(cfg), seeds=SEEDS)
156        idea_results.append((cfg, r))
157    best_cfg, idea = min(idea_results, key=lambda cr: cr[1]["mean"])
158    sig = {"prediction": "identity insertion and O(eta) untied response",
159           "confirmed": bool(SIGNATURES and np.mean([s["insertion_relative_output_change"] for s in SIGNATURES]) < 1e-5 and
160                              abs(np.mean([s["trained_loglog_slope"] for s in SIGNATURES if np.isfinite(s["trained_loglog_slope"])]) - 1.0) < 0.15),
161           "idea_best_cfg": best_cfg, "per_seed": SIGNATURES,
162           "mean_insertion_relative_output_change": float(np.mean([s["insertion_relative_output_change"] for s in SIGNATURES])) if SIGNATURES else None,
163           "mean_trained_pair_residual": float(np.mean([s["trained_pair_residual"] for s in SIGNATURES])) if SIGNATURES else None}
164    report = make_report(TRACK, MODEL, base, idea, {"mechanism_signature": sig})
165    # Required custom_track is not applicable: sequence is a built-in structural match.
166    Path("bench_report.json").write_text(json.dumps(report, indent=2))
167    print(json.dumps(report, indent=2))
168
169if __name__ == "__main__":
170    main()