File size: 8,183 Bytes
9c6799a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | """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 # per domain
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): # many random fillings of every template
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 # about as many sentences as ssv2
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()}) # task strings are the index
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) # [n_seq, SEQ]
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) # one document per row, starting at its first token
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) # 514 tokens
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) # Olsson et al.: loss at token 500 minus token 50
r["icl_score"] = float(ld[:, 500:550].mean() - ld[:, 50:100].mean()) # windows reduce per-token noise
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()
|