#!/usr/bin/env python3 """ scale_experiment.py — Scale comparison: KAIZEN 42M vs GPT-2 117M. Research question: does semantic structure recoverable by a linear head from frozen LM hidden states improve with model scale? Protocol (identical to Phase 7.0 / semantic_probe.py): - Same 50 paraphrase pairs from probe_paraphrase.json - Same 70/30 task-level split (35 train, 15 held-out) - Same LinearHead architecture (d_model→128, InfoNCE) - Same recall@1 metric on held-out tasks Output: - Raw recall@1 (no head) at each scale = lower bound - Recall@1 learning curve across epochs [0,10,25,50,100,150,200,300] - Final recall@1 = publishable single-number comparison KAIZEN 42M: embed_task via canonical prompt [BOS,USER]+q+[ASST], mean-pool, 512-dim GPT-2 117M: plain question text tokenized with GPT-2 BPE, mean-pool last hidden, 768-dim """ import json, random, sys, time import torch import torch.nn as nn import torch.nn.functional as F from tokenizers import Tokenizer from transformers import GPT2Model, GPT2Tokenizer from huggingface_hub import hf_hub_download from lora import KaizenWithLoRA from eval_benchmark import BENCHMARK_TASKS, build_prompt_ids, BLOCK_SIZE, MAX_GEN from online_learner import HF_TOKEN, HF_TOK_REPO # ── Config ──────────────────────────────────────────────────────────────────── def _get_default_ckpt(): from online_learner import HF_TOKEN from huggingface_hub import hf_hub_download return hf_hub_download('qoa/kaizen-42m', 'phase4_latest.pt', token=HF_TOKEN) DEFAULT_CKPT = _get_default_ckpt() PROBE_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'probe_paraphrase.json') SEED = 42 TRAIN_FRAC = 0.70 HEAD_DIM = 128 HEAD_LR = 1e-3 HEAD_TEMP = 0.07 CHECKPOINTS = [0, 10, 25, 50, 100, 150, 200, 300] # epochs at which to record recall@1 # ── LinearHead + InfoNCE ────────────────────────────────────────────────────── class LinearHead(nn.Module): def __init__(self, d_in: int, d_out: int = HEAD_DIM): super().__init__() self.proj = nn.Linear(d_in, d_out, bias=True) def forward(self, x: torch.Tensor) -> torch.Tensor: return F.normalize(self.proj(x), dim=-1) def infonce_loss(a: torch.Tensor, b: torch.Tensor, temp: float = HEAD_TEMP) -> torch.Tensor: """Bidirectional InfoNCE on (anchor, positive) pairs. a, b: [B, D].""" sim = a @ b.T / temp # [B, B] labels = torch.arange(sim.shape[0]) return (F.cross_entropy(sim, labels) + F.cross_entropy(sim.T, labels)) / 2 # ── Metrics ─────────────────────────────────────────────────────────────────── def recall_at_1_head(head: LinearHead, anc_embs: torch.Tensor, pos_embs: torch.Tensor, held_out_idx: list) -> float: """Recall@1 in g-space. For each held-out anchor i: query = head(anc_embs[i]). Gallery = head(pos_embs) over ALL 50 tasks (train + held-out). Correct if nearest gallery entry == i (the true paraphrase). """ with torch.no_grad(): g_pos = head(pos_embs) # [50, 128] g_anc = head(anc_embs[held_out_idx]) # [N_held, 128] correct = 0 for k, i in enumerate(held_out_idx): dists = ((g_pos - g_anc[k].unsqueeze(0)) ** 2).sum(dim=1) # [50] nearest = int(dists.argmin().item()) if nearest == i: correct += 1 return correct / len(held_out_idx) def recall_at_1_raw(anc_embs: torch.Tensor, pos_embs: torch.Tensor, held_out_idx: list) -> float: """Recall@1 in raw embedding space (L2-normalized, cosine-equivalent).""" with torch.no_grad(): a_n = F.normalize(anc_embs, dim=-1) # [50, d] p_n = F.normalize(pos_embs, dim=-1) # [50, d] correct = 0 for i in held_out_idx: dists = ((p_n - a_n[i].unsqueeze(0)) ** 2).sum(dim=1) nearest = int(dists.argmin().item()) if nearest == i: correct += 1 return correct / len(held_out_idx) # ── Training loop ───────────────────────────────────────────────────────────── def train_with_checkpoints(d_in: int, anc_train: torch.Tensor, pos_train: torch.Tensor, anc_all: torch.Tensor, pos_all: torch.Tensor, held_out_idx: list) -> dict: """Train LinearHead; return {epoch: recall@1} at each CHECKPOINT epoch.""" head = LinearHead(d_in=d_in) optimizer = torch.optim.Adam(head.parameters(), lr=HEAD_LR) N = anc_train.shape[0] results = {} if 0 in CHECKPOINTS: results[0] = recall_at_1_head(head, anc_all, pos_all, held_out_idx) max_ep = max(CHECKPOINTS) for ep in range(1, max_ep + 1): perm = torch.randperm(N) optimizer.zero_grad() loss = infonce_loss(head(anc_train[perm]), head(pos_train[perm])) loss.backward() optimizer.step() if ep in CHECKPOINTS: results[ep] = recall_at_1_head(head, anc_all, pos_all, held_out_idx) return results # ── Embedding extractors ────────────────────────────────────────────────────── def embed_kaizen(model: KaizenWithLoRA, tokenizer, questions: list) -> torch.Tensor: """Canonical prompt embed_task for all questions. Returns [N, 512].""" vecs = [] with torch.no_grad(): for q in questions: pids = build_prompt_ids(tokenizer, q)[:BLOCK_SIZE - MAX_GEN] x = torch.tensor([pids], dtype=torch.long) vecs.append(model.embed_task(x, adapter=None)) # [512] return torch.stack(vecs) def embed_gpt2(model: GPT2Model, tokenizer: GPT2Tokenizer, questions: list) -> torch.Tensor: """Mean-pool GPT-2 last hidden state for all questions. Returns [N, 768].""" vecs = [] with torch.no_grad(): for q in questions: inputs = tokenizer(q, return_tensors='pt', truncation=True, max_length=512) out = model(**inputs) # last_hidden_state: [1, seq_len, 768] → mean over seq_len vecs.append(out.last_hidden_state[0].mean(dim=0)) # [768] return torch.stack(vecs) # ── Main ────────────────────────────────────────────────────────────────────── def main(): t0 = time.time() rng = random.Random(SEED) print('Scale experiment: KAIZEN 42M vs GPT-2 117M') print('Metric: recall@1 on held-out paraphrase pairs') print(f'Checkpoints: {CHECKPOINTS}') print('=' * 64) # ── Probe data ─────────────────────────────────────────────────────────── with open(PROBE_PATH) as f: paraphrases = json.load(f) # list of [para_q, answer, category] orig_qs = [t[0] for t in BENCHMARK_TASKS] # 50 original questions para_qs = [p[0] for p in paraphrases] # 50 paraphrase questions # 70/30 task-level split — same seed as Phase 7.0 all_idx = list(range(50)) rng.shuffle(all_idx) n_train = int(TRAIN_FRAC * 50) # 35 train_idx = all_idx[:n_train] held_out_idx = all_idx[n_train:] # 15 print(f'Train tasks: {n_train} | Held-out tasks: {len(held_out_idx)}') # ── KAIZEN 42M embeddings ───────────────────────────────────────────────── print('\n[1/4] Loading KAIZEN 42M and embedding probe questions...') tok_file = hf_hub_download(HF_TOK_REPO, 'tokenizer.json', token=HF_TOKEN, cache_dir=None) kz_tok = Tokenizer.from_file(tok_file) kz_model = KaizenWithLoRA() kz_model.load_base(DEFAULT_CKPT) kz_model.eval() kz_params = sum(p.numel() for p in kz_model.parameters()) print(f' KAIZEN: {kz_params/1e6:.1f}M params') kz_anc = embed_kaizen(kz_model, kz_tok, orig_qs) # [50, 512] kz_pos = embed_kaizen(kz_model, kz_tok, para_qs) # [50, 512] del kz_model # free model weights; tensors remain valid print(f' Embeddings: {list(kz_anc.shape)} (elapsed {time.time()-t0:.0f}s)') # ── GPT-2 117M embeddings ───────────────────────────────────────────────── print('\n[2/4] Loading GPT-2 117M and embedding probe questions...') gpt2_tok = GPT2Tokenizer.from_pretrained('gpt2', cache_dir='/tmp/gpt2_cache') gpt2_model = GPT2Model.from_pretrained('gpt2', cache_dir='/tmp/gpt2_cache') gpt2_model.eval() gpt2_params = sum(p.numel() for p in gpt2_model.parameters()) print(f' GPT-2: {gpt2_params/1e6:.1f}M params') g2_anc = embed_gpt2(gpt2_model, gpt2_tok, orig_qs) # [50, 768] g2_pos = embed_gpt2(gpt2_model, gpt2_tok, para_qs) # [50, 768] del gpt2_model # free model weights print(f' Embeddings: {list(g2_anc.shape)} (elapsed {time.time()-t0:.0f}s)') # ── Raw recall@1 (no head) ──────────────────────────────────────────────── kz_raw = recall_at_1_raw(kz_anc, kz_pos, held_out_idx) g2_raw = recall_at_1_raw(g2_anc, g2_pos, held_out_idx) print(f'\nRaw recall@1 (frozen embeddings, no head, cosine NN):') print(f' KAIZEN 42M : {kz_raw:.4f}') print(f' GPT-2 117M : {g2_raw:.4f}') print(f' Δ (117M−42M): {g2_raw - kz_raw:+.4f}') # ── Head training: build train tensors ──────────────────────────────────── kz_anc_tr = kz_anc[train_idx] # [35, 512] kz_pos_tr = kz_pos[train_idx] # [35, 512] g2_anc_tr = g2_anc[train_idx] # [35, 768] g2_pos_tr = g2_pos[train_idx] # [35, 768] # ── Train KAIZEN head ───────────────────────────────────────────────────── print(f'\n[3/4] Training LinearHead for KAIZEN 42M (512→{HEAD_DIM}, InfoNCE)...') kz_curve = train_with_checkpoints( d_in=512, anc_train=kz_anc_tr, pos_train=kz_pos_tr, anc_all=kz_anc, pos_all=kz_pos, held_out_idx=held_out_idx, ) print(f' Done (elapsed {time.time()-t0:.0f}s)') # ── Train GPT-2 head ────────────────────────────────────────────────────── print(f'[4/4] Training LinearHead for GPT-2 117M (768→{HEAD_DIM}, InfoNCE)...') g2_curve = train_with_checkpoints( d_in=768, anc_train=g2_anc_tr, pos_train=g2_pos_tr, anc_all=g2_anc, pos_all=g2_pos, held_out_idx=held_out_idx, ) print(f' Done (elapsed {time.time()-t0:.0f}s)') # ── Results table ───────────────────────────────────────────────────────── print() print('=' * 64) print('Recall@1 on held-out paraphrase pairs (linear head, InfoNCE, 35 train)') print(f'{"Epoch":>8s} {"KAIZEN 42M":>12s} {"GPT-2 117M":>12s} {"Δ(117M-42M)":>12s}') print('-' * 52) print(f'{"raw":>8s} {kz_raw:12.4f} {g2_raw:12.4f} {g2_raw-kz_raw:+12.4f}') for ep in CHECKPOINTS: kz_r = kz_curve.get(ep, float('nan')) g2_r = g2_curve.get(ep, float('nan')) delta = g2_r - kz_r if (kz_r == kz_r and g2_r == g2_r) else float('nan') print(f'{ep:8d} {kz_r:12.4f} {g2_r:12.4f} {delta:+12.4f}') kz_final = kz_curve.get(max(CHECKPOINTS), 0.0) g2_final = g2_curve.get(max(CHECKPOINTS), 0.0) print() print(f'Summary (epoch {max(CHECKPOINTS)}):') print(f' KAIZEN 42M raw={kz_raw:.4f} → head={kz_final:.4f} ' f'(gain={kz_final-kz_raw:+.4f})') print(f' GPT-2 117M raw={g2_raw:.4f} → head={g2_final:.4f} ' f'(gain={g2_final-g2_raw:+.4f})') print(f' Scale benefit at epoch 300: {g2_final - kz_final:+.4f} ' f'({"117M > 42M" if g2_final > kz_final else "42M >= 117M"})') print(f'\nRuntime: {time.time()-t0:.0f}s') if __name__ == '__main__': main()