1 files changed, 224 insertions, 0 deletions
diff --git a/src/synthetic_pair.py b/src/synthetic_pair.py
new file mode 100644
index 0000000..7614fbe
--- /dev/null
+++ b/src/synthetic_pair.py
@@ -0,0 +1,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)