import json import math import numpy as np def tc_gaussian(cov, eps=1e-10): cov = np.asarray(cov, float) d = np.diag(cov) sign, ld = np.linalg.slogdet(cov + eps * np.eye(len(d))) if sign <= 0: return float("nan") return float(0.5 * (np.log(d + eps).sum() - ld)) def shrink_cov(x, alpha=0.0): s = np.cov(x, rowvar=False, ddof=1) return (1-alpha) * s + alpha * np.diag(np.diag(s)) def merge_score(cov, a, b, eps=1e-10): ab = list(a) + list(b) return tc_gaussian(cov[np.ix_(ab, ab)], eps) - tc_gaussian(cov[np.ix_(list(a), list(a))], eps) - tc_gaussian(cov[np.ix_(list(b), list(b))], eps) def agglomerative(cov, k, eps=1e-10): groups = [[i] for i in range(cov.shape[0])] while len(groups) > k: best = None for i in range(len(groups)): for j in range(i+1, len(groups)): val = merge_score(cov, groups[i], groups[j], eps) if best is None or val > best[0]: best = (val, i, j) _, i, j = best groups[i] = groups[i] + groups[j] del groups[j] return sorted([sorted(g) for g in groups]) def equicorr(d, rho): return (1-rho)*np.eye(d) + rho*np.ones((d,d)) def exact_equicorr_tc(d, rho): # det(Sigma)=(1-rho)^(d-1)*(1+(d-1)rho), unit variances. return -0.5*((d-1)*math.log(1-rho) + math.log(1+(d-1)*rho)) def fit_shared_lowrank(x, y, q=2, ridge=1e-3, iters=80): # Shared adapter basis U with task-specific heads V: Y ~= X U V. n, p = x.shape; t = y.shape[1] U = np.random.default_rng(777).normal(size=(p, q)) * .1 V = np.random.default_rng(778).normal(size=(q, t)) * .1 for _ in range(iters): xu = x @ U V = np.linalg.solve(xu.T @ xu + ridge*np.eye(q), xu.T @ y) xv = y @ V.T U = np.linalg.solve(x.T @ x + ridge*np.eye(p), x.T @ xv) return U, V def predict_shared_lowrank(x, model): U, V = model return x @ U @ V def fit_cluster_models(xtr, ytr, groups, ridge=1e-3): # Frozen base is zero; each routed adapter is a linear weight vector. p = xtr.shape[1] ws = [] for g in groups: yy = ytr[:, g].mean(axis=1) a = xtr.T @ xtr + ridge*np.eye(p) w = np.linalg.solve(a, xtr.T @ yy) ws.append(w) return ws def predict_cluster(x, ws, groups): out = np.zeros((len(x), sum(len(g) for g in groups))) # groups are contiguous in this experiment after canonical sorting. for w, g in zip(ws, groups): out[:, g] = x @ w[:, None] return out def regression_demo(rng): t, p, ntr, nv = 6, 12, 700, 1000 true_groups = [[0,1], [2,3], [4,5]] xtr, xv = rng.normal(size=(ntr,p)), rng.normal(size=(nv,p)) base_w = rng.normal(size=(3,p)) w = np.zeros((t,p)) for j,g in enumerate(true_groups): for task in g: w[task] = base_w[j] + .10*rng.normal(size=p) # Correlated task noise makes residual dependence operationally visible. noise_cov = np.eye(t)*.35 for g in true_groups: for i in g: for j in g: if i != j: noise_cov[i,j] = .28 etr = rng.multivariate_normal(np.zeros(t), noise_cov, ntr) ev = rng.multivariate_normal(np.zeros(t), noise_cov, nv) ytr, yv = xtr@w.T + etr, xv@w.T + ev # Warm-up predictor used only to collect held-out residuals. warm_model = fit_shared_lowrank(xtr, ytr, q=2) warm = predict_shared_lowrank(xv, warm_model) residual = yv - warm cov = shrink_cov(residual, alpha=.10) discovered = agglomerative(cov, 3) random_groups = [[0,2], [1,5], [3,4]] results = {} # Equal nominal budget: rank-2 shared factor has p*q + q*t = 36 parameters; # each of three rank-1 routed adapters has 3*p = 36 parameters. shared_model = fit_shared_lowrank(xtr, ytr, q=2) results['shared'] = {'mse': float(np.mean((predict_shared_lowrank(xv, shared_model)-yv)**2)), 'groups': [list(range(t))], 'params': p*2+2*t} for name, groups in [('random', random_groups), ('residual_tc', discovered), ('oracle', true_groups)]: ws = fit_cluster_models(xtr, ytr, groups) pred = predict_cluster(xv, ws, groups) results[name] = {'mse': float(np.mean((pred-yv)**2)), 'groups': groups, 'params': len(groups)*p} return results def main(): rng = np.random.default_rng(1015) # Prediction 1: exact TC rises according to closed form with rho. d = 4 rhos = [0.0, .1, .2, .4, .6, .8] tc_curve = [] for rho in rhos: x = rng.multivariate_normal(np.zeros(d), equicorr(d,rho), 30000) observed = tc_gaussian(np.cov(x, rowvar=False)) predicted = exact_equicorr_tc(d,rho) tc_curve.append({'rho':rho, 'predicted':predicted, 'observed':observed, 'abs_error':abs(observed-predicted)}) # Prediction 2: near independence TC is quadratic in rho, not linear. small = [] for rho in [.02,.04,.08,.12,.16]: pred = exact_equicorr_tc(d,rho) quadratic = d*(d-1)/4 * rho*rho small.append({'rho':rho, 'observed_exact':pred, 'quadratic_prediction':quadratic, 'ratio':pred/quadratic}) # Prediction 3: pairwise Gaussian TC is -1/2 log(1-rho^2), and # block clustering should recover the population pairs once signal exceeds sampling noise. pair_tc = [] for rho in [.0,.2,.4,.6,.8]: x = rng.multivariate_normal(np.zeros(2), equicorr(2,rho), 12000) pair_tc.append({'rho': rho, 'predicted': -0.5*math.log(max(1-rho*rho, 1e-12)), 'observed': tc_gaussian(np.cov(x, rowvar=False))}) recovery = [] for rho in [.0,.2,.4,.6,.8]: cov = np.eye(6) for g in [[0,1],[2,3],[4,5]]: i,j=g; cov[i,j]=cov[j,i]=rho x = rng.multivariate_normal(np.zeros(6), cov, 5000) est = shrink_cov(x, .05) groups = agglomerative(est, 3) recovery.append({'rho':rho, 'groups':groups, 'exact_recovery':groups==[[0,1],[2,3],[4,5]], 'within_tc':float(np.mean([tc_gaussian(est[np.ix_(g,g)]) for g in [[0,1],[2,3],[4,5]]]))}) regression = regression_demo(rng) output = {'tc_rho_sweep':tc_curve, 'small_rho_quadratic_sweep':small, 'pair_tc_sweep': pair_tc, 'block_clustering_sweep':recovery, 'regression_demo':regression} with open('results.json','w') as f: json.dump(output,f,indent=2) print(json.dumps(output, indent=2)) if __name__ == '__main__': main()