""" 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' REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) 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(clean=False): 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. Default: uniform random (slices inside words, # ~95% of the time — the confound the clean-boundary control fixes). # clean=True: always at a word start (immediately after a space), so it # sits at a clean boundary exactly like T_struct after "the ", but is # still unpredictable (uniform over words). if clean: starts = [i + 1 for i, c in enumerate(text) if c == ' ' and i + 1 < len(text)] avail = [p for p in starts if p not in set(struct_pos)] else: avail = [p for p in range(len(text)) if p not in set(struct_pos)] noise_pos = sorted(rng.choice(avail, size=n_target, replace=False).tolist()) # sanity: fraction of noise insertions that slice inside a word in_word = sum(1 for p in noise_pos if 0 < p < len(text) and (text[p - 1].isalnum() and text[p].isalnum())) # 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 mode = "clean-boundary" if clean else "random" print(f"prep({mode}): '{T_STRUCT}' x{n_target} after '{TRIGGER.strip()}', " f"'{T_NOISE}' x{n_target} {mode}, " f"freq each = {n_target/len(modified):.4%}, " f"in-word '#' = {in_word}/{n_target} ({in_word/n_target:.1%})") # 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=REPO_ROOT) 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) ap.add_argument('--clean', action='store_true', help='prep: place the noise token at word starts (after ' 'random spaces) instead of uniform random positions') ap.add_argument('--data_dir', default=DATA_DIR) ap.add_argument('--out_root', default=OUT_ROOT) ap.add_argument('--jlens_out', default=JLENS_OUT) a = ap.parse_args() if a.data_dir != DATA_DIR or a.out_root != OUT_ROOT or a.jlens_out != JLENS_OUT: DATA_DIR, OUT_ROOT, JLENS_OUT = a.data_dir, a.out_root, a.jlens_out if a.step == 'prep': prep(clean=a.clean) 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)