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
97 lines
4.6 KiB
Python
97 lines
4.6 KiB
Python
"""llm_moe figure — union-preserving recombination (route / max-merge) vs fusion (soup / ties).
|
|
|
|
The real-weight image of E8's *max*: keep every specialist intact and *select* (route per prompt, or
|
|
per module) instead of averaging the deltas. Two panels: (A) per-family accuracy for the base, each
|
|
specialist, the fusion merges, and the union operators — the union operators should match the best
|
|
specialist on every family (they *are* that specialist there), while fusion may dilute or compose;
|
|
(B) overall vs worst-family, fusion vs union, with the routing ceiling (moe_oracle) marked. The
|
|
suptitle reports whether union beats fusion (dilution regime) or they converge (composition regime).
|
|
Reads only the committed bundle.
|
|
|
|
Usage: python figures/plot_llm_moe.py [results/llm_moe]
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
|
|
sys.path.insert(0, str(Path(__file__).parent))
|
|
from _figlib import load_seed_bundles, savefig # noqa: E402
|
|
|
|
_FAMS = ["lists", "strings", "arith"]
|
|
_FUSION = {"merge_soup": "fuse:soup", "merge_ties": "fuse:ties"}
|
|
_UNION = {"moe_oracle": "route:oracle", "moe_learned": "route:learned", "max_merge": "max-merge"}
|
|
|
|
|
|
def _acc(df, model, metric):
|
|
r = df[(df["model"] == model) & (df["metric"] == metric)]["accuracy"]
|
|
return float(r.mean()) if len(r) else float("nan")
|
|
|
|
|
|
def main(results_dir: str = "results/llm_moe") -> None:
|
|
df, cfg = load_seed_bundles(results_dir) # seed-mean when the bundle has s{seed}/ sub-bundles
|
|
present = set(df["model"].unique())
|
|
specialists = sorted(m for m in present if m.startswith("spec_"))
|
|
fusion = [m for m in _FUSION if m in present]
|
|
union = [m for m in _UNION if m in present]
|
|
models = ["base"] + specialists + fusion + union
|
|
labels = {"base": "base", **{s: s.replace("spec_", "spec:") for s in specialists},
|
|
**_FUSION, **_UNION}
|
|
colors = {"base": "#7f7f7f", **{s: "#1f77b4" for s in specialists},
|
|
**{m: "#ff7f0e" for m in fusion}, **{m: "#2ca02c" for m in union}}
|
|
|
|
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
|
|
|
|
# Panel A: per-family accuracy, grouped by model.
|
|
ax = axes[0]
|
|
x = np.arange(len(_FAMS))
|
|
w = 0.8 / len(models)
|
|
for i, mdl in enumerate(models):
|
|
vals = [_acc(df, mdl, f) for f in _FAMS]
|
|
ax.bar(x + (i - (len(models) - 1) / 2) * w, vals, w, label=labels[mdl], color=colors[mdl],
|
|
alpha=0.9 if (mdl in union or mdl in fusion) else 0.65)
|
|
ax.set_xticks(x); ax.set_xticklabels(_FAMS)
|
|
ax.set(ylabel="accuracy", title="Per-family: fusion (orange) blends the deltas; union (green)\n"
|
|
"keeps each specialist intact and selects — no dilution")
|
|
ax.legend(frameon=False, fontsize=8, ncol=2)
|
|
|
|
# Panel B: overall vs worst-family, per model.
|
|
ax = axes[1]
|
|
x2 = np.arange(len(models))
|
|
for off, metric, hatch, lab in [(-0.2, "overall", "", "overall"),
|
|
(0.2, "worst_family", "//", "worst family")]:
|
|
ax.bar(x2 + off, [_acc(df, m, metric) for m in models], 0.38,
|
|
color=[colors[m] for m in models], hatch=hatch, alpha=0.85, label=lab,
|
|
edgecolor="white")
|
|
ax.set_xticks(x2); ax.set_xticklabels([labels[m] for m in models], rotation=25, ha="right",
|
|
fontsize=8)
|
|
ax.set(ylabel="accuracy", title="Overall (solid) vs worst-family (hatched):\nfusion vs union "
|
|
"recombination")
|
|
# Mark the routing ceiling (oracle) if present.
|
|
if "moe_oracle" in present:
|
|
ceil = _acc(df, "moe_oracle", "overall")
|
|
ax.axhline(ceil, ls=":", c="#2ca02c", lw=1, alpha=0.7)
|
|
ax.legend(frameon=False, fontsize=9)
|
|
|
|
best_fuse = max([_acc(df, m, "overall") for m in fusion], default=float("nan"))
|
|
best_union = max([_acc(df, m, "overall") for m in union], default=float("nan"))
|
|
router = _acc(df, "moe_learned", "router_acc") if "moe_learned" in present else float("nan")
|
|
if best_union > best_fuse + 0.01:
|
|
verdict = f"union {best_union:.2f} > fusion {best_fuse:.2f} overall (fusion dilutes)"
|
|
elif best_fuse > best_union + 0.01:
|
|
verdict = f"fusion {best_fuse:.2f} > union {best_union:.2f} overall (strong base composes)"
|
|
else:
|
|
verdict = f"union ≈ fusion ({best_union:.2f} vs {best_fuse:.2f}) overall"
|
|
rtxt = f"; learned router {router:.2f}" if router == router else ""
|
|
fig.suptitle(f"llm_moe — module-level union vs fusion recombination: {verdict}{rtxt} "
|
|
f"({cfg['base_model'].split('/')[-1]})", y=1.0, fontsize=12)
|
|
fig.tight_layout()
|
|
savefig(fig, results_dir, "llm_moe")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main(*sys.argv[1:])
|