File size: 5,476 Bytes
fc32096 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | import argparse
import json
import numpy as np
from common import ROOT, inside, save_json
from weights import load
def mean_window(directory, metadata, text, signal, window):
total, count = None, 0
for file in metadata['files']:
if file['text'] != text or file['signal'] != signal:
continue
start, end = file['token_start'], file['token_end']
left, right = max(start, window[0]), min(end, window[1])
if right <= left:
continue
array = np.load(directory / file['name'], mmap_mode='r')[0, left-start:right-start]
value = array.sum(0, dtype=np.float64)
total = value if total is None else total + value
count += right-left
if count != window[1]-window[0]:
raise ValueError(f'Incomplete {signal} window for text {text}: {count}')
return (total/count).astype(np.float32)
def branch_logps(directory, metadata, labels):
result = []
for row in labels:
values = []
for file in metadata['files']:
if file['text'] != row['text'] or file['signal'] != 'token_logprobs':
continue
array = np.load(directory / file['name'])
start = max(0, row['branch_start']-1-file['token_start'])
end = min(len(array), row.get('branch_end',file['token_start']+len(array)+1)-1-file['token_start'])
if end > start:
values.extend(array[start:end].tolist())
if not values:
raise ValueError('Missing branch scores')
result.append(dict(row, mean_logp=float(np.mean(values)), logp_sum=float(np.sum(values)), scored_tokens=len(values)))
return result
def main():
p = argparse.ArgumentParser()
p.add_argument('--probe', required=True)
p.add_argument('--labels', default=str(ROOT/'configs/matched_probe_001_labels.json'))
p.add_argument('--out', required=True)
p.add_argument('--model', default=str(ROOT/'models/original'))
p.add_argument('--logp-only', action='store_true')
p.add_argument('--protect-productive', action='store_true')
args = p.parse_args()
directory, out = inside(args.probe), inside(args.out)
metadata = json.loads((directory/'meta.json').read_text())
if not metadata.get('complete'):
raise ValueError('Probe is incomplete')
labels = json.loads(inside(args.labels).read_text())
out.mkdir(parents=True, exist_ok=True)
save_json(out/'branch_logps.json', branch_logps(directory, metadata, labels))
if args.logp_only:
return
arrays, ranking = {}, []
discovery = [r for r in labels if r['split']=='discovery']
pairs = sorted({r['pair'] for r in discovery})
for layer in range(24):
hidden, mlp = {}, {}
for row in discovery:
key = (row['pair'], row['condition'])
hidden[key] = mean_window(directory, metadata, row['text'], f'hidden_{layer+1}', row['mean_window'])
mlp[key] = mean_window(directory, metadata, row['text'], f'mlp_act_{layer}', row['mean_window'])
hp = np.stack([hidden[(pair,'productive')] for pair in pairs])
hr = np.stack([hidden[(pair,'repetitive')] for pair in pairs])
dh = hr-hp
dm = np.stack([mlp[(pair,'repetitive')]-mlp[(pair,'productive')] for pair in pairs])
direction = dh.mean(0)
if args.protect_productive:
_, singular, basis = np.linalg.svd(hp,full_matrices=False)
basis = basis[singular>singular[0]*1e-6]
direction = direction-basis.T@(basis@direction)
magnitude = np.linalg.norm(direction)
direction = direction/magnitude
weight = load(args.model, f'model.language_model.layers.{layer}.mlp.down_proj.weight').float().numpy()
projection = direction @ weight
scores = dm * projection
mean = scores.mean(0)
columns = np.argsort(mean)[-64:][::-1].copy()
arrays[f'u_{layer}'] = direction
arrays[f'columns_{layer}'] = columns
arrays[f'scores_{layer}'] = mean
arrays[f'pair_delta_mlp_{layer}'] = dm
arrays[f'productive_mlp_{layer}'] = np.stack([mlp[(pair,'productive')] for pair in pairs])
arrays[f'repetitive_mlp_{layer}'] = np.stack([mlp[(pair,'repetitive')] for pair in pairs])
arrays[f'pair_delta_hidden_{layer}'] = dh
arrays[f'productive_hidden_{layer}'] = hp
arrays[f'repetitive_hidden_{layer}'] = hr
ranking.append({'layer':layer,'direction_norm':float(magnitude),'top64_writer_contrast':float(mean[columns].sum()),
'top64_fraction_positive_contrast':float(mean[columns].sum()/np.maximum(mean,0).sum()),
'direction_pair_cosines':(dh@direction/np.linalg.norm(dh,axis=1)).tolist(),
'top_neurons':[{'neuron':int(j),'contrast':float(mean[j]),'positive_pairs':int((scores[:,j]>0).sum())} for j in columns[:12]],
'top64_consistent_pairs_mean':float((scores[:,columns]>0).mean()),
'productive_protection':args.protect_productive,
'max_productive_projection':float(np.max(np.abs(hp@direction)))})
np.savez(out/'directions_and_writers.npz', **arrays)
save_json(out/'layer_ranking.json', sorted(ranking,key=lambda x:x['top64_writer_contrast'],reverse=True))
print(json.dumps(sorted(ranking,key=lambda x:x['top64_writer_contrast'],reverse=True)[:5],indent=2))
if __name__ == '__main__':
main()
|