Giorgio Gilestro
84124de143
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
232 lines
13 KiB
Python
232 lines
13 KiB
Python
"""The composed society at LLM scale (C3) — E11 re-instantiated in a population of LoRA agents.
|
|
|
|
A population of `N` agents (LoRA adapters on a shared frozen base) evolves for `G` non-overlapping
|
|
generations under the four operators the paper composes: grounded evaluation, directed recombination
|
|
(sex), diversity-preserving selection, and retraining (mutation). The single grounding knob acts in
|
|
the *evaluation* channel, exactly as in E11: selection scores each agent by
|
|
``g * verifier_fitness + (1 - g) * conformity``, where conformity is agreement with the population's
|
|
own modal answer. The inheritance channel is identical in every arm and deliberately ungrounded —
|
|
each child is a fresh LoRA distilled from its source model's *own answers* (self-consumption made
|
|
literal), so knowledge survives only through the data channel.
|
|
|
|
Arms (four-arm ablation, mirroring E11): ``full`` / ``no_grounding`` (g=0; the verifier never enters
|
|
that arm's loop — offspring screening also falls back to conformity) / ``no_sex`` (children are
|
|
redistilled copies of selected parents) / ``no_diversity`` (plain top-P selection).
|
|
|
|
Pure, testable pieces live at module top (consensus, conformity, behavioural distance,
|
|
quality-diversity selection, complementary pairing); the GPU loop is
|
|
:func:`run_society_experiment`. ``python -m llm.experiment configs/llm/society_smoke.yaml``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import shutil
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
|
|
from .tasks import Task, make_tasks, _normalise
|
|
|
|
|
|
# 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:
|
|
"""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 run_society_experiment(cfg: dict) -> pd.DataFrame:
|
|
"""Run the society loop for every configured arm; return tidy long-form rows.
|
|
|
|
Config keys (with defaults): ``agents`` (6), ``generations`` (8), ``arms`` (all four),
|
|
``g`` (0.5), ``lam`` (0.3), ``n_test``/``n_val``/``n_conf``/``n_inherit`` pool sizes,
|
|
``n_candidates`` (6) offspring screened per pair, ``epochs`` (2) child SFT epochs,
|
|
``spec_train``/``spec_epochs`` gen-0 specialist budget, ``hard`` (False),
|
|
``keep_all_adapters`` (False).
|
|
"""
|
|
import torch
|
|
|
|
from .evaluate import generate, load_model
|
|
from .merge import load_specialists
|
|
from .specialise import train_lora_on_tasks, train_specialist
|
|
|
|
name = cfg["experiment"]
|
|
base = cfg["base_model"]
|
|
fams = list(cfg.get("families", ["lists", "strings", "arith"]))
|
|
N = int(cfg.get("agents", 6))
|
|
G = int(cfg.get("generations", 8))
|
|
arms = list(cfg.get("arms", ["full", "no_grounding", "no_sex", "no_diversity"]))
|
|
g_val = float(cfg.get("g", 0.5))
|
|
lam = float(cfg.get("lam", 0.3))
|
|
n_test = int(cfg.get("n_test", 40))
|
|
n_val = int(cfg.get("n_val", 30))
|
|
n_conf = int(cfg.get("n_conf", 60))
|
|
n_inherit = int(cfg.get("n_inherit", 240))
|
|
n_cand = int(cfg.get("n_candidates", 6))
|
|
epochs = int(cfg.get("epochs", 2))
|
|
elitism = int(cfg.get("elitism", 0))
|
|
n_parents = int(cfg.get("n_parents", max(2, N // 2)))
|
|
spec_train = int(cfg.get("spec_train", 300))
|
|
spec_epochs = int(cfg.get("spec_epochs", 2))
|
|
hard = bool(cfg.get("hard", False))
|
|
keep_all = bool(cfg.get("keep_all_adapters", False))
|
|
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" / f"{name}_s{seed}"
|
|
|
|
# fixed pools: test (reporting only), val (grounded selection signal + offspring screening)
|
|
test = sum([make_tasks(f, n_test, seed=1000 + i, hard=hard) for i, f in enumerate(fams)], [])
|
|
val = sum([make_tasks(f, n_val, seed=3000 + i, hard=hard) for i, f in enumerate(fams)], [])
|
|
rng = np.random.default_rng(seed)
|
|
|
|
# gen-0 founders: light per-family specialists, shared across arms (same starting population)
|
|
founders = []
|
|
for i in range(N):
|
|
fam = fams[i % len(fams)]
|
|
d = root / "founders" / f"agent{i}_{fam}"
|
|
if not (d / "adapter_config.json").exists():
|
|
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))
|
|
|
|
rows: list[dict] = []
|
|
for arm in arms:
|
|
s = arm_settings(arm, g_val)
|
|
agents = list(founders)
|
|
for t in range(G):
|
|
conf_pool = sum([make_tasks(f, n_conf // len(fams), seed=seed * 7919 + t * 13 + i,
|
|
hard=hard) for i, f in enumerate(fams)], [])
|
|
conf_prompts = [x.prompt for x in conf_pool]
|
|
|
|
# ---- produce & score: one base, all agents attached as adapters
|
|
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(generate(model, tok, [x.prompt for x in test]), test, fams))
|
|
fit_val.append(_fitness(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)
|
|
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
|
|
|
|
parents = select_parents(scores, dist, n_parents, diversity=s["diversity"], lam=lam)
|
|
|
|
# ---- rows (reporting uses the verifier in every arm; the loop does not)
|
|
for i in range(N):
|
|
base_row = {"experiment": name, "arm": arm, "seed": seed, "generation": t,
|
|
"agent": i, "selected": i in parents,
|
|
"conformity": float(conf[i]), "score": float(scores[i])}
|
|
for k, v in fit_test[i].items():
|
|
rows.append({**base_row, "metric": f"test_{k}", "value": v})
|
|
rows.append({**base_row, "metric": "val_overall", "value": float(fitness[i])})
|
|
rows.append({"experiment": name, "arm": arm, "seed": seed, "generation": t,
|
|
"agent": -1, "selected": False, "conformity": float("nan"),
|
|
"score": float("nan"), "metric": "consensus_acc", "value": cons_acc})
|
|
rows.append({"experiment": name, "arm": arm, "seed": seed, "generation": t,
|
|
"agent": -1, "selected": False, "conformity": float("nan"),
|
|
"score": float("nan"), "metric": "diversity_behav",
|
|
"value": float(dist[np.triu_indices(N, 1)].mean())})
|
|
|
|
if t == G - 1:
|
|
del model
|
|
torch.cuda.empty_cache()
|
|
break
|
|
|
|
# ---- breed: pick each child's source (merged offspring, or a copied parent)
|
|
inherit_pool = sum([make_tasks(f, n_inherit // len(fams),
|
|
seed=seed * 104729 + t * 17 + i, hard=hard)
|
|
for i, f in enumerate(fams)], [])
|
|
n_bred = N - elitism
|
|
child_sources: list[dict] = [] # per child: parent names + weights
|
|
if s["sex"]:
|
|
from .directed import sample_merge_weights
|
|
pairs = complementary_pairs(parents, dist, n_bred)
|
|
for c, (pa, pb) in enumerate(pairs):
|
|
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: # grounded screening: verifier on val
|
|
v = _fitness(generate(model, tok, [x.prompt for x in val]),
|
|
val, fams)["overall"]
|
|
else: # ungrounded screening: conformity only
|
|
outs_c = generate(model, tok, conf_prompts)
|
|
v = float(np.mean([_normalise(o) == cc
|
|
for o, cc in zip(outs_c, consensus)]))
|
|
if v > best_v:
|
|
best_i, best_v = ci, v
|
|
model.set_adapter(f"a{pa}") # never delete the active adapter
|
|
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, [x.prompt for x in inherit_pool])
|
|
model.set_adapter(f"a{pa}")
|
|
model.delete_adapter(cname)
|
|
child_sources.append({"parents": (pa, pb), "answers": answers})
|
|
else:
|
|
for c in range(n_bred):
|
|
p = parents[c % len(parents)]
|
|
model.set_adapter(f"a{p}")
|
|
answers = generate(model, tok, [x.prompt for x in inherit_pool])
|
|
child_sources.append({"parents": (p, p), "answers": answers})
|
|
del model
|
|
torch.cuda.empty_cache()
|
|
|
|
# ---- reproduce: fresh LoRA per child, trained on its source's own answers.
|
|
# Elites (overlapping generations): the top-scoring parents survive as unmodified
|
|
# copies, applied identically in every arm — reproduction here is lossy distillation,
|
|
# so without a survivor the mutation load erases the champion each generation.
|
|
new_agents = []
|
|
for e in range(elitism):
|
|
d = root / arm / f"gen{t + 1}" / f"elite{e}"
|
|
shutil.copytree(agents[parents[e]], d, dirs_exist_ok=True)
|
|
new_agents.append(str(d))
|
|
rows.append({"experiment": name, "arm": arm, "seed": seed, "generation": t + 1,
|
|
"agent": e, "selected": False, "conformity": float("nan"),
|
|
"score": float("nan"), "metric": "parents",
|
|
"value": float(parents[e] * 100 + parents[e])})
|
|
for c, src in enumerate(child_sources, start=elitism):
|
|
data = [Task(x.family, x.prompt, _answer_of(a))
|
|
for x, a in zip(inherit_pool, src["answers"]) if _answer_of(a)]
|
|
d = root / arm / f"gen{t + 1}" / f"agent{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: copy the source
|
|
# Reason: a fully-degenerate source emits no usable answers; crashing would kill
|
|
# a long sweep (cf. the neural terminal-collapse sentinel) — inherit unchanged.
|
|
shutil.copytree(agents[src["parents"][0]], d, dirs_exist_ok=True)
|
|
new_agents.append(str(d))
|
|
rows.append({"experiment": name, "arm": arm, "seed": seed, "generation": t + 1,
|
|
"agent": c, "selected": False, "conformity": float("nan"),
|
|
"score": float("nan"), "metric": "parents",
|
|
"value": float(src["parents"][0] * 100 + src["parents"][1])})
|
|
if not keep_all and t > 0: # disk hygiene: drop generation t
|
|
for d in agents: # (founders at t=0 are kept — shared)
|
|
shutil.rmtree(d, ignore_errors=True)
|
|
agents = new_agents
|
|
return pd.DataFrame(rows)
|