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
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
|
"""
GPT-2 Small J-lens — efficient sampling approach.
Instead of computing J-lens for all 50257 tokens, sample:
- 100 most common + 100 rarest = 200 tokens for frequency test
- 1000 random tokens for effective rank estimation
Each backward pass takes ~50ms on K2200.
200 tokens × 3 batches × 3 layers = 1800 backward passes ≈ 90s
"""
import sys, os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import torch, numpy as np, pickle, time, subprocess
try:
import transformers
except ImportError:
subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "transformers==4.44.0", "accelerate", "requests"])
import transformers
import requests
device = 'cuda'
def main():
print("=" * 60)
print("GPT-2 SMALL J-LENS — Frequency + Dimensionality Tests")
print("=" * 60)
# Load model
model = transformers.GPT2LMHeadModel.from_pretrained(
"openai-community/gpt2", torch_dtype=torch.float32).to(device)
model.eval()
tokenizer = transformers.GPT2Tokenizer.from_pretrained("openai-community/gpt2")
d_model = model.config.n_embd
n_layers = model.config.n_layer
vocab_size = model.config.vocab_size
print(f"GPT-2 Small: {n_layers} layers, d={d_model}, V={vocab_size}, V/d={vocab_size/d_model:.1f}x")
# Load corpus
text = requests.get(
"https://raw.githubusercontent.com/karpathy/nanoGPT/master/data/shakespeare_char/input.txt"
).text[:200000]
tokens = tokenizer(text, return_tensors='np', truncation=True, max_length=2000)['input_ids'][0]
print(f"Corpus: {len(tokens)} tokens")
# Estimate token frequencies in our corpus
from collections import Counter
freq = Counter(tokens.tolist())
total = len(tokens)
sorted_by_freq = sorted(freq.items(), key=lambda x: x[1], reverse=True)
# Pick: 50 most common, 50 rarest
sample_tokens = [t for t, _ in sorted_by_freq[:50]] + [t for t, _ in sorted_by_freq[-50:]]
sample_tokens = list(set(sample_tokens)) # dedupe
print(f"Sampling {len(sample_tokens)} tokens for frequency test")
# J-lens computation
n_batches = 3
seq_len = 32
layers_to_test = [3, 6, 9] # early, middle, late
# Accumulators: {layer: {token_id: [norm_sum, count]}}
accum = {layer: {tid: [0.0, 0] for tid in sample_tokens} for layer in layers_to_test}
for batch_i in range(n_batches):
ix = torch.randint(0, len(tokens) - seq_len - 1, (1,))
x = torch.from_numpy(tokens[ix[0]:ix[0]+seq_len].astype(np.int64)).unsqueeze(0).to(device)
for layer_idx in layers_to_test:
# Capture residual at this layer
resid_captured = {}
def hook(module, inp, out):
h = out[0] if isinstance(out, tuple) else out
resid_captured['val'] = h
target = model.transformer.h[layer_idx]
handle = target.register_forward_hook(hook)
# Forward (no no_grad — we need gradients)
result = model(x)
handle.remove()
logits = result.logits # (1, seq_len, V)
resid = resid_captured['val'] # (1, seq_len, d)
for tid in sample_tokens:
token_logprob = torch.nn.functional.log_softmax(logits, dim=-1)[:, :, tid].sum()
try:
grad = torch.autograd.grad(token_logprob, resid, retain_graph=True)[0]
norm = grad.norm().item()
accum[layer_idx][tid][0] += norm
accum[layer_idx][tid][1] += 1
except:
pass
del logits, result, resid
torch.cuda.empty_cache()
print(f" Batch {batch_i+1}/{n_batches} done")
# Results
print(f"\n{'='*60}")
print("FREQUENCY vs J-LENS NORM (GPT-2 Small)")
print(f"{'='*60}")
for layer_idx in layers_to_test:
print(f"\n--- Layer {layer_idx} ---")
norms = {}
for tid in sample_tokens:
s, c = accum[layer_idx][tid]
if c > 0:
norms[tid] = s / c
# Token frequencies
freqs = {tid: freq.get(tid, 0)/total*100 for tid in norms}
# Correlation
n_arr = np.array(list(norms.values()))
f_arr = np.array([freqs[t] for t in norms])
corr = np.corrcoef(n_arr, f_arr)[0, 1]
# Top/bottom by norm
sorted_tokens = sorted(norms.items(), key=lambda x: x[1], reverse=True)
print(f" Top 5 by J-lens norm:")
for tid, n in sorted_tokens[:5]:
tok_str = tokenizer.decode([tid]).replace('\n', '\\n')
print(f" '{tok_str}' (freq={freqs[tid]:.3f}%): norm={n:.4f}")
print(f" Bottom 5 by J-lens norm:")
for tid, n in sorted_tokens[-5:]:
tok_str = tokenizer.decode([tid]).replace('\n', '\\n')
print(f" '{tok_str}' (freq={freqs[tid]:.3f}%): norm={n:.4f}")
print(f" Pearson r(norm, freq): {corr:.3f}")
print(f"\n{'='*60}")
print("RESULTS SUMMARY")
print(f"{'='*60}")
print(f" nanoGPT (V/d=0.2x): full-rank J-space, r=-0.65 freq correlation")
print(f" GPT-2 (V/d=65x): {'r=' + str(corr)}")
print(f" HYPOTHESIS: frequency anti-correlation persists at scale")
if __name__ == '__main__':
main()
|
