Task-Tangent Capture Pruning / experiment.py
Mechanism confirmed, baseline not beaten
1import json, math, random
2from pathlib import Path
3import numpy as np
4
5SEED = 2772
6rng = np.random.default_rng(SEED)
7
8
9def tangent_sweep():
10 # T maps parameter perturbations to outputs; columns are candidate parameters.
11 n_out, n_param, n_task = 18, 160, 8
12 T = rng.normal(size=(n_out, n_param)) / math.sqrt(n_param)
13 # Task directions are deliberately anisotropic and not identical to output axes.
14 U = rng.normal(size=(n_out, n_task))
15 U[:, 0] *= 5.0
16 U[:, 1] *= 2.0
17 G = T.T @ U # columns T^* u_i
18 energy = np.sum(G * G, axis=1)
19 order_tangent = np.argsort(energy) # remove low task tangent energy first
20 order_mag = np.argsort(np.abs(rng.normal(size=n_param))) # replaced below by weight proxy
21 # A fixed unrelated parameter/weight vector makes magnitude pruning genuinely different.
22 weights = rng.normal(size=n_param)
23 order_mag = np.argsort(weights * weights)
24 total = float(np.sum(energy))
25
26 rows = []
27 for eps in [0.05, 0.10, 0.20, 0.30, 0.50]:
28 target = eps * eps
29 removed = []
30 cum = 0.0
31 for j in order_tangent:
32 if (cum + energy[j]) / total <= target:
33 removed.append(int(j)); cum += float(energy[j])
34 keep = np.ones(n_param, dtype=bool); keep[removed] = False
35 r = math.sqrt(np.sum(G[~keep] ** 2) / total)
36 # For each u, u^T(G-GP)u = ||(I-P)T^*u||^2 exactly.
37 d_quad = np.sum(G[~keep] * G[~keep], axis=0)
38 full_quad = np.sum(G * G, axis=0)
39 quad_ratio = float(np.sum(d_quad) / np.sum(full_quad))
40 # First-order tangent update: output discrepancy is T(I-P)T^*u;
41 # report the energy identity, which is the directly predicted quantity.
42 rows.append(dict(epsilon=eps, retained=int(keep.sum()), observed_ratio=r,
43 predicted_upper_bound=eps, observed_squared_ratio=quad_ratio,
44 predicted_squared_ratio=r*r))
45
46 # Sweep arbitrary masks and verify exact identity ratio^2 = relative lost quadratic task energy.
47 identity_err = []
48 for frac in [0.1, 0.3, 0.5, 0.7, 0.9]:
49 keep = rng.random(n_param) > frac
50 lost = np.sum(G[~keep] ** 2)
51 ratio2 = lost / total
52 quad = np.sum(np.sum(G[~keep] ** 2, axis=0)) / np.sum(np.sum(G ** 2, axis=0))
53 identity_err.append(abs(ratio2 - quad))
54
55 # Compare task-aware and magnitude masks at matched sparsity.
56 comparisons = []
57 for sparsity in [0.25, 0.50, 0.80]:
58 k = int(round(sparsity * n_param))
59 kt = np.ones(n_param, bool); kt[order_tangent[:k]] = False
60 km = np.ones(n_param, bool); km[order_mag[:k]] = False
61 rt = math.sqrt(np.sum(G[~kt] ** 2) / total)
62 rm = math.sqrt(np.sum(G[~km] ** 2) / total)
63 comparisons.append(dict(sparsity=sparsity, tangent_ratio=rt, magnitude_ratio=rm,
64 tangent_better=rt < rm))
65 return dict(rows=rows, max_identity_error=max(identity_err), comparisons=comparisons)
66
67
68def torch_mini_experiment():
69 # Small nonlinear regression, using per-example output residual VJPs as task tangents.
70 try:
71 import torch
72 torch.manual_seed(SEED)
73 device = 'cuda' if torch.cuda.is_available() else 'cpu'
74 try:
75 torch.zeros(1, device=device)
76 except Exception:
77 device = 'cpu'
78 n, d, h = 192, 10, 24
79 x = torch.randn(n, d, device=device)
80 true_w = torch.randn(d, 1, device=device)
81 y = torch.sin(x @ true_w) + 0.08 * torch.randn(n, 1, device=device)
82 model0 = torch.nn.Sequential(torch.nn.Linear(d,h), torch.nn.Tanh(), torch.nn.Linear(h,1)).to(device)
83 # Common initialization for both methods.
84 state = {k:v.detach().clone() for k,v in model0.state_dict().items()}
85 calib = slice(0, 64)
86 def make_model():
87 m = torch.nn.Sequential(torch.nn.Linear(d,h), torch.nn.Tanh(), torch.nn.Linear(h,1)).to(device)
88 m.load_state_dict(state); return m
89 def masks_for(m, sparsity):
90 # Collect one tangent g=T^*u per calibration example, where u is residual.
91 params = [p for p in m.parameters() if p.requires_grad]
92 scores = [torch.zeros_like(p) for p in params]
93 for j in range(calib.start, calib.stop):
94 out = m(x[j:j+1]); residual = out - y[j:j+1]
95 gs = torch.autograd.grad((residual*residual).sum()/2, params, retain_graph=False)
96 for s,g in zip(scores,gs): s += g.detach() ** 2
97 flat_score = torch.cat([s.flatten() for s in scores])
98 flat_mag = torch.cat([p.detach().abs().flatten() for p in params])
99 k = int(sparsity * flat_score.numel())
100 mt = torch.ones_like(flat_score); mm = torch.ones_like(flat_score)
101 mt[torch.argsort(flat_score)[:k]] = 0
102 mm[torch.argsort(flat_mag)[:k]] = 0
103 out=[]; a=0
104 for p in params:
105 z=p.numel(); out.append((mt[a:a+z].reshape_as(p), mm[a:a+z].reshape_as(p))); a += z
106 return out
107 def run(kind, masks, steps=30):
108 m=make_model(); params=list(m.parameters());
109 with torch.no_grad():
110 for p,(mt,mm) in zip(params,masks): p.mul_(mt if kind=='tangent' else mm)
111 def loss(): return ((m(x)-y)**2).mean()
112 initial=float(loss().detach().cpu()); opt=torch.optim.SGD(m.parameters(),lr=0.08)
113 vals=[]
114 for _ in range(steps):
115 opt.zero_grad(); z=loss(); z.backward()
116 with torch.no_grad():
117 for p,(mt,mm) in zip(params,masks):
118 p.grad.mul_(mt if kind=='tangent' else mm)
119 p.mul_(mt if kind=='tangent' else mm)
120 opt.step()
121 with torch.no_grad():
122 for p,(mt,mm) in zip(params,masks): p.mul_(mt if kind=='tangent' else mm)
123 vals.append(float(loss().detach().cpu()))
124 return initial, vals[-1]
125 result={"device":device,"sparsities":{}}
126 for sp in [0.5,0.8]:
127 masks=masks_for(model0,sp)
128 t=run('tangent',masks); mag=run('magnitude',masks)
129 result['sparsities'][str(sp)]={"tangent_initial":t[0],"magnitude_initial":mag[0],"tangent_final":t[1],"magnitude_final":mag[1]}
130 return result
131 except Exception as e:
132 return {"error": type(e).__name__ + ': ' + str(e)}
133
134if __name__ == '__main__':
135 result={'seed':SEED,'math_check':tangent_sweep(),'mini_experiment':torch_mini_experiment()}
136 Path('results.json').write_text(json.dumps(result, indent=2))
137 print(json.dumps(result, indent=2))