Channel-aware attention-head pruning / experiment.py
Mechanism failed
1import json
2from pathlib import Path
3import numpy as np
4
5SEED = 2304
6EPS = 1e-12
7rng = np.random.default_rng(SEED)
8
9
10def clr(x):
11 x = np.asarray(x, float)
12 lx = np.log(x)
13 return lx - np.mean(lx, axis=-1, keepdims=True)
14
15
16def row(s, c):
17 c = np.asarray(c, float)
18 c = c / c.sum()
19 return np.r_[s, (1.0 - s) * c]
20
21
22def split(p):
23 s = float(p[0])
24 c = np.maximum(np.asarray(p[1:], float) / (1.0 - s), EPS)
25 return s, c / c.sum()
26
27
28def distances(a, b):
29 sa, ca = split(a)
30 sb, cb = split(b)
31 q = np.log(ca) - np.log(cb)
32 g = float(q.mean())
33 z = q - g
34 dc2 = float(np.mean(z * z))
35 ua = np.log(sa / (1.0 - sa))
36 ub = np.log(sb / (1.0 - sb))
37 ds2 = 0.5 * (ua - ub) ** 2
38 full = float(np.sum((clr(a) - clr(b)) ** 2))
39 m = len(ca)
40 D = m + 1
41 # Exact CLR decomposition. k is the relative scaling of the content
42 # block; it is not present in the idea's stated sink-only coordinate.
43 k = np.log((1.0 - sa) / (1.0 - sb))
44 corrected = m * dc2 + (m / D) * (ua - ub - g) ** 2
45 stated = (m / D) * (dc2 + ds2)
46 return dict(full=full, content=dc2, sink=ds2, corrected=corrected,
47 stated=stated, content_logratio_mean=g, content_scale=k)
48
49
50def math_sweep():
51 invariant = []
52 for m in [2, 3, 7, 15]:
53 c1 = rng.dirichlet(np.ones(m) * .7)
54 c2 = rng.dirichlet(np.ones(m) * .7)
55 vals = [distances(row(s1, c1), row(s2, c2))['content']
56 for s1, s2 in [(.1, .2), (.5, .8), (.9, .97)]]
57 invariant.append({'m': m, 'predicted_range': 0.0,
58 'observed_range': float(max(vals) - min(vals))})
59
60 sink_scaling = []
61 for m in [2, 3, 7, 15, 31]:
62 c = rng.dirichlet(np.ones(m) * .7)
63 vals = []
64 for _ in range(100):
65 s1, s2 = rng.uniform(.05, .95, 2)
66 d = distances(row(s1, c), row(s2, c))
67 vals.append(d['full'] / d['sink'])
68 sink_scaling.append({'m': m, 'predicted': 2*m/(m+1),
69 'observed_mean': float(np.mean(vals)),
70 'observed_sd': float(np.std(vals))})
71
72 check = []
73 for m in [2, 3, 7, 15, 31]:
74 corrected_err, stated_err = [], []
75 for _ in range(300):
76 a = rng.dirichlet(np.ones(m+1) * .7)
77 b = rng.dirichlet(np.ones(m+1) * .7)
78 d = distances(a, b)
79 corrected_err.append(abs(d['full'] - d['corrected']) / (d['full'] + EPS))
80 stated_err.append(abs(d['full'] - d['stated']) / (d['full'] + EPS))
81 check.append({'m': m, 'corrected_max_relative_error': float(max(corrected_err)),
82 'stated_median_relative_error': float(np.median(stated_err)),
83 'stated_mean_relative_error': float(np.mean(stated_err))})
84 return {'content_sink_invariance': invariant,
85 'equal_content_sink_scaling': sink_scaling,
86 'decomposition_check': check}
87
88
89def cosine_distance(a, b):
90 return 1 - float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))
91
92
93def js(a, b):
94 q = .5 * (a + b)
95 return float(.5*np.sum(a*np.log((a+EPS)/(q+EPS))) +
96 .5*np.sum(b*np.log((b+EPS)/(q+EPS))))
97
98
99def pruning_toy():
100 m = 5
101 ca = np.array([.55, .25, .12, .06, .02])
102 cb = np.array([.02, .06, .12, .25, .55])
103 heads = [row(.95, ca), row(.95, cb), row(.70, ca),
104 row(.95, ca*np.array([1.01,.99,1,1,1])), row(.50, cb)]
105 names = ['reference', 'same_sink_diff_content', 'same_content_diff_sink',
106 'duplicate', 'unrelated']
107 pairs = []
108 for i in range(len(heads)):
109 for j in range(i+1, len(heads)):
110 d = distances(heads[i], heads[j])
111 pairs.append({'pair': [i,j], 'names': [names[i],names[j]],
112 'cosine': cosine_distance(heads[i],heads[j]),
113 'js': js(heads[i],heads[j]), 'content': d['content'],
114 'sink': d['sink'], 'full': d['full']})
115 raw = sorted(pairs, key=lambda x: x['cosine'])
116 admissible = [p for p in pairs if p['content'] < .02 and p['sink'] < .05]
117 return {'pairs': pairs, 'raw_cosine_nearest': raw[:4],
118 'channel_admissible': admissible,
119 'thresholds': {'content': .02, 'sink': .05},
120 'intended_redundant_pair': [0,3]}
121
122
123def main():
124 out = {'seed': SEED, 'math': math_sweep(), 'pruning_toy': pruning_toy()}
125 Path('results.json').write_text(json.dumps(out, indent=2))
126 print(json.dumps(out, indent=2))
127
128
129if __name__ == '__main__':
130 main()