Add mixture-of-experts routing prototype for Stage 1 and Stage 2

Both stages can now route through a pluggable Router (EnergyRouter as the first implementation, a soft turn-on gate over pre-step log-energy) into several small ExpertTrunks instead of one monolithic trunk. Trains as a differentiable soft mixture and dispatches to a single expert per row at eval time, which is the source of the per-call speedup this prototype is after (issue #5's ~10x native-Geant4 budget). Disabled by default, so existing configs/checkpoints are unaffected; build_models() centralizes routed-vs-monolith construction across train/predict/rollout. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-08 15:17:19 +02:00
parent af8dce53a7
commit bac541240f
6 changed files with 815 additions and 74 deletions
@@ -0,0 +1,296 @@
+"""Tests for the mixture-of-experts routing prototype (giant/model/network.py)."""
+
+import torch
+
+from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM
+from giant.model.network import (
+    DenoisingMLP,
+    EnergyRouter,
+    ROUTER_REGISTRY,
+    RoutedDenoisingMLP,
+    RoutedSecondaryDecoder,
+    SecondaryDecoder,
+    build_models,
+    build_router,
+)
+
+
+def _cond(B=8, pdg=3, mat=2):
+    cond_cont = torch.randn(B, COND_DIM)
+    cond_cat = torch.stack(
+        [torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1
+    )
+    return cond_cont, cond_cat
+
+
+def _routed_stage1(n_experts=4, pdg=3, mat=2, **router_kwargs):
+    router = build_router("energy", n_experts, **router_kwargs)
+    return RoutedDenoisingMLP(
+        pdg_vocab=pdg,
+        mat_vocab=mat,
+        router=router,
+        expert_hidden_dim=16,
+        expert_n_blocks=2,
+    )
+
+
+def _routed_sec_decoder(n_experts=4, pdg=3, mat=2, **router_kwargs):
+    router = build_router("energy", n_experts, **router_kwargs)
+    return RoutedSecondaryDecoder(
+        pdg_vocab=pdg,
+        mat_vocab=mat,
+        router=router,
+        expert_hidden_dim=16,
+        expert_n_blocks=2,
+    )
+
+
+# ── Router / EnergyRouter contract ──────────────────────────────────────────
+
+
+def test_energy_router_registered():
+    assert ROUTER_REGISTRY["energy"] is EnergyRouter
+
+
+def test_energy_router_gate_partition_of_unity():
+    router = EnergyRouter(n_experts=4)
+    cond_cont, cond_cat = _cond(16)
+    g = router.gate(cond_cont, cond_cat)
+    assert g.shape == (16, 4)
+    torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0)
+
+
+def test_energy_router_top1_matches_gate_argmax():
+    router = EnergyRouter(n_experts=4)
+    cond_cont, cond_cat = _cond(16)
+    assert torch.equal(
+        router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1)
+    )
+
+
+def test_energy_router_hardens_as_temperature_shrinks():
+    """As tau -> 0 the soft gate should converge to a one-hot at the argmax."""
+    router = EnergyRouter(n_experts=4, temperature=1e-4)
+    cond_cont, cond_cat = _cond(16)
+    g = router.gate(cond_cont, cond_cat)
+    top1 = router.top1(cond_cont, cond_cat)
+    onehot = torch.nn.functional.one_hot(top1, num_classes=4).float()
+    torch.testing.assert_close(g, onehot, atol=1e-3, rtol=0)
+
+
+def test_energy_router_balance_loss_is_nonnegative_scalar():
+    router = EnergyRouter(n_experts=4)
+    cond_cont, cond_cat = _cond(16)
+    loss = router.balance_loss(cond_cont, cond_cat)
+    assert loss.shape == ()
+    assert loss.item() >= 0.0
+
+
+def test_build_router_ignores_unrecognized_kwargs():
+    # lambda_balance is a model_config.router key but not an EnergyRouter kwarg
+    router = build_router("energy", 4, temperature=0.3, lambda_balance=0.5)
+    assert isinstance(router, EnergyRouter)
+    assert router.temperature == 0.3
+
+
+def test_build_router_unknown_type_raises():
+    try:
+        build_router("nonexistent", 4)
+    except ValueError:
+        return
+    raise AssertionError("expected ValueError for unknown router type")
+
+
+# ── RoutedDenoisingMLP ───────────────────────────────────────────────────────
+
+
+def test_routed_denoising_mlp_output_shape_train_and_eval():
+    B = 8
+    model = _routed_stage1()
+    x_t = torch.randn(B, X_DIM)
+    t = torch.rand(B)
+    cond_cont, cond_cat = _cond(B)
+
+    model.train()
+    out_train = model(x_t, t, cond_cont, cond_cat)
+    assert out_train.shape == (B, X_DIM)
+
+    model.eval()
+    with torch.no_grad():
+        out_eval = model(x_t, t, cond_cont, cond_cat)
+    assert out_eval.shape == (B, X_DIM)
+
+
+def test_routed_denoising_mlp_gradients_flow_in_train_mode():
+    """Soft mixture in train mode should touch every expert's parameters."""
+    B = 8
+    model = _routed_stage1(n_experts=3)
+    x_t = torch.randn(B, X_DIM)
+    t = torch.rand(B)
+    cond_cont, cond_cat = _cond(B)
+    model.train()
+    flow_loss = model(x_t, t, cond_cont, cond_cat).sum()
+    nsec_loss = model.predict_n_sec(cond_cont, cond_cat).sum()
+    (flow_loss + nsec_loss).backward()
+    for name, p in model.named_parameters():
+        assert p.grad is not None, f"no grad for {name}"
+
+
+def test_routed_denoising_mlp_eval_dispatch_matches_manual_grouping():
+    """Eval-mode grouped top-1 dispatch must equal running each row through
+    its assigned expert individually (batch order shouldn't matter)."""
+    B = 12
+    model = _routed_stage1(n_experts=4)
+    model.eval()
+    x_t = torch.randn(B, X_DIM)
+    t = torch.rand(B)
+    cond_cont, cond_cat = _cond(B)
+
+    with torch.no_grad():
+        batched = model(x_t, t, cond_cont, cond_cat)
+
+        t_emb = model.time_emb(t)
+        c_emb = model.cond_enc(cond_cont, cond_cat)
+        cond = torch.cat([t_emb, c_emb], dim=-1)
+        idx = model.router.top1(cond_cont, cond_cat)
+        manual = torch.zeros_like(x_t)
+        for i in range(B):
+            manual[i] = model.experts[int(idx[i])](x_t[i : i + 1], cond[i : i + 1])[0]
+
+    torch.testing.assert_close(batched, manual, atol=1e-5, rtol=1e-4)
+
+
+def test_routed_denoising_mlp_predict_n_sec_shape():
+    B = 6
+    model = _routed_stage1()
+    cond_cont, cond_cat = _cond(B)
+    logits = model.predict_n_sec(cond_cont, cond_cat)
+    assert logits.shape == (B, K_MAX + 1)
+
+
+def test_routed_denoising_mlp_pdg_embedding_weight_shape():
+    model = _routed_stage1(pdg=5, mat=2)
+    from giant.constants import EMB_DIM
+
+    assert model.pdg_embedding_weight().shape == (5, EMB_DIM)
+
+
+# ── RoutedSecondaryDecoder ───────────────────────────────────────────────────
+
+
+def test_routed_secondary_decoder_output_shape_train_and_eval():
+    B = 8
+    decoder = _routed_sec_decoder()
+    x_t = torch.randn(B, SEC_DIM)
+    t = torch.rand(B)
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+
+    decoder.train()
+    out_train = decoder(x_t, t, cond_cont, cond_cat, stage1_out)
+    assert out_train.shape == (B, SEC_DIM)
+
+    decoder.eval()
+    with torch.no_grad():
+        out_eval = decoder(x_t, t, cond_cont, cond_cat, stage1_out)
+    assert out_eval.shape == (B, SEC_DIM)
+
+
+def test_routed_secondary_decoder_gradients_flow():
+    B = 4
+    decoder = _routed_sec_decoder(n_experts=3)
+    x_t = torch.randn(B, SEC_DIM)
+    t = torch.rand(B)
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    decoder.train()
+    decoder(x_t, t, cond_cont, cond_cat, stage1_out).sum().backward()
+    for name, p in decoder.named_parameters():
+        assert p.grad is not None, f"no grad for {name}"
+
+
+# ── build_models dispatch ────────────────────────────────────────────────────
+
+
+def test_build_models_monolith_when_router_absent():
+    model_config = dict(
+        pdg_vocab=4,
+        mat_vocab=2,
+        hidden_dim=32,
+        n_blocks=2,
+        emb_dim=16,
+        dropout=0.1,
+        k_max=K_MAX,
+    )
+    stage1, sec_decoder = build_models(model_config)
+    assert isinstance(stage1, DenoisingMLP)
+    assert isinstance(sec_decoder, SecondaryDecoder)
+
+
+def test_build_models_monolith_when_router_disabled():
+    model_config = dict(
+        pdg_vocab=4,
+        mat_vocab=2,
+        hidden_dim=32,
+        n_blocks=2,
+        emb_dim=16,
+        dropout=0.1,
+        k_max=K_MAX,
+        router={"enabled": False, "type": "energy", "n_experts": 4},
+    )
+    stage1, sec_decoder = build_models(model_config)
+    assert isinstance(stage1, DenoisingMLP)
+    assert isinstance(sec_decoder, SecondaryDecoder)
+
+
+def test_build_models_routed_when_enabled():
+    model_config = dict(
+        pdg_vocab=4,
+        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": "energy",
+            "n_experts": 4,
+            "temperature": 0.5,
+            "learn_centers": True,
+            "lambda_balance": 0.0,
+        },
+    )
+    stage1, sec_decoder = build_models(model_config)
+    assert isinstance(stage1, RoutedDenoisingMLP)
+    assert isinstance(sec_decoder, RoutedSecondaryDecoder)
+    assert len(stage1.experts) == 4
+    assert len(sec_decoder.experts) == 4
+
+
+def test_build_models_routed_pair_is_drop_in_for_sample_flow():
+    """Exercise the exact calling convention giant/sample.py uses."""
+    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": "energy", "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)