llm_moe: the union operator (route/max-merge) vs fusion — and the regime flips at scale

Adds the union-preserving recombination operator that llm_merge lacked (E8's max,
not mean): keep each specialist LoRA intact and SELECT the right one per prompt
(MoE router: oracle, or training-free nearest-centroid over base embeddings) or
per module (max_merge = winner-take-all by delta norm). src/llm/moe.py, kind
llm_moe, reuses the cached specialists.

Result — a clean regime boundary for "merge, don't average":
- 0.5B: union wins. Routing 0.74 / worst-family 0.43 > soup 0.64 / 0.26, with no
  dilution (recovers each specialist's own-family peak). E8's max > mean in real
  weights, because at a weak base averaging dilutes.
- 7B (Imperial CX3, L40S, 9 min): the ordering INVERTS. Fusion wins — soup 0.87 >
  routing 0.84 > max_merge 0.78. Routing is capped at the best parent per family;
  fusion blends and, given a capable base, COMPOSES beyond any parent (soup lists
  0.62 > spec 0.57). Selection can't synthesise better than its best component;
  averaging-that-composes can.

So "merge, don't average" (E4/E8) is a weak-parent / small-model law, not
universal: union wins under dilution, fusion wins under composition. Refines E8
(its additive-landscape max>mean assumed no compositional headroom). The operator
to want is fusion-that-composes + offspring selection = the directed-sex ideal
(E10) — the natural next experiment.

Honest riders: the learned router is trivially perfect (lexically-distinct
families), and router-free max_merge is the weakest union (not input-adaptive).
+2 router unit tests (127 green). Results in results/llm_moe{,_hpc}/ (parquet
gitignored per the reproducibility contract).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

src/llm/experiment.py

@@ -24,12 +24,13 @@ from knowledge.experiment import save_artifacts
 
 from .evaluate import evaluate, generate, load_model
 from .merge import load_specialists, make_merge
+from .moe import build_max_merge, embed_prompts, learned_routes, moe_generate
 from .specialise import train_specialist
 from .tasks import FAMILIES, make_tasks, verify
 
 
 def _rows(experiment: str, model: str, kind: str, acc: dict) -> list[dict]:
-    worst = min(acc[f] for f in FAMILIES)
+    worst = min(acc[f] for f in acc if f in FAMILIES)
     out = [{"experiment": experiment, "model": model, "kind": kind, "metric": k, "accuracy": v}
            for k, v in acc.items()]
     out.append({"experiment": experiment, "model": model, "kind": kind,
@@ -85,16 +86,122 @@ def run_merge_experiment(cfg: dict) -> pd.DataFrame:
     return pd.DataFrame(rows)
 
 
+def _load_or_train_specialists(cfg: dict, base: str, fams: list[str], name: str,
+                               rows: list[dict]) -> list[str]:
+    """Reuse cached specialist adapters if present, else train one per family. Appends spec rows."""
+    import torch
+
+    n_train = int(cfg.get("n_train", 700))
+    epochs = int(cfg.get("epochs", 3))
+    lora = cfg.get("lora", {})
+    seed = int(cfg["seed"])
+    adapters_root = Path(cfg.get("adapters_dir", "models/llm"))
+    dirs: list[str] = []
+    for i, f in enumerate(fams):
+        d = str(adapters_root / f"spec_{f}")
+        if not (Path(d) / "adapter_config.json").exists():   # reuse across llm_merge / llm_moe runs
+            train_specialist(base, f, d, n_train=n_train, epochs=epochs, seed=seed + i,
+                             r=int(lora.get("r", 16)), alpha=int(lora.get("alpha", 32)))
+        dirs.append(d)
+        m, tok = load_model(base, adapter_dir=d)
+        rows += _rows(name, f"spec_{f}", "specialist", evaluate(m, tok, test_of(cfg, fams)))
+        del m; torch.cuda.empty_cache()
+    return dirs
+
+
+def test_of(cfg: dict, fams: list[str]) -> list:
+    """The held-out mixed test set (one seed offset per family), shared by both LLM runners."""
+    n_test = int(cfg.get("n_test", 100))
+    return sum([make_tasks(f, n_test, seed=1000 + i) for i, f in enumerate(fams)], [])
+
+
+def run_moe_experiment(cfg: dict) -> pd.DataFrame:
+    """Contrast union-preserving recombination (route / max-merge) with fusion (soup / ties).
+
+    Reuses the trained specialist adapters and evaluates, per operator, on the held-out mixed test
+    set. Union operators keep each specialist intact and *select* (per prompt via a router, or per
+    module via winner-take-all); fusion operators blend the deltas. The real-weight image of E8's
+    ``max`` vs the ``mean`` baseline. Returns long-form accuracies (+ a ``router_acc`` fidelity row
+    for each routed operator).
+    """
+    import torch
+
+    name = cfg["experiment"]
+    base = cfg["base_model"]
+    fams = list(cfg.get("families", list(FAMILIES)))
+    ops = list(cfg.get("operators", ["soup", "ties", "moe_oracle", "moe_learned", "max_merge"]))
+    n_route = int(cfg.get("n_route", 32))
+    rows: list[dict] = []
+
+    test = test_of(cfg, fams)
+    prompts = [t.prompt for t in test]
+    fam = np.array([t.family for t in test])
+    true_idx = np.array([fams.index(t.family) for t in test])
+
+    # base + specialists (parents)
+    m, tok = load_model(base)
+    rows += _rows(name, "base", "base", evaluate(m, tok, test))
+    del m; torch.cuda.empty_cache()
+    dirs = _load_or_train_specialists(cfg, base, fams, name, rows)
+    k = len(dirs)
+
+    # all specialists on one base for the recombination operators
+    model, tok = load_specialists(base, dirs)
+
+    def _score(outs: list[str]) -> dict:
+        corr = np.array([verify(o, t) for o, t in zip(outs, test)])
+        acc = {"overall": float(corr.mean())}
+        acc.update({f: float(corr[fam == f].mean()) for f in fams})
+        return acc
+
+    # precompute router assignments once, from the base model's own prompt embeddings
+    routes: dict[str, np.ndarray] = {}
+    if "moe_oracle" in ops:
+        routes["moe_oracle"] = true_idx
+    if "moe_learned" in ops:
+        route_train = sum([make_tasks(f, n_route, seed=2000 + i) for i, f in enumerate(fams)], [])
+        tr_emb = embed_prompts(model, tok, [t.prompt for t in route_train])
+        te_emb = embed_prompts(model, tok, prompts)
+        tr_fam = np.array([t.family for t in route_train])
+        routes["moe_learned"] = learned_routes(tr_emb, tr_fam, te_emb, fams)
+
+    for op in ops:
+        if op in ("soup", "ties", "task_arith"):
+            make_merge(model, k, op, op)
+            rows += _rows(name, f"merge_{op}", "merge", _score(generate(model, tok, prompts)))
+        elif op in ("moe_oracle", "moe_learned"):
+            r = routes[op]
+            rows += _rows(name, op, "moe", _score(moe_generate(model, tok, prompts, r, k)))
+            rows.append({"experiment": name, "model": op, "kind": "moe",
+                         "metric": "router_acc", "accuracy": float((r == true_idx).mean())})
+        elif op == "max_merge":
+            build_max_merge(model, k, "max_merge")
+            rows += _rows(name, "max_merge", "moe", _score(generate(model, tok, prompts)))
+        else:
+            raise ValueError(f"unknown operator {op!r} (soup|ties|task_arith|moe_oracle|"
+                             f"moe_learned|max_merge)")
+
+    return pd.DataFrame(rows)
+
+
+_RUNNERS = {"llm_merge": run_merge_experiment, "llm_moe": run_moe_experiment}
+
+
 def run_and_save(config_path: str | Path) -> Path:
-    """Load an LLM experiment YAML, run it, and write the artifact triple."""
+    """Load an LLM experiment YAML, run it (dispatch on ``kind``), and write the artifact triple."""
     config_path = Path(config_path)
     cfg = yaml.safe_load(config_path.read_text())
     out_dir = Path(cfg.get("output", {}).get("dir", f"results/{cfg['experiment']}"))
     cfg.setdefault("n_replicates", 1)
-    df = run_merge_experiment(cfg)
+    kind = cfg.get("kind", "llm_merge")
+    if kind not in _RUNNERS:
+        raise ValueError(f"unknown LLM experiment kind {kind!r} (expected one of {list(_RUNNERS)})")
+    df = _RUNNERS[kind](cfg)
+    extra = {"layer": "2", "tier": "llm", "base_model": cfg["base_model"]}
+    if kind == "llm_moe":
+        extra["operators"] = list(cfg.get("operators", []))
     save_artifacts(cfg, df, out_dir, extra_libs=("torch", "transformers", "peft"),
-                   extra_manifest={"layer": "2", "tier": "llm", "base_model": cfg["base_model"]},
-                   grid=None)
+                   extra_manifest=extra, grid=None)
     return out_dir
 
 

src/llm/moe.py

@@ -0,0 +1,154 @@
+"""Module-level, union-preserving recombination of specialists — the real-weight image of E8's *max*.
+
+`merge.py` fuses specialist adapters into one set of weights (``soup`` averages the deltas, ``ties``
+sign-reconciles them). Fusion can *dilute*: at a weak base, averaging a specialist's delta by ``1/K``
+erased its peak (the 0.5B ``llm_merge`` result). The alternative recombination operator — the one E8
+calls *max* / union-preserving — never averages the parents at all. It keeps every specialist adapter
+**intact** and, per input (or per module), **selects** the parent that owns that skill. Nothing is
+diluted because nothing is blended; the child is the *union* of the parents' capabilities.
+
+Two selection operators, both reusing the already-trained specialist adapters (no retraining):
+
+* **route** — a Mixture-of-Experts over the specialists: a cheap router assigns each prompt to one
+  specialist adapter, which then answers it. Router variants: ``oracle`` (route by the known task
+  family — the ceiling of routing) and ``learned`` (nearest-centroid over the *base* model's own
+  prompt embeddings — an honest, training-free router; its accuracy is reported).
+* **max_merge** — a router-free static union: build a single adapter that, per LoRA module, copies the
+  delta from the specialist whose delta has the largest norm there (winner-take-all per module). The
+  literal weight-space image of E8's element-wise ``max`` over teachers.
+
+The experiment (``kind: llm_moe``) contrasts these union operators against the fusion baselines
+(soup/ties) at both scales — the prediction being that union wins where fusion dilutes (0.5B) and the
+gap narrows once a capable base lets fusion *compose* rather than dilute (7B): the regime boundary of
+"merge, don't average" in real LLM weights.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from .evaluate import generate
+
+
+def embed_prompts(model, tok, prompts: list[str], batch_size: int = 32,
+                  device: str = "cuda") -> np.ndarray:
+    """Mean last-hidden-state embedding of each prompt from the **base** model (adapters disabled).
+
+    The router must see a parent-agnostic representation, so generation happens with
+    ``model.disable_adapter()`` — the embedding is the frozen base model's, not any specialist's.
+
+    Args:
+        model: a PEFT model with specialist adapters loaded.
+        tok: the matching tokenizer (left-padded).
+        prompts (list[str]): raw user prompts.
+        batch_size (int): forward-pass batch size.
+        device (str): compute device.
+
+    Returns:
+        np.ndarray: ``(len(prompts), hidden)`` float32 embeddings.
+    """
+    import torch
+
+    embs: list[np.ndarray] = []
+    with model.disable_adapter():                          # parent-agnostic base representation
+        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)
+            with torch.no_grad():
+                out = model(**enc, output_hidden_states=True)
+            h = out.hidden_states[-1]                       # (B, T, H)
+            mask = enc["attention_mask"].unsqueeze(-1)      # (B, T, 1)
+            mean = (h * mask).sum(1) / mask.sum(1).clamp(min=1)
+            embs.append(mean.float().cpu().numpy())
+    return np.concatenate(embs, axis=0)
+
+
+def learned_routes(train_emb: np.ndarray, train_fam: np.ndarray, test_emb: np.ndarray,
+                   fams: list[str]) -> np.ndarray:
+    """Nearest-centroid router: assign each test prompt to the family whose train centroid is closest.
+
+    A training-free expert selector — the "regulatory gene" that decides which specialist to express.
+
+    Args:
+        train_emb (np.ndarray): ``(n_train, H)`` base embeddings of labelled train prompts.
+        train_fam (np.ndarray): family label per train prompt.
+        test_emb (np.ndarray): ``(n_test, H)`` base embeddings of test prompts.
+        fams (list[str]): family order (index i ↔ adapter ``a{i}``).
+
+    Returns:
+        np.ndarray: ``(n_test,)`` int expert index in ``[0, len(fams))``.
+    """
+    centroids = np.stack([train_emb[train_fam == f].mean(0) for f in fams])   # (K, H)
+    # Cosine distance is scale-robust for hidden states.
+    c = centroids / (np.linalg.norm(centroids, axis=1, keepdims=True) + 1e-8)
+    t = test_emb / (np.linalg.norm(test_emb, axis=1, keepdims=True) + 1e-8)
+    sims = t @ c.T                                          # (n_test, K)
+    return sims.argmax(1)
+
+
+def moe_generate(model, tok, prompts: list[str], routes: np.ndarray, k: int,
+                 **gen_kw) -> list[str]:
+    """Generate each prompt with its routed specialist adapter active (grouped by expert for speed).
+
+    Args:
+        model: PEFT model with adapters ``a0..a{k-1}`` loaded.
+        tok: tokenizer.
+        prompts (list[str]): prompts to answer.
+        routes (np.ndarray): expert index per prompt (from an oracle or learned router).
+        k (int): number of experts.
+        **gen_kw: forwarded to :func:`llm.evaluate.generate`.
+
+    Returns:
+        list[str]: completions, in the original prompt order.
+    """
+    outs: list[str] = [""] * len(prompts)
+    for expert in range(k):
+        idx = np.nonzero(routes == expert)[0]
+        if len(idx) == 0:
+            continue
+        model.set_adapter(f"a{expert}")
+        completions = generate(model, tok, [prompts[i] for i in idx], **gen_kw)
+        for j, i in enumerate(idx):
+            outs[i] = completions[j]
+    return outs
+
+
+def build_max_merge(model, k: int, name: str):
+    """Add a router-free union adapter ``name``: per LoRA module, keep the largest-norm specialist delta.
+
+    For each LoRA-adapted module, the effective delta of specialist ``a_i`` is ``B_i @ A_i`` (scaled).
+    We select, per module, the specialist whose delta has the largest Frobenius norm and copy its
+    ``A``/``B`` into a fresh adapter — a deterministic winner-take-all union (E8's element-wise max at
+    the granularity of a module). No averaging, so no dilution.
+
+    Args:
+        model: PEFT model with adapters ``a0..a{k-1}``.
+        k (int): number of specialists.
+        name (str): name of the new merged adapter (created as a copy of ``a0`` then overwritten).
+
+    Returns:
+        The model, with adapter ``name`` added and set active.
+    """
+    import torch
+    from peft.tuners.lora import LoraLayer
+
+    # Seed the new adapter from a0's config, then overwrite its weights per module.
+    model.add_weighted_adapter([f"a{i}" for i in range(k)], [1.0] + [0.0] * (k - 1), name,
+                               combination_type="linear")
+    with torch.no_grad():
+        for module in model.modules():
+            if not isinstance(module, LoraLayer) or name not in module.lora_A:
+                continue
+            norms = []
+            for i in range(k):
+                a, b = f"a{i}", f"a{i}"
+                delta = module.lora_B[a].weight @ module.lora_A[a].weight
+                norms.append(float(delta.norm()) * module.scaling.get(f"a{i}", 1.0))
+            win = int(np.argmax(norms))
+            module.lora_A[name].weight.copy_(module.lora_A[f"a{win}"].weight)
+            module.lora_B[name].weight.copy_(module.lora_B[f"a{win}"].weight)
+            module.scaling[name] = module.scaling.get(f"a{win}", 1.0)
+    model.set_adapter(name)
+    return model