import json, math, os, sys from pathlib import Path import numpy as np import torch from torch import nn sys.path.insert(0, "/home/maxwelhelp/all/math2nn") from bench import (get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report) # Sequence is the matched track: the task requires multi-token temporal # correlations, which is the structure affected by progressive depth. TRACK = "sequence" MODEL = "transformer_tiny" EPOCHS = 8 BATCH = 128 # The union of idea and baseline learning rates is explicitly shared. LR_GRID = [0.0015, 0.003, 0.006] WD_GRID = [0.0] SEEDS = tuple(range(8)) N_TRAIN, N_TEST = 400, 200 class DiagonalCoupling(nn.Module): """Exact additive coupling with channelwise (diagonal) f and g.""" def __init__(self, dim, scale=0.10): super().__init__() assert dim % 2 == 0 h = dim // 2 self.f = nn.Parameter(torch.randn(h) * scale) self.g = nn.Parameter(torch.randn(h) * scale) def forward(self, x): a, b = x.chunk(2, dim=-1) ap = a + b * self.f bp = b + ap * self.g return torch.cat((ap, bp), dim=-1) def inverse(self, y): ya, yb = y.chunk(2, dim=-1) b = yb - ya * self.g a = ya - b * self.f return torch.cat((a, b), dim=-1) class IdentityPair(nn.Module): """C_theta followed by C_theta^{-1}, then untied during optimization.""" def __init__(self, dim, scale=0.10): super().__init__() self.forward_block = DiagonalCoupling(dim, scale) self.inverse_block = DiagonalCoupling(dim, scale) self.inverse_block.load_state_dict(self.forward_block.state_dict()) def forward(self, x): return self.inverse_block.inverse(self.forward_block(x)) class PairedTransformer(nn.Module): """Bench transformer structure with one cheap identity-paired block.""" def __init__(self, input_dim, out_dim): super().__init__() d, depth = 64, 2 self.win = input_dim self.inp = nn.Linear(1, d) self.pos = nn.Parameter(torch.zeros(1, input_dim, d)) nn.init.normal_(self.pos, std=.02) layer = nn.TransformerEncoderLayer(d, nhead=2, dim_feedforward=128, batch_first=True, dropout=0.0) self.enc = nn.TransformerEncoder(layer, depth) self.pair = IdentityPair(input_dim * d) self.head = nn.Linear(input_dim * d, out_dim) def forward(self, x): h = self.inp(x.unsqueeze(-1)) + self.pos[:, :x.shape[1]] z = self.enc(h).reshape(x.shape[0], -1) return self.head(self.pair(z)) def seed_all(seed): np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): try: torch.cuda.manual_seed_all(seed) except Exception: pass def ds(seed): return get_dataset(TRACK, seed, n_train=N_TRAIN, n_test=N_TEST) def base_fn(cfg): def run(seed): seed_all(seed + 10000) net = make_model(MODEL, ds(seed)["input_shape"], 1) _, metric, _ = train_model(net, ds(seed), epochs=EPOCHS, lr=cfg["lr"], batch=BATCH, weight_decay=cfg["weight_decay"], log=lambda *_: None) return float(metric) return run SIGNATURES = [] def idea_fn(cfg): def run(seed): seed_all(seed + 10000) d = ds(seed) net = PairedTransformer(d["input_shape"][0], 1) # Quantify insertion behavior before any optimizer update. net.eval() with torch.no_grad(): x = d["xte"][:min(64, len(d["xte"]))] h = net.inp(x.unsqueeze(-1)) + net.pos[:, :x.shape[1]] z = net.enc(h).reshape(x.shape[0], -1) insertion_rel = float((net.pair(z) - z).norm() / z.norm().clamp_min(1e-12)) trained, metric, _ = train_model(net, d, epochs=EPOCHS, lr=cfg["lr"], batch=BATCH, weight_decay=cfg["weight_decay"], log=lambda *_: None) if trained is None: return float("nan") trained.eval() with torch.no_grad(): x = d["xte"][:min(64, len(d["xte"]))].to(next(trained.parameters()).device) h = trained.inp(x.unsqueeze(-1)) + trained.pos[:, :x.shape[1]] z = trained.enc(h).reshape(x.shape[0], -1) residual = float((trained.pair(z) - z).norm() / z.norm().clamp_min(1e-12)) # Retest the predicted O(eta) untied response on the trained model. trained.zero_grad(set_to_none=True) out = trained.head(trained.pair(z)) loss = (out ** 2).mean() loss.backward() grads = [p.grad.detach().clone() for p in trained.pair.parameters()] original = [p.detach().clone() for p in trained.pair.parameters()] vals = [] for eta in (1e-5, 1e-4, 1e-3): with torch.no_grad(): for p, q, g in zip(trained.pair.parameters(), original, grads): p.copy_(q - eta * g) changed = (trained.head(trained.pair(z)) - out.detach()).norm() / out.detach().norm().clamp_min(1e-12) vals.append((eta, float(changed))) with torch.no_grad(): for p, q in zip(trained.pair.parameters(), original): p.copy_(q) valid = [(a,b) for a,b in vals if b > 0] 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") SIGNATURES.append({"seed": seed, "insertion_relative_output_change": insertion_rel, "trained_pair_residual": residual, "untied_eta_response": vals, "trained_loglog_slope": slope}) return float(metric) return run def main(): # Baseline sweep includes every idea learning rate, satisfying search parity. grid = [{"lr": lr, "weight_decay": wd} for lr in LR_GRID for wd in WD_GRID] base = sweep_baseline(base_fn, grid) # Evaluate all three idea settings on all eight paired seeds, then report best. idea_results = [] for cfg in grid: r = evaluate(idea_fn(cfg), seeds=SEEDS) idea_results.append((cfg, r)) best_cfg, idea = min(idea_results, key=lambda cr: cr[1]["mean"]) sig = {"prediction": "identity insertion and O(eta) untied response", "confirmed": bool(SIGNATURES and np.mean([s["insertion_relative_output_change"] for s in SIGNATURES]) < 1e-5 and abs(np.mean([s["trained_loglog_slope"] for s in SIGNATURES if np.isfinite(s["trained_loglog_slope"])]) - 1.0) < 0.15), "idea_best_cfg": best_cfg, "per_seed": SIGNATURES, "mean_insertion_relative_output_change": float(np.mean([s["insertion_relative_output_change"] for s in SIGNATURES])) if SIGNATURES else None, "mean_trained_pair_residual": float(np.mean([s["trained_pair_residual"] for s in SIGNATURES])) if SIGNATURES else None} report = make_report(TRACK, MODEL, base, idea, {"mechanism_signature": sig}) # Required custom_track is not applicable: sequence is a built-in structural match. Path("bench_report.json").write_text(json.dumps(report, indent=2)) print(json.dumps(report, indent=2)) if __name__ == "__main__": main()