import json, math, random from pathlib import Path import numpy as np import torch import torch.nn as nn import sys sys.path.insert(0, "/home/maxwelhelp/all/math2nn") from bench import get_dataset, make_model, evaluate, sweep_baseline, make_report SEEDS = tuple(range(8)) SWEEP_SEEDS = tuple(range(4)) EPOCHS = 12 BATCH = 128 # This is the only new solver/readout knob; it is fixed before running. GAP = 0.25 K = 7 def fejer_coeffs(k): return np.ones(k + 1, dtype=np.float64) / (k + 1) def jackson_coeffs(k): n = k // 2 + 1 return np.convolve(np.ones(n) / n, np.ones(n) / n) def math_checks(): out = {} for k in (3, 7, 15): f, j = fejer_coeffs(k), jackson_coeffs(k) out[str(k)] = { "fejer_p1_error": float(abs(f.sum() - 1)), "jackson_p1_error": float(abs(j.sum() - 1)), "jackson_degree": int(len(j) - 1), "nonnegative": bool(j.min() >= 0), } rng = np.random.default_rng(11) a = rng.normal(size=20) # Direct polynomial versus the explicit reflection iterates. c = fejer_coeffs(15) direct = np.zeros_like(a) z = a.copy() for x in c: direct += x * z z = (2 * 0.75 - 1) * z horner = np.zeros_like(a) t = 2 * .75 - 1 for x in c[::-1]: horner = horner * t + x * a out["iterate_identity_error"] = float(np.linalg.norm(direct - horner)) return out def filter_output(raw, method="idea", k=K, gap=GAP): """Apply p(2F-I) to y0=0, where F(y)=(1-gap)y+gap*raw. This is an end-to-end differentiable filtered prediction. The baseline uses the same rnn_small and optimizer but returns raw predictions. """ if method == "baseline": return raw # Spectral-gap selector: remain conservative at the critical scale. if gap * k < 2.0: coeff_np = fejer_coeffs(k) else: coeff_np = jackson_coeffs(k) coeff = torch.as_tensor(coeff_np, dtype=raw.dtype, device=raw.device) # Reflection map applied to a state, with raw as the fixed-point forcing. z = torch.zeros_like(raw) total = torch.zeros_like(raw) for c in coeff: total = total + c * z z = 2 * ((1 - gap) * z + gap * raw) - z # Safety check analogous to the proposal, evaluated per batch. fejer = torch.zeros_like(raw) zf = torch.zeros_like(raw) for _ in range(k + 1): fejer = fejer + zf / (k + 1) zf = 2 * ((1 - gap) * zf + gap * raw) - zf if torch.mean(total.square()) > 1.10 * torch.mean(fejer.square()): return fejer return total def seed_all(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def train_one(seed, lr, method): seed_all(seed) d = get_dataset("dynamics", seed, n_train=400, n_test=200) model = make_model("rnn_small", d["input_shape"], d["out_dim"]) lossf = nn.MSELoss() device = "cuda" if torch.cuda.is_available() else "cpu" try: model.to(device) xtr, ytr = d["xtr"].to(device), d["ytr"].to(device) opt = torch.optim.Adam(model.parameters(), lr=lr) for _ in range(EPOCHS): model.train() perm = torch.randperm(len(xtr), device=device) for i in range(0, len(xtr), BATCH): ix = perm[i:i+BATCH] raw = model(xtr[ix]) pred = filter_output(raw, method) loss = lossf(pred, ytr[ix]) opt.zero_grad(); loss.backward(); opt.step() model.eval() with torch.no_grad(): raw = model(d["xte"].to(device)) pred = filter_output(raw, method) metric = float(lossf(pred, d["yte"].to(device))) return metric except RuntimeError: # Explicit CPU fallback, including shared-GPU OOM/cuDNN failures. seed_all(seed) d = get_dataset("dynamics", seed, n_train=400, n_test=200) model = make_model("rnn_small", d["input_shape"], d["out_dim"]) opt = torch.optim.Adam(model.parameters(), lr=lr) for _ in range(EPOCHS): perm = torch.randperm(len(d["xtr"])) for i in range(0, len(perm), BATCH): ix = perm[i:i+BATCH]; raw = model(d["xtr"][ix]) loss = lossf(filter_output(raw, method), d["ytr"][ix]) opt.zero_grad(); loss.backward(); opt.step() with torch.no_grad(): raw = model(d["xte"]); return float(lossf(filter_output(raw, method), d["yte"])) def train_fn(cfg, method): return lambda seed: train_one(seed, float(cfg["lr"]), method) def signature(seed=0, lr=0.003): """Measure behavior on a trained idea model, not a synthetic matrix.""" seed_all(seed) d = get_dataset("dynamics", seed, n_train=400, n_test=200) model = make_model("rnn_small", d["input_shape"], d["out_dim"]) opt = torch.optim.Adam(model.parameters(), lr=lr); lossf = nn.MSELoss() for _ in range(EPOCHS): p = torch.randperm(len(d["xtr"])) for i in range(0, len(p), BATCH): ix=p[i:i+BATCH]; raw=model(d["xtr"][ix]); loss=lossf(filter_output(raw,"idea"),d["ytr"][ix]) opt.zero_grad(); loss.backward(); opt.step() model.eval() with torch.no_grad(): x=d["xte"]; raw=model(x); filt=filter_output(raw,"idea") # observed suppression measured on actual trained predictions observed=float(torch.linalg.vector_norm(filt)/ (torch.linalg.vector_norm(raw)+1e-12)) t=2*GAP-1 f_pred=float(sum(fejer_coeffs(K)[j] * (1-t**(j+1))/(1-t) * GAP for j in range(K+1))) j_pred=float(sum(jackson_coeffs(K)[j] * (1-t**(j+1))/(1-t) * GAP for j in range(len(jackson_coeffs(K))))) return {"gap_hat": GAP, "K": K, "sK": GAP*K, "predicted_fejer_gain": f_pred, "predicted_jackson_gain": j_pred, "observed_filtered_over_raw_norm": observed, "prediction_error_abs": abs(observed-j_pred), "confirmed": abs(observed-j_pred) < .08} def main(): checks=math_checks() assert checks["iterate_identity_error"] < 1e-10 grid=[{"lr":1e-3},{"lr":3e-3},{"lr":1e-2}] base=sweep_baseline(lambda cfg: train_fn(cfg,"baseline"), grid, seeds=SWEEP_SEEDS) idea_trials=[] for cfg in grid: r=evaluate(train_fn(cfg,"idea"), seeds=SEEDS) idea_trials.append({"cfg":cfg,"result":r}) best=min(idea_trials,key=lambda q:q["result"]["mean"]) report=make_report("dynamics","rnn_small",base,best["result"],{ "prediction": "Jackson filtering suppresses the trained model output by p(2*gap-1) around its fixed point", "signature": signature(0, float(best["cfg"]["lr"])), "idea_sweep": idea_trials, "math_checks": checks, "protocol": {"epochs":EPOCHS,"n_train":400,"n_test":200,"gap":GAP,"K":K, "matched_architecture":True,"lr_union":grid} }) Path("bench_report.json").write_text(json.dumps(report,indent=2)) print(json.dumps(report,indent=2)) if __name__ == "__main__": main()