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>
This commit is contained in:
Giorgio Gilestro 2026-07-05 17:53:47 +01:00
parent 585264d0b4
commit 8da0dac007
18 changed files with 647 additions and 6 deletions

View file

@ -8,6 +8,9 @@ accepts correct answers (including verbose model phrasings) and rejects wrong on
from __future__ import annotations
import numpy as np
from llm.moe import learned_routes
from llm.tasks import FAMILIES, make_tasks, verify
@ -37,3 +40,29 @@ def test_verifier_rejects_wrong():
assert not verify(wrong, t)
lt = next(x for x in make_tasks("lists", 30, seed=6) if x.answer.startswith("["))
assert not verify("[9, 9, 9]", lt) or lt.answer == "[9, 9, 9]"
def test_learned_router_assigns_nearest_centroid():
# Three well-separated families in a 4-D "embedding" space; the nearest-centroid router
# (the MoE expert-selection gene) must route each test prompt to its own family's specialist.
rng = np.random.default_rng(0)
fams = ["lists", "strings", "arith"]
anchors = {"lists": [5, 0, 0, 0], "strings": [0, 5, 0, 0], "arith": [0, 0, 5, 0]}
train_emb = np.array([anchors[f] for f in fams for _ in range(8)], dtype=float)
train_emb += rng.normal(scale=0.1, size=train_emb.shape)
train_fam = np.array([f for f in fams for _ in range(8)])
test_fam = np.array(["arith", "lists", "strings", "arith"])
test_emb = np.array([anchors[f] for f in test_fam], dtype=float) + rng.normal(scale=0.1, size=(4, 4))
routes = learned_routes(train_emb, train_fam, test_emb, fams)
assert [fams[r] for r in routes] == list(test_fam) # each routed to its own family
def test_learned_router_is_cosine_scale_invariant():
# Cosine routing must ignore prompt-embedding magnitude (long vs short prompts): a test point on a
# family's ray routes there regardless of its norm.
fams = ["a", "b"]
train_emb = np.array([[1.0, 0.0], [1.0, 0.0], [0.0, 1.0], [0.0, 1.0]])
train_fam = np.array(["a", "a", "b", "b"])
test_emb = np.array([[10.0, 0.0], [0.0, 0.01]]) # very different magnitudes
routes = learned_routes(train_emb, train_fam, test_emb, fams)
assert [fams[r] for r in routes] == ["a", "b"]