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
111
112
113
114
115
116
|
"""Decomposition + robustness statistics for the committed results (results.md).
Computes, from the saved jlens_v3 artifacts and the base checkpoint:
1. W_U row-norm decomposition: r(||W_U[k]||, freq), r vs log-frequency,
Spearman, and per-layer partial correlation of faithful norm with
frequency after regressing out ||W_U[k]||.
2. Pearson p-values (normal-approx two-sided) and Spearman per layer.
3. Synthetic-pair ratio with across-seed mean/SD and bootstrap CI.
Run inside the meru container where artifacts exist:
python3 src/stats_decomp.py
"""
import os
import pickle
import sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import numpy as np
import torch
V = 65
TRAIN_BIN = 'data/shakespeare_char/train.bin'
CKPT = 'out-shakespeare-char/ckpt.pt'
def rankdata(a):
"""Average ranks (ties averaged) — minimal scipy-free rankdata."""
sorter = np.argsort(a, kind='mergesort')
ranks = np.empty(len(a), dtype=float)
ranks[sorter] = np.arange(1, len(a) + 1)
i = 0
while i < len(a):
j = i
while j + 1 < len(a) and a[sorter[j + 1]] == a[sorter[i]]:
j += 1
if j > i:
ranks[sorter[i:j + 1]] = (i + j) / 2.0 + 1
i = j + 1
return ranks
def spearman(x, y):
return np.corrcoef(rankdata(x), rankdata(y))[0, 1]
def pearson_p(r, n):
"""Two-sided p for Pearson r via normal approx on Fisher's z (fine for n=65)."""
import math
from math import erf, sqrt
z = abs(r) * math.sqrt((n - 2) / max(1 - r * r, 1e-12))
return 2 * (1 - 0.5 * (1 + erf(z / sqrt(2))))
def load_layer(l):
for p in (f'outputs/jlens_v3/layer{l}.pt', f'outputs/jlens_v3_layer{l}.pt'):
if os.path.exists(p):
return torch.load(p, map_location='cpu')
raise FileNotFoundError(f'no jlens artifact for layer {l}')
def main():
train = np.memmap(TRAIN_BIN, dtype=np.uint16, mode='r')
counts = np.bincount(train, minlength=V).astype(float)
freq = counts / counts.sum() * 100
logf = np.log10(freq)
ckpt = torch.load(CKPT, map_location='cpu')
wu = ckpt['model']['lm_head.weight'].float().numpy()
wu_norms = np.linalg.norm(wu, axis=1)
print("=== W_U row-norm decomposition ===")
print(f"r(||W_U[k]||, freq) = {np.corrcoef(wu_norms, freq)[0, 1]:+.3f}")
print(f"r(||W_U[k]||, log10 freq) = {np.corrcoef(wu_norms, logf)[0, 1]:+.3f}")
print(f"Spearman(||W_U[k]||, freq) = {spearman(wu_norms, freq):+.3f}")
print("\n=== per-layer decomposition ===")
print(f"{'layer':<6}{'r(faith,freq)':>14}{'p':>10}{'r(faith,WU)':>13}"
f"{'Spear(faith)':>14}{'partial|WU':>13}")
for l in range(6):
d = load_layer(l)
fn = np.array([d['faithful_norms'][k] for k in range(V)])
rf = np.corrcoef(fn, freq)[0, 1]
rfw = np.corrcoef(fn, wu_norms)[0, 1]
sf = spearman(fn, freq)
A = np.vstack([wu_norms, np.ones(V)]).T
resid = fn - A @ np.linalg.lstsq(A, fn, rcond=None)[0]
r_part = np.corrcoef(resid, freq)[0, 1]
print(f"L{l:<5}{rf:>+14.3f}{pearson_p(rf, V):>10.1e}{rfw:>+13.3f}"
f"{sf:>+14.3f}{r_part:>+13.3f}")
print("\n=== synthetic pair ratio ===")
with open('data/synth_pair/meta.pkl', 'rb') as f:
meta = pickle.load(f)
iid_s = meta['stoi']['@']
iid_n = meta['stoi']['#']
ratios = []
for s in range(3):
vals_s, vals_n = [], []
for l in (2, 3, 4):
d = torch.load(f'outputs/synth_pair/seed{s}/layer{l}.pt',
map_location='cpu')
fn = d['faithful_norms']
vals_s.append(fn[iid_s])
vals_n.append(fn[iid_n])
ratios.append(np.mean(vals_s) / max(np.mean(vals_n), 1e-9))
r_arr = np.array(ratios)
rng = np.random.RandomState(0)
boot = [np.mean(rng.choice(r_arr, 3, replace=True)) for _ in range(10000)]
print(f"ratio per seed: {[f'{r:.3f}' for r in ratios]}")
print(f"mean +/- SD: {r_arr.mean():.3f} +/- {r_arr.std(ddof=1):.3f}")
print(f"bootstrap 95% CI: [{np.percentile(boot, 2.5):.3f}, "
f"{np.percentile(boot, 97.5):.3f}]")
if __name__ == '__main__':
main()
|
