Something went wrong. Try again.
This repository has no description
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174#!/usr/bin/env python3"""Epistemic dashboard for a trained colony: is it a discussion or a monologue?
A colony that scores like a solo unit could be failing in very different ways,and they call for opposite fixes:
(a) the units are clones in practice — identical proposals, so aggregation is a no-op and "K units" is just one unit at K times the cost; (b) the units differ but aggregation destroys the difference — the mean washes out minority-but-correct proposals; (c) the units differ and aggregation helps, but every member is so weak that the ceiling is low anyway.
Metrics reported: - solo accuracy of each unit's own proposal (before aggregation) - best-single-unit vs aggregate accuracy -> does combining add anything? - oracle accuracy (any unit right) -> headroom a better rule could capture - inter-unit disagreement rate -> are they saying different things? - per-round evolution of all of the above -> does discussion converge or collapse?
Usage: python scripts/analyze_colony.py --checkpoint <ckpt> --config <all_config.yaml> \ --data data/sudoku-testsub-12k --batches 4"""
import argparseimport jsonimport osimport sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import torch # noqa: E402import yaml # noqa: E402
from dsem.device import get_device # noqa: E402from dsem.pretrain import PretrainConfig, create_dataloader, create_model # noqa: E402
IGNORE = -100
def main() -> None: ap = argparse.ArgumentParser() ap.add_argument("--checkpoint", required=True) ap.add_argument("--config", required=True) ap.add_argument("--data", required=True) ap.add_argument("--batches", type=int, default=4) ap.add_argument("--out", default=None) args = ap.parse_args()
device = get_device() with open(args.config) as f: cfg_dict = yaml.unsafe_load(f) config = PretrainConfig(**cfg_dict) config.data_paths_test = [args.data] config.evaluators = []
loader, metadata = create_dataloader(config, "test", test_set_mode=True, epochs_per_iter=1, global_batch_size=config.global_batch_size, rank=0, world_size=1) model, _, _ = create_model(config, metadata, rank=0, world_size=1) model.load_state_dict(torch.load(args.checkpoint, map_location=device), assign=True) model.eval()
inner = model.model.inner if hasattr(model, "model") else model.inner if hasattr(inner, "_orig_mod"): inner = inner._orig_mod K = inner.config.num_agents T = inner.config.H_cycles
acc = {"per_agent": torch.zeros(T, K), "best": torch.zeros(T), "agg": torch.zeros(T), "oracle": torch.zeros(T), "disagree": torch.zeros(T)} alone = {"per_unit": torch.zeros(K), "any": torch.zeros(()), "disagree": torch.zeros(())} seen = 0
for bi, (_set, batch, _gbs) in enumerate(loader): if bi >= args.batches: break batch = {k: v.to(device) for k, v in batch.items()} labels = batch["labels"] mask = labels != IGNORE B_now = labels.shape[0]
with torch.device(device): carry = model.initial_carry(batch) inner_carry = inner.reset_carry(torch.ones(labels.shape[0], dtype=torch.bool, device=device), carry.inner_carry) # Eval runs the full ACT budget (halt_max_steps supervision steps), each of # which is T rounds of discussion. Measuring one supervision step would # describe a model 1/16th of the way through its inference. from dsem.models.recursive_reasoning.colony import ColonyInnerCarry first = None for si in range(inner.config.halt_max_steps): out = inner.forward_instrumented(inner_carry, batch) if si == 0: first = out["rounds"][0] # blackboard still at init: units on their own inner_carry = ColonyInnerCarry(y=out["y"], z=out["z"])
# "alone" = each unit's very first proposal, before any information has # reached it through the blackboard. This is what a member knows by itself. pa0 = first["per_agent"] solo = ((pa0 == labels.unsqueeze(1)) | ~mask.unsqueeze(1)).all(-1) # [B, K] alone["per_unit"] += solo.sum(0).float().cpu() alone["any"] += solo.any(-1).sum().float().cpu() unan0 = (pa0 == pa0[:, :1]).all(1) alone["disagree"] += (~unan0 & mask).sum().float().cpu() / max(mask.sum().item() / B_now, 1)
B = labels.shape[0] seen += B for t, r in enumerate(out["rounds"]): pa, ag = r["per_agent"], r["aggregate"] solved_a = ((pa == labels.unsqueeze(1)) | ~mask.unsqueeze(1)).all(-1) # [B, K] solved_g = ((ag == labels) | ~mask).all(-1) # [B] acc["per_agent"][t] += solved_a.sum(0).float().cpu() acc["best"][t] += solved_a.any(-1).sum().float().cpu() # oracle over units == best possible vote acc["oracle"][t] += solved_a.any(-1).sum().float().cpu() acc["agg"][t] += solved_g.sum().float().cpu() # cellwise disagreement: fraction of cells where units are not unanimous unan = (pa == pa[:, :1]).all(1) # [B, seq] acc["disagree"][t] += (~unan & mask).sum().float().cpu() / max(mask.sum().item() / B, 1)
n = max(seen, 1) print(f"colony: K={K}, T={T} rounds, {seen} puzzles from {args.data}\n") print(f"{'round':>5} {'agg%':>7} {'best-unit%':>11} {'mean-unit%':>11} {'disagree%':>10}") results = [] for t in range(T): per = acc["per_agent"][t] / n * 100 row = dict( round=t + 1, aggregate=float(acc["agg"][t] / n * 100), oracle_any_unit=float(acc["oracle"][t] / n * 100), mean_unit=float(per.mean()), per_unit=[round(float(v), 2) for v in per], disagreement=float(acc["disagree"][t] / n * 100), ) results.append(row) print(f"{t+1:>5} {row['aggregate']:>7.2f} {row['oracle_any_unit']:>11.2f}" f" {row['mean_unit']:>11.2f} {row['disagreement']:>10.2f}")
last = results[-1] print(f"\nper-unit solve rates (final round): {last['per_unit']}")
solo_mean = float((alone["per_unit"] / n * 100).mean()) solo_any = float(alone["any"] / n * 100) solo_dis = float(alone["disagree"] / n * 100) print(f"\nBEFORE any discussion (first round, blackboard empty):") print(f" mean unit alone: {solo_mean:.2f}% any unit alone: {solo_any:.2f}%" f" disagreement: {solo_dis:.2f}%") print(f"AFTER discussion: colony {last['aggregate']:.2f}%") lift = last["aggregate"] - solo_mean print(f" => collective lift: {lift:+.2f} points over the average member acting alone")
print("\ndiagnosis:") if last["disagreement"] < 1.0: print(" (a) UNITS ARE EFFECTIVELY CLONES — disagreement <1% of cells.") print(" Aggregation has nothing to combine; standpoints are not producing diversity.") elif last["aggregate"] < last["oracle_any_unit"] - 2: print(f" (b) AGGREGATION LOSES INFORMATION — some unit solves {last['oracle_any_unit']:.1f}%") print(f" but the colony only reports {last['aggregate']:.1f}%. A better decision rule") print(" (vote/confidence) or a way for minorities to persuade could capture that gap.") elif last["aggregate"] > last["mean_unit"] + 2: print(" (c) AGGREGATION HELPS — the colony beats its average member; the ceiling is") print(" member quality, so bigger/heterogeneous members are the lever.") else: print(" (d) NO CLEAR MECHANISM — units differ but combining neither helps nor hurts.")
if args.out: os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True) with open(args.out, "w") as f: json.dump({"K": K, "T": T, "puzzles": seen, "rounds": results}, f, indent=2)
if __name__ == "__main__": main()