import json, random, sys from pathlib import Path import numpy as np import torch import torch.nn as nn from torch.distributions import Normal sys.path.insert(0, "/home/maxwelhelp/all/math2nn") from bench import get_dataset, make_report, sweep_baseline from bench.protocol import evaluate def load_ds(seed): ds = get_dataset(TRACK, seed, n_train=NTRAIN, n_test=NTEST) ds["ytr"] = ds["ytr"].reshape(NTRAIN, D) ds["yte"] = ds["yte"].reshape(NTEST, D) return ds TRACK = "correlated_multitask_regression" MODEL = "mlp_tiny" SEEDS = tuple(range(8)) SWEEP_SEEDS = tuple(range(4)) EPOCHS = 18 BATCH = 128 NTRAIN, NTEST = 1200, 500 LRS = [1e-3, 3e-3, 6e-3] D = 6 EPS = 1e-5 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 device(): return torch.device("cuda" if torch.cuda.is_available() else "cpu") def trunk(): return nn.Sequential(nn.Linear(12, 64), nn.ReLU(), nn.Linear(64, 64), nn.ReLU()) class DiagonalGaussian(nn.Module): def __init__(self): super().__init__() self.base = trunk() self.head = nn.Linear(64, 12) def forward(self, x): a = self.head(self.base(x)) return a[:, :D], a[:, D:].clamp(-4.0, 3.0) class ConditionalCopula(nn.Module): def __init__(self): super().__init__() self.base = trunk() # Marginal parameters, then autoregressive Gaussian-copula Cholesky factors. self.head = nn.Linear(64, 2 * D + D * (D - 1) // 2) def forward(self, x): a = self.head(self.base(x)) mu = a[:, :D] log_scale = a[:, D:2 * D].clamp(-4.0, 3.0) raw = a[:, 2 * D:] L = torch.zeros(x.shape[0], D, D, device=x.device, dtype=x.dtype) k = 0 for i in range(D): L[:, i, i] = 1.0 for j in range(i): L[:, i, j] = raw[:, k].tanh() * 0.8 k += 1 # Row normalization makes L L^T a correlation matrix. R0 = L @ L.transpose(1, 2) s = torch.sqrt(torch.diagonal(R0, dim1=1, dim2=2).clamp_min(EPS)) L = L / s.unsqueeze(-1) return mu, log_scale, L def train_baseline(ds, lr, seed): seed_all(seed) net = DiagonalGaussian().to(device()) opt = torch.optim.Adam(net.parameters(), lr=lr) x, y = ds["xtr"].to(device()), ds["ytr"].to(device()) for _ in range(EPOCHS): net.train() perm = torch.randperm(len(x), device=x.device) for i in range(0, len(x), BATCH): q = perm[i:i + BATCH] mu, ls = net(x[q]) # Standard independent Gaussian NLL, the baseline being replaced. loss = (0.5 * (((y[q] - mu) / ls.exp()) ** 2 + 2 * ls)).mean() opt.zero_grad(); loss.backward(); opt.step() net.eval() with torch.no_grad(): pred, _ = net(ds["xte"].to(device())) mse = ((pred - ds["yte"].to(device())) ** 2).mean().item() return mse, net.cpu() def train_idea(ds, lr, seed): seed_all(seed) net = ConditionalCopula().to(device()) opt = torch.optim.Adam(net.parameters(), lr=lr) x, y = ds["xtr"].to(device()), ds["ytr"].to(device()) normal = Normal(torch.tensor(0., device=x.device), torch.tensor(1., device=x.device)) for _ in range(EPOCHS): net.train() perm = torch.randperm(len(x), device=x.device) for i in range(0, len(x), BATCH): q = perm[i:i + BATCH] mu, ls, L = net(x[q]) scale = ls.exp() z = (y[q] - mu) / scale # p(y|z_context) = product marginal densities times copula density. whiten = torch.linalg.solve_triangular(L, z.unsqueeze(-1), upper=False).squeeze(-1) log_joint_std = -0.5 * (whiten ** 2).sum(1) - torch.log(torch.diagonal(L, dim1=1, dim2=2)).sum(1) - D * 0.5 * np.log(2 * np.pi) loss = -(log_joint_std - ls.sum(1)).mean() opt.zero_grad(); loss.backward(); opt.step() net.eval() with torch.no_grad(): mu, _, _ = net(ds["xte"].to(device())) mse = ((mu - ds["yte"].to(device())) ** 2).mean().item() return mse, net.cpu() def run(): # The union of idea and baseline learning rates is identical; baseline sweep is fair. ds0 = load_ds(0) grid = [{"lr": x, "epochs": EPOCHS} for x in LRS] base_block = sweep_baseline( lambda cfg: lambda seed: train_baseline(load_ds(seed), cfg["lr"], seed)[0], grid, seeds=SWEEP_SEEDS) best_lr = float(base_block["best_cfg"]["lr"]) idea_grid = LRS idea_cfg_results = [] for lr in idea_grid: r = evaluate(lambda seed, lr=lr: train_idea(load_ds(seed), lr, seed)[0], seeds=SEEDS) idea_cfg_results.append({"lr": lr, "result": r}) best_idea = min(idea_cfg_results, key=lambda z: z["result"]["mean"]) idea_res = best_idea["result"] # Signature is measured from trained models: predicted residual rank correlation vs observed. sig_rows = [] for seed in SEEDS: ds = load_ds(seed) bm, bn = train_baseline(ds, best_lr, seed) im, inn = train_idea(ds, float(best_idea["lr"]), seed) with torch.no_grad(): xb = ds["xte"] bmu, _ = bn(xb) imu, _, L = inn(xb) residual = ds["yte"] - imu obs = np.corrcoef(residual.numpy(), rowvar=False) pred = (L @ L.transpose(1, 2)).mean(0).numpy() off = np.triu_indices(D, 1) sig_rows.append({"seed": seed, "observed_residual_corr_mean": float(obs[off].mean()), "predicted_copula_corr_mean": float(pred[off].mean()), "baseline_mse": bm, "idea_mse": im}) pred_mean = float(np.mean([r["predicted_copula_corr_mean"] for r in sig_rows])) obs_mean = float(np.mean([r["observed_residual_corr_mean"] for r in sig_rows])) signature = {"prediction": "copula dependence on uniform/standardized scale captures positive cross-output residual dependence", "predicted_mean_offdiag_corr": pred_mean, "observed_mean_offdiag_residual_corr": obs_mean, "abs_error": abs(pred_mean - obs_mean), "confirmed": bool(abs(pred_mean - obs_mean) < 0.12), "per_seed": sig_rows} report = make_report(TRACK, MODEL, base_block, idea_res, {"signature": signature, "idea_sweep": idea_cfg_results, "custom_track": {"name": TRACK, "file": "/home/maxwelhelp/all/math2nn/bench/custom_tracks/correlated_multitask_regression.py", "domain": "multi_task_learning"}}) Path("bench_report.json").write_text(json.dumps(report, indent=2)) print(json.dumps(report, indent=2)) if __name__ == "__main__": run()