Channel-aware attention-head pruning / experiment.py

Mechanism failed

Raw ⬇ ZIP
  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()