Manuscript revision and pending experiment work, snapshot before restructuring
Clarity pass over the main text (36-item audit), Discussion rewrite and cut, acknowledgements, Souly et al. as ref 62, lettered SI panels, model section moved under Results; plus the untracked curriculum/society/compose/smol configs, runners, figures, stats and tests that the SI already cites. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Y64o8FKP7rCuXzC48pxpMm
This commit is contained in:
parent
e4804adabc
commit
84124de143
450 changed files with 52813 additions and 1202 deletions
348
src/llm/calibrate.py
Normal file
348
src/llm/calibrate.py
Normal file
|
|
@ -0,0 +1,348 @@
|
|||
"""Calibration gates for the v2 society (``kind: llm_society_calib``; prereg §4).
|
||||
|
||||
Four stages, each a config ``stage:``, each printing its gate table and writing the usual artifact
|
||||
triple. Nothing in the campaign is chosen by feel: the family set, the inheritance pool size, and
|
||||
the recombination operator are all fixed here, by measurement, before any campaign job is submitted.
|
||||
|
||||
* ``families`` (C1a–c, C4): base and specialist accuracy per candidate family; specialist confidence
|
||||
AUC (own- vs off-family prompts — the routing precondition); pairwise confidence-weighted functional
|
||||
conflict between specialists (the ``llm_epistasis`` measure; the grid's no-conflict axis sits at
|
||||
≈0.20–0.26 and its conflict axis at ≈0.46–0.52, so the gate is < 0.35); gen-0 behavioural distance.
|
||||
* ``transmission`` (C2): retention of a founder's own-family skill in a child distilled from the
|
||||
founder's *own answers*, as a function of examples-per-family ``k`` and epochs. Sets ``k_inherit``.
|
||||
* ``cross`` (C3): one two-founder cross, union-distil vs best-of-6 linear-merge-distil; both
|
||||
families' accuracy in each child. Tests the operator choice for twenty minutes, not eighty hours.
|
||||
* ``consensus`` (C5): consensus accuracy over the founders at gen 0 (must be low, else conformity
|
||||
is a truth proxy and the ``no_grounding`` arm cannot fail by the predicted mechanism).
|
||||
|
||||
Specialists are cached under ``models/llm/society_v2/calib_s{seed}/`` and shared across stages.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
from . import families as _families # noqa: F401
|
||||
from .society_ops import behavioural_distance, consensus_answers, conformity_scores, fitness_of, route_union
|
||||
from .tasks import Task, make_tasks, verify, _normalise
|
||||
|
||||
|
||||
def _auc(pos: np.ndarray, neg: np.ndarray) -> float:
|
||||
"""Rank AUC: P(confidence on own-family prompt > confidence on off-family prompt)."""
|
||||
if len(pos) == 0 or len(neg) == 0:
|
||||
return float("nan")
|
||||
return float(np.mean([(p > n) + 0.5 * (p == n) for p in pos for n in neg]))
|
||||
|
||||
|
||||
def _specialists(cfg, fams, root, base):
|
||||
from .specialise import train_specialist
|
||||
dirs = []
|
||||
n_tr, n_ep = int(cfg.get("spec_train", 600)), int(cfg.get("spec_epochs", 3))
|
||||
r_ = int(cfg.get("lora", {}).get("r", 16))
|
||||
for i, f in enumerate(fams):
|
||||
d = root / (f"spec_{f}_n{n_tr}e{n_ep}" + ("" if r_ == 16 else f"_r{r_}")) # keyed by budget + rank
|
||||
if not (d / "adapter_config.json").exists():
|
||||
train_specialist(base, f, str(d), n_train=int(cfg.get("spec_train", 600)),
|
||||
epochs=int(cfg.get("spec_epochs", 3)), seed=int(cfg["seed"]) * 100 + i,
|
||||
r=int(cfg.get("lora", {}).get("r", 16)),
|
||||
alpha=int(cfg.get("lora", {}).get("alpha", 32)))
|
||||
dirs.append(str(d))
|
||||
return dirs
|
||||
|
||||
|
||||
def _stage_families(cfg, base, fams, root) -> list[dict]:
|
||||
import torch
|
||||
from .epistasis import generate_with_confidence
|
||||
from .evaluate import generate, load_model
|
||||
from .merge import load_specialists
|
||||
|
||||
n_test, n_probe = int(cfg.get("n_test", 100)), int(cfg.get("n_probe", 10))
|
||||
spec_hi = float(cfg.get("spec_hi", 0.90)) # prereg §4 amendment 1: 1.0 in stage A2
|
||||
tests = {f: make_tasks(f, n_test, seed=1000 + i) for i, f in enumerate(fams)}
|
||||
probe = sum([make_tasks(f, n_probe, seed=5000 + i) for i, f in enumerate(fams)], [])
|
||||
p_prompts = [x.prompt for x in probe]
|
||||
p_fam = np.array([x.family for x in probe])
|
||||
rows = []
|
||||
|
||||
model, tok = load_model(base) # base alone
|
||||
base_acc = {f: fitness_of(generate(model, tok, [x.prompt for x in tests[f]]), tests[f], [f])[f]
|
||||
for f in fams}
|
||||
del model; torch.cuda.empty_cache()
|
||||
|
||||
dirs = _specialists(cfg, fams, root, base)
|
||||
model, tok = load_specialists(base, dirs)
|
||||
spec_acc, ans, conf = {}, {}, {}
|
||||
for i, f in enumerate(fams):
|
||||
model.set_adapter(f"a{i}")
|
||||
spec_acc[f] = fitness_of(generate(model, tok, [x.prompt for x in tests[f]]), tests[f], [f])[f]
|
||||
ans[f], conf[f] = generate_with_confidence(model, tok, p_prompts)
|
||||
del model; torch.cuda.empty_cache()
|
||||
|
||||
for f in fams:
|
||||
uniq = len({t.prompt for t in make_tasks(f, 600, seed=7)})
|
||||
auc = _auc(conf[f][p_fam == f], conf[f][p_fam != f])
|
||||
in_band = 0.05 <= base_acc[f] <= 0.40 and 0.60 <= spec_acc[f] <= spec_hi
|
||||
for m, v in (("base_acc", base_acc[f]), ("spec_acc", spec_acc[f]), ("conf_auc", auc),
|
||||
("unique_of_600", uniq), ("in_band", float(in_band))):
|
||||
rows.append({"stage": "families", "family": f, "other": "", "metric": m, "value": float(v)})
|
||||
norm = {f: [_normalise(a) for a in ans[f]] for f in fams}
|
||||
dist = behavioural_distance([ans[f] for f in fams])
|
||||
for i, fa in enumerate(fams):
|
||||
for j, fb in enumerate(fams):
|
||||
if j <= i:
|
||||
continue
|
||||
dis = np.array([a != b for a, b in zip(norm[fa], norm[fb])], dtype=float)
|
||||
epi = float(np.mean(conf[fa] * conf[fb] * dis))
|
||||
rows.append({"stage": "families", "family": fa, "other": fb, "metric": "epi_conf", "value": epi})
|
||||
rows.append({"stage": "families", "family": fa, "other": fb, "metric": "dist", "value": float(dist[i, j])})
|
||||
|
||||
print(f"\n== C1 family band (base ∈ [0.05,0.40], specialist ∈ [0.60,{spec_hi:.2f}]) ==")
|
||||
for f in fams:
|
||||
flag = "OK " if 0.05 <= base_acc[f] <= 0.40 and 0.60 <= spec_acc[f] <= spec_hi else "-- "
|
||||
print(f" {flag}{f:12s} base {base_acc[f]:.2f} spec {spec_acc[f]:.2f} "
|
||||
f"confAUC {_auc(conf[f][p_fam == f], conf[f][p_fam != f]):.2f} "
|
||||
f"unique/600 {len({t.prompt for t in make_tasks(f, 600, seed=7)})}")
|
||||
epis = [r["value"] for r in rows if r["metric"] == "epi_conf"]
|
||||
dists = [r["value"] for r in rows if r["metric"] == "dist"]
|
||||
print(f"== C1b pairwise epi_conf: max {max(epis):.3f} (gate < 0.35) "
|
||||
f"C1c min pairwise distance {min(dists):.2f} (gate ≥ 0.5)")
|
||||
return rows
|
||||
|
||||
|
||||
def _stage_transmission(cfg, base, fams, root) -> list[dict]:
|
||||
import torch
|
||||
from .evaluate import generate
|
||||
from .merge import load_specialists
|
||||
from .specialise import train_lora_on_tasks
|
||||
|
||||
probe_f = list(cfg.get("probe_families", fams[:3]))
|
||||
ks, eps = list(cfg.get("ks", [25, 50, 100, 150])), list(cfg.get("epochs_grid", [2, 3]))
|
||||
n_test = int(cfg.get("n_test", 100))
|
||||
dirs = _specialists(cfg, fams, root, base)
|
||||
rows = []
|
||||
for f in probe_f:
|
||||
i = fams.index(f)
|
||||
test = make_tasks(f, n_test, seed=1000 + i)
|
||||
model, tok = load_specialists(base, [dirs[i]])
|
||||
founder_acc = fitness_of(generate(model, tok, [x.prompt for x in test]), test, [f])[f]
|
||||
pools = {k: sum([make_tasks(g, k, seed=9000 + k * 31 + j) for j, g in enumerate(fams)], [])
|
||||
for k in ks}
|
||||
answers = {k: generate(model, tok, [x.prompt for x in pools[k]]) for k in ks}
|
||||
del model; torch.cuda.empty_cache()
|
||||
for k in ks:
|
||||
supplied = fitness_of(answers[k], pools[k], [f])[f]
|
||||
data = [Task(x.family, x.prompt, a.strip().split("\n")[0][:64]) for x, a in zip(pools[k], answers[k])
|
||||
if a.strip()]
|
||||
for ep in eps:
|
||||
d = root / "transmission" / f"{f}_k{k}_e{ep}"
|
||||
train_lora_on_tasks(base, data, str(d), epochs=ep, seed=int(cfg["seed"]) * 7 + k + ep)
|
||||
m, t2 = load_specialists(base, [str(d)])
|
||||
child_acc = fitness_of(generate(m, t2, [x.prompt for x in test]), test, [f])[f]
|
||||
del m; torch.cuda.empty_cache()
|
||||
ret = child_acc / founder_acc if founder_acc > 0 else float("nan")
|
||||
for name, v in (("founder_acc", founder_acc), ("supplied_acc", supplied),
|
||||
("child_acc", child_acc), ("retention", ret)):
|
||||
rows.append({"stage": "transmission", "family": f, "other": f"k{k}_e{ep}",
|
||||
"metric": name, "value": float(v), "k": k, "epochs": ep})
|
||||
print(f" {f:12s} k={k:3d} ep={ep}: founder {founder_acc:.2f} supplied {supplied:.2f} "
|
||||
f"child {child_acc:.2f} retention {ret:.2f}", flush=True)
|
||||
df = pd.DataFrame(rows)
|
||||
ret = df[df.metric == "retention"].groupby(["k", "epochs"]).value.mean().unstack()
|
||||
print("\n== C2 mean retention (gate: choose k* = min k with retention ≥ 0.85) ==")
|
||||
print(ret.round(2).to_string())
|
||||
return rows
|
||||
|
||||
|
||||
def _stage_transmission_conf(cfg, base, fams, root) -> list[dict]:
|
||||
"""C2b — confidence-gated inheritance: retention when the child learns only the prompts its source
|
||||
is confident on (verifier-free). Also reports the τ that separates own- from off-family confidence
|
||||
(Youden's J on the founder's own confidences, using family labels for calibration only) and the
|
||||
off-family harm of ungated inheritance (child off-family accuracy vs base)."""
|
||||
import torch
|
||||
from .epistasis import generate_with_confidence
|
||||
from .evaluate import generate, load_model
|
||||
from .merge import load_specialists
|
||||
from .specialise import train_lora_on_tasks
|
||||
|
||||
probe_f = list(cfg.get("probe_families", fams[:3]))
|
||||
k, ep = int(cfg.get("k_pool", 300)), int(cfg.get("epochs", 3))
|
||||
taus = list(cfg.get("taus", [0.5, 0.7, 0.85]))
|
||||
n_test = int(cfg.get("n_test", 100))
|
||||
dirs = _specialists(cfg, fams, root, base)
|
||||
tests = {f: make_tasks(f, n_test, seed=1000 + i) for i, f in enumerate(fams)}
|
||||
pool = sum([make_tasks(g, k, seed=9300 + j) for j, g in enumerate(fams)], [])
|
||||
p_fam = np.array([x.family for x in pool])
|
||||
|
||||
model, tok = load_model(base)
|
||||
base_acc = {f: fitness_of(generate(model, tok, [x.prompt for x in tests[f]]), tests[f], [f])[f] for f in fams}
|
||||
del model; torch.cuda.empty_cache()
|
||||
|
||||
rows = []
|
||||
for f in probe_f:
|
||||
i = fams.index(f)
|
||||
model, tok = load_specialists(base, [dirs[i]])
|
||||
founder_acc = fitness_of(generate(model, tok, [x.prompt for x in tests[f]]), tests[f], [f])[f]
|
||||
ans, conf = generate_with_confidence(model, tok, [x.prompt for x in pool])
|
||||
del model; torch.cuda.empty_cache()
|
||||
own, off = conf[p_fam == f], conf[p_fam != f]
|
||||
# Youden-optimal τ over a grid — calibration only (uses family labels)
|
||||
grid = np.linspace(0.05, 0.99, 95)
|
||||
J = [np.mean(own >= t) - np.mean(off >= t) for t in grid]
|
||||
tau_star = float(grid[int(np.argmax(J))])
|
||||
rows.append({"stage": "transmission_conf", "family": f, "other": "", "metric": "tau_youden", "value": tau_star})
|
||||
rows.append({"stage": "transmission_conf", "family": f, "other": "", "metric": "own_conf_median", "value": float(np.median(own))})
|
||||
rows.append({"stage": "transmission_conf", "family": f, "other": "", "metric": "off_conf_median", "value": float(np.median(off))})
|
||||
print(f" {f:12s} founder {founder_acc:.2f} own-conf median {np.median(own):.2f} off-conf median "
|
||||
f"{np.median(off):.2f} τ* {tau_star:.2f} (keeps {np.mean(own >= tau_star):.0%} own, "
|
||||
f"{np.mean(off >= tau_star):.0%} off)", flush=True)
|
||||
for tau in [None] + taus:
|
||||
keep = np.ones(len(pool), bool) if tau is None else conf >= tau
|
||||
data = [Task(x.family, x.prompt, a.strip().split("\n")[0][:64])
|
||||
for x, a, kp in zip(pool, ans, keep) if kp and a.strip()]
|
||||
if len(data) < 8:
|
||||
print(f" {f:12s} τ={tau}: only {len(data)} prompts kept — skipped"); continue
|
||||
d = root / "transmission_conf" / f"{f}_tau{tau}"
|
||||
train_lora_on_tasks(base, data, str(d), epochs=ep, seed=int(cfg["seed"]) * 13 + int((tau or 0) * 100))
|
||||
m, t2 = load_specialists(base, [str(d)])
|
||||
child_own = fitness_of(generate(m, t2, [x.prompt for x in tests[f]]), tests[f], [f])[f]
|
||||
off_f = [g for g in fams if g != f][:3] # off-family harm on three other families
|
||||
child_off = float(np.mean([fitness_of(generate(m, t2, [x.prompt for x in tests[g]]), tests[g], [g])[g] for g in off_f]))
|
||||
base_off = float(np.mean([base_acc[g] for g in off_f]))
|
||||
del m; torch.cuda.empty_cache()
|
||||
ret = child_own / founder_acc if founder_acc > 0 else float("nan")
|
||||
lab = "none" if tau is None else f"{tau:.2f}"
|
||||
for name, v in (("retention", ret), ("child_own_acc", child_own), ("n_kept", len(data)),
|
||||
("child_off_acc", child_off), ("base_off_acc", base_off)):
|
||||
rows.append({"stage": "transmission_conf", "family": f, "other": f"tau{lab}", "metric": name, "value": float(v)})
|
||||
print(f" {f:12s} τ={lab}: kept {len(data):4d}/{len(pool)} child own {child_own:.2f} "
|
||||
f"retention {ret:.2f} off-family child {child_off:.2f} vs base {base_off:.2f}", flush=True)
|
||||
df = pd.DataFrame(rows)
|
||||
r = df[df.metric == "retention"].groupby("other").value.mean()
|
||||
print("\n== C2b mean retention by gate (target ≥ 0.85) ==\n" + r.round(2).to_string())
|
||||
return rows
|
||||
|
||||
|
||||
def _stage_cross(cfg, base, fams, root) -> list[dict]:
|
||||
import torch
|
||||
from .directed import sample_merge_weights
|
||||
from .epistasis import generate_with_confidence
|
||||
from .evaluate import generate
|
||||
from .merge import load_specialists
|
||||
from .specialise import train_lora_on_tasks
|
||||
|
||||
fa, fb = cfg["cross"]
|
||||
k, ep = int(cfg.get("k_inherit", 100)), int(cfg.get("epochs", 3))
|
||||
n_cand = int(cfg.get("n_candidates", 6))
|
||||
ia, ib = fams.index(fa), fams.index(fb)
|
||||
tests = {f: make_tasks(f, int(cfg.get("n_test", 100)), seed=1000 + fams.index(f)) for f in (fa, fb)}
|
||||
val = sum([make_tasks(g, 10, seed=3000 + j) for j, g in enumerate(fams)], [])
|
||||
inherit = sum([make_tasks(g, k, seed=9100 + j) for j, g in enumerate(fams)], [])
|
||||
inh_p = [x.prompt for x in inherit]
|
||||
dirs = _specialists(cfg, fams, root, base)
|
||||
rows = []
|
||||
|
||||
model, tok = load_specialists(base, [dirs[ia], dirs[ib]])
|
||||
parent_acc = {}
|
||||
for j, f in enumerate((fa, fb)):
|
||||
model.set_adapter(f"a{j}")
|
||||
for g in (fa, fb):
|
||||
parent_acc[(f, g)] = fitness_of(generate(model, tok, [x.prompt for x in tests[g]]), tests[g], [g])[g]
|
||||
model.set_adapter("a0"); ans_a, cf_a = generate_with_confidence(model, tok, inh_p)
|
||||
model.set_adapter("a1"); ans_b, cf_b = generate_with_confidence(model, tok, inh_p)
|
||||
union_ans, src = route_union(ans_a, cf_a, ans_b, cf_b)
|
||||
|
||||
rng = np.random.default_rng(int(cfg["seed"]))
|
||||
w = sample_merge_weights(2, n_cand, rng)
|
||||
best_i, best_v = 0, -np.inf
|
||||
for ci in range(n_cand):
|
||||
model.add_weighted_adapter(["a0", "a1"], w[ci].tolist(), f"c{ci}", combination_type="linear")
|
||||
model.set_adapter(f"c{ci}")
|
||||
v = fitness_of(generate(model, tok, [x.prompt for x in val]), val, fams)["overall"]
|
||||
if v > best_v:
|
||||
best_i, best_v = ci, v
|
||||
model.set_adapter("a0"); model.delete_adapter(f"c{ci}")
|
||||
model.add_weighted_adapter(["a0", "a1"], w[best_i].tolist(), "win", combination_type="linear")
|
||||
model.set_adapter("win"); linear_ans, cf_lin = generate_with_confidence(model, tok, inh_p)
|
||||
del model; torch.cuda.empty_cache()
|
||||
|
||||
# optional confidence gate (C2b amendment): each child keeps only prompts its own source is
|
||||
# confident on — the union by max(parent confidences), the linear merge by its own confidence.
|
||||
gate = cfg.get("conf_gate")
|
||||
gate = None if gate is None else float(gate)
|
||||
keep_u = np.ones(len(inherit), bool) if gate is None else np.maximum(cf_a, cf_b) >= gate
|
||||
keep_l = np.ones(len(inherit), bool) if gate is None else cf_lin >= gate
|
||||
|
||||
out = {}
|
||||
for op, answers, keep in (("union", union_ans, keep_u), ("linear", linear_ans, keep_l)):
|
||||
supplied = fitness_of(answers, inherit, [fa, fb])
|
||||
data = [Task(x.family, x.prompt, a.strip().split("\n")[0][:64])
|
||||
for x, a, kp in zip(inherit, answers, keep) if kp and a.strip()]
|
||||
rows.append({"stage": "cross", "family": "", "other": op, "metric": "n_kept", "value": float(len(data))})
|
||||
r_, a_ = int(cfg.get("lora", {}).get("r", 16)), int(cfg.get("lora", {}).get("alpha", 32))
|
||||
d = root / "cross" / f"{fa}_{fb}_{op}_g{gate}_e{ep}_r{r_}"
|
||||
train_lora_on_tasks(base, data, str(d), epochs=ep, seed=int(cfg["seed"]) * 11, r=r_, alpha=a_)
|
||||
m, t2 = load_specialists(base, [str(d)])
|
||||
for g in (fa, fb):
|
||||
acc = fitness_of(generate(m, t2, [x.prompt for x in tests[g]]), tests[g], [g])[g]
|
||||
out[(op, g)] = acc
|
||||
rows.append({"stage": "cross", "family": g, "other": op, "metric": "child_acc", "value": float(acc)})
|
||||
rows.append({"stage": "cross", "family": g, "other": op, "metric": "supplied_acc", "value": float(supplied[g])})
|
||||
del m; torch.cuda.empty_cache()
|
||||
for f in (fa, fb):
|
||||
rows.append({"stage": "cross", "family": f, "other": "parent", "metric": "parent_own_acc",
|
||||
"value": float(parent_acc[(f, f)])})
|
||||
rows.append({"stage": "cross", "family": fa, "other": fb, "metric": "union_share_b", "value": float(src.mean())})
|
||||
|
||||
print(f"\n== C3 cross {fa} × {fb} (gate: union ≥ 0.85×parent on each; union ≥ linear on the min) ==")
|
||||
for g in (fa, fb):
|
||||
print(f" {g:12s} parent {parent_acc[(g, g)]:.2f} union child {out[('union', g)]:.2f} "
|
||||
f"linear child {out[('linear', g)]:.2f}")
|
||||
print(f" union routed {src.mean():.2f} of prompts to {fb}")
|
||||
return rows
|
||||
|
||||
|
||||
def _stage_consensus(cfg, base, fams, root) -> list[dict]:
|
||||
import torch
|
||||
from .evaluate import generate
|
||||
from .merge import load_specialists
|
||||
|
||||
probe = sum([make_tasks(f, int(cfg.get("n_probe", 10)), seed=5000 + i) for i, f in enumerate(fams)], [])
|
||||
dirs = _specialists(cfg, fams, root, base)
|
||||
model, tok = load_specialists(base, dirs)
|
||||
outs = []
|
||||
for i in range(len(fams)):
|
||||
model.set_adapter(f"a{i}"); outs.append(generate(model, tok, [x.prompt for x in probe]))
|
||||
del model; torch.cuda.empty_cache()
|
||||
cons = consensus_answers(outs)
|
||||
cons_acc = float(np.mean([verify(c, x) for c, x in zip(cons, probe)]))
|
||||
conf = conformity_scores(outs, cons)
|
||||
dist = behavioural_distance(outs)
|
||||
fit = [fitness_of(o, probe, fams)["overall"] for o in outs]
|
||||
rows = [{"stage": "consensus", "family": "", "other": "", "metric": "consensus_acc", "value": cons_acc},
|
||||
{"stage": "consensus", "family": "", "other": "", "metric": "min_pair_dist",
|
||||
"value": float(dist[np.triu_indices(len(fams), 1)].min())}]
|
||||
for i, f in enumerate(fams):
|
||||
rows.append({"stage": "consensus", "family": f, "other": "", "metric": "conformity", "value": float(conf[i])})
|
||||
rows.append({"stage": "consensus", "family": f, "other": "", "metric": "probe_acc", "value": float(fit[i])})
|
||||
print(f"\n== C5 consensus accuracy at gen 0: {cons_acc:.2f} (gate < 0.35) "
|
||||
f"corr(conformity, accuracy) = {np.corrcoef(conf, fit)[0, 1]:+.2f} "
|
||||
f"min pairwise distance {rows[1]['value']:.2f}")
|
||||
return rows
|
||||
|
||||
|
||||
def run_calibration(cfg: dict) -> pd.DataFrame:
|
||||
"""Dispatch on ``stage`` and return the stage's rows as a DataFrame."""
|
||||
base = cfg["base_model"]
|
||||
fams = list(cfg["families"])
|
||||
root = Path(cfg.get("adapters_dir", "models/llm")) / "society_v2" / f"calib_s{int(cfg['seed'])}"
|
||||
stage = cfg["stage"]
|
||||
fn = {"families": _stage_families, "transmission": _stage_transmission,
|
||||
"transmission_conf": _stage_transmission_conf,
|
||||
"cross": _stage_cross, "consensus": _stage_consensus}[stage]
|
||||
rows = fn(cfg, base, fams, root)
|
||||
df = pd.DataFrame(rows)
|
||||
df["experiment"] = cfg["experiment"]
|
||||
return df
|
||||
348
src/llm/compose.py
Normal file
348
src/llm/compose.py
Normal file
|
|
@ -0,0 +1,348 @@
|
|||
"""`llm_compose` — does a composed capability survive inheritance? (prereg v3)
|
||||
|
||||
Two single-skill LoRA lineages on a shared frozen base — **math** and **code**. Each generation both
|
||||
lineages reproduce by self-consumption (a fresh LoRA distilled from their own confidence-gated answers
|
||||
on fresh prompts, optionally mixed with a fraction ``g`` of verified real examples — E2's immigration
|
||||
in the training mix). Each generation the *current* two parents are merged and evaluated on the
|
||||
held-out composed target, GSM8k-Hard, program-aided with an execution verifier.
|
||||
|
||||
The composed model is a **measurement, not a lineage**: it is re-formed each generation from whatever
|
||||
the parents currently are, which separates "does composition survive parental drift?" from "does the
|
||||
composed model itself drift?". The ``composed`` arm makes it a lineage as well.
|
||||
|
||||
Measured every generation (§1.6): ``q_math`` (GSM8K), ``q_code`` (MBPP, execution-verified),
|
||||
``rho`` (behavioural agreement on a shared probe + LoRA-delta cosine), composed accuracy for the
|
||||
merge *and each parent alone*, hence the **surplus** (merge − best parent) and the
|
||||
**union-exceedance** (composed-solved items neither parent solves — the super-linear signature).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
from .compose_data import (ProgTask, code_train, composed_target, gsm8k_probe, gsm_hard,
|
||||
math_train, mbpp_probe, probe_pool)
|
||||
from .execute import extract_code, numeric_match, run_solution, verify_program
|
||||
from .tasks import Task, _normalise
|
||||
|
||||
LINEAGES = ("math", "code")
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ scoring helpers (pure-ish)
|
||||
|
||||
|
||||
def score_composed(completions: list[str], tasks: list[ProgTask]) -> tuple[float, np.ndarray, float]:
|
||||
"""Program-aided accuracy over the composed target.
|
||||
|
||||
Returns:
|
||||
(accuracy, per-item correctness, executable-rate) — the last is the fraction whose code ran
|
||||
at all, which is the code sub-skill isolated from the maths.
|
||||
"""
|
||||
ok, ran = [], []
|
||||
for c, t in zip(completions, tasks):
|
||||
good, res = verify_program(c, t.target)
|
||||
ok.append(bool(good)); ran.append(bool(res.ok))
|
||||
return float(np.mean(ok)), np.array(ok, dtype=bool), float(np.mean(ran))
|
||||
|
||||
|
||||
_NUM = __import__("re").compile(r"-?\d[\d,]*\.?\d*")
|
||||
|
||||
|
||||
def _last_number(text: str) -> float | None:
|
||||
"""The last number in a completion, commas stripped — the conventional GSM8K readout."""
|
||||
m = _NUM.findall(text)
|
||||
for tok in reversed(m):
|
||||
try:
|
||||
return float(tok.replace(",", "").rstrip("."))
|
||||
except ValueError:
|
||||
continue
|
||||
return None
|
||||
|
||||
|
||||
def score_gsm8k(completions: list[str], tasks: list[Task]) -> float:
|
||||
"""Math own-skill: the number after '####' if the model used it, else the last number.
|
||||
|
||||
Graded numerically rather than by string equality, so "18" and "18.00" both count and a model
|
||||
that reasons past its answer is not punished for formatting.
|
||||
"""
|
||||
hit = []
|
||||
for c, t in zip(completions, tasks):
|
||||
tail = c.split("####")[-1] if "####" in c else c
|
||||
got, want = _last_number(tail), _last_number(t.answer)
|
||||
hit.append(got is not None and want is not None and numeric_match(got, want))
|
||||
return float(np.mean(hit))
|
||||
|
||||
|
||||
def score_mbpp(completions: list[str], items: list[dict], timeout_s: float = 6.0) -> float:
|
||||
"""Code own-skill: the emitted function must pass MBPP's reference asserts."""
|
||||
hit = []
|
||||
for c, it in zip(completions, items):
|
||||
prog = "\n".join(list(it["imports"]) + [extract_code(c)] + list(it["tests"])
|
||||
+ ["def solution():\n return 1"])
|
||||
hit.append(run_solution(prog, timeout_s=timeout_s).ok)
|
||||
return float(np.mean(hit))
|
||||
|
||||
|
||||
def union_exceedance(merged_ok: np.ndarray, parent_ok: list[np.ndarray]) -> float:
|
||||
"""Fraction of all items that the merge solves and *no* parent solves (super-linear signature)."""
|
||||
any_parent = np.zeros_like(merged_ok)
|
||||
for p in parent_ok:
|
||||
any_parent |= p
|
||||
return float(np.mean(merged_ok & ~any_parent))
|
||||
|
||||
|
||||
def mean_logprob(model, tok, prompts: list[str], completions: list[str], *,
|
||||
batch_size: int = 8, device: str = "cuda") -> np.ndarray:
|
||||
"""Mean token log-probability of each completion under the model, teacher-forced.
|
||||
|
||||
``generate(..., output_scores=True)`` keeps one (batch × vocab) tensor per step, which at 320 new
|
||||
tokens is gigabytes; one extra forward pass over prompt+completion costs a fraction of that and
|
||||
gives the same quantity. Returns exp(mean logprob) ∈ (0, 1] — the same self-certainty scale the
|
||||
epistasis module uses.
|
||||
"""
|
||||
import torch
|
||||
|
||||
out = np.zeros(len(prompts))
|
||||
for i in range(0, len(prompts), batch_size):
|
||||
p_chunk, c_chunk = prompts[i:i + batch_size], completions[i:i + batch_size]
|
||||
ids, labels = [], []
|
||||
from .evaluate import format_prompt
|
||||
for p, c in zip(p_chunk, c_chunk):
|
||||
pi = tok(format_prompt(tok, p), add_special_tokens=False).input_ids
|
||||
ci = tok(c if c.strip() else " ", add_special_tokens=False).input_ids[:512]
|
||||
ids.append(pi + ci); labels.append([-100] * len(pi) + ci)
|
||||
L = max(len(x) for x in ids)
|
||||
pad = tok.pad_token_id
|
||||
inp = torch.full((len(ids), L), pad, dtype=torch.long)
|
||||
lab = torch.full((len(ids), L), -100, dtype=torch.long)
|
||||
att = torch.zeros((len(ids), L), dtype=torch.long)
|
||||
for b, (x, y) in enumerate(zip(ids, labels)): # right-pad (scoring, not generation)
|
||||
inp[b, :len(x)] = torch.tensor(x); lab[b, :len(y)] = torch.tensor(y)
|
||||
att[b, :len(x)] = 1
|
||||
with torch.no_grad():
|
||||
logits = model(input_ids=inp.to(device), attention_mask=att.to(device)).logits[:, :-1]
|
||||
tgt = lab[:, 1:].to(device)
|
||||
mask = tgt != -100
|
||||
# cross_entropy streams the log-partition internally: no float32 copy of the logits and
|
||||
# no second tensor for log_softmax, which together were ~8 GB at batch 8 x 800 x 152k.
|
||||
nll = torch.nn.functional.cross_entropy(
|
||||
logits.reshape(-1, logits.size(-1)).float(),
|
||||
tgt.reshape(-1).clamp(min=0), reduction="none").view(tgt.shape)
|
||||
n = mask.sum(-1).clamp(min=1)
|
||||
out[i:i + len(p_chunk)] = torch.exp(-(nll * mask).sum(-1) / n).cpu().numpy()
|
||||
del logits, nll
|
||||
return out
|
||||
|
||||
|
||||
def predicted_composition(q_math: np.ndarray, q_code: np.ndarray, rho: np.ndarray,
|
||||
observed0: float) -> np.ndarray:
|
||||
"""The framework's forecast Ĉ_t (prereg §2), one free scale fixed at generation 0.
|
||||
|
||||
Ĉ_t = c0 · q_math_t · q_code_t · (1 − rho_t)/(1 − rho_0), c0 chosen so Ĉ_0 = observed_0.
|
||||
"""
|
||||
q_math, q_code, rho = map(np.asarray, (q_math, q_code, rho))
|
||||
shape = q_math * q_code * (1.0 - rho) / max(1e-9, 1.0 - rho[0])
|
||||
c0 = observed0 / shape[0] if shape[0] > 1e-9 else 0.0
|
||||
return c0 * shape
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ the GPU loop
|
||||
|
||||
|
||||
def _state_path(root: Path, arm: str) -> Path:
|
||||
return root / arm / "state.json"
|
||||
|
||||
|
||||
def run_compose(cfg: dict) -> pd.DataFrame:
|
||||
"""Run every configured arm; return tidy long-form rows (resumable per arm)."""
|
||||
import torch
|
||||
|
||||
from .epistasis import delta_geometry
|
||||
from .evaluate import generate
|
||||
from .merge import load_specialists
|
||||
from .specialise import train_lora_on_tasks
|
||||
|
||||
name, base, seed = cfg["experiment"], cfg["base_model"], int(cfg["seed"])
|
||||
G = int(cfg.get("generations", 6))
|
||||
arms = list(cfg.get("arms", ["dry", "grounded", "dry_linear"]))
|
||||
n_hard = int(cfg.get("n_hard", 200))
|
||||
n_hard_val = int(cfg.get("n_hard_val", 0)) # >0 enables directed merge-weight selection
|
||||
n_gsm = int(cfg.get("n_gsm8k", 150))
|
||||
n_mbpp = int(cfg.get("n_mbpp", 100))
|
||||
n_probe = int(cfg.get("n_probe", 60))
|
||||
k_inh = int(cfg.get("k_inherit", 300))
|
||||
epochs = int(cfg.get("epochs", 3))
|
||||
conf_gate = cfg.get("conf_gate")
|
||||
conf_gate = None if conf_gate is None else float(conf_gate)
|
||||
spec_train = int(cfg.get("spec_train", 1200))
|
||||
spec_epochs = int(cfg.get("spec_epochs", 3))
|
||||
max_new = int(cfg.get("max_new_tokens", 320))
|
||||
tb = int(cfg.get("train_batch_size", 2)) # the causal-LM loss upcasts logits to fp32:
|
||||
tlen = int(cfg.get("train_max_len", 448)) # memory ~ batch * len * vocab(152k) * 4 B
|
||||
r, alpha = (int(cfg.get("lora", {}).get(k, v)) for k, v in (("r", 16), ("alpha", 32)))
|
||||
root = Path(cfg.get("adapters_dir", "models/llm")) / "compose" / f"{name}_s{seed}"
|
||||
out_dir = Path(cfg.get("output", {}).get("dir", f"results/{name}"))
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# fixed evaluation sets (identical across arms, generations and seeds-within-a-config)
|
||||
target_kind = str(cfg.get("target", "gsm_hard"))
|
||||
hard = composed_target(target_kind, n_hard, 1000)
|
||||
hard_prompts = [t.prompt for t in hard]
|
||||
# Directed recombination (E10): candidate merge weights are screened on a DISJOINT validation
|
||||
# split and only the winner is reported on the test split, so nothing is selected on the numbers
|
||||
# we report. n_hard_val = 0 keeps the fixed 0.5/0.5 blend of the original design.
|
||||
hard_val = composed_target(target_kind, n_hard_val, 5000, split="val") if n_hard_val else []
|
||||
val_prompts = [t.prompt for t in hard_val]
|
||||
weight_grid = [list(w) for w in cfg.get("merge_weights", [[0.5, 0.5]])]
|
||||
gsm = gsm8k_probe(n_gsm, 2000)
|
||||
mbpp = mbpp_probe(n_mbpp, 3000)
|
||||
probes = probe_pool(n_probe, 4000)
|
||||
|
||||
def gen(model, tok, prompts, **kw):
|
||||
return generate(model, tok, prompts, max_new_tokens=max_new,
|
||||
batch_size=int(cfg.get("batch_size", 16)), **kw)
|
||||
|
||||
# ---- founders: one specialist per lineage, cached and shared across arms
|
||||
founders = {}
|
||||
for lin in LINEAGES:
|
||||
d = root / "founders" / lin
|
||||
if not (d / "adapter_config.json").exists():
|
||||
data = (math_train(spec_train, seed * 100 + 1, with_answers=True) if lin == "math"
|
||||
else code_train(spec_train, seed * 100 + 2, with_answers=True))
|
||||
train_lora_on_tasks(base, data, str(d), epochs=spec_epochs, seed=seed * 100 + 7,
|
||||
r=r, alpha=alpha, batch_size=tb, max_len=tlen)
|
||||
founders[lin] = str(d)
|
||||
|
||||
frames = []
|
||||
for a_idx, arm in enumerate(arms):
|
||||
# Operator per arm, explicit: the gen-0 sweep showed the blend ratio dominates the operator,
|
||||
# so which operator an arm uses is a stated choice, not an inference from its name.
|
||||
op = str(cfg.get("arm_ops", {}).get(arm, "linear" if arm.endswith("_linear") else "cat"))
|
||||
g_frac = float(cfg.get("g", 0.10)) if arm.startswith("grounded") else 0.0
|
||||
rows: list[dict] = []
|
||||
parents, t0 = dict(founders), 0
|
||||
st_path = _state_path(root, arm)
|
||||
partial = out_dir / f"partial_{arm}_s{seed}.parquet"
|
||||
if bool(cfg.get("resume", True)) and st_path.exists():
|
||||
st = json.loads(st_path.read_text())
|
||||
if all(Path(p, "adapter_config.json").exists() for p in st["parents"].values()):
|
||||
parents, t0 = dict(st["parents"]), int(st["generation"])
|
||||
if partial.exists():
|
||||
rows = pd.read_parquet(partial).to_dict("records")
|
||||
print(f"[{arm}] resuming at generation {t0}", flush=True)
|
||||
|
||||
for t in range(t0, G + 1):
|
||||
rng = np.random.default_rng([seed, a_idx, t])
|
||||
tag = {"experiment": name, "arm": arm, "seed": seed, "generation": t, "operator": op}
|
||||
|
||||
model, tok = load_specialists(base, [parents["math"], parents["code"]])
|
||||
|
||||
# ---- own-skill retention q_t, and each parent alone on the composed target
|
||||
model.set_adapter("a0")
|
||||
q_math = score_gsm8k(gen(model, tok, [x.prompt for x in gsm]), gsm)
|
||||
math_comp = gen(model, tok, hard_prompts)
|
||||
_, math_ok, math_exec = score_composed(math_comp, hard)
|
||||
probe_math = gen(model, tok, probes)
|
||||
model.set_adapter("a1")
|
||||
q_code = score_mbpp(gen(model, tok, [x["prompt"] for x in mbpp]), mbpp)
|
||||
code_comp = gen(model, tok, hard_prompts)
|
||||
_, code_ok, code_exec = score_composed(code_comp, hard)
|
||||
probe_code = gen(model, tok, probes)
|
||||
|
||||
rho_behav = float(np.mean([_normalise(a) == _normalise(b)
|
||||
for a, b in zip(probe_math, probe_code)]))
|
||||
rho_geom = float(delta_geometry(parents["math"], parents["code"])["delta_cos"])
|
||||
|
||||
# ---- merge and measure the composition (optionally selecting the weights on val)
|
||||
chosen = weight_grid[0]
|
||||
if hard_val and len(weight_grid) > 1:
|
||||
best_v = -np.inf
|
||||
for wi, w in enumerate(weight_grid):
|
||||
cn = f"cand{wi}"
|
||||
model.add_weighted_adapter(["a0", "a1"], w, cn, combination_type=op)
|
||||
model.set_adapter(cn)
|
||||
v, _, _ = score_composed(gen(model, tok, val_prompts), hard_val)
|
||||
if v > best_v:
|
||||
best_v, chosen = v, w
|
||||
model.set_adapter("a0"); model.delete_adapter(cn)
|
||||
rows.append({**tag, "metric": "chosen_weight_math", "value": float(chosen[0])})
|
||||
model.add_weighted_adapter(["a0", "a1"], chosen, "merged", combination_type=op)
|
||||
model.set_adapter("merged")
|
||||
merged_comp = gen(model, tok, hard_prompts)
|
||||
comp_acc, merged_ok, merged_exec = score_composed(merged_comp, hard)
|
||||
model.set_adapter("a0"); model.delete_adapter("merged")
|
||||
|
||||
best_parent = max(float(math_ok.mean()), float(code_ok.mean()))
|
||||
surplus = comp_acc - best_parent
|
||||
uex = union_exceedance(merged_ok, [math_ok, code_ok])
|
||||
for metric, value in (
|
||||
("q_math", q_math), ("q_code", q_code),
|
||||
("rho_behav", rho_behav), ("rho_geom", rho_geom),
|
||||
("composed_acc", comp_acc), ("composed_exec", merged_exec),
|
||||
("parent_math_composed", float(math_ok.mean())),
|
||||
("parent_code_composed", float(code_ok.mean())),
|
||||
("parent_math_exec", math_exec), ("parent_code_exec", code_exec),
|
||||
("best_parent_composed", best_parent), ("surplus", surplus),
|
||||
("union_exceedance", uex),
|
||||
):
|
||||
rows.append({**tag, "metric": metric, "value": float(value)})
|
||||
print(f"[{arm}] gen {t}/{G}: composed {comp_acc:.3f} (best parent {best_parent:.3f}, "
|
||||
f"surplus {surplus:+.3f}, uex {uex:.3f}) q_math {q_math:.2f} q_code {q_code:.2f} "
|
||||
f"rho {rho_behav:.2f}", flush=True)
|
||||
|
||||
if t == G:
|
||||
del model; torch.cuda.empty_cache()
|
||||
break
|
||||
|
||||
# ---- reproduce: each lineage distils from its own gated answers (+ g real examples)
|
||||
new_parents = {}
|
||||
for li, lin in enumerate(LINEAGES):
|
||||
pool = (math_train(k_inh, seed * 9001 + t * 31 + li, with_answers=False)
|
||||
if lin == "math" else
|
||||
code_train(k_inh, seed * 9001 + t * 31 + li, with_answers=False))
|
||||
model.set_adapter(f"a{li}")
|
||||
pool_prompts = [x.prompt for x in pool]
|
||||
answers = gen(model, tok, pool_prompts)
|
||||
if conf_gate is None:
|
||||
keep = np.ones(len(pool), bool)
|
||||
else:
|
||||
conf = mean_logprob(model, tok, pool_prompts, answers,
|
||||
batch_size=int(cfg.get("score_batch_size", 8)))
|
||||
keep = conf >= conf_gate
|
||||
data = [Task(lin, x.prompt, a.strip()[:512])
|
||||
for x, a, kp in zip(pool, answers, keep) if kp and a.strip()]
|
||||
n_self = len(data)
|
||||
if g_frac > 0: # immigration: verified real examples, fresh each gen
|
||||
n_real = max(1, int(round(g_frac / (1 - g_frac) * n_self)))
|
||||
real = (math_train(n_real, seed * 7717 + t * 13 + li, with_answers=True)
|
||||
if lin == "math" else
|
||||
code_train(n_real, seed * 7717 + t * 13 + li, with_answers=True))
|
||||
data += [x for x in real if x.answer.strip()]
|
||||
rows.append({**tag, "metric": f"n_inherit_{lin}", "value": float(n_self)})
|
||||
rows.append({**tag, "metric": f"n_real_{lin}", "value": float(len(data) - n_self)})
|
||||
d = root / arm / f"gen{t + 1}" / lin
|
||||
if len(data) >= 8:
|
||||
train_lora_on_tasks(base, data, str(d), epochs=epochs,
|
||||
seed=seed * 31 + t * 7 + li, r=r, alpha=alpha,
|
||||
batch_size=tb, max_len=tlen)
|
||||
else: # terminal degeneration: inherit unchanged
|
||||
shutil.copytree(parents[lin], d, dirs_exist_ok=True)
|
||||
rows.append({**tag, "metric": f"degenerate_{lin}", "value": 1.0})
|
||||
new_parents[lin] = str(d)
|
||||
del model; torch.cuda.empty_cache()
|
||||
|
||||
for lin in LINEAGES: # disk hygiene: drop the superseded generation
|
||||
if parents[lin] != founders[lin]:
|
||||
shutil.rmtree(parents[lin], ignore_errors=True)
|
||||
parents = new_parents
|
||||
st_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
st_path.write_text(json.dumps({"generation": t + 1, "parents": parents}))
|
||||
pd.DataFrame(rows).to_parquet(partial, index=False)
|
||||
|
||||
frames.append(pd.DataFrame(rows))
|
||||
return pd.concat(frames, ignore_index=True)
|
||||
170
src/llm/compose_data.py
Normal file
170
src/llm/compose_data.py
Normal file
|
|
@ -0,0 +1,170 @@
|
|||
"""Datasets and prompts for the composition experiment (prereg v3 §1.1).
|
||||
|
||||
Two single-skill training corpora and three evaluation sets, all pinned to fixed subsets by seed so
|
||||
every generation sees *fresh* prompts from a *fixed* pool and no split ever leaks into another:
|
||||
|
||||
* **math** lineage — trains on MetaMathQA (natural-language chain-of-thought); own-skill probe is
|
||||
GSM8K test (numeric answer).
|
||||
* **code** lineage — trains on CodeAlpaca-20k (instruction → code); own-skill probe is MBPP
|
||||
(sanitized) with execution against the reference tests.
|
||||
* **composed** target — GSM8k-Hard, program-aided: emit ``solution()``, execute, compare numerically.
|
||||
Out-of-domain for both lineages, which is the point (LoRA Soups, COLING 2025).
|
||||
|
||||
The composed prompt is fixed here rather than tuned per model, so the operator contrast is never
|
||||
confounded with prompt search.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from functools import lru_cache
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .tasks import Task
|
||||
|
||||
COMPOSED_PROMPT = (
|
||||
"{question}\n\n"
|
||||
"Write a Python function `solution()` that takes no arguments and returns the numeric answer. "
|
||||
"Reply with only the code, inside a ```python code block."
|
||||
)
|
||||
MATH_PROMPT = "{question}\n\nSolve this step by step, then give the final numeric answer after '####'."
|
||||
CODE_PROMPT = "{instruction}\n\nReply with only the Python code, inside a ```python code block."
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ProgTask:
|
||||
"""A program-aided task: prompt, numeric target, and its source id."""
|
||||
|
||||
prompt: str
|
||||
target: float
|
||||
idx: int
|
||||
|
||||
|
||||
@lru_cache(maxsize=8)
|
||||
def _load(name: str, split: str, config: str | None = None):
|
||||
from datasets import load_dataset
|
||||
return load_dataset(name, config, split=split) if config else load_dataset(name, split=split)
|
||||
|
||||
|
||||
def _pick(n_total: int, n: int, seed: int) -> np.ndarray:
|
||||
return np.random.default_rng(seed).choice(n_total, size=min(n, n_total), replace=False)
|
||||
|
||||
|
||||
def _split_pool(n_total: int, split: str) -> np.ndarray:
|
||||
"""Deterministic disjoint halves of an index pool.
|
||||
|
||||
The composed target's val split screens merge weights and the test split reports them; with a
|
||||
pool as small as MATH-500's 271 usable items, sampling both by seed alone would overlap and leak
|
||||
selection into the reported number. Splitting first makes the disjointness structural.
|
||||
"""
|
||||
idx = np.random.default_rng(20260907).permutation(n_total)
|
||||
cut = int(0.7 * n_total)
|
||||
return idx[:cut] if split == "test" else idx[cut:]
|
||||
|
||||
|
||||
def gsm_hard(n: int, seed: int, split: str = "test") -> list[ProgTask]:
|
||||
"""The composed target: GSM8k-Hard, program-aided. ``split`` picks a disjoint val/test half."""
|
||||
d = _load("reasoning-machines/gsm-hard", "train")
|
||||
pool = _split_pool(len(d), split)
|
||||
return [ProgTask(COMPOSED_PROMPT.format(question=d[int(pool[j])]["input"]),
|
||||
float(d[int(pool[j])]["target"]), int(pool[j]))
|
||||
for j in _pick(len(pool), n, seed)]
|
||||
|
||||
|
||||
_FRAC = __import__("re").compile(r"-?\\d?frac\{(-?\d+)\}\{(-?\d+)\}")
|
||||
|
||||
|
||||
def _numeric_answer(a: str) -> float | None:
|
||||
"""Parse a MATH-500 answer to a float, or None if it is not a plain number/simple fraction."""
|
||||
import re
|
||||
a = (a.strip().replace("\\!", "").replace("{,}", "").replace(",", "")
|
||||
.replace("\\%", "").replace("$", "").replace("\\dfrac", "\\frac"))
|
||||
if re.fullmatch(r"-?\d+(\.\d+)?", a):
|
||||
return float(a)
|
||||
f = re.fullmatch(r"-?\\frac\{(-?\d+)\}\{(-?\d+)\}", a)
|
||||
if f and int(f.group(2)) != 0:
|
||||
v = int(f.group(1)) / int(f.group(2))
|
||||
return -v if a.startswith("-") else v
|
||||
return None
|
||||
|
||||
|
||||
def math500(n: int, seed: int, split: str = "test", min_level: int = 3) -> list[ProgTask]:
|
||||
"""Harder composed target: MATH-500 (competition maths), program-aided, numeric answers only.
|
||||
|
||||
The alternative target for when the base can already do the *reasoning* in GSM8k-Hard once code
|
||||
removes the arithmetic burden (measured 2026-09-07: code-only 0.427 there). MetaMathQA is built
|
||||
from GSM8K **and** MATH, so the math specialist is trained on exactly this reasoning while the
|
||||
base is weak at it — which restores E8's premise that each parent supplies something scarce.
|
||||
``min_level`` filters MATH's 1–5 difficulty scale.
|
||||
"""
|
||||
d = _load("HuggingFaceH4/MATH-500", "test")
|
||||
pool = [(i, _numeric_answer(d[i]["answer"])) for i in range(len(d))]
|
||||
pool = [(i, v) for i, v in pool if v is not None and int(d[i]["level"]) >= min_level]
|
||||
half = _split_pool(len(pool), split)
|
||||
idx = [int(half[j]) for j in _pick(len(half), n, seed)]
|
||||
return [ProgTask(COMPOSED_PROMPT.format(question=d[pool[j][0]]["problem"]),
|
||||
float(pool[j][1]), int(pool[j][0])) for j in idx]
|
||||
|
||||
|
||||
def composed_target(kind: str, n: int, seed: int, split: str = "test", **kw) -> list[ProgTask]:
|
||||
"""Dispatch the composed target by name: ``gsm_hard`` (default) or ``math500``."""
|
||||
return {"gsm_hard": gsm_hard, "math500": math500}[kind](n, seed, split=split, **kw)
|
||||
|
||||
|
||||
def gsm8k_probe(n: int, seed: int) -> list[Task]:
|
||||
"""Math own-skill probe: GSM8K test, answer after '####'."""
|
||||
d = _load("openai/gsm8k", "test", "main")
|
||||
out = []
|
||||
for i in _pick(len(d), n, seed):
|
||||
r = d[int(i)]
|
||||
out.append(Task("math", MATH_PROMPT.format(question=r["question"]),
|
||||
r["answer"].split("####")[-1].strip().replace(",", "")))
|
||||
return out
|
||||
|
||||
|
||||
def mbpp_probe(n: int, seed: int) -> list[dict]:
|
||||
"""Code own-skill probe: MBPP sanitized, with the reference asserts as the verifier."""
|
||||
d = _load("google-research-datasets/mbpp", "test", "sanitized")
|
||||
out = []
|
||||
for i in _pick(len(d), n, seed):
|
||||
r = d[int(i)]
|
||||
out.append({"prompt": CODE_PROMPT.format(instruction=r["prompt"]),
|
||||
"tests": list(r["test_list"]), "imports": list(r.get("test_imports") or []),
|
||||
"idx": int(i)})
|
||||
return out
|
||||
|
||||
|
||||
def math_train(n: int, seed: int, *, with_answers: bool) -> list[Task]:
|
||||
"""MetaMathQA examples. ``with_answers=False`` returns prompts only (self-consumption pool)."""
|
||||
d = _load("meta-math/MetaMathQA", "train")
|
||||
out = []
|
||||
for i in _pick(len(d), n, seed):
|
||||
r = d[int(i)]
|
||||
# Reason: MetaMathQA responses end "The answer is: X"; rewrite to the '####' convention the
|
||||
# probe grades on, so training format and evaluation format agree.
|
||||
body = r["response"].strip()
|
||||
ans = body.split("The answer is:")[-1].strip()
|
||||
target = (body[:480] + f"\n#### {ans}") if with_answers else ""
|
||||
out.append(Task("math", MATH_PROMPT.format(question=r["query"]), target))
|
||||
return out
|
||||
|
||||
|
||||
def code_train(n: int, seed: int, *, with_answers: bool) -> list[Task]:
|
||||
"""CodeAlpaca examples (input-free subset, so the prompt is self-contained)."""
|
||||
d = _load("sahil2801/CodeAlpaca-20k", "train")
|
||||
idx = [int(i) for i in _pick(len(d), n * 3, seed) if not d[int(i)]["input"].strip()][:n]
|
||||
return [Task("code", CODE_PROMPT.format(instruction=d[i]["instruction"]),
|
||||
d[i]["output"].strip()[:512] if with_answers else "") for i in idx]
|
||||
|
||||
|
||||
def probe_pool(n: int, seed: int) -> list[str]:
|
||||
"""Shared prompts both lineages answer, for the behavioural correlation rho_t (§1.6).
|
||||
|
||||
Half composed-task prompts, half a mix of each lineage's own domain — so rho reflects agreement
|
||||
on ground both lineages actually walk on, not on prompts only one has ever seen.
|
||||
"""
|
||||
half = n // 2
|
||||
return ([t.prompt for t in gsm_hard(half, seed + 71)]
|
||||
+ [t.prompt for t in gsm8k_probe(half // 2, seed + 72)]
|
||||
+ [t.prompt for t in code_train(n - half - half // 2, seed + 73, with_answers=False)])
|
||||
494
src/llm/curriculum.py
Normal file
494
src/llm/curriculum.py
Normal file
|
|
@ -0,0 +1,494 @@
|
|||
"""`llm_curriculum` — does a society accumulate more than its members? (prereg v4)
|
||||
|
||||
Continual learning in a population. ``L`` lineages of LoRA adapters on a shared frozen base each meet
|
||||
a **new task family every generation** and learn it on verified real data, inheriting their parent's
|
||||
adapter rather than restarting from the base — so what a lineage acquires is passed on as *structure*
|
||||
(the Lamarckian channel v3 lacked). Old families are kept alive by replay, and lineages periodically
|
||||
**recombine**.
|
||||
|
||||
The curriculum is a cyclic Latin square: every lineage sees all ``F`` families but in a different
|
||||
order, so at generation ``t`` each has met ``t`` families and (early on) *different* ones.
|
||||
Complementarity is therefore a known function of generation — maximal at ``t = F/L``, zero at
|
||||
``t = F`` — which lets the framework predict the *shape* of any recombination advantage, not just its
|
||||
sign (prereg §4, H6).
|
||||
|
||||
Arms differ only in who a lineage recombines with and how old skills are maintained:
|
||||
``isolated`` (nobody) · ``society`` (a decorrelated contemporary) · ``society_dry`` (contemporary,
|
||||
self-generated replay) · ``seed_bank`` (its own ancestor at ``t-3`` — temporal rather than spatial
|
||||
complementarity). Three single-shot baselines at matched budget make the multigenerational claim
|
||||
falsifiable: ``sequential``, ``single_shot_merge``, ``joint``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
from . import families as _families # noqa: F401 (registers the extra task families)
|
||||
from .tasks import Task, make_tasks, verify as _verify_proc
|
||||
from .curriculum_data import REAL_FAMILIES, real_tasks, verify_real
|
||||
|
||||
|
||||
def tasks_of(family: str, n: int, seed: int, split: str = "train") -> list[Task]:
|
||||
"""Dataset-backed families draw from a real split; procedural families are generated by seed."""
|
||||
if family in REAL_FAMILIES:
|
||||
return real_tasks(family, n, seed, split)
|
||||
return make_tasks(family, n, seed=seed)
|
||||
|
||||
|
||||
def verify(output: str, task: Task) -> bool:
|
||||
"""Family-appropriate verifier: dataset-backed families carry their own; the rest exact-match."""
|
||||
return verify_real(output, task) if task.family in REAL_FAMILIES else _verify_proc(output, task)
|
||||
|
||||
ARMS = ("isolated", "society", "society_dry", "seed_bank")
|
||||
BASELINES = ("sequential", "single_shot_merge", "joint")
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ pure schedule / bookkeeping
|
||||
|
||||
|
||||
def latin_square(n_lineages: int, n_families: int) -> list[list[int]]:
|
||||
"""Curriculum orders: lineage ``i`` starts at family ``i * F / L`` and cycles.
|
||||
|
||||
With L=3, F=9 the lineages hold disjoint family sets at generation 3 ({0,1,2}, {3,4,5}, {6,7,8})
|
||||
and identical sets at generation 9 — complementarity by construction, decaying on a known
|
||||
schedule.
|
||||
"""
|
||||
if n_families % n_lineages:
|
||||
raise ValueError(f"{n_families} families must divide evenly among {n_lineages} lineages")
|
||||
step = n_families // n_lineages
|
||||
return [[(i * step + t) % n_families for t in range(n_families)] for i in range(n_lineages)]
|
||||
|
||||
|
||||
def complementarity(orders: list[list[int]], t: int) -> float:
|
||||
"""Mean pairwise Jaccard *distance* between lineages' seen-family sets after ``t`` generations.
|
||||
|
||||
1.0 = fully disjoint curricula, 0.0 = identical. The predictor H6 tests the advantage against.
|
||||
"""
|
||||
seen = [set(o[:t]) for o in orders]
|
||||
if t == 0 or len(seen) < 2:
|
||||
return 0.0
|
||||
ds = [1.0 - len(a & b) / len(a | b) for i, a in enumerate(seen) for b in seen[i + 1:]]
|
||||
return float(np.mean(ds))
|
||||
|
||||
|
||||
def resolve_orders(cfg: dict, fams: list[str], n_lineages: int, generations: int) -> list[list[int]]:
|
||||
"""Per-lineage family orders (indices into ``fams``).
|
||||
|
||||
``cfg["orders"]`` (a list of family-name lists, one per lineage) overrides the Latin square, so a
|
||||
curriculum can decouple partner complementarity from generation number (the Latin square makes
|
||||
them collinear). Every name must be a configured family and every order must cover the run.
|
||||
"""
|
||||
if not cfg.get("orders"):
|
||||
return latin_square(n_lineages, len(fams))
|
||||
orders = cfg["orders"]
|
||||
if len(orders) != n_lineages:
|
||||
raise ValueError(f"orders has {len(orders)} lineages, config has {n_lineages}")
|
||||
out = []
|
||||
for o in orders:
|
||||
bad = [f for f in o if f not in fams]
|
||||
if bad:
|
||||
raise ValueError(f"unknown families in orders: {bad}")
|
||||
if len(o) < generations:
|
||||
raise ValueError(f"order {o} shorter than {generations} generations")
|
||||
out.append([fams.index(f) for f in o])
|
||||
return out
|
||||
|
||||
|
||||
def replay_split(seen: list[int], n_replay: int) -> dict[int, int]:
|
||||
"""Split a fixed replay budget evenly across every family seen so far.
|
||||
|
||||
Fixed *total* (not per-family), so protection per family thins as the curriculum grows and
|
||||
forgetting stays a live pressure — the realistic regime, and the one where recombination has
|
||||
something to fix (prereg §8 decision 2).
|
||||
"""
|
||||
if not seen or n_replay <= 0:
|
||||
return {}
|
||||
per = n_replay // len(seen)
|
||||
rem = n_replay - per * len(seen)
|
||||
return {f: per + (1 if k < rem else 0) for k, f in enumerate(seen)}
|
||||
|
||||
|
||||
def partner_for(arm: str, lineage: int, t: int, n_lineages: int, ancestor_depth: int = 3,
|
||||
merge_until: int | None = None):
|
||||
"""Who does ``lineage`` recombine with this generation?
|
||||
|
||||
Returns ``("contemporary", j)``, ``("ancestor", t - depth)``, or ``None``. ``merge_until`` is
|
||||
the forced-stop control: recombination is allowed only at generations ``t < merge_until``
|
||||
(the fixed "merge early, then stop" schedule the declinable merge is compared against).
|
||||
"""
|
||||
if arm == "isolated":
|
||||
return None
|
||||
if merge_until is not None and t >= merge_until:
|
||||
return None
|
||||
if arm == "seed_bank":
|
||||
return ("ancestor", t - ancestor_depth) if t >= ancestor_depth else None
|
||||
return ("contemporary", (lineage + 1) % n_lineages) if n_lineages > 1 else None
|
||||
|
||||
|
||||
def cull_step(scores: list[float]) -> tuple[int, int] | None:
|
||||
"""Truncation selection at fixed population size: ``(culled_slot, source_slot)``.
|
||||
|
||||
The lowest-scoring lineage is replaced by a copy of the highest-scoring one (differential
|
||||
reproduction: the fittest genome leaves two descendants, the least fit none). Ties are left
|
||||
alone, so a population of equals is never reshuffled. NaN scores are treated as the worst.
|
||||
|
||||
Args:
|
||||
scores (list[float]): one fitness per lineage slot (all-families accuracy this generation).
|
||||
|
||||
Returns:
|
||||
tuple[int, int] | None: (slot to overwrite, slot to copy from), or None when best == worst.
|
||||
"""
|
||||
s = [(-np.inf if np.isnan(x) else float(x)) for x in scores]
|
||||
worst, best = int(np.argmin(s)), int(np.argmax(s))
|
||||
return None if s[best] <= s[worst] else (worst, best)
|
||||
|
||||
|
||||
def inherit_slot(i: int, j: int, adapters: list, history: list[list[str]], budget: list[int],
|
||||
archive: dict[tuple[int, int], str], t: int) -> None:
|
||||
"""Make slot ``i`` carry lineage ``j``'s genome: adapter path, taught families, example budget,
|
||||
and the ancestry archive up to generation ``t`` (so a seed-bank partner follows the genome, not
|
||||
the slot). Paths are aliased, never copied: every generation trains into a fresh ``gen{t}/lin{i}``
|
||||
directory, so a shared path is read-only from here on."""
|
||||
adapters[i] = adapters[j]
|
||||
history[i] = list(history[j])
|
||||
budget[i] = budget[j]
|
||||
for k in range(t + 1):
|
||||
if (j, k) in archive:
|
||||
archive[(i, k)] = archive[(j, k)]
|
||||
|
||||
|
||||
def cumulative_accuracy(per_family: dict[str, float], seen: list[str]) -> float:
|
||||
"""Mean accuracy over the families a model is supposed to know (the primary outcome)."""
|
||||
vals = [per_family[f] for f in seen if f in per_family and not np.isnan(per_family[f])]
|
||||
return float(np.mean(vals)) if vals else float("nan")
|
||||
|
||||
|
||||
def forgetting(history: list[dict[str, float]], fam: str, learned_at: int, now: int) -> float:
|
||||
"""Drop in ``fam`` from the generation it was learned to ``now`` (positive = forgotten)."""
|
||||
if learned_at >= len(history) or now >= len(history):
|
||||
return float("nan")
|
||||
a, b = history[learned_at].get(fam), history[now].get(fam)
|
||||
return float("nan") if a is None or b is None else float(a - b)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ the GPU loop
|
||||
|
||||
|
||||
def _eval_all(model, tok, gen_fn, tests: dict[str, list[Task]]) -> dict[str, float]:
|
||||
"""Per-family exact-match accuracy over every family in the curriculum."""
|
||||
out = {}
|
||||
for fam, ts in tests.items():
|
||||
outs = gen_fn(model, tok, [x.prompt for x in ts])
|
||||
out[fam] = float(np.mean([verify(o, x) for o, x in zip(outs, ts)]))
|
||||
return out
|
||||
|
||||
|
||||
def run_curriculum(cfg: dict) -> pd.DataFrame:
|
||||
"""Run every configured arm and baseline; return tidy long-form rows (resumable per arm)."""
|
||||
import torch
|
||||
|
||||
from .evaluate import generate
|
||||
from .merge import load_specialists
|
||||
from .specialise import continue_lora_training, train_lora_on_tasks
|
||||
|
||||
name, base, seed = cfg["experiment"], cfg["base_model"], int(cfg["seed"])
|
||||
fams = list(cfg["families"])
|
||||
F, L = len(fams), int(cfg.get("lineages", 3))
|
||||
G = int(cfg.get("generations", F))
|
||||
arms = list(cfg.get("arms", ARMS))
|
||||
baselines = list(cfg.get("baselines", []))
|
||||
n_new = int(cfg.get("n_new", 300))
|
||||
n_replay = int(cfg.get("n_replay", 150))
|
||||
n_test = int(cfg.get("n_test", 60))
|
||||
n_val = int(cfg.get("n_val", 30))
|
||||
epochs = int(cfg.get("epochs", 3))
|
||||
lr_new = float(cfg.get("lr", 1e-4))
|
||||
depth = int(cfg.get("ancestor_depth", 3))
|
||||
merge_weights = [list(w) for w in cfg.get("merge_weights", [[0.5, 0.5], [0.3, 0.7], [0.7, 0.3]])]
|
||||
op = str(cfg.get("operator", "linear"))
|
||||
allow_veto = bool(cfg.get("allow_veto", False))
|
||||
cull = bool(cfg.get("cull", False))
|
||||
merge_until = cfg.get("merge_until")
|
||||
merge_until = None if merge_until is None else int(merge_until)
|
||||
max_new = int(cfg.get("max_new_tokens", 32))
|
||||
bs = int(cfg.get("batch_size", 32))
|
||||
tb, tlen = int(cfg.get("train_batch_size", 2)), int(cfg.get("train_max_len", 448))
|
||||
r, alpha = int(cfg.get("lora", {}).get("r", 16)), int(cfg.get("lora", {}).get("alpha", 32))
|
||||
root = Path(cfg.get("adapters_dir", "models/llm")) / "curriculum" / f"{name}_s{seed}"
|
||||
out_dir = Path(cfg.get("output", {}).get("dir", f"results/{name}"))
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
orders = resolve_orders(cfg, fams, L, G)
|
||||
tests = {f: tasks_of(f, n_test, 1000 + i, "test") for i, f in enumerate(fams)}
|
||||
vals = {f: tasks_of(f, n_val, 2000 + i, "train") for i, f in enumerate(fams)}
|
||||
|
||||
def gen(model, tok, prompts):
|
||||
return generate(model, tok, prompts, max_new_tokens=max_new, batch_size=bs)
|
||||
|
||||
def train_data(fam_idx: int, seen_idx: list[int], t: int, lineage: int) -> tuple[list[Task], int]:
|
||||
"""New-family examples plus a replay split over families already seen. Returns (data, n)."""
|
||||
data = tasks_of(fams[fam_idx], n_new, seed * 977 + t * 31 + fam_idx, "train")
|
||||
for f, k in replay_split(seen_idx, n_replay).items():
|
||||
data += tasks_of(fams[f], k, seed * 613 + t * 17 + f + lineage, "train")
|
||||
return data, len(data)
|
||||
|
||||
rows: list[dict] = [] # every arm, returned at the end
|
||||
arm_rows: list[dict] = [] # the arm currently running — what its partial file holds
|
||||
|
||||
def measure_base():
|
||||
"""Base model on every family, before any adapter — the format-transfer reference."""
|
||||
from .evaluate import load_model
|
||||
m, tk = load_model(base)
|
||||
acc = _eval_all(m, tk, gen, tests)
|
||||
del m
|
||||
torch.cuda.empty_cache()
|
||||
return acc
|
||||
|
||||
def record(arm: str, t: int, who: str, acc: dict, seen: list[str], budget: int, **extra):
|
||||
tag = {"experiment": name, "arm": arm, "seed": seed, "generation": t, "model": who}
|
||||
for f, v in acc.items():
|
||||
arm_rows.append({**tag, "metric": f"acc_{f}", "value": v})
|
||||
# Two different questions, both recorded, easily confused:
|
||||
# retention_seen — mean over families this model was TAUGHT: "of what you learned, how much
|
||||
# do you still know?" It starts at ~1 family freshly learned and can only
|
||||
# fall as the curriculum grows. This is the forgetting readout.
|
||||
# all_families — mean over ALL F families in the curriculum: total capability, which is
|
||||
# what ACCUMULATES. This is the pre-registered primary outcome (§4).
|
||||
arm_rows.append({**tag, "metric": "retention_seen", "value": cumulative_accuracy(acc, seen)})
|
||||
arm_rows.append({**tag, "metric": "all_families", "value": float(np.mean(list(acc.values())))})
|
||||
arm_rows.append({**tag, "metric": "n_seen", "value": float(len(seen))})
|
||||
arm_rows.append({**tag, "metric": "budget_examples", "value": float(budget)})
|
||||
for k, v in extra.items():
|
||||
arm_rows.append({**tag, "metric": k, "value": float(v)})
|
||||
|
||||
base_acc = measure_base()
|
||||
for f, v in base_acc.items():
|
||||
rows.append({"experiment": name, "arm": "base", "seed": seed, "generation": -1,
|
||||
"model": "base", "metric": f"acc_{f}", "value": v})
|
||||
rows.append({"experiment": name, "arm": "base", "seed": seed, "generation": -1, "model": "base",
|
||||
"metric": "all_families", "value": float(np.mean(list(base_acc.values())))})
|
||||
print(f"[base] all-families {np.mean(list(base_acc.values())):.3f} "
|
||||
+ " ".join(f"{k}={v:.2f}" for k, v in base_acc.items()), flush=True)
|
||||
|
||||
# ---------------------------------------------------------------- society arms
|
||||
for arm in arms:
|
||||
arm_rows = [] # this arm's rows only
|
||||
state = root / arm / "state.json"
|
||||
partial = out_dir / f"partial_{arm}_s{seed}.parquet"
|
||||
adapters: list[str | None] = [None] * L # current adapter per lineage
|
||||
history: list[list[str]] = [[] for _ in range(L)]
|
||||
archive: dict[tuple[int, int], str] = {} # (lineage, generation) -> adapter dir
|
||||
budget = [0] * L
|
||||
t0 = 0
|
||||
if bool(cfg.get("resume", True)) and state.exists():
|
||||
st = json.loads(state.read_text())
|
||||
if all(a is None or Path(a, "adapter_config.json").exists() for a in st["adapters"]):
|
||||
adapters, history, budget, t0 = st["adapters"], st["history"], st["budget"], st["t"]
|
||||
archive = {tuple(map(int, k.split(","))): v for k, v in st["archive"].items()}
|
||||
if partial.exists():
|
||||
arm_rows = pd.read_parquet(partial).to_dict("records")
|
||||
print(f"[{arm}] resuming at generation {t0}", flush=True)
|
||||
|
||||
for t in range(t0, G):
|
||||
for i in range(L):
|
||||
fam_idx = orders[i][t]
|
||||
seen_idx = list(dict.fromkeys(orders[i][k] for k in range(t)))
|
||||
data, n = train_data(fam_idx, seen_idx, t, i)
|
||||
budget[i] += n
|
||||
d = root / arm / f"gen{t}" / f"lin{i}"
|
||||
if adapters[i] is None: # generation 0: the founding adapter
|
||||
train_lora_on_tasks(base, data, str(d), epochs=epochs, r=r, alpha=alpha,
|
||||
seed=seed * 41 + i, batch_size=tb, max_len=tlen)
|
||||
else: # inherit the parent's weights, then add
|
||||
continue_lora_training(base, adapters[i], data, str(d), epochs=epochs,
|
||||
lr=lr_new, seed=seed * 41 + t * L + i, batch_size=tb,
|
||||
max_len=tlen)
|
||||
adapters[i] = str(d)
|
||||
# Union with what the slot already carries: after a cull the slot holds a genome
|
||||
# taught elsewhere, and those families stay "taught" (no-op for uncull arms).
|
||||
history[i] = list(dict.fromkeys(history[i] + [fams[k] for k in seen_idx]
|
||||
+ [fams[fam_idx]]))
|
||||
archive[(i, t)] = str(d)
|
||||
|
||||
# ---- recombination (after everyone has learned this generation's family)
|
||||
merged: list[str | None] = [None] * L
|
||||
model, tok = load_specialists(base, [a for a in adapters])
|
||||
for i in range(L):
|
||||
p = partner_for(arm, i, t, L, depth, merge_until)
|
||||
if p is None:
|
||||
continue
|
||||
kind, j = p
|
||||
if kind == "ancestor":
|
||||
anc = archive.get((i, j))
|
||||
if anc is None:
|
||||
continue
|
||||
model.load_adapter(anc, adapter_name=f"anc{i}")
|
||||
names = [f"a{i}", f"anc{i}"]
|
||||
else:
|
||||
names = [f"a{i}", f"a{j}"]
|
||||
# directed recombination: weights screened on a held-out validation split.
|
||||
# With `allow_veto`, "keep the parent unchanged" is itself a candidate offspring —
|
||||
# a lineage may decline a merge that no weighting makes worthwhile. Without it the
|
||||
# merge is obligate, which is what v5 tested and nobody would build.
|
||||
seen_f = history[i]
|
||||
|
||||
def _val_score(adapter_name):
|
||||
model.set_adapter(adapter_name)
|
||||
acc = {f: float(np.mean([verify(o, x) for o, x in
|
||||
zip(gen(model, tok, [x.prompt for x in vals[f]]), vals[f])]))
|
||||
for f in seen_f}
|
||||
return cumulative_accuracy(acc, seen_f)
|
||||
|
||||
veto_score = _val_score(f"a{i}") if allow_veto else -np.inf
|
||||
best_w, best_v = merge_weights[0], -np.inf
|
||||
for wi, w in enumerate(merge_weights):
|
||||
cn = f"c{i}_{wi}"
|
||||
model.add_weighted_adapter(names, w, cn, combination_type=op)
|
||||
model.set_adapter(cn)
|
||||
acc = {f: float(np.mean([verify(o, x) for o, x in
|
||||
zip(gen(model, tok, [x.prompt for x in vals[f]]), vals[f])]))
|
||||
for f in seen_f}
|
||||
v = cumulative_accuracy(acc, seen_f)
|
||||
if v > best_v:
|
||||
best_v, best_w = v, w
|
||||
model.set_adapter(f"a{i}"); model.delete_adapter(cn)
|
||||
if allow_veto:
|
||||
declined = best_v <= veto_score
|
||||
arm_rows.append({"experiment": name, "arm": arm, "seed": seed, "generation": t,
|
||||
"model": f"lineage{i}", "metric": "veto_used",
|
||||
"value": float(declined)})
|
||||
arm_rows.append({"experiment": name, "arm": arm, "seed": seed, "generation": t,
|
||||
"model": f"lineage{i}", "metric": "veto_margin",
|
||||
"value": float(best_v - veto_score)})
|
||||
if declined: # no merge improves on the parent — keep it
|
||||
if kind == "ancestor":
|
||||
model.delete_adapter(f"anc{i}")
|
||||
continue
|
||||
mn = f"m{i}"
|
||||
model.add_weighted_adapter(names, best_w, mn, combination_type=op)
|
||||
model.set_adapter(mn)
|
||||
md = root / arm / f"gen{t}" / f"merged{i}"
|
||||
model.save_pretrained(str(md.parent), selected_adapters=[mn])
|
||||
shutil.rmtree(md, ignore_errors=True)
|
||||
shutil.move(str(md.parent / mn), str(md))
|
||||
merged[i] = str(md)
|
||||
model.set_adapter(f"a{i}"); model.delete_adapter(mn)
|
||||
if kind == "ancestor":
|
||||
model.delete_adapter(f"anc{i}")
|
||||
|
||||
# ---- measure every lineage's deployed model on ALL families
|
||||
for i in range(L):
|
||||
use = merged[i] or adapters[i]
|
||||
if merged[i]:
|
||||
model.load_adapter(use, adapter_name=f"dep{i}")
|
||||
model.set_adapter(f"dep{i}")
|
||||
else:
|
||||
model.set_adapter(f"a{i}")
|
||||
acc = _eval_all(model, tok, gen, tests)
|
||||
record(arm, t, f"lineage{i}", acc, history[i], budget[i],
|
||||
complementarity=complementarity(orders, t + 1))
|
||||
if merged[i]:
|
||||
model.set_adapter(f"a{i}"); model.delete_adapter(f"dep{i}")
|
||||
adapters[i] = use # the merged offspring continues the lineage
|
||||
del model; torch.cuda.empty_cache()
|
||||
|
||||
def _acc_of(i):
|
||||
return {r["metric"][4:]: r["value"] for r in arm_rows
|
||||
if r["arm"] == arm and r["generation"] == t and r["model"] == f"lineage{i}"
|
||||
and r["metric"].startswith("acc_")}
|
||||
|
||||
# ---- differential reproduction: the least fit slot is re-founded from the fittest
|
||||
if cull:
|
||||
scores = [float(np.mean(list(_acc_of(i).values()))) for i in range(L)]
|
||||
step = cull_step(scores)
|
||||
for i in range(L):
|
||||
culled = step is not None and i == step[0]
|
||||
tag = {"experiment": name, "arm": arm, "seed": seed, "generation": t,
|
||||
"model": f"lineage{i}"}
|
||||
arm_rows.append({**tag, "metric": "culled", "value": float(culled)})
|
||||
arm_rows.append({**tag, "metric": "cull_source",
|
||||
"value": float(step[1]) if culled else -1.0})
|
||||
if step is not None:
|
||||
inherit_slot(step[0], step[1], adapters, history, budget, archive, t)
|
||||
print(f"[{arm}] gen {t}: culled lineage{step[0]} ({scores[step[0]]:.3f}), "
|
||||
f"re-founded from lineage{step[1]} ({scores[step[1]]:.3f})", flush=True)
|
||||
|
||||
state.parent.mkdir(parents=True, exist_ok=True)
|
||||
state.write_text(json.dumps({"adapters": adapters, "history": history, "budget": budget,
|
||||
"t": t + 1,
|
||||
"archive": {f"{a},{b}": v for (a, b), v in archive.items()}}))
|
||||
pd.DataFrame(arm_rows).to_parquet(partial, index=False)
|
||||
|
||||
best = max(float(np.mean(list(_acc_of(i).values()))) for i in range(L))
|
||||
ret = max(cumulative_accuracy(_acc_of(i), history[i]) for i in range(L))
|
||||
print(f"[{arm}] gen {t}/{G - 1}: ALL-FAMILIES {best:.3f} (retention-of-taught {ret:.3f}) "
|
||||
f"complementarity {complementarity(orders, t + 1):.2f} budget {budget[0]}",
|
||||
flush=True)
|
||||
|
||||
rows += arm_rows # fold the finished arm into the run
|
||||
|
||||
# ---------------------------------------------------------------- single-shot baselines
|
||||
arm_rows = rows # baselines append straight to the run
|
||||
for bl in baselines:
|
||||
d = root / "baselines" / bl
|
||||
chunk = F // L
|
||||
if bl == "joint":
|
||||
data = sum([tasks_of(f, (n_new + n_replay) * G // F, seed * 811 + i, "train")
|
||||
for i, f in enumerate(fams)], [])
|
||||
train_lora_on_tasks(base, data, str(d), epochs=epochs, r=r, alpha=alpha,
|
||||
seed=seed * 53, batch_size=tb, max_len=tlen)
|
||||
model, tok = load_specialists(base, [str(d)])
|
||||
record(bl, G - 1, "baseline", _eval_all(model, tok, gen, tests), fams, len(data))
|
||||
del model; torch.cuda.empty_cache()
|
||||
elif bl == "sequential":
|
||||
cur, used = None, 0
|
||||
for t, f in enumerate(fams):
|
||||
data, n = train_data(t, list(range(t)), t, 0); used += n
|
||||
dd = d / f"step{t}"
|
||||
if cur is None:
|
||||
train_lora_on_tasks(base, data, str(dd), epochs=epochs, r=r, alpha=alpha,
|
||||
seed=seed * 59, batch_size=tb, max_len=tlen)
|
||||
else:
|
||||
continue_lora_training(base, cur, data, str(dd), epochs=epochs, lr=lr_new,
|
||||
seed=seed * 59 + t, batch_size=tb, max_len=tlen)
|
||||
cur = str(dd)
|
||||
model, tok = load_specialists(base, [cur])
|
||||
record(bl, G - 1, "baseline", _eval_all(model, tok, gen, tests), fams, used)
|
||||
del model; torch.cuda.empty_cache()
|
||||
elif bl == "single_shot_merge":
|
||||
specs, used = [], 0
|
||||
for i in range(L):
|
||||
mine = [orders[i][k] for k in range(F)][:chunk] if chunk else []
|
||||
data = sum([tasks_of(fams[k], (n_new + n_replay) * G // F,
|
||||
seed * 733 + k, "train") for k in mine], [])
|
||||
used += len(data)
|
||||
dd = d / f"spec{i}"
|
||||
train_lora_on_tasks(base, data, str(dd), epochs=epochs, r=r, alpha=alpha,
|
||||
seed=seed * 67 + i, batch_size=tb, max_len=tlen)
|
||||
specs.append(str(dd))
|
||||
model, tok = load_specialists(base, specs)
|
||||
best_w, best_v, wgrid = None, -np.inf, cfg.get("baseline_weights", [[1 / L] * L])
|
||||
for wi, w in enumerate(wgrid): # same directed selection the arms get
|
||||
cn = f"bl{wi}"
|
||||
model.add_weighted_adapter([f"a{i}" for i in range(L)], list(w), cn,
|
||||
combination_type=op)
|
||||
model.set_adapter(cn)
|
||||
acc = {f: float(np.mean([verify(o, x) for o, x in
|
||||
zip(gen(model, tok, [x.prompt for x in vals[f]]), vals[f])]))
|
||||
for f in fams}
|
||||
v = cumulative_accuracy(acc, fams)
|
||||
if v > best_v:
|
||||
best_v, best_w = v, list(w)
|
||||
model.set_adapter("a0"); model.delete_adapter(cn)
|
||||
model.add_weighted_adapter([f"a{i}" for i in range(L)], best_w, "blm",
|
||||
combination_type=op)
|
||||
model.set_adapter("blm")
|
||||
record(bl, G - 1, "baseline", _eval_all(model, tok, gen, tests), fams, used)
|
||||
del model; torch.cuda.empty_cache()
|
||||
|
||||
return pd.DataFrame(rows)
|
||||
187
src/llm/curriculum_data.py
Normal file
187
src/llm/curriculum_data.py
Normal file
|
|
@ -0,0 +1,187 @@
|
|||
"""A real curriculum: naturally heterogeneous tasks with distinct answer formats (prereg v4/v5).
|
||||
|
||||
Why this exists. v2's nine procedural families could not serve as a curriculum: they share one
|
||||
answer convention (so the first family teaches the format for all nine — the base's 0.094 became 0.417
|
||||
after *one* family), were calibrated for low mutual conflict (so learning one never damaged another
|
||||
and there was nothing to forget except one confusable pair), and the base was merely unformatted
|
||||
rather than incapable. Continual learning needs the opposite on all three counts, which is why the
|
||||
literature's benchmarks use tasks like these.
|
||||
|
||||
Every family here is drawn from a public dataset with a **train** split (acquisition and replay) and a
|
||||
disjoint **test** split (evaluation), and each has its own answer format and verifier:
|
||||
|
||||
| family | source | format | verifier |
|
||||
|---|---|---|---|
|
||||
| gsm8k | openai/gsm8k | number | last number, numeric tolerance |
|
||||
| mbpp | google-research-datasets/mbpp (full) | Python function | execute against reference asserts |
|
||||
| boolq | google/boolq | yes / no | label |
|
||||
| mnli | nyu-mll/glue mnli | entailment / neutral / contradiction | label |
|
||||
| sst2 | nyu-mll/glue sst2 | positive / negative | label |
|
||||
| csqa | tau/commonsense_qa | letter A–E | label |
|
||||
| arc | allenai/ai2_arc ARC-Easy | letter A–D | label |
|
||||
| winogrande | allenai/winogrande xl | option 1 / 2 | label |
|
||||
| squad | rajpurkar/squad | extractive span | normalised exact match against gold aliases |
|
||||
| nq_open | google-research-datasets/nq_open | short free text | normalised exact match against aliases |
|
||||
| hellaswag | Rowan/hellaswag | letter A–D | label |
|
||||
|
||||
Selection into the curriculum is by calibration (base ≤ 0.40, specialist ≥ 0.60, and the
|
||||
single-lineage zero-replay probe must show mean forgetting ≥ 0.15), not by preference.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
import string
|
||||
from functools import lru_cache
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .execute import extract_code, numeric_match, run_solution
|
||||
from .tasks import Task, _normalise
|
||||
|
||||
# ------------------------------------------------------------------ registry
|
||||
|
||||
FORMATS = {
|
||||
"gsm8k": "number", "mbpp": "code", "boolq": "label", "mnli": "label", "sst2": "label",
|
||||
"csqa": "label", "arc": "label", "winogrande": "label", "squad": "span", "nq_open": "text",
|
||||
"hellaswag": "label",
|
||||
}
|
||||
REAL_FAMILIES = tuple(FORMATS)
|
||||
|
||||
_SRC = {
|
||||
"gsm8k": ("openai/gsm8k", "main", "train", "test"),
|
||||
"mbpp": ("google-research-datasets/mbpp", "full", "train", "test"),
|
||||
"boolq": ("google/boolq", None, "train", "validation"),
|
||||
"mnli": ("nyu-mll/glue", "mnli", "train", "validation_matched"),
|
||||
"sst2": ("nyu-mll/glue", "sst2", "train", "validation"),
|
||||
"csqa": ("tau/commonsense_qa", None, "train", "validation"),
|
||||
"arc": ("allenai/ai2_arc", "ARC-Easy", "train", "test"),
|
||||
"winogrande": ("allenai/winogrande", "winogrande_xl", "train", "validation"),
|
||||
"squad": ("rajpurkar/squad", None, "train", "validation"),
|
||||
"nq_open": ("google-research-datasets/nq_open", None, "train", "validation"),
|
||||
"hellaswag": ("Rowan/hellaswag", None, "train", "validation"),
|
||||
}
|
||||
|
||||
|
||||
@lru_cache(maxsize=32)
|
||||
def _load(family: str, split: str):
|
||||
from datasets import load_dataset
|
||||
name, cfg, tr, te = _SRC[family]
|
||||
sp = tr if split == "train" else te
|
||||
return load_dataset(name, cfg, split=sp) if cfg else load_dataset(name, split=sp)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ per-family formatting
|
||||
|
||||
_L = "ABCDE"
|
||||
|
||||
|
||||
def _fmt(family: str, r: dict) -> Task:
|
||||
if family == "gsm8k":
|
||||
ans = r["answer"].strip()
|
||||
return Task(family, f"{r['question']}\n\nSolve step by step, then give the final numeric "
|
||||
f"answer after '####'.", ans[:480] if "####" in ans[-12:] else
|
||||
ans[:440] + "\n#### " + ans.split("####")[-1].strip())
|
||||
if family == "mbpp":
|
||||
tests = list(r["test_list"])
|
||||
sig = tests[0].split("assert ")[-1].split("(")[0].strip() if tests else "solution"
|
||||
prompt = (f"{r['text']}\n\nWrite the Python function `{sig}`. Reply with only the code, "
|
||||
f"inside a ```python code block.")
|
||||
return Task(family, prompt, "```python\n" + r["code"].strip() + "\n```",
|
||||
json.dumps({"tests": tests, "setup": r.get("test_setup_code", "") or ""}))
|
||||
if family == "boolq":
|
||||
return Task(family, f"Passage: {r['passage'][:900]}\n\nQuestion: {r['question']}?\n\n"
|
||||
f"Answer yes or no.", "yes" if r["answer"] else "no")
|
||||
if family == "mnli":
|
||||
lab = ["entailment", "neutral", "contradiction"][int(r["label"])]
|
||||
return Task(family, f"Premise: {r['premise']}\nHypothesis: {r['hypothesis']}\n\nDoes the "
|
||||
f"premise entail the hypothesis? Answer entailment, neutral, or contradiction.", lab)
|
||||
if family == "sst2":
|
||||
return Task(family, f"Review: {r['sentence'].strip()}\n\nIs the sentiment positive or "
|
||||
f"negative? Answer with one word.", "positive" if r["label"] == 1 else "negative")
|
||||
if family == "csqa":
|
||||
ch = r["choices"]
|
||||
opts = "\n".join(f"{l}. {t}" for l, t in zip(ch["label"], ch["text"]))
|
||||
return Task(family, f"{r['question']}\n{opts}\n\nAnswer with the letter only.", r["answerKey"])
|
||||
if family == "arc":
|
||||
ch = r["choices"]
|
||||
labels = [("ABCD"[i] if l.isdigit() else l) for i, l in enumerate(ch["label"])]
|
||||
key = r["answerKey"]; key = "ABCD"[int(key) - 1] if key.isdigit() else key
|
||||
opts = "\n".join(f"{l}. {t}" for l, t in zip(labels, ch["text"]))
|
||||
return Task(family, f"{r['question']}\n{opts}\n\nAnswer with the letter only.", key)
|
||||
if family == "winogrande":
|
||||
return Task(family, f"{r['sentence']}\n\nWhat does the blank refer to?\n1. {r['option1']}\n"
|
||||
f"2. {r['option2']}\n\nAnswer 1 or 2.", str(r["answer"]))
|
||||
if family == "squad":
|
||||
golds = list(dict.fromkeys(r["answers"]["text"]))
|
||||
return Task(family, f"Context: {r['context'][:1200]}\n\nQuestion: {r['question']}\n\n"
|
||||
f"Answer with the exact phrase from the context.", golds[0],
|
||||
json.dumps({"aliases": golds}))
|
||||
if family == "nq_open":
|
||||
golds = list(dict.fromkeys(r["answer"]))
|
||||
return Task(family, f"Question: {r['question']}?\n\nAnswer in a few words.", golds[0],
|
||||
json.dumps({"aliases": golds}))
|
||||
if family == "hellaswag":
|
||||
ends = list(r["endings"])
|
||||
opts = "\n".join(f"{_L[i]}. {e}" for i, e in enumerate(ends))
|
||||
return Task(family, f"{r['ctx']}\n\nWhich ending is most plausible?\n{opts}\n\nAnswer with "
|
||||
f"the letter only.", _L[int(r["label"])])
|
||||
raise KeyError(family)
|
||||
|
||||
|
||||
def real_tasks(family: str, n: int, seed: int, split: str = "train") -> list[Task]:
|
||||
"""``n`` tasks from ``family``'s ``split``, chosen deterministically by ``seed``."""
|
||||
d = _load(family, split)
|
||||
idx = np.random.default_rng(seed).choice(len(d), size=min(n, len(d)), replace=False)
|
||||
return [_fmt(family, d[int(i)]) for i in idx]
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ verifiers
|
||||
|
||||
_ART = re.compile(r"\b(a|an|the)\b")
|
||||
_NUM = re.compile(r"-?\d[\d,]*\.?\d*")
|
||||
|
||||
|
||||
def _squad_norm(s: str) -> str:
|
||||
s = s.lower()
|
||||
s = "".join(c for c in s if c not in string.punctuation)
|
||||
s = _ART.sub(" ", s)
|
||||
return " ".join(s.split())
|
||||
|
||||
|
||||
def _last_number(text: str) -> float | None:
|
||||
for tok in reversed(_NUM.findall(text)):
|
||||
try:
|
||||
return float(tok.replace(",", "").rstrip("."))
|
||||
except ValueError:
|
||||
continue
|
||||
return None
|
||||
|
||||
|
||||
def _first_line_answer(output: str) -> str:
|
||||
"""Model outputs for short-answer formats: take the first non-empty line, strip a leading label."""
|
||||
for line in output.strip().splitlines():
|
||||
line = line.strip()
|
||||
if line:
|
||||
return re.sub(r"^(answer|final answer)\s*[:\-]\s*", "", line, flags=re.I).strip()
|
||||
return ""
|
||||
|
||||
|
||||
def verify_real(output: str, task: Task) -> bool:
|
||||
"""Family-appropriate correctness for a dataset-backed task."""
|
||||
kind = FORMATS[task.family]
|
||||
if kind == "number":
|
||||
got = _last_number(output.split("####")[-1] if "####" in output else output)
|
||||
want = _last_number(task.answer.split("####")[-1])
|
||||
return got is not None and want is not None and numeric_match(got, want)
|
||||
if kind == "code":
|
||||
meta = json.loads(task.meta or "{}")
|
||||
prog = "\n".join([meta.get("setup", ""), extract_code(output)] + list(meta.get("tests", []))
|
||||
+ ["def solution():\n return 1"])
|
||||
return run_solution(prog, timeout_s=6.0).ok
|
||||
if kind == "label":
|
||||
return _normalise(_first_line_answer(output)) == _normalise(task.answer)
|
||||
aliases = json.loads(task.meta or "{}").get("aliases", [task.answer])
|
||||
got = _squad_norm(_first_line_answer(output))
|
||||
return bool(got) and any(got == _squad_norm(a) for a in aliases)
|
||||
|
|
@ -29,6 +29,20 @@ def load_model(name: str, device: str = "cuda", adapter_dir: str | None = None):
|
|||
return model, tok
|
||||
|
||||
|
||||
def format_prompt(tok, prompt: str) -> str:
|
||||
"""Render one user prompt to text, identically for training and inference.
|
||||
|
||||
Instruct checkpoints get their own chat template. Base checkpoints (no template — the regime the
|
||||
composition experiment needs, since an instruction-tuned base already has the skills the
|
||||
specialists are supposed to supply) get a plain instruction/response format. Training and
|
||||
generation must agree on this or the adapter learns a format it is never evaluated in.
|
||||
"""
|
||||
if getattr(tok, "chat_template", None):
|
||||
return tok.apply_chat_template([{"role": "user", "content": prompt}],
|
||||
add_generation_prompt=True, tokenize=False)
|
||||
return f"### Instruction:\n{prompt}\n\n### Response:\n"
|
||||
|
||||
|
||||
def generate(model, tok, prompts: list[str], max_new_tokens: int = 24,
|
||||
batch_size: int = 32, device: str = "cuda") -> list[str]:
|
||||
"""Greedy chat-formatted batched generation; returns the decoded completions."""
|
||||
|
|
@ -37,9 +51,8 @@ def generate(model, tok, prompts: list[str], max_new_tokens: int = 24,
|
|||
outs: list[str] = []
|
||||
for i in range(0, len(prompts), batch_size):
|
||||
chunk = prompts[i:i + batch_size]
|
||||
msgs = [[{"role": "user", "content": p}] for p in chunk]
|
||||
enc = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt",
|
||||
return_dict=True, padding=True).to(device)
|
||||
enc = tok([format_prompt(tok, p) for p in chunk], return_tensors="pt", padding=True,
|
||||
add_special_tokens=False).to(device)
|
||||
with torch.no_grad():
|
||||
gen = model.generate(**enc, max_new_tokens=max_new_tokens, do_sample=False,
|
||||
pad_token_id=tok.pad_token_id)
|
||||
|
|
|
|||
164
src/llm/execute.py
Normal file
164
src/llm/execute.py
Normal file
|
|
@ -0,0 +1,164 @@
|
|||
"""Sandboxed execution verifier for program-aided answers (prereg v3 §1.1, blueprint §3.6).
|
||||
|
||||
The composed task is *program-aided* math: the model emits Python defining ``solution()``, the code is
|
||||
run, and its return value is compared numerically to the reference. Execution is the verifier — the
|
||||
one place in this tier where something outside the model population decides what is true.
|
||||
|
||||
**Threat model, stated honestly.** The code comes from a small instruction-tuned model answering maths
|
||||
questions, so the realistic hazards are accidental: runaway loops, unbounded allocation, stray file
|
||||
writes, a fork bomb. Those are contained by running each snippet in a fresh subprocess with CPU,
|
||||
address-space, file-size and process-count rlimits, a wall-clock timeout, an empty temporary working
|
||||
directory, and a stripped environment, launched with ``-I -S`` (isolated, no site imports). This is
|
||||
*not* a security boundary against an adversary who controls the model — it does not use namespaces,
|
||||
seccomp, or a container, and generated code could still open a socket. Do not point this at
|
||||
untrusted model output without adding one of those.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
from dataclasses import dataclass
|
||||
|
||||
_FENCE = re.compile(r"```(?:python|py)?\s*(.*?)```", re.S | re.I)
|
||||
|
||||
_DRIVER = '''
|
||||
import json, sys, os, resource
|
||||
resource.setrlimit(resource.RLIMIT_AS, ({mem}, {mem}))
|
||||
resource.setrlimit(resource.RLIMIT_CPU, ({cpu}, {cpu}))
|
||||
resource.setrlimit(resource.RLIMIT_FSIZE, (1 << 20, 1 << 20))
|
||||
try:
|
||||
resource.setrlimit(resource.RLIMIT_NPROC, (64, 64))
|
||||
except (ValueError, OSError):
|
||||
pass
|
||||
|
||||
_SRC = open({path!r}).read() # read before the hook goes up; hooks cannot be removed
|
||||
_JAIL = os.path.realpath({jail!r})
|
||||
_BLOCKED = ("socket.socket", "socket.connect", "socket.getaddrinfo", "subprocess.Popen",
|
||||
"os.system", "os.exec", "os.fork", "os.posix_spawn", "shutil.rmtree", "ftplib.connect",
|
||||
"urllib.Request", "webbrowser.open")
|
||||
|
||||
def _guard(event, args):
|
||||
if event == "open":
|
||||
mode = str(args[1]) if len(args) > 1 and args[1] else "r"
|
||||
if any(c in mode for c in "wxa+"):
|
||||
p = os.path.realpath(str(args[0]))
|
||||
if not p.startswith(_JAIL):
|
||||
raise PermissionError("sandbox: write outside jail: " + p)
|
||||
elif event.startswith(_BLOCKED):
|
||||
raise PermissionError("sandbox: blocked " + event)
|
||||
|
||||
sys.addaudithook(_guard)
|
||||
|
||||
_out = {{"status": "ok", "value": None}}
|
||||
try:
|
||||
_ns = {{}}
|
||||
exec(compile(_SRC, "solution.py", "exec"), _ns)
|
||||
fn = _ns.get("solution")
|
||||
if fn is None:
|
||||
_out = {{"status": "no_solution", "value": None}}
|
||||
else:
|
||||
v = fn()
|
||||
_out = {{"status": "ok", "value": v if isinstance(v, (int, float, bool)) else str(v)}}
|
||||
except MemoryError:
|
||||
_out = {{"status": "memory", "value": None}}
|
||||
except RecursionError:
|
||||
_out = {{"status": "recursion", "value": None}}
|
||||
except BaseException as e:
|
||||
_out = {{"status": "error:" + type(e).__name__, "value": None}}
|
||||
sys.stdout.write("\\x00RESULT\\x00" + json.dumps(_out, default=str))
|
||||
'''
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ExecResult:
|
||||
"""Outcome of running one generated snippet."""
|
||||
|
||||
ok: bool # solution() ran and returned a value
|
||||
value: float | None # numeric return value, if it could be coerced
|
||||
status: str # ok | no_code | no_solution | timeout | memory | recursion | error:Type
|
||||
|
||||
|
||||
def extract_code(text: str) -> str:
|
||||
"""Pull the Python out of a completion: fenced block if present, else the raw text.
|
||||
|
||||
Keeps only up to the last line that still looks like code, so trailing prose ("This function
|
||||
computes…") does not become a SyntaxError.
|
||||
"""
|
||||
m = _FENCE.findall(text)
|
||||
body = m[-1] if m else text
|
||||
lines = body.splitlines()
|
||||
while lines and not lines[-1].strip():
|
||||
lines.pop()
|
||||
return "\n".join(lines).strip()
|
||||
|
||||
|
||||
def run_solution(code: str, *, timeout_s: float = 6.0, mem_mb: int = 768,
|
||||
cpu_s: int = 5) -> ExecResult:
|
||||
"""Run ``code`` (which should define ``solution()``) in a sandboxed subprocess.
|
||||
|
||||
Args:
|
||||
code (str): the snippet, already extracted from any markdown fence.
|
||||
timeout_s (float): wall-clock limit; the process is killed past it.
|
||||
mem_mb (int): address-space limit inside the child.
|
||||
cpu_s (int): CPU-time limit inside the child.
|
||||
|
||||
Returns:
|
||||
ExecResult: ``ok`` iff ``solution()`` returned; ``value`` coerced to float where possible.
|
||||
"""
|
||||
if not code.strip():
|
||||
return ExecResult(False, None, "no_code")
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
src = os.path.join(tmp, "solution.py")
|
||||
with open(src, "w") as fh:
|
||||
fh.write(code)
|
||||
drv = os.path.join(tmp, "_driver.py")
|
||||
with open(drv, "w") as fh:
|
||||
fh.write(_DRIVER.format(mem=mem_mb * 1024 * 1024, cpu=cpu_s, path=src, jail=tmp))
|
||||
try:
|
||||
proc = subprocess.run(
|
||||
[sys.executable, "-I", "-S", drv],
|
||||
cwd=tmp, capture_output=True, text=True, timeout=timeout_s,
|
||||
env={"PATH": "/usr/bin:/bin", "HOME": tmp, "TMPDIR": tmp,
|
||||
"PYTHONHASHSEED": "0", "OPENBLAS_NUM_THREADS": "1"},
|
||||
)
|
||||
except subprocess.TimeoutExpired:
|
||||
return ExecResult(False, None, "timeout")
|
||||
except OSError as e: # spawn failure — infrastructure, not model
|
||||
return ExecResult(False, None, f"error:{type(e).__name__}")
|
||||
marker = proc.stdout.rfind("\x00RESULT\x00")
|
||||
if marker < 0:
|
||||
# Killed by a signal with no result: the CPU rlimit (SIGXCPU) fires before the wall
|
||||
# clock, so a runaway loop lands here — report it as the timeout it is.
|
||||
return ExecResult(False, None, "timeout" if proc.returncode < 0 else "error:NoResult")
|
||||
try:
|
||||
out = json.loads(proc.stdout[marker + len("\x00RESULT\x00"):])
|
||||
except json.JSONDecodeError:
|
||||
return ExecResult(False, None, "error:BadJSON")
|
||||
if out["status"] != "ok":
|
||||
return ExecResult(False, None, out["status"])
|
||||
try:
|
||||
return ExecResult(True, float(out["value"]), "ok")
|
||||
except (TypeError, ValueError):
|
||||
return ExecResult(True, None, "ok")
|
||||
|
||||
|
||||
def numeric_match(value: float | None, target: float, *, rel: float = 1e-4,
|
||||
abs_: float = 1e-6) -> bool:
|
||||
"""Compare a returned value to the reference with a relative-or-absolute tolerance."""
|
||||
if value is None:
|
||||
return False
|
||||
try:
|
||||
return abs(value - target) <= max(abs_, rel * abs(target))
|
||||
except (TypeError, ValueError, OverflowError):
|
||||
return False
|
||||
|
||||
|
||||
def verify_program(completion: str, target: float, **kw) -> tuple[bool, ExecResult]:
|
||||
"""Full pipeline for one completion: extract → execute → compare. Returns (correct, result)."""
|
||||
res = run_solution(extract_code(completion), **kw)
|
||||
return (res.ok and numeric_match(res.value, target)), res
|
||||
|
|
@ -273,9 +273,32 @@ def run_society_dispatch(cfg: dict) -> pd.DataFrame:
|
|||
return run_society_experiment(cfg)
|
||||
|
||||
|
||||
def run_society_v2_dispatch(cfg: dict) -> pd.DataFrame:
|
||||
from .society_v2 import run_society_v2 # local import: torch-heavy
|
||||
return run_society_v2(cfg)
|
||||
|
||||
|
||||
def run_compose_dispatch(cfg: dict) -> pd.DataFrame:
|
||||
from .compose import run_compose # local import: torch-heavy
|
||||
return run_compose(cfg)
|
||||
|
||||
|
||||
def run_curriculum_dispatch(cfg: dict) -> pd.DataFrame:
|
||||
from .curriculum import run_curriculum # local import: torch-heavy
|
||||
return run_curriculum(cfg)
|
||||
|
||||
|
||||
def run_calib_dispatch(cfg: dict) -> pd.DataFrame:
|
||||
from .calibrate import run_calibration # local import: torch-heavy
|
||||
return run_calibration(cfg)
|
||||
|
||||
|
||||
_RUNNERS = {"llm_merge": run_merge_experiment, "llm_moe": run_moe_experiment,
|
||||
"llm_directed": run_directed_experiment, "llm_speciation": run_speciation_dispatch,
|
||||
"llm_epistasis": run_epistasis_dispatch, "llm_society": run_society_dispatch}
|
||||
"llm_epistasis": run_epistasis_dispatch, "llm_society": run_society_dispatch,
|
||||
"llm_society_v2": run_society_v2_dispatch, "llm_society_calib": run_calib_dispatch,
|
||||
"llm_compose": run_compose_dispatch,
|
||||
"llm_curriculum": run_curriculum_dispatch}
|
||||
|
||||
|
||||
def run_and_save(config_path: str | Path) -> Path:
|
||||
|
|
@ -312,6 +335,21 @@ def run_and_save(config_path: str | Path) -> Path:
|
|||
if kind == "llm_society":
|
||||
extra["society"] = {k: cfg.get(k) for k in
|
||||
("agents", "generations", "arms", "g", "lam", "n_candidates")}
|
||||
if kind == "llm_society_v2":
|
||||
extra["society_v2"] = {k: cfg.get(k) for k in
|
||||
("families", "agents", "generations", "arms", "g", "lam",
|
||||
"k_inherit", "epochs", "n_test", "n_val", "n_conf")}
|
||||
if kind == "llm_compose":
|
||||
extra["compose"] = {k: cfg.get(k) for k in
|
||||
("arms", "generations", "g", "k_inherit", "conf_gate", "epochs",
|
||||
"n_hard", "n_gsm8k", "n_mbpp", "lora")}
|
||||
if kind == "llm_curriculum":
|
||||
extra["curriculum"] = {k: cfg.get(k) for k in
|
||||
("families", "lineages", "generations", "arms", "baselines",
|
||||
"n_new", "n_replay", "operator", "ancestor_depth", "lora",
|
||||
"allow_veto", "merge_until", "orders", "cull")}
|
||||
if kind == "llm_society_calib":
|
||||
extra["calibration"] = {"stage": cfg.get("stage"), "families": cfg.get("families")}
|
||||
if kind == "llm_directed":
|
||||
extra["directed"] = {"n_candidates": int(cfg.get("n_candidates", 16)),
|
||||
"concentration": float(cfg.get("concentration", 0.5)),
|
||||
|
|
|
|||
259
src/llm/families.py
Normal file
259
src/llm/families.py
Normal file
|
|
@ -0,0 +1,259 @@
|
|||
"""Candidate task families for the L=12 society (pre-registration §3.2, §4 C1).
|
||||
|
||||
The v1 society had three families, so eight agents were near-clones and the competence space had
|
||||
2^3 = 8 states (E8 needs room for a child to hold a combination no parent had). This module adds
|
||||
fourteen further *disjoint* procedural families with exact-match answers; calibration gate C1 keeps
|
||||
the twelve whose base / specialist accuracies fall in the pre-set band, so membership is measured,
|
||||
not chosen.
|
||||
|
||||
Verifier safety (``tasks._normalise``): every answer is an integer, a bracketed integer list, or a
|
||||
single lowercase alphabetic token — the three forms the normaliser canonicalises. Families whose
|
||||
natural answer would be multi-word or alphanumeric (run-length codes, hex strings, word-order
|
||||
reversal) are deliberately absent because the verifier could not score them.
|
||||
|
||||
The generators register into :data:`llm.tasks._GEN` on import, so ``make_tasks(family, ...)`` works
|
||||
unchanged for the new names; :data:`llm.tasks.FAMILIES` (the original three) is untouched because the
|
||||
merge / routing / directed experiments depend on it.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .tasks import Task, _GEN, _WORDS, _fmt_list, _caesar
|
||||
|
||||
_TAIL = " Output only the answer and nothing else."
|
||||
_ROMAN = (("M", 1000), ("CM", 900), ("D", 500), ("CD", 400), ("C", 100), ("XC", 90), ("L", 50),
|
||||
("XL", 40), ("X", 10), ("IX", 9), ("V", 5), ("IV", 4), ("I", 1))
|
||||
|
||||
|
||||
def _pseudo(rng, lo: int = 5, hi: int = 8) -> str:
|
||||
"""A random lowercase pseudo-word — a large prompt space so training cannot cover the test set."""
|
||||
letters = "abcdefghijklmnopqrstuvwxyz"
|
||||
return "".join(letters[i] for i in rng.integers(0, 26, size=int(rng.integers(lo, hi))))
|
||||
|
||||
|
||||
def _to_roman(n: int) -> str:
|
||||
out = ""
|
||||
for sym, val in _ROMAN:
|
||||
while n >= val:
|
||||
out += sym; n -= val
|
||||
return out
|
||||
|
||||
|
||||
def _roman(rng, hard=False) -> Task:
|
||||
n = int(rng.integers(4, 1000)) # 1-999: a 600-item train set cannot cover it
|
||||
if rng.random() < 0.5:
|
||||
return Task("roman", f"Write the number {n} as a Roman numeral." + _TAIL, _to_roman(n))
|
||||
return Task("roman", f"Convert the Roman numeral {_to_roman(n)} to an ordinary number." + _TAIL,
|
||||
str(n))
|
||||
|
||||
|
||||
def _binary(rng, hard=False) -> Task:
|
||||
n = int(rng.integers(5, 512))
|
||||
if rng.random() < 0.5:
|
||||
return Task("binary", f"Write the decimal number {n} in binary." + _TAIL, format(n, "b"))
|
||||
return Task("binary", f"Convert the binary number {format(n, 'b')} to decimal." + _TAIL, str(n))
|
||||
|
||||
|
||||
def _sortletters(rng, hard=False) -> Task:
|
||||
w = _pseudo(rng)
|
||||
if rng.random() < 0.5:
|
||||
return Task("sortletters", f"Sort the letters of the word \"{w}\" into alphabetical order "
|
||||
f"and write them as one word." + _TAIL, "".join(sorted(w)))
|
||||
return Task("sortletters", f"Sort the letters of the word \"{w}\" into reverse alphabetical "
|
||||
f"order and write them as one word." + _TAIL, "".join(sorted(w, reverse=True)))
|
||||
|
||||
|
||||
def _setops(rng, hard=False) -> Task:
|
||||
a = sorted(set(rng.integers(1, 15, size=int(rng.integers(4, 7))).tolist()))
|
||||
b = sorted(set(rng.integers(1, 15, size=int(rng.integers(4, 7))).tolist()))
|
||||
op = rng.choice(["intersection", "union", "difference"])
|
||||
if op == "intersection":
|
||||
ans = sorted(set(a) & set(b)); what = "the elements present in both lists"
|
||||
elif op == "union":
|
||||
ans = sorted(set(a) | set(b)); what = "the elements present in either list, without repeats"
|
||||
else:
|
||||
ans = sorted(set(a) - set(b)); what = "the elements of the first list that are not in the second"
|
||||
return Task("setops", f"Given the lists {_fmt_list(a)} and {_fmt_list(b)}, list {what}, sorted "
|
||||
f"ascending, as a bracketed list (write [] if there are none)." + _TAIL, _fmt_list(ans))
|
||||
|
||||
|
||||
def _numtheory(rng, hard=False) -> Task:
|
||||
kind = rng.choice(["gcd", "lcm", "mod", "parity"])
|
||||
if kind == "gcd":
|
||||
a, b = rng.integers(6, 80, size=2).tolist()
|
||||
return Task("numtheory", f"What is the greatest common divisor of {a} and {b}?" + _TAIL,
|
||||
str(math.gcd(a, b)))
|
||||
if kind == "lcm":
|
||||
a, b = rng.integers(2, 13, size=2).tolist()
|
||||
return Task("numtheory", f"What is the least common multiple of {a} and {b}?" + _TAIL,
|
||||
str(a * b // math.gcd(a, b)))
|
||||
if kind == "mod":
|
||||
a = int(rng.integers(20, 300)); b = int(rng.integers(3, 12))
|
||||
return Task("numtheory", f"What is the remainder when {a} is divided by {b}?" + _TAIL,
|
||||
str(a % b))
|
||||
n = int(rng.integers(10, 999))
|
||||
return Task("numtheory", f"Is the number {n} even or odd? Answer with one word." + _TAIL,
|
||||
"even" if n % 2 == 0 else "odd")
|
||||
|
||||
|
||||
def _mixedtoken(rng, hard=False) -> Task:
|
||||
letters = "abcdefghijklmnopqrstuvwxyz"
|
||||
n = int(rng.integers(6, 10))
|
||||
chars = []
|
||||
for _ in range(n):
|
||||
r = rng.random()
|
||||
if r < 0.35:
|
||||
chars.append(str(rng.integers(0, 10)))
|
||||
elif r < 0.65:
|
||||
chars.append(letters[rng.integers(0, 26)].upper())
|
||||
else:
|
||||
chars.append(letters[rng.integers(0, 26)])
|
||||
tokn = "".join(chars)
|
||||
if not any(c.isdigit() for c in tokn): # guarantee at least one digit
|
||||
tokn = tokn[:-1] + str(rng.integers(1, 10))
|
||||
kind = rng.choice(["upper", "digits", "digitsum", "extract"])
|
||||
if kind == "upper":
|
||||
return Task("mixedtoken", f"How many uppercase letters are in \"{tokn}\"?" + _TAIL,
|
||||
str(sum(c.isupper() for c in tokn)))
|
||||
if kind == "digits":
|
||||
return Task("mixedtoken", f"How many digits are in \"{tokn}\"?" + _TAIL,
|
||||
str(sum(c.isdigit() for c in tokn)))
|
||||
ds = [int(c) for c in tokn if c.isdigit()]
|
||||
if kind == "digitsum":
|
||||
return Task("mixedtoken", f"What is the sum of the digits that appear in \"{tokn}\"?" + _TAIL,
|
||||
str(sum(ds)))
|
||||
return Task("mixedtoken", f"Write the digits that appear in \"{tokn}\", in order, as a single "
|
||||
f"number." + _TAIL, str(int("".join(map(str, ds)))))
|
||||
|
||||
|
||||
def _caesar_task(rng, hard=False) -> Task:
|
||||
w = _pseudo(rng); k = int(rng.integers(1, 6))
|
||||
return Task("caesar", f"Shift every letter of the word \"{w}\" forward by {k} places in the "
|
||||
f"alphabet, wrapping around from z to a." + _TAIL, _caesar(w, k))
|
||||
|
||||
|
||||
def _vectors(rng, hard=False) -> Task:
|
||||
a = rng.integers(-5, 9, size=3).tolist(); b = rng.integers(-5, 9, size=3).tolist()
|
||||
if rng.random() < 0.5:
|
||||
return Task("vectors", f"What is the dot product of the vectors {_fmt_list(a)} and "
|
||||
f"{_fmt_list(b)}?" + _TAIL, str(sum(x * y for x, y in zip(a, b))))
|
||||
return Task("vectors", f"Add the vectors {_fmt_list(a)} and {_fmt_list(b)} element by element "
|
||||
f"and write the result as a bracketed list." + _TAIL,
|
||||
_fmt_list([x + y for x, y in zip(a, b)]))
|
||||
|
||||
|
||||
def _progression(rng, hard=False) -> Task:
|
||||
a = int(rng.integers(1, 20)); d = int(rng.integers(2, 9)); n = int(rng.integers(5, 15))
|
||||
if rng.random() < 0.5:
|
||||
return Task("progression", f"An arithmetic sequence starts at {a} and increases by {d} each "
|
||||
f"step. What is its {n}th term (the first term is term 1)?" + _TAIL,
|
||||
str(a + (n - 1) * d))
|
||||
return Task("progression", f"An arithmetic sequence starts at {a} and increases by {d} each step. "
|
||||
f"What is the sum of its first {n} terms?" + _TAIL, str(n * (2 * a + (n - 1) * d) // 2))
|
||||
|
||||
|
||||
def _charfreq(rng, hard=False) -> Task:
|
||||
w = _pseudo(rng, 7, 11)
|
||||
if rng.random() < 0.5:
|
||||
ans = min(set(w), key=lambda c: (-w.count(c), c))
|
||||
return Task("charfreq", f"Which letter occurs most often in \"{w}\"? If several tie, give the "
|
||||
f"one earliest in the alphabet." + _TAIL, ans)
|
||||
return Task("charfreq", f"How many distinct letters does \"{w}\" contain?" + _TAIL,
|
||||
str(len(set(w))))
|
||||
|
||||
|
||||
def _digits(rng, hard=False) -> Task:
|
||||
n = int(rng.integers(100, 9999))
|
||||
kind = rng.choice(["sum", "reverse", "count"])
|
||||
if kind == "sum":
|
||||
return Task("digits", f"What is the sum of the digits of {n}?" + _TAIL,
|
||||
str(sum(int(c) for c in str(n))))
|
||||
if kind == "reverse":
|
||||
return Task("digits", f"Write the digits of {n} in reverse order as a number." + _TAIL,
|
||||
str(int(str(n)[::-1])))
|
||||
m = n * int(rng.integers(1, 12))
|
||||
return Task("digits", f"How many digits does the number {m} have?" + _TAIL, str(len(str(m))))
|
||||
|
||||
|
||||
def _liststats(rng, hard=False) -> Task:
|
||||
n = int(rng.choice([5, 7]))
|
||||
xs = rng.integers(0, 40, size=n).tolist()
|
||||
kind = rng.choice(["median", "range", "evens", "argmax"])
|
||||
if kind == "median":
|
||||
return Task("liststats", f"What is the median of {_fmt_list(xs)}?" + _TAIL,
|
||||
str(sorted(xs)[n // 2]))
|
||||
if kind == "range":
|
||||
return Task("liststats", f"What is the range (largest minus smallest) of {_fmt_list(xs)}?"
|
||||
+ _TAIL, str(max(xs) - min(xs)))
|
||||
if kind == "evens":
|
||||
return Task("liststats", f"How many even numbers are in {_fmt_list(xs)}?" + _TAIL,
|
||||
str(sum(x % 2 == 0 for x in xs)))
|
||||
return Task("liststats", f"At which position (counting from 0) is the largest value in "
|
||||
f"{_fmt_list(xs)}? If it appears more than once, give the first position." + _TAIL,
|
||||
str(int(np.argmax(xs))))
|
||||
|
||||
|
||||
def _alphabet(rng, hard=False) -> Task:
|
||||
letters = "abcdefghijklmnopqrstuvwxyz"
|
||||
kind = rng.choice(["pos", "at", "next"])
|
||||
if kind == "pos":
|
||||
c = letters[rng.integers(0, 26)]
|
||||
return Task("alphabet", f"What is the position of the letter \"{c}\" in the alphabet "
|
||||
f"(a is 1)?" + _TAIL, str(letters.index(c) + 1))
|
||||
if kind == "at":
|
||||
k = int(rng.integers(1, 27))
|
||||
return Task("alphabet", f"Which letter is at position {k} of the alphabet (a is 1)?" + _TAIL,
|
||||
letters[k - 1])
|
||||
c = letters[rng.integers(0, 25)]; k = int(rng.integers(1, 5))
|
||||
return Task("alphabet", f"Which letter comes {k} places after \"{c}\" in the alphabet?" + _TAIL,
|
||||
letters[(letters.index(c) + k) % 26])
|
||||
|
||||
|
||||
def _wordlen(rng, hard=False) -> Task:
|
||||
w = _pseudo(rng, 4, 12)
|
||||
return Task("wordlen", f"How many letters are in the word \"{w}\"?" + _TAIL, str(len(w)))
|
||||
|
||||
|
||||
def _lettercount(rng, hard=False) -> Task:
|
||||
w = _pseudo(rng, 6, 12)
|
||||
c = w[rng.integers(0, len(w))] # a letter that occurs at least once
|
||||
return Task("lettercount", f"How many times does the letter \"{c}\" occur in \"{w}\"?" + _TAIL,
|
||||
str(w.count(c)))
|
||||
|
||||
|
||||
def _sumeven(rng, hard=False) -> Task:
|
||||
xs = rng.integers(1, 30, size=int(rng.integers(5, 8))).tolist()
|
||||
if rng.random() < 0.5:
|
||||
return Task("sumeven", f"What is the sum of the even numbers in {_fmt_list(xs)}?" + _TAIL,
|
||||
str(sum(x for x in xs if x % 2 == 0)))
|
||||
return Task("sumeven", f"What is the sum of the odd numbers in {_fmt_list(xs)}?" + _TAIL,
|
||||
str(sum(x for x in xs if x % 2 == 1)))
|
||||
|
||||
|
||||
def _prime(rng, hard=False) -> Task:
|
||||
n = int(rng.integers(4, 400))
|
||||
if rng.random() < 0.5:
|
||||
isp = n > 1 and all(n % d for d in range(2, int(n ** 0.5) + 1))
|
||||
return Task("prime", f"Is {n} a prime number? Answer yes or no." + _TAIL, "yes" if isp else "no")
|
||||
spf = next(d for d in range(2, n + 1) if n % d == 0)
|
||||
return Task("prime", f"What is the smallest prime factor of {n}?" + _TAIL, str(spf))
|
||||
|
||||
|
||||
EXTRA_FAMILIES = ("roman", "binary", "sortletters", "setops", "numtheory", "mixedtoken", "caesar",
|
||||
"vectors", "progression", "charfreq", "digits", "liststats", "alphabet", "prime",
|
||||
"wordlen", "lettercount", "sumeven")
|
||||
"""The seventeen candidate families added for the society; gate C1 selects twelve from these plus
|
||||
the original three. ``wordlen``, ``lettercount``, ``sumeven`` were added after the first C1 table
|
||||
(2026-09-07) left only six in band."""
|
||||
|
||||
_GEN.update({"roman": _roman, "binary": _binary, "sortletters": _sortletters, "setops": _setops,
|
||||
"numtheory": _numtheory, "mixedtoken": _mixedtoken, "caesar": _caesar_task,
|
||||
"vectors": _vectors, "progression": _progression, "charfreq": _charfreq,
|
||||
"digits": _digits, "liststats": _liststats, "alphabet": _alphabet, "prime": _prime,
|
||||
"wordlen": _wordlen, "lettercount": _lettercount, "sumeven": _sumeven})
|
||||
|
||||
ALL_CANDIDATES = ("lists", "strings", "arith") + EXTRA_FAMILIES
|
||||
|
|
@ -26,129 +26,29 @@ from pathlib import Path
|
|||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
from .tasks import Task, make_tasks, verify, _normalise
|
||||
from .tasks import Task, make_tasks, _normalise
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------- pure operators
|
||||
|
||||
|
||||
def consensus_answers(outputs: list[list[str]]) -> list[str]:
|
||||
"""The population's modal (normalised) answer per prompt.
|
||||
|
||||
Args:
|
||||
outputs (list[list[str]]): ``outputs[i][p]`` = agent i's raw answer to prompt p.
|
||||
|
||||
Returns:
|
||||
list[str]: per-prompt modal normalised answer (ties broken lexicographically —
|
||||
deterministic).
|
||||
"""
|
||||
n_prompts = len(outputs[0])
|
||||
cons = []
|
||||
for p in range(n_prompts):
|
||||
votes: dict[str, int] = {}
|
||||
for out in outputs:
|
||||
key = _normalise(out[p])
|
||||
votes[key] = votes.get(key, 0) + 1
|
||||
# Reason: max count, then lexicographically smallest key, so consensus is deterministic.
|
||||
top = max(votes.values())
|
||||
cons.append(min(k for k, v in votes.items() if v == top))
|
||||
return cons
|
||||
|
||||
|
||||
def conformity_scores(outputs: list[list[str]], consensus: list[str]) -> np.ndarray:
|
||||
"""Each agent's agreement rate with the population consensus (the E11 conformity signal)."""
|
||||
return np.array([np.mean([_normalise(o) == c for o, c in zip(out, consensus)])
|
||||
for out in outputs], dtype=float)
|
||||
|
||||
|
||||
def behavioural_distance(outputs: list[list[str]]) -> np.ndarray:
|
||||
"""Pairwise disagreement rate between agents' normalised answers (verifier-free diversity)."""
|
||||
n = len(outputs)
|
||||
norm = [[_normalise(o) for o in out] for out in outputs]
|
||||
d = np.zeros((n, n))
|
||||
for i in range(n):
|
||||
for j in range(i + 1, n):
|
||||
d[i, j] = d[j, i] = float(np.mean([a != b for a, b in zip(norm[i], norm[j])]))
|
||||
return d
|
||||
|
||||
|
||||
def select_parents(scores: np.ndarray, dist: np.ndarray, k: int, *, diversity: bool,
|
||||
lam: float = 0.3) -> list[int]:
|
||||
"""Select ``k`` parents by score, optionally with quality-diversity preservation.
|
||||
|
||||
Diversity ON: greedy — take the top-scoring agent, then repeatedly add the agent maximising
|
||||
``score + lam * (mean behavioural distance to those already selected)``. Diversity OFF: plain
|
||||
top-k by score.
|
||||
|
||||
Args:
|
||||
scores (np.ndarray): per-agent selection score ``g*fitness + (1-g)*conformity``.
|
||||
dist (np.ndarray): pairwise behavioural-distance matrix.
|
||||
k (int): number of parents.
|
||||
diversity (bool): quality-diversity (True) or greedy (False).
|
||||
lam (float): diversity weight.
|
||||
|
||||
Returns:
|
||||
list[int]: selected agent indices (deterministic).
|
||||
"""
|
||||
if not diversity:
|
||||
return list(np.argsort(-scores, kind="stable")[:k])
|
||||
chosen = [int(np.argmax(scores))]
|
||||
while len(chosen) < k:
|
||||
best_i, best_v = -1, -np.inf
|
||||
for i in range(len(scores)):
|
||||
if i in chosen:
|
||||
continue
|
||||
v = scores[i] + lam * float(np.mean([dist[i, j] for j in chosen]))
|
||||
if v > best_v:
|
||||
best_i, best_v = i, v
|
||||
chosen.append(best_i)
|
||||
return chosen
|
||||
|
||||
|
||||
def complementary_pairs(parents: list[int], dist: np.ndarray, n_children: int) -> list[tuple[int, int]]:
|
||||
"""Mating plan: pairs of parents in descending behavioural distance, cycled to fill the slots.
|
||||
|
||||
Directed mate choice (E10): the most complementary (most-disagreeing) pairs breed first; the pair
|
||||
list cycles until ``n_children`` slots are filled.
|
||||
"""
|
||||
pairs = sorted(((a, b) for i, a in enumerate(parents) for b in parents[i + 1:]),
|
||||
key=lambda ab: -dist[ab[0], ab[1]])
|
||||
if not pairs: # single parent: self-pair (degenerate)
|
||||
pairs = [(parents[0], parents[0])]
|
||||
return [pairs[i % len(pairs)] for i in range(n_children)]
|
||||
# Pure operators live in society_ops (shared with the v2 loop); re-exported here so the v1 code
|
||||
# path and its tests are unchanged.
|
||||
from .society_ops import ( # noqa: E402,F401
|
||||
answer_of as _answer_of, arm_settings as _arm_settings_v2, behavioural_distance,
|
||||
complementary_pairs, conformity_scores, consensus_answers, fitness_of as _fitness,
|
||||
select_parents,
|
||||
)
|
||||
|
||||
|
||||
def arm_settings(arm: str, g: float) -> dict:
|
||||
"""Resolve an arm name to its operator switches (the four-arm ablation)."""
|
||||
if arm == "full":
|
||||
return {"g": g, "sex": True, "diversity": True}
|
||||
if arm == "no_grounding":
|
||||
return {"g": 0.0, "sex": True, "diversity": True}
|
||||
if arm == "no_sex":
|
||||
return {"g": g, "sex": False, "diversity": True}
|
||||
if arm == "no_diversity":
|
||||
return {"g": g, "sex": True, "diversity": False}
|
||||
raise ValueError(f"unknown arm {arm!r} (full|no_grounding|no_sex|no_diversity)")
|
||||
"""v1 four-arm switches (boolean ``sex``), resolved from the shared table."""
|
||||
s = _arm_settings_v2(arm, g)
|
||||
if s["sex"] == "linear":
|
||||
raise ValueError("sex_linear is a v2 arm (kind: llm_society_v2)")
|
||||
return {"g": s["g"], "sex": s["sex"] is not None, "diversity": s["diversity"]}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------- the GPU loop
|
||||
|
||||
|
||||
def _fitness(outs: list[str], tasks: list[Task], fams: list[str]) -> dict:
|
||||
"""Overall / per-family / worst-family verifier accuracy from precomputed outputs."""
|
||||
corr = np.array([verify(o, t) for o, t in zip(outs, tasks)])
|
||||
fam = np.array([t.family for t in tasks])
|
||||
acc = {"overall": float(corr.mean())}
|
||||
acc.update({f: float(corr[fam == f].mean()) for f in fams})
|
||||
acc["worst_family"] = min(acc[f] for f in fams)
|
||||
return acc
|
||||
|
||||
|
||||
def _answer_of(raw: str) -> str:
|
||||
"""Trim a raw completion to a single short answer line for the inheritance data."""
|
||||
return raw.strip().split("\n")[0][:64].strip()
|
||||
|
||||
|
||||
def run_society_experiment(cfg: dict) -> pd.DataFrame:
|
||||
"""Run the society loop for every configured arm; return tidy long-form rows.
|
||||
|
||||
|
|
|
|||
209
src/llm/society_ops.py
Normal file
209
src/llm/society_ops.py
Normal file
|
|
@ -0,0 +1,209 @@
|
|||
"""Pure, testable operators for the LLM society (v1 and v2 — see ``tasks/prereg-llm-society-v2.md``).
|
||||
|
||||
Everything here is NumPy on strings and floats; nothing touches a model. The v2 additions are the
|
||||
E11-faithful pieces the v1 design lacked:
|
||||
|
||||
* :func:`pooled_survival` — selection acts on *survival over parents + offspring*, keeping the top N
|
||||
by ``score + λ·novelty`` (E11's rule), instead of truncating parents before breeding, which in v1
|
||||
discarded half the families at generation 1 with no operator able to restore them (E6).
|
||||
* :func:`mating_plan` — complementary pairing over the whole population with a per-agent use cap, so
|
||||
every founder's family can reach the next generation.
|
||||
* :func:`route_union` — the union-preserving recombination operator in the *inheritance data*: per
|
||||
prompt, the child learns the answer of the more confident parent. This is E4's ``max`` per item,
|
||||
verifier-free (legal in ``no_grounding``), and directed in E10's sense.
|
||||
* :func:`choose_single_parent` — score-proportional parent sampling for the ``no_sex`` arm, so that
|
||||
arm still has selection without recombination.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .tasks import Task, _normalise, verify
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ population readouts (v1)
|
||||
|
||||
|
||||
def consensus_answers(outputs: list[list[str]]) -> list[str]:
|
||||
"""The population's modal (normalised) answer per prompt; ties broken lexicographically."""
|
||||
n_prompts = len(outputs[0])
|
||||
cons = []
|
||||
for p in range(n_prompts):
|
||||
votes: dict[str, int] = {}
|
||||
for out in outputs:
|
||||
key = _normalise(out[p])
|
||||
votes[key] = votes.get(key, 0) + 1
|
||||
top = max(votes.values())
|
||||
cons.append(min(k for k, v in votes.items() if v == top))
|
||||
return cons
|
||||
|
||||
|
||||
def conformity_scores(outputs: list[list[str]], consensus: list[str]) -> np.ndarray:
|
||||
"""Each agent's agreement rate with the population consensus (the E11 conformity signal)."""
|
||||
return np.array([np.mean([_normalise(o) == c for o, c in zip(out, consensus)])
|
||||
for out in outputs], dtype=float)
|
||||
|
||||
|
||||
def behavioural_distance(outputs: list[list[str]]) -> np.ndarray:
|
||||
"""Pairwise disagreement rate between agents' normalised answers (verifier-free diversity)."""
|
||||
n = len(outputs)
|
||||
norm = [[_normalise(o) for o in out] for out in outputs]
|
||||
d = np.zeros((n, n))
|
||||
for i in range(n):
|
||||
for j in range(i + 1, n):
|
||||
d[i, j] = d[j, i] = float(np.mean([a != b for a, b in zip(norm[i], norm[j])]))
|
||||
return d
|
||||
|
||||
|
||||
def novelty(dist: np.ndarray) -> np.ndarray:
|
||||
"""Per-agent mean behavioural distance to the rest of the pool (E11's ``_novelty``)."""
|
||||
n = dist.shape[0]
|
||||
if n < 2:
|
||||
return np.zeros(n)
|
||||
return dist.sum(axis=1) / (n - 1)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ selection
|
||||
|
||||
|
||||
def select_parents(scores: np.ndarray, dist: np.ndarray, k: int, *, diversity: bool,
|
||||
lam: float = 0.3) -> list[int]:
|
||||
"""v1 parent truncation (kept for the v1 code path and its tests): greedy QD or plain top-k."""
|
||||
if not diversity:
|
||||
return list(np.argsort(-scores, kind="stable")[:k])
|
||||
chosen = [int(np.argmax(scores))]
|
||||
while len(chosen) < k:
|
||||
best_i, best_v = -1, -np.inf
|
||||
for i in range(len(scores)):
|
||||
if i in chosen:
|
||||
continue
|
||||
v = scores[i] + lam * float(np.mean([dist[i, j] for j in chosen]))
|
||||
if v > best_v:
|
||||
best_i, best_v = i, v
|
||||
chosen.append(best_i)
|
||||
return chosen
|
||||
|
||||
|
||||
def pooled_survival(scores: np.ndarray, dist: np.ndarray, n_keep: int, *, lam: float) -> list[int]:
|
||||
"""E11 survival: keep the top ``n_keep`` of the pool by ``score + lam * novelty``.
|
||||
|
||||
``lam = 0`` is the greedy (``no_diversity``) arm. Parents and children compete on equal terms, so
|
||||
a strong parent survives by out-scoring its children (elitism is emergent, not a knob) and no
|
||||
lineage is excluded before it has bred.
|
||||
|
||||
Args:
|
||||
scores (np.ndarray): per-member selection score over the pooled parents + children.
|
||||
dist (np.ndarray): pairwise behavioural distance over the pool.
|
||||
n_keep (int): population size to keep.
|
||||
lam (float): novelty weight (E11 ``novelty``).
|
||||
|
||||
Returns:
|
||||
list[int]: indices into the pool, best merit first (deterministic).
|
||||
"""
|
||||
merit = np.asarray(scores, dtype=float) + lam * novelty(dist)
|
||||
order = np.argsort(-merit, kind="stable")
|
||||
return [int(i) for i in order[:n_keep]]
|
||||
|
||||
|
||||
def choose_single_parent(scores: np.ndarray, rng: np.random.Generator) -> int:
|
||||
"""Score-proportional parent sampling (``no_sex``): selection without recombination."""
|
||||
s = np.asarray(scores, dtype=float)
|
||||
w = s - s.min() + 1e-6
|
||||
return int(rng.choice(len(s), p=w / w.sum()))
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ mating and recombination
|
||||
|
||||
|
||||
def mating_plan(dist: np.ndarray, n_pairs: int, *, max_use: int = 2) -> list[tuple[int, int]]:
|
||||
"""Complementary pairing over the whole population with a per-agent use cap.
|
||||
|
||||
Pairs are taken in descending behavioural distance (directed mate choice, E10) subject to each
|
||||
agent appearing in at most ``max_use`` pairs, so no single agent monopolises reproduction and
|
||||
every agent's knowledge has a route to the next generation. If the cap exhausts the candidates
|
||||
before ``n_pairs`` is reached, the remaining slots cycle the most-distant pairs (degenerate but
|
||||
never empty).
|
||||
"""
|
||||
n = dist.shape[0]
|
||||
cands = sorted(((i, j) for i in range(n) for j in range(i + 1, n)), key=lambda ij: -dist[ij])
|
||||
if not cands:
|
||||
return [(0, 0)] * n_pairs
|
||||
use = np.zeros(n, dtype=int)
|
||||
plan: list[tuple[int, int]] = []
|
||||
for i, j in cands:
|
||||
if len(plan) >= n_pairs:
|
||||
break
|
||||
if use[i] < max_use and use[j] < max_use:
|
||||
plan.append((i, j)); use[i] += 1; use[j] += 1
|
||||
k = 0
|
||||
while len(plan) < n_pairs: # cap exhausted: cycle the best pairs
|
||||
plan.append(cands[k % len(cands)]); k += 1
|
||||
return plan
|
||||
|
||||
|
||||
def route_union(ans_a: list[str], conf_a: np.ndarray, ans_b: list[str],
|
||||
conf_b: np.ndarray) -> tuple[list[str], np.ndarray]:
|
||||
"""Union-preserving recombination of two parents' answer sets: per prompt, the more confident wins.
|
||||
|
||||
Confidence is the parent's own self-certainty (exp mean token log-prob) — no verifier. Ties go to
|
||||
parent A (deterministic).
|
||||
|
||||
Returns:
|
||||
(answers, source): the child's inheritance answers and a 0/1 array naming the parent per
|
||||
prompt (for the H6 diagnostic of *where* a skill is lost).
|
||||
"""
|
||||
src = (np.asarray(conf_b) > np.asarray(conf_a)).astype(int)
|
||||
out = [b if s else a for a, b, s in zip(ans_a, ans_b, src)]
|
||||
return out, src
|
||||
|
||||
|
||||
def complementary_pairs(parents: list[int], dist: np.ndarray, n_children: int) -> list[tuple[int, int]]:
|
||||
"""v1 mating plan over a parent subset (kept for the v1 code path and its tests)."""
|
||||
pairs = sorted(((a, b) for i, a in enumerate(parents) for b in parents[i + 1:]),
|
||||
key=lambda ab: -dist[ab[0], ab[1]])
|
||||
if not pairs:
|
||||
pairs = [(parents[0], parents[0])]
|
||||
return [pairs[i % len(pairs)] for i in range(n_children)]
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ arms and readouts
|
||||
|
||||
|
||||
def arm_settings(arm: str, g: float) -> dict:
|
||||
"""Resolve an arm name to its operator switches.
|
||||
|
||||
``sex`` is ``"union"`` (confidence-routed union inheritance, the v2 default), ``"linear"`` (v1's
|
||||
screened 2-parent LoRA blend — the H2 control arm ``sex_linear``), or ``None``.
|
||||
"""
|
||||
table = {
|
||||
"full": {"g": g, "sex": "union", "diversity": True},
|
||||
"no_grounding": {"g": 0.0, "sex": "union", "diversity": True},
|
||||
"no_sex": {"g": g, "sex": None, "diversity": True},
|
||||
"no_diversity": {"g": g, "sex": "union", "diversity": False},
|
||||
"sex_linear": {"g": g, "sex": "linear", "diversity": True},
|
||||
}
|
||||
if arm not in table:
|
||||
raise ValueError(f"unknown arm {arm!r} (expected one of {list(table)})")
|
||||
return dict(table[arm])
|
||||
|
||||
|
||||
def fitness_of(outs: list[str], tasks: list[Task], fams: list[str]) -> dict:
|
||||
"""Overall / per-family / worst-family verifier accuracy from precomputed outputs."""
|
||||
corr = np.array([verify(o, t) for o, t in zip(outs, tasks)])
|
||||
fam = np.array([t.family for t in tasks])
|
||||
acc = {"overall": float(corr.mean())}
|
||||
acc.update({f: float(corr[fam == f].mean()) if (fam == f).any() else float("nan") for f in fams})
|
||||
acc["worst_family"] = min(acc[f] for f in fams)
|
||||
return acc
|
||||
|
||||
|
||||
def families_alive(per_family_acc: list[dict], fams: list[str], threshold: float = 0.6) -> int:
|
||||
"""Number of families on which at least one agent is competent (≥ ``threshold``) — the count of
|
||||
'alleles' still present in the population; loss is permanent (E6)."""
|
||||
return int(sum(any(a.get(f, 0.0) >= threshold for a in per_family_acc) for f in fams))
|
||||
|
||||
|
||||
def answer_of(raw: str) -> str:
|
||||
"""Trim a raw completion to a single short answer line for the inheritance data."""
|
||||
return raw.strip().split("\n")[0][:64].strip()
|
||||
319
src/llm/society_v2.py
Normal file
319
src/llm/society_v2.py
Normal file
|
|
@ -0,0 +1,319 @@
|
|||
"""The composed society at LLM scale, v2 — E11 re-instantiated faithfully (``kind: llm_society_v2``).
|
||||
|
||||
Pre-registered in ``tasks/prereg-llm-society-v2.md``; this module is its §3. What changed from v1 and
|
||||
why is in that document's §1. In one paragraph: ``L`` disjoint task families and **one founder per
|
||||
family** (ρ = 0 by construction); recombination is a **confidence-routed union** of two parents'
|
||||
inheritance answers (E4's ``max`` per item, verifier-free) rather than a linear LoRA blend; selection
|
||||
acts on **survival over the pooled parents + children** by ``score + λ·novelty`` (E11's rule) rather
|
||||
than on breeding eligibility; the run **checkpoints every generation** and resumes.
|
||||
|
||||
One generation (§3.4): produce & score the population → mating plan by complementarity → each child's
|
||||
inheritance data from its parents' own answers (union / linear / single) → source diagnostics →
|
||||
train N fresh LoRAs → score the children against the *parents'* consensus → keep the top N of the
|
||||
2N pool → checkpoint.
|
||||
|
||||
Row schema (long form): ``experiment, arm, seed, generation, agent, role, parents, selected,
|
||||
conformity, score, metric, value``. ``role`` ∈ {population, child, child_source, summary}.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
from . import families as _families # noqa: F401 (registers the extra task families)
|
||||
from .society_ops import (answer_of, arm_settings, behavioural_distance, choose_single_parent,
|
||||
conformity_scores, consensus_answers, families_alive, fitness_of,
|
||||
mating_plan, novelty, pooled_survival, route_union)
|
||||
from .tasks import Task, make_tasks, verify, _normalise
|
||||
|
||||
|
||||
def _pool(fams: list[str], per_family: int, seed: int, hard: bool) -> list[Task]:
|
||||
return sum([make_tasks(f, per_family, seed=seed + i, hard=hard) for i, f in enumerate(fams)], [])
|
||||
|
||||
|
||||
def _row(base: dict, **kw) -> dict:
|
||||
r = dict(base); r.update(kw); return r
|
||||
|
||||
|
||||
def _train_founder_locked(d: Path, train, timeout_s: int = 3600) -> None:
|
||||
"""Train the founder at ``d`` exactly once across concurrent jobs (lock file, O_EXCL)."""
|
||||
import os, time
|
||||
done = d / "adapter_config.json"
|
||||
if done.exists():
|
||||
return
|
||||
d.parent.mkdir(parents=True, exist_ok=True)
|
||||
lock = d.parent / (d.name + ".lock")
|
||||
try:
|
||||
fd = os.open(lock, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
|
||||
except FileExistsError:
|
||||
t0 = time.time()
|
||||
while not done.exists(): # another job is training it
|
||||
if time.time() - t0 > timeout_s:
|
||||
raise TimeoutError(f"waited {timeout_s}s for founder {d}")
|
||||
time.sleep(15)
|
||||
return
|
||||
try:
|
||||
os.write(fd, str(os.getpid()).encode()); os.close(fd)
|
||||
train()
|
||||
finally:
|
||||
lock.unlink(missing_ok=True)
|
||||
|
||||
|
||||
class _State:
|
||||
"""Per-arm checkpoint: the current population's adapter dirs and the generation reached."""
|
||||
|
||||
def __init__(self, path: Path):
|
||||
self.path = path
|
||||
|
||||
def load(self) -> dict | None:
|
||||
return json.loads(self.path.read_text()) if self.path.exists() else None
|
||||
|
||||
def save(self, generation: int, agents: list[str]) -> None:
|
||||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self.path.write_text(json.dumps({"generation": generation, "agents": agents}))
|
||||
|
||||
|
||||
def run_society_v2(cfg: dict) -> pd.DataFrame:
|
||||
"""Run every configured arm of the v2 society; return tidy rows (resumable per arm)."""
|
||||
import torch
|
||||
|
||||
from .epistasis import generate_with_confidence
|
||||
from .evaluate import generate
|
||||
from .merge import load_specialists
|
||||
from .specialise import train_lora_on_tasks, train_specialist
|
||||
|
||||
name = cfg["experiment"]
|
||||
base = cfg["base_model"]
|
||||
fams = list(cfg["families"])
|
||||
L = len(fams)
|
||||
N = int(cfg.get("agents", L))
|
||||
G = int(cfg.get("generations", 12))
|
||||
arms = list(cfg.get("arms", ["full", "no_grounding", "no_sex", "no_diversity"]))
|
||||
g_val = float(cfg.get("g", 0.85))
|
||||
lam = float(cfg.get("lam", 0.3))
|
||||
n_test = int(cfg.get("n_test", 20)) # per family
|
||||
n_val = int(cfg.get("n_val", 10)) # per family
|
||||
n_conf = int(cfg.get("n_conf", 10)) # per family, fresh each generation
|
||||
k_inh = int(cfg.get("k_inherit", 100)) # per family, fresh each generation (gate C2)
|
||||
epochs = int(cfg.get("epochs", 3))
|
||||
n_cand = int(cfg.get("n_candidates", 6)) # sex_linear only
|
||||
spec_train = int(cfg.get("spec_train", 600))
|
||||
spec_epochs = int(cfg.get("spec_epochs", 3))
|
||||
hard = bool(cfg.get("hard", False))
|
||||
resume = bool(cfg.get("resume", True))
|
||||
seed = int(cfg["seed"])
|
||||
lora = cfg.get("lora", {})
|
||||
r, alpha = int(lora.get("r", 16)), int(lora.get("alpha", 32))
|
||||
root = Path(cfg.get("adapters_dir", "models/llm")) / "society_v2" / f"{name}_s{seed}"
|
||||
out_dir = Path(cfg.get("output", {}).get("dir", f"results/{name}"))
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
max_use = int(cfg.get("max_mate_use", 2))
|
||||
# Confidence-gated inheritance (prereg §4a, C2b): a child learns only the prompts its source is
|
||||
# confident on (exp mean token log-prob ≥ conf_gate). Verifier-free; identical in every arm;
|
||||
# None = ungated (the pre-registered v2 default, which C2 showed loses 20-40%/generation).
|
||||
conf_gate = cfg.get("conf_gate")
|
||||
conf_gate = None if conf_gate is None else float(conf_gate)
|
||||
|
||||
test = _pool(fams, n_test, 1000, hard)
|
||||
val = _pool(fams, n_val, 3000, hard)
|
||||
|
||||
# ---- founders: one specialist per family (agent i -> family i mod L), cached and shared by arms.
|
||||
# Arm-jobs of one seed may start together on the cluster: an O_EXCL lock makes the first train
|
||||
# and the rest wait on the finished adapter, so no two jobs write the same founder.
|
||||
founders = []
|
||||
for i in range(N):
|
||||
fam = fams[i % L]
|
||||
d = root / "founders" / f"agent{i}_{fam}"
|
||||
_train_founder_locked(d, lambda d=d, fam=fam, i=i: train_specialist(
|
||||
base, fam, str(d), n_train=spec_train, epochs=spec_epochs, seed=seed * 100 + i,
|
||||
hard=hard, r=r, alpha=alpha))
|
||||
founders.append(str(d))
|
||||
|
||||
frames: list[pd.DataFrame] = []
|
||||
for a_idx, arm in enumerate(arms):
|
||||
s = arm_settings(arm, g_val)
|
||||
state = _State(root / arm / "state.json")
|
||||
partial = out_dir / f"partial_{arm}_s{seed}.parquet"
|
||||
rows: list[dict] = []
|
||||
agents, t0 = list(founders), 0
|
||||
st = state.load() if resume else None
|
||||
if st and all(Path(d, "adapter_config.json").exists() for d in st["agents"]):
|
||||
agents, t0 = list(st["agents"]), int(st["generation"])
|
||||
if partial.exists():
|
||||
rows = pd.read_parquet(partial).to_dict("records")
|
||||
print(f"[{arm}] resuming at generation {t0}")
|
||||
|
||||
for t in range(t0, G + 1):
|
||||
rng = np.random.default_rng([seed, a_idx, t]) # resume-safe per-generation RNG
|
||||
conf_pool = _pool(fams, n_conf, seed * 7919 + t * 13, hard)
|
||||
conf_prompts = [x.prompt for x in conf_pool]
|
||||
tag = {"experiment": name, "arm": arm, "seed": seed, "generation": t}
|
||||
|
||||
# ---- (1) produce & score the population
|
||||
model, tok = load_specialists(base, agents)
|
||||
fit_test, fit_val, conf_outs = [], [], []
|
||||
for i in range(N):
|
||||
model.set_adapter(f"a{i}")
|
||||
fit_test.append(fitness_of(generate(model, tok, [x.prompt for x in test]), test, fams))
|
||||
fit_val.append(fitness_of(generate(model, tok, [x.prompt for x in val]), val, fams))
|
||||
conf_outs.append(generate(model, tok, conf_prompts))
|
||||
consensus = consensus_answers(conf_outs)
|
||||
conf = conformity_scores(conf_outs, consensus)
|
||||
dist = behavioural_distance(conf_outs)
|
||||
nov = novelty(dist)
|
||||
cons_acc = float(np.mean([verify(c, x) for c, x in zip(consensus, conf_pool)]))
|
||||
fitness = np.array([fv["overall"] for fv in fit_val])
|
||||
scores = s["g"] * fitness + (1.0 - s["g"]) * conf
|
||||
|
||||
for i in range(N):
|
||||
b = _row(tag, agent=i, role="population", parents=Path(agents[i]).name,
|
||||
selected=True, conformity=float(conf[i]), score=float(scores[i]))
|
||||
for k, v in fit_test[i].items():
|
||||
rows.append(_row(b, metric=f"test_{k}", value=v))
|
||||
rows.append(_row(b, metric="val_overall", value=float(fitness[i])))
|
||||
rows.append(_row(b, metric="novelty", value=float(nov[i])))
|
||||
summ = _row(tag, agent=-1, role="summary", parents="", selected=False,
|
||||
conformity=float("nan"), score=float("nan"))
|
||||
rows.append(_row(summ, metric="consensus_acc", value=cons_acc))
|
||||
rows.append(_row(summ, metric="diversity_behav", value=float(nov.mean())))
|
||||
rows.append(_row(summ, metric="families_alive", value=float(families_alive(fit_test, fams))))
|
||||
rows.append(_row(summ, metric="gap_conformity_minus_truth",
|
||||
value=float(conf.mean() - np.mean([f["overall"] for f in fit_test]))))
|
||||
if t == G:
|
||||
del model; torch.cuda.empty_cache()
|
||||
break
|
||||
|
||||
# ---- (2)-(3) mating plan and each child's inheritance data
|
||||
inherit = _pool(fams, k_inh, seed * 104729 + t * 17, hard)
|
||||
inh_prompts = [x.prompt for x in inherit]
|
||||
child_src: list[dict] = []
|
||||
if s["sex"] == "union":
|
||||
plan = mating_plan(dist, N, max_use=max_use)
|
||||
needed = sorted({p for ab in plan for p in ab})
|
||||
ans, cf = {}, {}
|
||||
for p in needed: # each parent answers the pool once
|
||||
model.set_adapter(f"a{p}")
|
||||
ans[p], cf[p] = generate_with_confidence(model, tok, inh_prompts)
|
||||
for pa, pb in plan:
|
||||
routed, src = route_union(ans[pa], cf[pa], ans[pb], cf[pb])
|
||||
keep = np.maximum(cf[pa], cf[pb]) >= conf_gate if conf_gate is not None else None
|
||||
child_src.append({"parents": (pa, pb), "answers": routed,
|
||||
"share_b": float(src.mean()), "keep": keep})
|
||||
elif s["sex"] == "linear":
|
||||
from .directed import sample_merge_weights
|
||||
plan = mating_plan(dist, N, max_use=max_use)
|
||||
for c, (pa, pb) in enumerate(plan):
|
||||
w = sample_merge_weights(2, n_cand, rng)
|
||||
best_i, best_v = 0, -np.inf
|
||||
for ci in range(n_cand):
|
||||
cname = f"g{t}c{c}k{ci}"
|
||||
model.add_weighted_adapter([f"a{pa}", f"a{pb}"], w[ci].tolist(), cname,
|
||||
combination_type="linear")
|
||||
model.set_adapter(cname)
|
||||
if s["g"] > 0:
|
||||
v = fitness_of(generate(model, tok, [x.prompt for x in val]), val,
|
||||
fams)["overall"]
|
||||
else:
|
||||
oc = generate(model, tok, conf_prompts)
|
||||
v = float(np.mean([_normalise(o) == cc for o, cc in zip(oc, consensus)]))
|
||||
if v > best_v:
|
||||
best_i, best_v = ci, v
|
||||
model.set_adapter(f"a{pa}"); model.delete_adapter(cname)
|
||||
cname = f"g{t}c{c}win"
|
||||
model.add_weighted_adapter([f"a{pa}", f"a{pb}"], w[best_i].tolist(), cname,
|
||||
combination_type="linear")
|
||||
model.set_adapter(cname)
|
||||
answers = generate(model, tok, inh_prompts)
|
||||
model.set_adapter(f"a{pa}"); model.delete_adapter(cname)
|
||||
child_src.append({"parents": (pa, pb), "answers": answers, "share_b": float("nan")})
|
||||
else: # no_sex: single parent, score-proportional
|
||||
picks = [choose_single_parent(scores, rng) for _ in range(N)]
|
||||
ans, cf = {}, {}
|
||||
for p in sorted(set(picks)):
|
||||
model.set_adapter(f"a{p}")
|
||||
ans[p], cf[p] = generate_with_confidence(model, tok, inh_prompts)
|
||||
for p in picks:
|
||||
keep = cf[p] >= conf_gate if conf_gate is not None else None
|
||||
child_src.append({"parents": (p, p), "answers": ans[p], "share_b": float("nan"),
|
||||
"keep": keep})
|
||||
|
||||
# ---- source diagnostics (H6): what each child was *supplied*, per family
|
||||
for c, src in enumerate(child_src):
|
||||
supplied = fitness_of(src["answers"], inherit, fams)
|
||||
b = _row(tag, agent=c, role="child_source", parents=f"{src['parents'][0]}+{src['parents'][1]}",
|
||||
selected=False, conformity=float("nan"), score=float("nan"))
|
||||
for k, v in supplied.items():
|
||||
rows.append(_row(b, metric=f"source_{k}", value=v))
|
||||
rows.append(_row(b, metric="source_share_b", value=src["share_b"]))
|
||||
del model; torch.cuda.empty_cache()
|
||||
|
||||
# ---- (4) inherit: a fresh LoRA per child on its source's own answers
|
||||
children, degenerate = [], 0
|
||||
for c, src in enumerate(child_src):
|
||||
keep = src.get("keep")
|
||||
if keep is None:
|
||||
keep = np.ones(len(inherit), dtype=bool)
|
||||
data = [Task(x.family, x.prompt, answer_of(a))
|
||||
for x, a, kp in zip(inherit, src["answers"], keep) if kp and answer_of(a)]
|
||||
rows.append(_row(tag, agent=c, role="child_source", parents=f"{src['parents'][0]}+{src['parents'][1]}",
|
||||
selected=False, conformity=float("nan"), score=float("nan"),
|
||||
metric="n_inherit_kept", value=float(len(data))))
|
||||
d = root / arm / f"gen{t + 1}" / f"child{c}"
|
||||
if len(data) >= 8:
|
||||
train_lora_on_tasks(base, data, str(d), epochs=epochs,
|
||||
seed=seed * 31 + t * N + c, r=r, alpha=alpha)
|
||||
else: # terminal degeneration: inherit unchanged
|
||||
shutil.copytree(agents[src["parents"][0]], d, dirs_exist_ok=True); degenerate += 1
|
||||
children.append(str(d))
|
||||
|
||||
# ---- (5) survive: score children against the PARENTS' consensus; keep top N of 2N
|
||||
model, tok = load_specialists(base, agents + children)
|
||||
c_test, c_val, c_conf = [], [], []
|
||||
for j in range(N):
|
||||
model.set_adapter(f"a{N + j}")
|
||||
c_test.append(fitness_of(generate(model, tok, [x.prompt for x in test]), test, fams))
|
||||
c_val.append(fitness_of(generate(model, tok, [x.prompt for x in val]), val, fams))
|
||||
c_conf.append(generate(model, tok, conf_prompts))
|
||||
del model; torch.cuda.empty_cache()
|
||||
conf_c = conformity_scores(c_conf, consensus)
|
||||
fit_c = np.array([fv["overall"] for fv in c_val])
|
||||
scores_c = s["g"] * fit_c + (1.0 - s["g"]) * conf_c
|
||||
pool_scores = np.concatenate([scores, scores_c])
|
||||
pool_dist = behavioural_distance(conf_outs + c_conf)
|
||||
keep = pooled_survival(pool_scores, pool_dist, N, lam=lam if s["diversity"] else 0.0)
|
||||
keep_set = set(keep)
|
||||
|
||||
for j in range(N):
|
||||
b = _row(tag, agent=j, role="child",
|
||||
parents=f"{child_src[j]['parents'][0]}+{child_src[j]['parents'][1]}",
|
||||
selected=(N + j) in keep_set, conformity=float(conf_c[j]),
|
||||
score=float(scores_c[j]))
|
||||
for k, v in c_test[j].items():
|
||||
rows.append(_row(b, metric=f"test_{k}", value=v))
|
||||
rows.append(_row(b, metric="val_overall", value=float(fit_c[j])))
|
||||
rows.append(_row(summ, metric="best_newborn_overall",
|
||||
value=float(max(f["overall"] for f in c_test))))
|
||||
rows.append(_row(summ, metric="n_degenerate", value=float(degenerate)))
|
||||
rows.append(_row(summ, metric="n_parents_survive",
|
||||
value=float(sum(1 for k in keep if k < N))))
|
||||
|
||||
pool_dirs = agents + children
|
||||
survivors = [pool_dirs[k] for k in keep]
|
||||
# disk hygiene: drop non-survivors (founders are shared across arms — keep them)
|
||||
for k, d in enumerate(pool_dirs):
|
||||
if k not in keep_set and not d.startswith(str(root / "founders")):
|
||||
shutil.rmtree(d, ignore_errors=True)
|
||||
agents = survivors
|
||||
state.save(t + 1, agents)
|
||||
pd.DataFrame(rows).to_parquet(partial, index=False)
|
||||
print(f"[{arm}] gen {t + 1}/{G}: best {max(f['overall'] for f in fit_test):.3f} "
|
||||
f"newborn {max(f['overall'] for f in c_test):.3f} cons {cons_acc:.2f} "
|
||||
f"alive {families_alive(fit_test, fams)} deg {degenerate}", flush=True)
|
||||
|
||||
frames.append(pd.DataFrame(rows))
|
||||
return pd.concat(frames, ignore_index=True)
|
||||
|
|
@ -18,9 +18,8 @@ from .tasks import make_tasks
|
|||
|
||||
def _encode(tok, task, device):
|
||||
"""Return (input_ids, labels) for one example with the prompt tokens masked out of the loss."""
|
||||
text = tok.apply_chat_template([{"role": "user", "content": task.prompt}],
|
||||
add_generation_prompt=True, tokenize=False)
|
||||
prompt_ids = tok(text, add_special_tokens=False).input_ids
|
||||
from .evaluate import format_prompt
|
||||
prompt_ids = tok(format_prompt(tok, task.prompt), add_special_tokens=False).input_ids
|
||||
answer_ids = tok(task.answer + tok.eos_token, add_special_tokens=False).input_ids
|
||||
ids = prompt_ids + answer_ids
|
||||
labels = [-100] * len(prompt_ids) + answer_ids
|
||||
|
|
@ -29,7 +28,7 @@ def _encode(tok, task, device):
|
|||
|
||||
def train_lora_on_tasks(base_name: str, tasks: list, out_dir: str, *, epochs: int = 3,
|
||||
lr: float = 2e-4, batch_size: int = 8, r: int = 16, alpha: int = 32,
|
||||
seed: int = 0, device: str = "cuda") -> str:
|
||||
seed: int = 0, device: str = "cuda", max_len: int | None = None) -> str:
|
||||
"""Fine-tune a fresh LoRA adapter on an arbitrary list of ``tasks`` and save it to ``out_dir``.
|
||||
|
||||
The reusable answer-only SFT primitive: each task supplies a ``.prompt`` and a ``.answer`` (which
|
||||
|
|
@ -60,9 +59,20 @@ def train_lora_on_tasks(base_name: str, tasks: list, out_dir: str, *, epochs: in
|
|||
target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
|
||||
"gate_proj", "up_proj", "down_proj"])
|
||||
model = get_peft_model(model, lora)
|
||||
return _sft(model, tok, tasks, epochs=epochs, lr=lr, batch_size=batch_size, seed=seed,
|
||||
device=device, max_len=max_len, out_dir=out_dir)
|
||||
|
||||
|
||||
def _sft(model, tok, tasks, *, epochs, lr, batch_size, seed, device, max_len, out_dir):
|
||||
"""The shared answer-only SFT loop. ``model`` is an already-prepared PEFT model."""
|
||||
import torch
|
||||
|
||||
model.train()
|
||||
|
||||
encoded = [_encode(tok, t, device) for t in tasks]
|
||||
if max_len: # cap sequence length: the causal-LM loss
|
||||
encoded = [(i[:max_len], l[:max_len]) for i, l in encoded] # upcasts logits to float32,
|
||||
# so memory ~ batch*len*vocab*4
|
||||
opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=lr)
|
||||
rng = np.random.default_rng(seed)
|
||||
pad = tok.pad_token_id
|
||||
|
|
@ -101,3 +111,36 @@ def train_specialist(base_name: str, family: str, out_dir: str, *, n_train: int
|
|||
tasks = make_tasks(family, n_train, seed=seed, hard=hard)
|
||||
return train_lora_on_tasks(base_name, tasks, out_dir, epochs=epochs, lr=lr,
|
||||
batch_size=batch_size, r=r, alpha=alpha, seed=seed, device=device)
|
||||
|
||||
|
||||
def continue_lora_training(base_name: str, adapter_dir: str, tasks: list, out_dir: str, *,
|
||||
epochs: int = 3, lr: float = 1e-4, batch_size: int = 2, seed: int = 0,
|
||||
device: str = "cuda", max_len: int | None = 448) -> str:
|
||||
"""Continue training an **existing** adapter on new data — the Lamarckian channel.
|
||||
|
||||
``train_lora_on_tasks`` builds a fresh adapter from the frozen base, so knowledge survives only
|
||||
through the data (Weismannian, and correct for the collapse experiments). Continual learning needs
|
||||
the opposite: a child starts from its parent's weights and adds to them, so what a lineage
|
||||
acquires is inherited as *structure*. The learning rate defaults lower than a fresh specialist's,
|
||||
because the adapter is already in a good region and full-rate updates overwrite it.
|
||||
|
||||
Args:
|
||||
base_name (str): HF id of the frozen base.
|
||||
adapter_dir (str): the parent adapter to start from.
|
||||
tasks (list): objects with ``.prompt`` and ``.answer``.
|
||||
out_dir (str): where to save the child adapter.
|
||||
|
||||
Returns:
|
||||
str: ``out_dir``.
|
||||
"""
|
||||
import torch
|
||||
from peft import PeftModel
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
tok = AutoTokenizer.from_pretrained(base_name)
|
||||
if tok.pad_token_id is None:
|
||||
tok.pad_token = tok.eos_token
|
||||
base = AutoModelForCausalLM.from_pretrained(base_name, dtype=torch.bfloat16).to(device)
|
||||
model = PeftModel.from_pretrained(base, adapter_dir, is_trainable=True)
|
||||
return _sft(model, tok, tasks, epochs=epochs, lr=lr, batch_size=batch_size, seed=seed,
|
||||
device=device, max_len=max_len, out_dir=out_dir)
|
||||
|
|
|
|||
|
|
@ -73,6 +73,15 @@ def _acc(model, tok, tasks: list[Task]) -> float:
|
|||
return float(np.mean([verify(o, t) for o, t in zip(outs, tasks)]))
|
||||
|
||||
|
||||
|
||||
def adapter_root(cfg: dict) -> Path:
|
||||
"""Seed-specific adapter directory for the speciation sweeps.
|
||||
|
||||
Children are retrained at every sweep point, so the directory is scratch; it is keyed by seed so
|
||||
that seeds running concurrently (an HPC array) never overwrite each other's adapters.
|
||||
"""
|
||||
return Path(cfg.get("adapters_dir", "models/llm")) / f"speciation_s{int(cfg['seed'])}"
|
||||
|
||||
def run_speciation_experiment(cfg: dict) -> pd.DataFrame:
|
||||
"""Run the conflict-cliff and/or duration sweeps; return long-form accuracies.
|
||||
|
||||
|
|
@ -91,7 +100,7 @@ def run_speciation_experiment(cfg: dict) -> pd.DataFrame:
|
|||
r, alpha = int(lora.get("r", 16)), int(lora.get("alpha", 32))
|
||||
seed = int(cfg["seed"])
|
||||
hard = bool(cfg.get("hard", False))
|
||||
root = Path(cfg.get("adapters_dir", "models/llm")) / "speciation"
|
||||
root = adapter_root(cfg)
|
||||
|
||||
# Fixed evaluation sets (identical across the sweep; convention pairs grade the SAME prompts).
|
||||
test_a = make_tasks(fam_a, n_test, seed=1000, hard=hard)
|
||||
|
|
|
|||
|
|
@ -31,11 +31,16 @@ _WORDS = ("apple", "table", "river", "cloud", "stone", "plant", "music", "green"
|
|||
|
||||
@dataclass(frozen=True)
|
||||
class Task:
|
||||
"""One verifiable task: a prompt, its canonical answer, and its family."""
|
||||
"""One verifiable task: a prompt, its canonical answer, and its family.
|
||||
|
||||
``meta`` is optional verifier context for dataset-backed families (JSON: unit tests for code,
|
||||
answer aliases for span/short-text QA). Procedural families leave it ``None``.
|
||||
"""
|
||||
|
||||
family: str
|
||||
prompt: str
|
||||
answer: str
|
||||
meta: str | None = None
|
||||
|
||||
|
||||
def _fmt_list(xs) -> str:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue