Add ComposedRouter for multi-axis MoE gating

Route on several independent axes at once (e.g. energy x pdg), each with its own expert count and hyperparameters. The joint gate is the outer product of per-axis softmax gates, so it stays a partition of unity and top1/balance_loss factor per-axis. Config uses flat axis{i}_{field} keys in model.router (TOML/CLI friendly), also settable via repeatable --router-axis flags. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-15 13:53:21 +02:00
parent 550dc679c7
commit f3fec8bcb3
4 changed files with 405 additions and 20 deletions
@@ -4,6 +4,7 @@ import torch

 from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM
 from giant.model.network import (
+    ComposedRouter,
    DenoisingMLP,
    EnergyRouter,
    PdgRouter,
@@ -12,6 +13,7 @@ from giant.model.network import (
    RoutedDenoisingMLP,
    RoutedSecondaryDecoder,
    SecondaryDecoder,
+    build_composed_router,
    build_models,
    build_router,
 )
@@ -289,6 +291,224 @@ def test_build_models_routed_with_process_router():
    assert stage1.router.mat_emb.num_embeddings == 2


+# ── ComposedRouter ───────────────────────────────────────────────────────────
+
+
+def test_composed_router_n_experts_is_product():
+    router = ComposedRouter(
+        [EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)]
+    )
+    assert router.n_experts == 12
+
+
+def test_composed_router_gate_partition_of_unity():
+    router = ComposedRouter(
+        [EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)]
+    )
+    cond_cont, cond_cat = _cond(16, pdg=5)
+    g = router.gate(cond_cont, cond_cat)
+    assert g.shape == (16, 12)
+    torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0)
+
+
+def test_composed_router_gate_is_outer_product_of_sub_gates():
+    energy_router = EnergyRouter(n_experts=4)
+    pdg_router = PdgRouter(n_experts=3, pdg_vocab=5)
+    router = ComposedRouter([energy_router, pdg_router])
+    cond_cont, cond_cat = _cond(16, pdg=5)
+
+    g_energy = energy_router.gate(cond_cont, cond_cat)  # (16, 4)
+    g_pdg = pdg_router.gate(cond_cont, cond_cat)  # (16, 3)
+    expected = (g_energy.unsqueeze(-1) * g_pdg.unsqueeze(1)).flatten(1)  # (16, 12)
+
+    torch.testing.assert_close(router.gate(cond_cont, cond_cat), expected)
+
+
+def test_composed_router_top1_factors_into_per_axis_argmax():
+    """Joint argmax over the outer product must equal the pair of per-axis
+    argmaxes, flattened with the same row-major index convention as gate()."""
+    energy_router = EnergyRouter(n_experts=4)
+    pdg_router = PdgRouter(n_experts=3, pdg_vocab=5)
+    router = ComposedRouter([energy_router, pdg_router])
+    cond_cont, cond_cat = _cond(16, pdg=5)
+
+    joint_idx = router.top1(cond_cont, cond_cat)
+    energy_idx = energy_router.top1(cond_cont, cond_cat)
+    pdg_idx = pdg_router.top1(cond_cont, cond_cat)
+    expected = energy_idx * pdg_router.n_experts + pdg_idx
+
+    assert torch.equal(joint_idx, expected)
+
+
+def test_composed_router_supports_different_expert_counts_per_axis():
+    router = ComposedRouter(
+        [EnergyRouter(n_experts=5), PdgRouter(n_experts=2, pdg_vocab=5)]
+    )
+    assert router.n_experts == 10
+    cond_cont, cond_cat = _cond(8, pdg=5)
+    assert router.gate(cond_cont, cond_cat).shape == (8, 10)
+
+
+def test_composed_router_classify_loss_sums_sub_router_losses():
+    """energy/pdg both default to zero, so the composed loss should too."""
+    router = ComposedRouter(
+        [EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)]
+    )
+    cond_cont, cond_cat = _cond(16, pdg=5)
+    labels = torch.randint(0, 4, (16,))
+    loss = router.classify_loss(cond_cont, cond_cat, labels)
+    assert loss.shape == ()
+    assert loss.item() == 0.0
+
+
+def test_composed_router_rejects_empty_router_list():
+    try:
+        ComposedRouter([])
+    except ValueError:
+        return
+    raise AssertionError("expected ValueError for empty router list")
+
+
+def test_composed_router_not_in_registry():
+    assert "composed" not in ROUTER_REGISTRY
+
+
+# ── _parse_composed_axes (axis{i}_{field} flat-key config convention) ───────
+
+
+def test_parse_composed_axes_groups_indexed_keys():
+    from giant.model.network import _parse_composed_axes
+
+    router_cfg = {
+        "enabled": True,
+        "type": "composed",
+        "axis0_type": "energy",
+        "axis0_n_experts": 4,
+        "axis0_temperature": 0.3,
+        "axis1_type": "pdg",
+        "axis1_n_experts": 3,
+        "axis1_emb_dim": 6,
+    }
+    axes = _parse_composed_axes(router_cfg)
+    assert axes == [
+        {"type": "energy", "n_experts": 4, "temperature": 0.3},
+        {"type": "pdg", "n_experts": 3, "emb_dim": 6},
+    ]
+
+
+def test_parse_composed_axes_ignores_unrelated_keys():
+    from giant.model.network import _parse_composed_axes
+
+    router_cfg = {
+        "enabled": True,
+        "type": "composed",
+        "lambda_balance": 0.0,
+        "axis0_type": "energy",
+        "axis0_n_experts": 4,
+    }
+    axes = _parse_composed_axes(router_cfg)
+    assert axes == [{"type": "energy", "n_experts": 4}]
+
+
+def test_parse_composed_axes_raises_on_index_gap():
+    from giant.model.network import _parse_composed_axes
+
+    router_cfg = {
+        "type": "composed",
+        "axis0_type": "energy",
+        "axis0_n_experts": 4,
+        # axis1 missing entirely
+        "axis2_type": "pdg",
+        "axis2_n_experts": 3,
+    }
+    try:
+        _parse_composed_axes(router_cfg)
+    except ValueError:
+        return
+    raise AssertionError("expected ValueError for a gap in axis indices")
+
+
+def test_build_composed_router_resolves_per_axis_specs():
+    router = build_composed_router(
+        [
+            {"type": "energy", "n_experts": 4, "temperature": 0.3},
+            {"type": "pdg", "n_experts": 3, "emb_dim": 6},
+        ],
+        pdg_vocab=5,
+        mat_vocab=2,
+    )
+    assert isinstance(router, ComposedRouter)
+    assert router.n_experts == 12
+    energy_router, pdg_router = router.routers
+    assert isinstance(energy_router, EnergyRouter)
+    assert energy_router.temperature == 0.3
+    assert isinstance(pdg_router, PdgRouter)
+    assert pdg_router.pdg_emb.num_embeddings == 5
+    assert pdg_router.pdg_emb.embedding_dim == 6
+
+
+def test_build_models_routed_with_composed_router():
+    model_config = dict(
+        pdg_vocab=5,
+        mat_vocab=2,
+        emb_dim=16,
+        dropout=0.1,
+        k_max=K_MAX,
+        expert_hidden_dim=16,
+        expert_n_blocks=2,
+        router={
+            "enabled": True,
+            "type": "composed",
+            "axis0_type": "energy",
+            "axis0_n_experts": 4,
+            "axis1_type": "pdg",
+            "axis1_n_experts": 3,
+        },
+    )
+    stage1, sec_decoder = build_models(model_config)
+    assert isinstance(stage1, RoutedDenoisingMLP)
+    assert isinstance(stage1.router, ComposedRouter)
+    assert len(stage1.experts) == 12
+    assert len(sec_decoder.experts) == 12
+    # stage1 and sec_decoder must not share router weights (same convention
+    # as the single-axis routers built by build_models).
+    assert stage1.router is not sec_decoder.router
+
+
+def test_build_models_routed_pair_composed_router_is_drop_in_for_sample_flow():
+    from giant.sample import sample_flow, sample_secondaries
+
+    model_config = dict(
+        pdg_vocab=3,
+        mat_vocab=2,
+        emb_dim=16,
+        dropout=0.1,
+        k_max=K_MAX,
+        expert_hidden_dim=8,
+        expert_n_blocks=1,
+        router={
+            "enabled": True,
+            "type": "composed",
+            "axis0_type": "energy",
+            "axis0_n_experts": 2,
+            "axis1_type": "pdg",
+            "axis1_n_experts": 2,
+        },
+    )
+    stage1, sec_decoder = build_models(model_config)
+    B = 5
+    cond_cont, cond_cat = _cond(B, pdg=3, mat=2)
+    stage1_norm, n_sec_pred = sample_flow(stage1, cond_cont, cond_cat, steps=2)
+    assert stage1_norm.shape == (B, X_DIM)
+    assert n_sec_pred.shape == (B,)
+
+    sec_cont, sec_type_emb, sec_valid = sample_secondaries(
+        sec_decoder, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2
+    )
+    assert sec_cont.shape == (B, K_MAX, 4)
+    assert sec_valid.shape == (B, K_MAX)
+
+
 # ── RoutedDenoisingMLP ───────────────────────────────────────────────────────