""" 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' REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) 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': # Same number of upweighted positions as the q-mode model, but on # random non-q targets. NOTE: gradient-magnitude distribution differs # from q-mode (random positions spread across the batch vs rare 'q' # positions); the design controls for "any 2x reweighting changes the # model", not for exact gradient-mass matching. g = torch.Generator().manual_seed(1000 + batch_k) # CPU generator (randperm) 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 CPU generator, fixed start # (torch.randint does not accept CUDA generators on torch 2.4) g = torch.Generator().manual_seed(20260731 + seed) gval = torch.Generator().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(): _, 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') logits = model(X, Y)[0] # pass Y: targets=None would give last-position logits only 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=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) 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)