Add causal experiments: synthetic frequency-matched pair (Gemini design) + loss-reweighting (Codex design); jlens_v3 output_dir + last-layer device fix

1 files changed, 197 insertions, 0 deletions

diff --git a/src/loss_reweight.py b/src/loss_reweight.py
new file mode 100644
index 0000000..3274807
--- /dev/null
+++ b/src/loss_reweight.py
@@ -0,0 +1,197 @@
+"""
+Loss-reweighting ablation (GPT-5.6-Terra's design, adapted to faithful J-lens).
+
+Tests whether increasing a token's EFFECTIVE frequency/importance reduces its
+faithful J-lens norm — without the confounds of the old random-insertion
+ablation (which shifted positions, destroyed n-grams, and used an unmatched
+control run).
+
+Design per seed (identical init + identical minibatch order for all three):
+  q-upweight:   cross-entropy terms whose target is 'q' are weighted x2.
+  control:      ordinary loss.
+  ctrl_random:  same-total-loss control: weight x2 on the SAME NUMBER of
+                randomly chosen non-'q' target positions (deterministic per
+                batch index, so all models share the same control positions).
+
+If increased effective frequency causally reduces the faithful J-lens norm of
+'q', the q-upweight model must show a lower norm than BOTH controls.
+
+Steps:
+  python3 src/loss_reweight.py --step train --mode q --seed 0 [--max_iters 3000]
+  python3 src/loss_reweight.py --step jlens --mode q --seed 0 [--layers 2,3,4]
+  python3 src/loss_reweight.py --step summary
+"""
+import sys, os, argparse, subprocess, pickle
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from typing import Any
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+DATA_DIR = 'data/shakespeare_char'
+OUT_ROOT = 'out-loss-reweight'
+JLENS_OUT = 'outputs/loss_reweight'
+TARGET = 'q'
+WEIGHT = 2.0
+MODEL_ARGS: dict[str, Any] = dict(n_layer=6, n_head=6, n_embd=384, block_size=128,
+                                  bias=False, dropout=0.2)
+LOSS_MODES = ('q', 'control', 'ctrl_random')
+
+
+def _batch(data, blk, bs, g, device):
+    ix = torch.randint(len(data) - blk, (bs,), generator=g)
+    x = torch.stack([torch.from_numpy(data[i:i+blk].astype(np.int64)) for i in ix])
+    y = torch.stack([torch.from_numpy(data[i+1:i+1+blk].astype(np.int64)) for i in ix])
+    return x.to(device), y.to(device)
+
+
+def _weighted_loss(logits, y, mode, q_id, batch_k, V, device):
+    """Per-token weighted CE. Returns scalar loss."""
+    logp = F.log_softmax(logits.view(-1, V), dim=-1)
+    nll = -logp.gather(1, y.view(-1, 1)).squeeze(1)  # (B*T,)
+    w = torch.ones_like(nll)
+    if mode == 'q':
+        w[y.view(-1) == q_id] = WEIGHT
+    elif mode == 'ctrl_random':
+        g = torch.Generator(device=device).manual_seed(1000 + batch_k)
+        n_q = int((y == q_id).sum().item())
+        flat = torch.arange(y.numel(), device=device)
+        non_q = flat[y.view(-1) != q_id]
+        if len(non_q) > 0 and n_q > 0:
+            pick = non_q[torch.randperm(len(non_q), generator=g)[:min(n_q, len(non_q))]]
+            w[pick] = WEIGHT
+    return (nll * w).mean()
+
+
+def train(mode, 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)
+    q_id = meta['stoi'][TARGET]
+    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"[{mode}] 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']
+    V = args['vocab_size']
+    # identical minibatch order for every model: per-seed generator, fixed start
+    g = torch.Generator(device=device).manual_seed(20260731 + seed)
+    gval = torch.Generator(device=device).manual_seed(777 + seed)
+
+    def get_batch(split):
+        d = train_data if split == 'train' else val_data
+        gg = g if split == 'train' else gval
+        return _batch(d, blk, bs, gg, device)
+
+    best_val = 1e9
+    out_dir = f'{OUT_ROOT}/{mode}/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():
+                    logits = model(X)[0]
+                lv.append(F.cross_entropy(logits.view(-1, V), Y.view(-1)).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')
+        logits = model(X)[0]
+        loss = _weighted_loss(logits, Y, mode, q_id, it, V, device)
+        loss.backward()
+        opt.step()
+        opt.zero_grad(set_to_none=True)
+    print(f"[{mode}] seed {seed} done. best_val={best_val:.4f}")
+
+
+def jlens(mode, seed, layers, n_prompts):
+    out_dir = f'{JLENS_OUT}/{mode}/seed{seed}'
+    cmd = ["python3", "-u", "src/jlens_v3.py",
+           "--checkpoint", f'{OUT_ROOT}/{mode}/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)
+    q_id = meta['stoi'][TARGET]
+    seeds = sorted(set(
+        d.split('seed')[1] for mode in LOSS_MODES
+        for d in os.listdir(f'{JLENS_OUT}/{mode}')
+        if d.startswith('seed')))
+    print(f"\n{'='*78}")
+    print(f"LOSS-REWEIGHTING: faithful J-lens norm of '{TARGET}' "
+          f"({WEIGHT}x CE) vs controls")
+    print(f"{'='*78}")
+    print(f"{'seed':<5}{'layer':<6}" + "".join(f"{m:>14}" for m in LOSS_MODES))
+    for s in seeds:
+        for l in map(int, layers.split(',')):
+            row = [s, str(l)]
+            for m in LOSS_MODES:
+                d = torch.load(f'{JLENS_OUT}/{m}/seed{s}/layer{l}.pt',
+                               map_location='cpu')
+                row.append(f"{d['faithful_norms'][q_id]:.4f}")
+            print(f"{row[0]:<5}{row[1]:<6}" + "".join(f"{v:>14}" for v in row[2:]))
+        # mean over middle layers per mode
+        mids = [l for l in map(int, layers.split(','))]
+        means = {}
+        for m in LOSS_MODES:
+            vals = []
+            for l in mids:
+                d = torch.load(f'{JLENS_OUT}/{m}/seed{s}/layer{l}.pt',
+                               map_location='cpu')
+                vals.append(d['faithful_norms'][q_id])
+            means[m] = np.mean(vals)
+        print(f"  -> mean over layers: q={means['q']:.4f} "
+              f"control={means['control']:.4f} ctrl_random={means['ctrl_random']:.4f}")
+        print(f"  -> q/control = {means['q']/max(means['control'],1e-9):.3f}  "
+              f"q/ctrl_random = {means['q']/max(means['ctrl_random'],1e-9):.3f}")
+
+
+if __name__ == '__main__':
+    ap = argparse.ArgumentParser()
+    ap.add_argument('--step', required=True, choices=['train', 'jlens', 'summary'])
+    ap.add_argument('--mode', choices=LOSS_MODES)
+    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='2,3,4')
+    ap.add_argument('--n_prompts', type=int, default=10)
+    a = ap.parse_args()
+    if a.step == 'train':
+        assert a.mode, "need --mode"
+        train(a.mode, a.seed, a.max_iters, a.batch_size)
+    elif a.step == 'jlens':
+        assert a.mode, "need --mode"
+        jlens(a.mode, a.seed, a.layers, a.n_prompts)
+    else:
+        summary(a.layers)