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
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
|
"""
Frequency-Matched Synthetic Pair Test (Gemini 3.1 Pro's design).
Two synthetic tokens at IDENTICAL unigram frequency in an otherwise-normal
Shakespeare corpus:
T_struct '@' : appears only after the trigger sequence "the " (high conditional
predictability, structured context).
T_noise '#' : injected at uniform random positions (zero conditional structure).
Hypotheses:
Frequency-only: J-lens norms of '@' and '#' are identical at every layer
(same unigram frequency, same rarity).
Structure/workspace (Anthropic): '@' maintains a higher faithful J-lens norm,
especially in intermediate layers (model tracks the trigger
context in the residual stream).
Uses the FAITHFUL J-lens (jlens_v3: rows of W_U * J_l) plus the old proxy.
Steps:
python3 src/synthetic_pair.py --step prep # build data/synth_pair
python3 src/synthetic_pair.py --step train --seed 0 [--max_iters 3000]
python3 src/synthetic_pair.py --step jlens --seed 0 [--layers 0,1,2,3,4,5]
python3 src/synthetic_pair.py --step summary # compare @ vs #
"""
import sys, os, argparse, subprocess, pickle, random
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from typing import Any
import numpy as np
import torch
DATA_SRC = 'data/shakespeare_char/input.txt'
DATA_DIR = 'data/synth_pair'
OUT_ROOT = 'out-synth-pair'
JLENS_OUT = 'outputs/synth_pair'
T_STRUCT = '@'
T_NOISE = '#'
TRIGGER = 'the '
TARGET_FREQ = 0.001 # 0.1%
MODEL_ARGS: dict[str, Any] = dict(n_layer=6, n_head=6, n_embd=384, block_size=128,
bias=False, dropout=0.2)
def prep():
with open(DATA_SRC) as f:
text = f.read()
n_target = max(1, int(TARGET_FREQ * len(text)))
# positions for T_struct: after occurrences of TRIGGER
trig_positions = []
start = 0
while True:
i = text.find(TRIGGER, start)
if i < 0:
break
trig_positions.append(i + len(TRIGGER))
start = i + len(TRIGGER)
assert len(trig_positions) >= n_target, f"only {len(trig_positions)} triggers"
rng = np.random.RandomState(42)
struct_pos = sorted(rng.choice(trig_positions, size=n_target, replace=False).tolist())
# positions for T_noise: uniform random, disjoint from struct_pos
noise_pos = sorted(rng.choice(
[p for p in range(len(text)) if p not in set(struct_pos)],
size=n_target, replace=False).tolist())
# insert with offset (both sets sorted -> single merge pass)
insertions = [(p, T_STRUCT) for p in struct_pos] + [(p, T_NOISE) for p in noise_pos]
insertions.sort()
out = []
prev = 0
for pos, ch in insertions:
out.append(text[prev:pos])
out.append(ch)
prev = pos
out.append(text[prev:])
modified = ''.join(out)
assert modified.count(T_STRUCT) == modified.count(T_NOISE) == n_target
print(f"prep: '{T_STRUCT}' x{n_target} after '{TRIGGER.strip()}', "
f"'{T_NOISE}' x{n_target} random, "
f"freq each = {n_target/len(modified):.4%}")
# build vocab (existing chars + the two synthetic)
chars = sorted(set(text)) + [T_STRUCT, T_NOISE]
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for i, c in enumerate(chars)}
data = np.array([stoi[c] for c in modified], dtype=np.uint16)
n = int(0.9 * len(data))
os.makedirs(DATA_DIR, exist_ok=True)
data[:n].tofile(os.path.join(DATA_DIR, 'train.bin'))
data[n:].tofile(os.path.join(DATA_DIR, 'val.bin'))
with open(os.path.join(DATA_DIR, 'meta.pkl'), 'wb') as f:
pickle.dump({'stoi': stoi, 'itos': itos, 'vocab_size': len(chars)}, f)
with open(os.path.join(DATA_DIR, 'input.txt'), 'w') as f:
f.write(modified)
print(f"prep: vocab={len(chars)}, train={n:,} val={len(data)-n:,} tokens")
def train(seed, max_iters, batch_size):
sys.path.insert(0, '.')
from model import GPT, GPTConfig
torch.manual_seed(seed)
np.random.seed(seed)
device = 'cuda'
train_data = np.memmap(f'{DATA_DIR}/train.bin', dtype=np.uint16, mode='r')
val_data = np.memmap(f'{DATA_DIR}/val.bin', dtype=np.uint16, mode='r')
with open(f'{DATA_DIR}/meta.pkl', 'rb') as f:
meta = pickle.load(f)
args = dict(
n_layer=int(MODEL_ARGS['n_layer']), n_head=int(MODEL_ARGS['n_head']),
n_embd=int(MODEL_ARGS['n_embd']), block_size=int(MODEL_ARGS['block_size']),
bias=bool(MODEL_ARGS['bias']), dropout=float(MODEL_ARGS['dropout']),
vocab_size=int(meta['vocab_size']),
)
model = GPT(GPTConfig(**args)).to(device)
print(f"seed {seed}: params={sum(p.numel() for p in model.parameters())/1e6:.2f}M")
opt = model.configure_optimizers(weight_decay=0.1, learning_rate=1e-3,
betas=(0.9, 0.99), device_type='cuda')
bs = batch_size
blk = args['block_size']
def get_batch(split):
d = train_data if split == 'train' else val_data
ix = torch.randint(len(d) - blk, (bs,))
x = torch.stack([torch.from_numpy(d[i:i+blk].astype(np.int64)) for i in ix])
y = torch.stack([torch.from_numpy(d[i+1:i+1+blk].astype(np.int64)) for i in ix])
return x.to(device), y.to(device)
best_val = 1e9
out_dir = f'{OUT_ROOT}/seed{seed}'
os.makedirs(out_dir, exist_ok=True)
for it in range(max_iters):
if it % 500 == 0:
model.eval()
lv = []
for _ in range(50):
X, Y = get_batch('val')
with torch.no_grad():
_, loss = model(X, Y)
lv.append(loss.item())
v = np.mean(lv)
model.train()
if v < best_val:
best_val = v
torch.save({'model': model.state_dict(), 'model_args': args,
'best_val_loss': best_val}, f'{out_dir}/ckpt.pt')
if it % 1000 == 0:
print(f" iter {it}: val={v:.4f}")
X, Y = get_batch('train')
_, loss = model(X, Y)
loss.backward()
opt.step()
opt.zero_grad(set_to_none=True)
print(f"seed {seed} done. best_val={best_val:.4f}")
def jlens(seed, layers, n_prompts):
out_dir = f'{JLENS_OUT}/seed{seed}'
cmd = ["python3", "-u", "src/jlens_v3.py",
"--checkpoint", f'{OUT_ROOT}/seed{seed}/ckpt.pt',
"--data_dir", DATA_DIR,
"--n_prompts", str(n_prompts),
"--layers", layers,
"--chunk", "16",
"--output_dir", out_dir]
print("running:", " ".join(cmd))
r = subprocess.run(cmd, cwd='/workspace/code')
assert r.returncode == 0, "jlens_v3 failed"
def summary(layers):
with open(f'{DATA_DIR}/meta.pkl', 'rb') as f:
meta = pickle.load(f)
sid = meta['stoi']
iid_s = sid[T_STRUCT]
iid_n = sid[T_NOISE]
train_data = np.memmap(f'{DATA_DIR}/train.bin', dtype=np.uint16, mode='r')
V = meta['vocab_size']
counts = np.bincount(train_data, minlength=V).astype(float)
freq = counts / counts.sum() * 100
seeds = sorted([d for d in os.listdir(JLENS_OUT) if d.startswith('seed')])
print(f"\n{'='*72}")
print("SYNTHETIC PAIR: faithful J-lens norms of T_struct '@' vs T_noise '#'")
print(f"('@' freq={freq[iid_s]:.3f}%, '#' freq={freq[iid_n]:.3f}%)")
print(f"{'='*72}")
print(f"{'seed':<5}{'layer':<6}{'@ norm':>9}{'# norm':>9}{'ratio':>8} freq-corr r")
for s in seeds:
for l in map(int, layers.split(',')):
d = torch.load(f'{JLENS_OUT}/{s}/layer{l}.pt', map_location='cpu')
fn = d['faithful_norms']
f_arr = np.array([freq[k] for k in range(V)])
rf = np.corrcoef(np.array([fn[k] for k in range(V)]), f_arr)[0, 1]
print(f"{s:<5}{l:<6}{fn[iid_s]:>9.4f}{fn[iid_n]:>9.4f}"
f"{fn[iid_s]/max(fn[iid_n],1e-9):>8.2f} {rf:+.3f}")
# per-seed ratio across middle layers
mids = [l for l in map(int, layers.split(',')) if l in (2, 3, 4)]
rs = []
rn = []
for l in mids:
d = torch.load(f'{JLENS_OUT}/{s}/layer{l}.pt', map_location='cpu')
fn = d['faithful_norms']
rs.append(fn[iid_s])
rn.append(fn[iid_n])
print(f" -> mean middle-layer ratio @/# = {np.mean(rs)/max(np.mean(rn),1e-9):.3f}")
if __name__ == '__main__':
ap = argparse.ArgumentParser()
ap.add_argument('--step', required=True, choices=['prep', 'train', 'jlens', 'summary'])
ap.add_argument('--seed', type=int, default=0)
ap.add_argument('--max_iters', type=int, default=3000)
ap.add_argument('--batch_size', type=int, default=16)
ap.add_argument('--layers', default='0,1,2,3,4,5')
ap.add_argument('--n_prompts', type=int, default=10)
a = ap.parse_args()
if a.step == 'prep':
prep()
elif a.step == 'train':
train(a.seed, a.max_iters, a.batch_size)
elif a.step == 'jlens':
jlens(a.seed, a.layers, a.n_prompts)
else:
summary(a.layers)
|
