import json, math, random from pathlib import Path import numpy as np SEED = 2772 rng = np.random.default_rng(SEED) def tangent_sweep(): # T maps parameter perturbations to outputs; columns are candidate parameters. n_out, n_param, n_task = 18, 160, 8 T = rng.normal(size=(n_out, n_param)) / math.sqrt(n_param) # Task directions are deliberately anisotropic and not identical to output axes. U = rng.normal(size=(n_out, n_task)) U[:, 0] *= 5.0 U[:, 1] *= 2.0 G = T.T @ U # columns T^* u_i energy = np.sum(G * G, axis=1) order_tangent = np.argsort(energy) # remove low task tangent energy first order_mag = np.argsort(np.abs(rng.normal(size=n_param))) # replaced below by weight proxy # A fixed unrelated parameter/weight vector makes magnitude pruning genuinely different. weights = rng.normal(size=n_param) order_mag = np.argsort(weights * weights) total = float(np.sum(energy)) rows = [] for eps in [0.05, 0.10, 0.20, 0.30, 0.50]: target = eps * eps removed = [] cum = 0.0 for j in order_tangent: if (cum + energy[j]) / total <= target: removed.append(int(j)); cum += float(energy[j]) keep = np.ones(n_param, dtype=bool); keep[removed] = False r = math.sqrt(np.sum(G[~keep] ** 2) / total) # For each u, u^T(G-GP)u = ||(I-P)T^*u||^2 exactly. d_quad = np.sum(G[~keep] * G[~keep], axis=0) full_quad = np.sum(G * G, axis=0) quad_ratio = float(np.sum(d_quad) / np.sum(full_quad)) # First-order tangent update: output discrepancy is T(I-P)T^*u; # report the energy identity, which is the directly predicted quantity. rows.append(dict(epsilon=eps, retained=int(keep.sum()), observed_ratio=r, predicted_upper_bound=eps, observed_squared_ratio=quad_ratio, predicted_squared_ratio=r*r)) # Sweep arbitrary masks and verify exact identity ratio^2 = relative lost quadratic task energy. identity_err = [] for frac in [0.1, 0.3, 0.5, 0.7, 0.9]: keep = rng.random(n_param) > frac lost = np.sum(G[~keep] ** 2) ratio2 = lost / total quad = np.sum(np.sum(G[~keep] ** 2, axis=0)) / np.sum(np.sum(G ** 2, axis=0)) identity_err.append(abs(ratio2 - quad)) # Compare task-aware and magnitude masks at matched sparsity. comparisons = [] for sparsity in [0.25, 0.50, 0.80]: k = int(round(sparsity * n_param)) kt = np.ones(n_param, bool); kt[order_tangent[:k]] = False km = np.ones(n_param, bool); km[order_mag[:k]] = False rt = math.sqrt(np.sum(G[~kt] ** 2) / total) rm = math.sqrt(np.sum(G[~km] ** 2) / total) comparisons.append(dict(sparsity=sparsity, tangent_ratio=rt, magnitude_ratio=rm, tangent_better=rt < rm)) return dict(rows=rows, max_identity_error=max(identity_err), comparisons=comparisons) def torch_mini_experiment(): # Small nonlinear regression, using per-example output residual VJPs as task tangents. try: import torch torch.manual_seed(SEED) device = 'cuda' if torch.cuda.is_available() else 'cpu' try: torch.zeros(1, device=device) except Exception: device = 'cpu' n, d, h = 192, 10, 24 x = torch.randn(n, d, device=device) true_w = torch.randn(d, 1, device=device) y = torch.sin(x @ true_w) + 0.08 * torch.randn(n, 1, device=device) model0 = torch.nn.Sequential(torch.nn.Linear(d,h), torch.nn.Tanh(), torch.nn.Linear(h,1)).to(device) # Common initialization for both methods. state = {k:v.detach().clone() for k,v in model0.state_dict().items()} calib = slice(0, 64) def make_model(): m = torch.nn.Sequential(torch.nn.Linear(d,h), torch.nn.Tanh(), torch.nn.Linear(h,1)).to(device) m.load_state_dict(state); return m def masks_for(m, sparsity): # Collect one tangent g=T^*u per calibration example, where u is residual. params = [p for p in m.parameters() if p.requires_grad] scores = [torch.zeros_like(p) for p in params] for j in range(calib.start, calib.stop): out = m(x[j:j+1]); residual = out - y[j:j+1] gs = torch.autograd.grad((residual*residual).sum()/2, params, retain_graph=False) for s,g in zip(scores,gs): s += g.detach() ** 2 flat_score = torch.cat([s.flatten() for s in scores]) flat_mag = torch.cat([p.detach().abs().flatten() for p in params]) k = int(sparsity * flat_score.numel()) mt = torch.ones_like(flat_score); mm = torch.ones_like(flat_score) mt[torch.argsort(flat_score)[:k]] = 0 mm[torch.argsort(flat_mag)[:k]] = 0 out=[]; a=0 for p in params: z=p.numel(); out.append((mt[a:a+z].reshape_as(p), mm[a:a+z].reshape_as(p))); a += z return out def run(kind, masks, steps=30): m=make_model(); params=list(m.parameters()); with torch.no_grad(): for p,(mt,mm) in zip(params,masks): p.mul_(mt if kind=='tangent' else mm) def loss(): return ((m(x)-y)**2).mean() initial=float(loss().detach().cpu()); opt=torch.optim.SGD(m.parameters(),lr=0.08) vals=[] for _ in range(steps): opt.zero_grad(); z=loss(); z.backward() with torch.no_grad(): for p,(mt,mm) in zip(params,masks): p.grad.mul_(mt if kind=='tangent' else mm) p.mul_(mt if kind=='tangent' else mm) opt.step() with torch.no_grad(): for p,(mt,mm) in zip(params,masks): p.mul_(mt if kind=='tangent' else mm) vals.append(float(loss().detach().cpu())) return initial, vals[-1] result={"device":device,"sparsities":{}} for sp in [0.5,0.8]: masks=masks_for(model0,sp) t=run('tangent',masks); mag=run('magnitude',masks) result['sparsities'][str(sp)]={"tangent_initial":t[0],"magnitude_initial":mag[0],"tangent_final":t[1],"magnitude_final":mag[1]} return result except Exception as e: return {"error": type(e).__name__ + ': ' + str(e)} if __name__ == '__main__': result={'seed':SEED,'math_check':tangent_sweep(),'mini_experiment':torch_mini_experiment()} Path('results.json').write_text(json.dumps(result, indent=2)) print(json.dumps(result, indent=2))