| """Post-hoc probes of trained models: per-domain loss, in-context copying, and the ICL score. |
| |
| Domains test whether gains concentrate on text about the physical world in iVideoGPT's robot videos: |
| ssv2 Something-Something-v2 action templates (hand and object motions like those in the robot videos) |
| ssv2_abs the same templates filled with abstract nouns (same style, no physical content) |
| abs_tmpl short non-physical templates filled with the same nouns (same style, no physical verbs) |
| droid DROID robot task instructions |
| piqa PIQA goals + correct solutions (physical commonsense) |
| fineweb held-out FineWeb-Edu (in-distribution control) |
| wiki WikiText-103 test (encyclopedic control) |
| gsm8k GSM8K questions + answers (math control) |
| Copying: loss on the second copy of a repeated random-token sequence (induction ability). |
| ICL score (Olsson et al. 2022): loss at token 500 minus loss at token 50 on FineWeb-Edu. |
| |
| Usage: python wp/probe.py --runs RUN [RUN ...] --out report/probe.json |
| """ |
|
|
| import argparse |
| import json |
| import random |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from datasets import load_dataset |
| from huggingface_hub import hf_hub_download, snapshot_download |
| from transformers import AutoTokenizer, LlamaForCausalLM |
|
|
| REPO = "Broyojo/world-pretraining-runs" |
| SEQ = 1024 |
| MAX_TOKENS = 120_000 |
|
|
|
|
| def domain_texts(): |
| random.seed(0) |
| cls = Path(hf_hub_download("HuggingFaceM4/something_something_v2", "classes.py", repo_type="dataset")).read_text() |
| templates = [l.strip().strip('",') for l in cls.splitlines() if l.strip().startswith('"') and l.strip().endswith('",')] |
| objects = ["a cup", "a pen", "a book", "a ball", "a phone", "a box", "a bottle", "a spoon", "a towel", "a remote", |
| "a mug", "a pencil", "a notebook", "an apple", "a banana", "a shoe", "a glass", "a plate", "a key", "a wallet", |
| "a candle", "a battery", "a coin", "a tissue", "a sock", "a lid", "a fork", "a brush", "a toy car", "a can", |
| "a jar", "a sponge", "a bowl", "a charger", "a marker", "a lighter", "a bag", "a hat", "a paper", "a stone"] |
| abstract = ["a theory", "an idea", "a budget", "a policy", "a contract", "a rumor", "a deadline", "a mortgage", |
| "a hypothesis", "a tax", "a debt", "a reputation", "a merger", "a doctrine", "a proposal", "a treaty", |
| "a statute", "a dividend", "a grievance", "a promise", "a salary", "a verdict", "a concept", "a motive", |
| "a schedule", "a pension", "a subsidy", "a loan", "a lawsuit", "a tariff", "a mandate", "an opinion", |
| "an argument", "an election", "an alliance", "an audit", "an estimate", "a belief", "a duty", "a right"] |
| abs_templates = ["Discussing something with a colleague", "Approving something before the meeting", |
| "Reviewing something for the committee", "Explaining something to the board", |
| "Questioning something during the debate", "Announcing something on Tuesday", |
| "Postponing something until next year", "Rejecting something after the vote", |
| "Comparing something with something", "Defending something in the report", |
| "Criticizing something in an editorial", "Renewing something for another term", |
| "Negotiating something with the union", "Revising something after the audit", |
| "Summarizing something for the director", "Considering something over the weekend", |
| "Supporting something at the conference", "Doubting something in private", |
| "Recommending something to the minister", "Mentioning something in the minutes"] |
|
|
| def fill(tmpl, nouns): |
| out = [] |
| for _ in range(60): |
| for t in tmpl: |
| s = t |
| while "something" in s: |
| s = s.replace("something", random.choice(nouns), 1) |
| out.append(s.replace("number of", str(random.randint(2, 5))) + ".") |
| return out |
|
|
| ssv2 = fill(templates, objects) |
| ssv2_abs = fill(templates, abstract) |
| abs_tmpl = fill(abs_templates, abstract) * 9 |
| import pandas as pd |
| tasks = pd.read_parquet(hf_hub_download("lerobot/droid_1.0.1", "meta/tasks.parquet", repo_type="dataset")) |
| droid = sorted({str(t).strip() for t in tasks.index if str(t).strip()}) |
| random.shuffle(droid) |
| droid = [t[0].upper() + t[1:] + ("" if t.endswith(".") else ".") for t in droid[:8000]] |
| piqa = load_dataset("baber/piqa", split="validation") |
| piqa = [f"{r['goal']} {r['sol1'] if r['label'] == 0 else r['sol2']}" for r in piqa] |
| fw = [r["text"] for _, r in zip(range(3000), load_dataset("HuggingFaceFW/fineweb-edu", "sample-10BT", split="train", streaming=True))] |
| wiki = [t for t in load_dataset("Salesforce/wikitext", "wikitext-103-raw-v1", split="test")["text"] if len(t) > 200] |
| gsm = [f"{r['question']} {r['answer']}" for r in load_dataset("openai/gsm8k", "main", split="test")] |
| return {"ssv2": ssv2, "ssv2_abs": ssv2_abs, "abs_tmpl": abs_tmpl, "droid": droid, "piqa": piqa, "fineweb": fw[:400], "wiki": wiki, "gsm8k": gsm, "_long": fw} |
|
|
|
|
| def pack(tok, texts): |
| ids = [] |
| for t in texts: |
| ids += [tok.eos_token_id] + tok(t)["input_ids"] |
| if len(ids) >= MAX_TOKENS: |
| break |
| ids = ids[: (len(ids) // (SEQ + 1)) * (SEQ + 1)] |
| return torch.tensor(ids).view(-1, SEQ + 1) |
|
|
|
|
| @torch.no_grad() |
| def token_losses(model, batch, device): |
| out = [] |
| for i in range(0, len(batch), 8): |
| x = batch[i:i + 8].to(device) |
| logits = model(input_ids=x[:, :-1]).logits.float() |
| out.append(torch.nn.functional.cross_entropy(logits.transpose(1, 2), x[:, 1:], reduction="none").cpu()) |
| return torch.cat(out) |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser() |
| p.add_argument("--runs", nargs="+", required=True) |
| p.add_argument("--out", default="report/probe.json") |
| p.add_argument("--cache", default="probe_cache") |
| args = p.parse_args() |
| device = "mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu") |
| tok = AutoTokenizer.from_pretrained("gpt2") |
| cache = Path(args.cache) |
| cache.mkdir(exist_ok=True) |
| data_f = cache / "domains.pt" |
| if not data_f.exists(): |
| texts = domain_texts() |
| long_docs = [ids[:SEQ + 1] for ids in (tok(t)["input_ids"] for t in texts.pop("_long")) if len(ids) > SEQ][:128] |
| data = {k: pack(tok, v) for k, v in texts.items()} |
| data["_long"] = torch.tensor(long_docs) |
| torch.save(data, data_f) |
| data = torch.load(data_f) |
| g = torch.Generator().manual_seed(0) |
| rand = torch.randint(1000, 40000, (64, 256), generator=g) |
| copy = torch.cat([torch.full((64, 1), tok.eos_token_id), rand, rand, rand[:, :1]], dim=1) |
| results = json.loads(Path(args.out).read_text()) if Path(args.out).exists() else {} |
| for run in args.runs: |
| if run in results: |
| continue |
| path = snapshot_download(REPO, allow_patterns=f"runs/{run}/*", local_dir=cache / "models") |
| model = LlamaForCausalLM.from_pretrained(f"{path}/runs/{run}", torch_dtype=torch.float32).to(device).eval() |
| r = {k: float(token_losses(model, v, device).mean()) for k, v in data.items() if not k.startswith("_")} |
| ld = token_losses(model, data["_long"], device) |
| r["icl_score"] = float(ld[:, 500:550].mean() - ld[:, 50:100].mean()) |
| cl = token_losses(model, copy, device) |
| r["copy_first"] = float(cl[:, 1:256].mean()) |
| r["copy_second"] = float(cl[:, 257:512].mean()) |
| results[run] = r |
| print(run, {k: round(v, 3) for k, v in r.items()}, flush=True) |
| Path(args.out).write_text(json.dumps(results, indent=1)) |
| del model |
| print({k: tuple(v.shape) for k, v in data.items()}) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
