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>

giant/cli.py

@@ -39,31 +39,30 @@ from giant.data.transforms import (
     Normalizer,
 )
 from giant.geometry import GeometryOracle
-from giant.model.network import DenoisingMLP, SecondaryDecoder
+from giant.model.network import build_models
 from giant.pipeline import run_train_job
 from giant.rollout import rollout as run_rollout
 from giant.sample import sample_flow, sample_secondaries, snap_type_to_pdg_idx
 
-_STAGE1_MODEL_KEYS = {
-    "pdg_vocab",
-    "mat_vocab",
-    "hidden_dim",
-    "n_blocks",
-    "emb_dim",
-    "dropout",
-    "k_max",
-}
-_SEC_DECODER_MODEL_KEYS = {
-    "pdg_vocab",
-    "mat_vocab",
-    "hidden_dim",
-    "n_blocks",
-    "emb_dim",
-    "dropout",
-}
-
 app = typer.Typer(no_args_is_help=True)
 
+
+def _batch_size_estimate_dims(model_cfg: dict) -> tuple[int, int]:
+    """Pick the (hidden_dim, n_blocks) that dominate per-call activation memory.
+
+    Routed models spend their FLOPs in the (smaller) expert trunks, not the
+    monolith's hidden_dim/n_blocks, so estimate_batch_size needs the expert
+    dims instead when routing is enabled.
+    """
+    router_cfg = model_cfg.get("router")
+    if router_cfg and router_cfg.get("enabled"):
+        return (
+            model_cfg.get("expert_hidden_dim", 128),
+            model_cfg.get("expert_n_blocks", 3),
+        )
+    return model_cfg["hidden_dim"], model_cfg["n_blocks"]
+
+
 _CEPH_PREDICTIONS = Path("/ceph/lbogner/geant_steps/predictions")
 
 
@@ -159,6 +158,23 @@ def train(
             "--dropout", "-d", help="Dropout probability in ResBlocks (default: 0.1)"
         ),
     ] = None,
+    router: Annotated[
+        Optional[bool],
+        typer.Option(
+            "--router/--no-router",
+            help="Route both stages through a mixture of small experts "
+            "instead of one monolithic trunk (see model.router in config.toml)",
+        ),
+    ] = None,
+    router_type: Annotated[
+        Optional[str],
+        typer.Option(
+            "--router-type", help="Router implementation name (see ROUTER_REGISTRY)"
+        ),
+    ] = None,
+    n_experts: Annotated[
+        Optional[int], typer.Option("--n-experts", help="Number of routed experts")
+    ] = None,
     val_fraction: Annotated[
         Optional[float], typer.Option("--val-fraction", "-f")
     ] = None,
@@ -238,7 +254,7 @@ def train(
         }.items()
         if v is not None
     }
-    cli_model = {
+    cli_model: dict[str, object] = {
         k: v
         for k, v in {
             "hidden_dim": hidden_dim,
@@ -248,6 +264,17 @@ def train(
         }.items()
         if v is not None
     }
+    cli_router = {
+        k: v
+        for k, v in {
+            "enabled": router,
+            "type": router_type,
+            "n_experts": n_experts,
+        }.items()
+        if v is not None
+    }
+    if cli_router:
+        cli_model["router"] = cli_router
     cfg = gconfig.merge_cli_overrides(
         gconfig.DEFAULT_CONFIG, config, cli_train, cli_model
     )
@@ -256,9 +283,10 @@ def train(
     _device = torch.device(device) if device else gconfig.auto_device()
 
     if batch_size_auto:
+        est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(m)
         try:
             t["batch_size"] = gconfig.estimate_batch_size(
-                m["hidden_dim"], m["n_blocks"], _device
+                est_hidden_dim, est_n_blocks, _device
             )
         except ValueError as exc:
             typer.echo(f"error: {exc}", err=True)
@@ -386,10 +414,11 @@ def predict(
     model_cfg = ckpt["model_config"]
 
     if batch_size_auto:
+        est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(model_cfg)
         try:
             batch_size_value = gconfig.estimate_batch_size(
-                model_cfg["hidden_dim"],
-                model_cfg["n_blocks"],
+                est_hidden_dim,
+                est_n_blocks,
                 _device,
                 training=False,
             )
@@ -408,15 +437,10 @@ def predict(
     cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"])
     tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"])
 
-    model = DenoisingMLP(
-        **{k: v for k, v in model_cfg.items() if k in _STAGE1_MODEL_KEYS}
-    )
+    model, sec_decoder = build_models(model_cfg)
     model.load_state_dict(ckpt["model"])
     model.to(_device).eval()
 
-    sec_decoder = SecondaryDecoder(
-        **{k: v for k, v in model_cfg.items() if k in _SEC_DECODER_MODEL_KEYS}
-    )
     sec_decoder.load_state_dict(ckpt["sec_decoder"])
     sec_decoder.to(_device).eval()
 
@@ -759,12 +783,9 @@ def rollout(
     cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"])
     tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"])
 
-    model = DenoisingMLP(**{k: v for k, v in model_cfg.items() if k in _STAGE1_MODEL_KEYS})
+    model, sec_decoder = build_models(model_cfg)
     model.load_state_dict(ckpt["model"])
     model.to(_device).eval()
-    sec_decoder = SecondaryDecoder(
-        **{k: v for k, v in model_cfg.items() if k in _SEC_DECODER_MODEL_KEYS}
-    )
     sec_decoder.load_state_dict(ckpt["sec_decoder"])
     sec_decoder.to(_device).eval()
     typer.echo(f"loaded checkpoint: {checkpoint}")

giant/config.py

@@ -28,6 +28,16 @@ DEFAULT_CONFIG: dict = {
         "n_blocks": 6,
         "emb_dim": 16,
         "dropout": 0.1,
+        "router": {
+            "enabled": False,
+            "type": "energy",  # selects the Router impl from ROUTER_REGISTRY
+            "n_experts": 4,
+            "expert_hidden_dim": 128,
+            "expert_n_blocks": 3,
+            "temperature": 0.5,  # energy-router kwarg
+            "learn_centers": True,  # energy-router kwarg
+            "lambda_balance": 0.0,  # optional load-balance aux loss weight
+        },
     },
 }
 
@@ -164,15 +174,29 @@ def merge_cli_overrides(
     train_overrides: dict,
     model_overrides: dict,
 ) -> dict:
-    """Resolve config as defaults -> TOML file -> explicit CLI flags."""
+    """Resolve config as defaults -> TOML file -> explicit CLI flags.
+
+    `model.router` is deep-merged one level (rather than replaced wholesale)
+    at each stage, so a TOML file or CLI flag only overriding e.g.
+    `router.enabled` doesn't drop the rest of the router defaults.
+    """
     cfg = {"train": dict(defaults["train"]), "model": dict(defaults["model"])}
+    cfg["model"]["router"] = dict(defaults["model"]["router"])
     if config_path is not None:
         file_cfg = load_toml(config_path)
-        for section in ("train", "model"):
-            cfg[section].update(file_cfg.get(section, {}))
+        cfg["train"].update(file_cfg.get("train", {}))
+        file_model = dict(file_cfg.get("model", {}))
+        file_router = file_model.pop("router", None)
+        cfg["model"].update(file_model)
+        if file_router:
+            cfg["model"]["router"].update(file_router)
         warn_if_git_hash_mismatch(file_cfg, config_path)
+    model_overrides = dict(model_overrides)
+    router_overrides = model_overrides.pop("router", None)
     cfg["train"].update(train_overrides)
     cfg["model"].update(model_overrides)
+    if router_overrides:
+        cfg["model"]["router"].update(router_overrides)
     return cfg
 
 
@@ -184,17 +208,38 @@ def seed_everything(seed: int) -> None:
         torch.cuda.manual_seed_all(seed)
 
 
+def _toml_value(v) -> str:
+    if isinstance(v, bool):
+        return "true" if v else "false"
+    if isinstance(v, str):
+        return repr(v)
+    return str(v)
+
+
 def save_config(cfg: dict, out_dir: Path, meta: dict) -> None:
     lines = []
+    # One-level-nested dict values (e.g. model.router) are rendered as their
+    # own [section.subsection] table after the parent section, since TOML
+    # doesn't accept a bare dict as a `key = value` scalar line.
+    nested_sections: list[tuple[str, dict]] = []
     for section, values in cfg.items():
         lines.append(f"[{section}]")
         for k, v in values.items():
-            lines.append(f"{k:<14} = {repr(v) if isinstance(v, str) else v}")
+            if isinstance(v, dict):
+                nested_sections.append((f"{section}.{k}", v))
+                continue
+            lines.append(f"{k:<14} = {_toml_value(v)}")
+        lines.append("")
+
+    for name, values in nested_sections:
+        lines.append(f"[{name}]")
+        for k, v in values.items():
+            lines.append(f"{k:<14} = {_toml_value(v)}")
         lines.append("")
 
     lines.append("[meta]")
     for k, v in meta.items():
-        lines.append(f"{k:<14} = {repr(v) if isinstance(v, str) else v}")
+        lines.append(f"{k:<14} = {_toml_value(v)}")
 
     (out_dir / "config.toml").write_text("\n".join(lines))
 

giant/model/network.py

@@ -1,9 +1,10 @@
+import inspect
 import math
 
 import torch
 import torch.nn as nn
 
-from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM
+from giant.constants import COND_DIM, EMB_DIM, K_MAX, SEC_DIM, X_DIM
 
 
 class SinusoidalEmbedding(nn.Module):
@@ -236,3 +237,366 @@ class SecondaryDecoder(nn.Module):
         for block in self.blocks:
             x = block(x, cond)
         return self.out_proj(x)
+
+
+class Router(nn.Module):
+    """Contract for a pluggable mixture-of-experts routing axis.
+
+    Subclasses implement `gate` (soft partition-of-unity weights over
+    experts, used in train mode for a fully differentiable mixture);
+    `top1` and `balance_loss` have working defaults so a new routing axis
+    is usually a one-method add. See `ROUTER_REGISTRY` / `build_router`.
+    """
+
+    def __init__(self, n_experts: int) -> None:
+        super().__init__()
+        self.n_experts = n_experts
+
+    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
+        """(B, n_experts) soft weights, rows summing to 1."""
+        raise NotImplementedError
+
+    def top1(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
+        """(B,) hard expert index, used for eval-time grouped dispatch."""
+        return self.gate(cond_cont, cond_cat).argmax(dim=-1)
+
+    def balance_loss(
+        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
+    ) -> torch.Tensor:
+        """Importance CV^2 load-balancing auxiliary loss (Shazeer et al. 2017)."""
+        importance = self.gate(cond_cont, cond_cat).sum(dim=0)  # (n_experts,)
+        return (importance.std() / (importance.mean() + 1e-8)) ** 2
+
+
+ROUTER_REGISTRY: dict[str, type[Router]] = {}
+
+
+def register_router(name: str):
+    def decorator(cls: type[Router]) -> type[Router]:
+        ROUTER_REGISTRY[name] = cls
+        return cls
+
+    return decorator
+
+
+def build_router(name: str, n_experts: int, **kwargs) -> Router:
+    """Factory: look up a `Router` subclass by name from the registry.
+
+    Every registered router type is fed the same `model.router` config
+    dict; kwargs not declared by that type's constructor are silently
+    dropped, so per-type hyperparameters (e.g. EnergyRouter's
+    `temperature`) can coexist in one config without special-casing.
+    """
+    if name not in ROUTER_REGISTRY:
+        raise ValueError(
+            f"unknown router type {name!r}; available: {sorted(ROUTER_REGISTRY)}"
+        )
+    cls = ROUTER_REGISTRY[name]
+    accepted = set(inspect.signature(cls.__init__).parameters) - {"self", "n_experts"}
+    filtered = {k: v for k, v in kwargs.items() if k in accepted}
+    return cls(n_experts=n_experts, **filtered)
+
+
+@register_router("energy")
+class EnergyRouter(Router):
+    """Soft turn-on gate over normalized pre-step log-energy.
+
+    Reads `cond_cont[:, energy_idx]` (ignores cond_cat). Learnable (or
+    fixed) 1-D centers, initialized spread across [-2, 2] — roughly the
+    z-normalized energy range. `gate(e) = softmax_i(-(e - c_i)^2 / tau)`,
+    differentiable in e; as tau -> 0 this hardens to nearest-center
+    (Voronoi) selection, which is exactly what `top1` uses at eval.
+    """
+
+    def __init__(
+        self,
+        n_experts: int = 4,
+        temperature: float = 0.5,
+        learn_centers: bool = True,
+        energy_idx: int = 3,
+    ) -> None:
+        super().__init__(n_experts)
+        self.temperature = temperature
+        self.energy_idx = energy_idx
+        centers = torch.linspace(-2.0, 2.0, n_experts)
+        if learn_centers:
+            self.centers = nn.Parameter(centers)
+        else:
+            self.register_buffer("centers", centers)
+
+    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
+        e = cond_cont[:, self.energy_idx].unsqueeze(-1)  # (B, 1)
+        d2 = (e - self.centers.unsqueeze(0)) ** 2  # (B, n_experts)
+        return torch.softmax(-d2 / self.temperature, dim=-1)
+
+
+class ExpertTrunk(nn.Module):
+    """One small expert: `input_proj -> ResBlock stack -> out_proj`.
+
+    Same shape as the monolithic DenoisingMLP/SecondaryDecoder trunk, but
+    intended to be narrower/shallower (per-call cost is the whole point).
+    """
+
+    def __init__(
+        self,
+        in_dim: int,
+        hidden_dim: int,
+        n_blocks: int,
+        merged_cond_dim: int,
+        dropout: float = 0.1,
+    ) -> None:
+        super().__init__()
+        self.input_proj = nn.Linear(in_dim, hidden_dim)
+        self.blocks = nn.ModuleList(
+            [
+                ResBlock(hidden_dim, merged_cond_dim, dropout=dropout)
+                for _ in range(n_blocks)
+            ]
+        )
+        self.out_proj = nn.Linear(hidden_dim, in_dim)
+
+    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+        x = self.input_proj(x)
+        for block in self.blocks:
+            x = block(x, cond)
+        return self.out_proj(x)
+
+
+def _route_forward(
+    experts: nn.ModuleList,
+    router: Router,
+    x: torch.Tensor,
+    cond: torch.Tensor,
+    cond_cont: torch.Tensor,
+    cond_cat: torch.Tensor,
+    training: bool,
+) -> torch.Tensor:
+    """Shared dispatch for both Routed* trunks.
+
+    Train mode: full soft mixture `sum_i gate_i * expert_i(x)` — fully
+    differentiable, N-expert compute. Eval mode: grouped top-1 dispatch —
+    each row runs exactly one (small) expert, which is the actual source
+    of the per-call speedup this architecture is for.
+    """
+    if training:
+        weights = router.gate(cond_cont, cond_cat)  # (B, n_experts)
+        out = torch.zeros_like(x)
+        for i, expert in enumerate(experts):
+            out = out + weights[:, i : i + 1] * expert(x, cond)
+        return out
+
+    idx = router.top1(cond_cont, cond_cat)  # (B,)
+    out = torch.zeros_like(x)
+    for i, expert in enumerate(experts):
+        mask = idx == i
+        if mask.any():
+            out[mask] = expert(x[mask], cond[mask])
+    return out
+
+
+class RoutedDenoisingMLP(nn.Module):
+    """Routed drop-in for `DenoisingMLP`.
+
+    Shares the time embedding, `ConditionEncoder`, and `n_sec_head` (all
+    tiny) across experts and routes only the trunk (where the FLOPs are).
+    Same `forward`/`predict_n_sec`/`pdg_embedding_weight` signatures as
+    `DenoisingMLP`, so sample.py/rollout.py/validate.py need no changes.
+    """
+
+    def __init__(
+        self,
+        pdg_vocab: int,
+        mat_vocab: int,
+        router: Router,
+        expert_hidden_dim: int = 128,
+        expert_n_blocks: int = 3,
+        emb_dim: int = EMB_DIM,
+        time_dim: int = 64,
+        cond_out_dim: int = 128,
+        x_dim: int = X_DIM,
+        dropout: float = 0.1,
+        k_max: int = K_MAX,
+    ) -> None:
+        super().__init__()
+        self.router = router
+        self.time_emb = SinusoidalEmbedding(time_dim)
+        self.cond_enc = ConditionEncoder(
+            pdg_vocab=pdg_vocab,
+            mat_vocab=mat_vocab,
+            emb_dim=emb_dim,
+            out_dim=cond_out_dim,
+        )
+        merged_cond_dim = time_dim + cond_out_dim
+        self.experts = nn.ModuleList(
+            [
+                ExpertTrunk(
+                    x_dim,
+                    expert_hidden_dim,
+                    expert_n_blocks,
+                    merged_cond_dim,
+                    dropout=dropout,
+                )
+                for _ in range(router.n_experts)
+            ]
+        )
+        self.n_sec_head = nn.Sequential(
+            nn.Linear(cond_out_dim, cond_out_dim),
+            nn.SiLU(),
+            nn.Linear(cond_out_dim, k_max + 1),
+        )
+
+    def forward(
+        self,
+        x_t: torch.Tensor,
+        t: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+    ) -> torch.Tensor:
+        t_emb = self.time_emb(t)
+        c_emb = self.cond_enc(cond_cont, cond_cat)
+        cond = torch.cat([t_emb, c_emb], dim=-1)
+        return _route_forward(
+            self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training
+        )
+
+    def predict_n_sec(
+        self,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+    ) -> torch.Tensor:
+        """Return n_sec logits (B, K_MAX+1) from conditioning alone."""
+        c_emb = self.cond_enc(cond_cont, cond_cat)
+        return self.n_sec_head(c_emb)
+
+    def pdg_embedding_weight(self) -> torch.Tensor:
+        """Return the PDG embedding table weights for secondary type targets."""
+        return self.cond_enc.pdg_emb.weight
+
+
+class RoutedSecondaryDecoder(nn.Module):
+    """Routed drop-in for `SecondaryDecoder`.
+
+    Shares the time embedding and `SecondaryConditionEncoder` across
+    experts and routes only the trunk. Same `forward` signature as
+    `SecondaryDecoder`.
+    """
+
+    def __init__(
+        self,
+        pdg_vocab: int,
+        mat_vocab: int,
+        router: Router,
+        expert_hidden_dim: int = 128,
+        expert_n_blocks: int = 3,
+        emb_dim: int = EMB_DIM,
+        time_dim: int = 64,
+        cond_out_dim: int = 128,
+        stage1_proj_dim: int = 64,
+        sec_dim: int = SEC_DIM,
+        dropout: float = 0.1,
+    ) -> None:
+        super().__init__()
+        self.router = router
+        self.time_emb = SinusoidalEmbedding(time_dim)
+        self.cond_enc = SecondaryConditionEncoder(
+            pdg_vocab=pdg_vocab,
+            mat_vocab=mat_vocab,
+            emb_dim=emb_dim,
+            cond_out_dim=cond_out_dim,
+            stage1_proj_dim=stage1_proj_dim,
+            out_dim=cond_out_dim,
+        )
+        merged_cond_dim = time_dim + cond_out_dim
+        self.experts = nn.ModuleList(
+            [
+                ExpertTrunk(
+                    sec_dim,
+                    expert_hidden_dim,
+                    expert_n_blocks,
+                    merged_cond_dim,
+                    dropout=dropout,
+                )
+                for _ in range(router.n_experts)
+            ]
+        )
+
+    def forward(
+        self,
+        x_t: torch.Tensor,
+        t: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_out: torch.Tensor,
+    ) -> torch.Tensor:
+        t_emb = self.time_emb(t)
+        c_emb = self.cond_enc(cond_cont, cond_cat, stage1_out)
+        cond = torch.cat([t_emb, c_emb], dim=-1)
+        return _route_forward(
+            self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training
+        )
+
+
+_STAGE1_MODEL_KEYS = {
+    "pdg_vocab",
+    "mat_vocab",
+    "hidden_dim",
+    "n_blocks",
+    "emb_dim",
+    "dropout",
+    "k_max",
+}
+_SEC_DECODER_MODEL_KEYS = {
+    "pdg_vocab",
+    "mat_vocab",
+    "hidden_dim",
+    "n_blocks",
+    "emb_dim",
+    "dropout",
+}
+
+
+def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]:
+    """Construct (stage1, sec_decoder) from a persisted/CLI model_config dict.
+
+    Dispatches to the routed pair when `model_config["router"]["enabled"]`
+    is truthy; a missing/absent "router" key (pre-routing checkpoints)
+    falls back to the monolithic pair unchanged, so this is a drop-in
+    replacement for the ad-hoc constructions it replaces.
+    """
+    router_cfg = model_config.get("router")
+    if router_cfg and router_cfg.get("enabled"):
+        router_kwargs = {
+            k: v
+            for k, v in router_cfg.items()
+            if k not in ("enabled", "type", "n_experts")
+        }
+        shared = dict(
+            pdg_vocab=model_config["pdg_vocab"],
+            mat_vocab=model_config["mat_vocab"],
+            expert_hidden_dim=model_config.get("expert_hidden_dim", 128),
+            expert_n_blocks=model_config.get("expert_n_blocks", 3),
+            emb_dim=model_config.get("emb_dim", EMB_DIM),
+            dropout=model_config.get("dropout", 0.1),
+        )
+        stage1 = RoutedDenoisingMLP(
+            router=build_router(
+                router_cfg["type"], router_cfg["n_experts"], **router_kwargs
+            ),
+            k_max=model_config.get("k_max", K_MAX),
+            **shared,
+        )
+        sec_decoder = RoutedSecondaryDecoder(
+            router=build_router(
+                router_cfg["type"], router_cfg["n_experts"], **router_kwargs
+            ),
+            **shared,
+        )
+        return stage1, sec_decoder
+
+    stage1 = DenoisingMLP(
+        **{k: v for k, v in model_config.items() if k in _STAGE1_MODEL_KEYS}
+    )
+    sec_decoder = SecondaryDecoder(
+        **{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS}
+    )
+    return stage1, sec_decoder

giant/pipeline.py

@@ -14,7 +14,7 @@ from giant.data.loader import (
 )
 from giant.data.transforms import build_features, _WelfordAccumulator
 from giant.data.dataset import make_event_split, StreamingStepsDataset
-from giant.model.network import DenoisingMLP, SecondaryDecoder
+from giant.model.network import build_models
 from giant.train import train as run_training
 
 
@@ -114,23 +114,22 @@ def run_train_job(
         "update giant/constants.py if emb_dim changed"
     )
 
-    stage1_model = DenoisingMLP(
-        pdg_vocab=len(pdg_map),
-        mat_vocab=len(mat_map),
-        hidden_dim=m["hidden_dim"],
-        n_blocks=m["n_blocks"],
-        emb_dim=emb_dim,
-        dropout=m["dropout"],
-        k_max=K_MAX,
-    )
-    sec_decoder = SecondaryDecoder(
-        pdg_vocab=len(pdg_map),
-        mat_vocab=len(mat_map),
-        hidden_dim=m["hidden_dim"],
-        n_blocks=m["n_blocks"],
-        emb_dim=emb_dim,
-        dropout=m["dropout"],
-    )
+    router_cfg = m["router"]
+    model_config = {
+        "pdg_vocab": len(pdg_map),
+        "mat_vocab": len(mat_map),
+        "hidden_dim": m["hidden_dim"],
+        "n_blocks": m["n_blocks"],
+        "emb_dim": emb_dim,
+        "dropout": m["dropout"],
+        "k_max": K_MAX,
+        "sec_slot_dim": SEC_SLOT_DIM,
+        "router": dict(router_cfg),
+        "expert_hidden_dim": router_cfg["expert_hidden_dim"],
+        "expert_n_blocks": router_cfg["expert_n_blocks"],
+    }
+
+    stage1_model, sec_decoder = build_models(model_config)
     echo(
         f"stage1: {sum(p.numel() for p in stage1_model.parameters()):,} parameters  |  "
         f"sec_decoder: {sum(p.numel() for p in sec_decoder.parameters()):,} parameters"
@@ -148,17 +147,6 @@ def run_train_job(
     )
     config.save_config(cfg, out_dir, meta)
 
-    model_config = {
-        "pdg_vocab": len(pdg_map),
-        "mat_vocab": len(mat_map),
-        "hidden_dim": m["hidden_dim"],
-        "n_blocks": m["n_blocks"],
-        "emb_dim": emb_dim,
-        "dropout": m["dropout"],
-        "k_max": K_MAX,
-        "sec_slot_dim": SEC_SLOT_DIM,
-    }
-
     run_training(
         stage1_model=stage1_model,
         sec_decoder=sec_decoder,
@@ -172,6 +160,7 @@ def run_train_job(
         out_dir=out_dir,
         lambda_nsec=t.get("lambda_nsec", 0.1),
         lambda_s2=t.get("lambda_s2", 1.0),
+        lambda_balance=router_cfg.get("lambda_balance", 0.0),
         normalizer_dict={"cond": cond_norm.to_dict(), "target": tgt_norm.to_dict()},
         pdg_map={str(k): v for k, v in pdg_map.items()},
         mat_map={str(k): v for k, v in mat_map.items()},

giant/train.py

@@ -26,10 +26,12 @@ _METRICS_FIELDS = [
     "train_loss_s1",
     "train_loss_nsec",
     "train_loss_s2",
+    "train_loss_balance",
     "val_loss",
     "val_loss_s1",
     "val_loss_nsec",
     "val_loss_s2",
+    "val_loss_balance",
     "lr",
     "epoch_time_s",
 ]
@@ -108,8 +110,9 @@ def _compute_losses(
     device: torch.device,
     lambda_nsec: float,
     lambda_s2: float,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    """Compute (total_loss, L_s1, L_nsec, L_s2) for one batch."""
+    lambda_balance: float = 0.0,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Compute (total_loss, L_s1, L_nsec, L_s2, L_balance) for one batch."""
     cond_cont, cond_cat, x1_s1, n_sec, sec_cont, sec_pdg_idx = batch
     cond_cont = cond_cont.to(device)
     cond_cat = cond_cat.to(device)
@@ -150,8 +153,21 @@ def _compute_losses(
         sec_mask,
     )
 
+    # Optional MoE load-balance auxiliary loss: only present when both stages
+    # are routed (RoutedDenoisingMLP/RoutedSecondaryDecoder carry `.router`,
+    # the monolith models don't), computed on cond_cont alone (cheap — no
+    # trunk compute) so it's reported even when lambda_balance == 0.
+    if hasattr(stage1_model, "router") and hasattr(sec_decoder, "router"):
+        l_balance = stage1_model.router.balance_loss(
+            cond_cont, cond_cat
+        ) + sec_decoder.router.balance_loss(cond_cont, cond_cat)
+    else:
+        l_balance = torch.zeros((), device=device)
+
     total = l_s1 + lambda_nsec * l_nsec + lambda_s2 * l_s2
-    return total, l_s1, l_nsec, l_s2
+    if lambda_balance > 0:
+        total = total + lambda_balance * l_balance
+    return total, l_s1, l_nsec, l_s2, l_balance
 
 
 def train(
@@ -167,6 +183,7 @@ def train(
     out_dir: str | Path,
     lambda_nsec: float = 0.1,
     lambda_s2: float = 1.0,
+    lambda_balance: float = 0.0,
     normalizer_dict: dict | None = None,
     pdg_map: dict | None = None,
     mat_map: dict | None = None,
@@ -245,6 +262,7 @@ def train(
             train_s1_sum = 0.0
             train_nsec_sum = 0.0
             train_s2_sum = 0.0
+            train_balance_sum = 0.0
             train_n = 0
             ema_loss = 0.0
             bar = tqdm(
@@ -256,7 +274,7 @@ def train(
                 dynamic_ncols=True,
             )
             for batch in bar:
-                loss, l_s1, l_nsec, l_s2 = _compute_losses(
+                loss, l_s1, l_nsec, l_s2, l_balance = _compute_losses(
                     stage1_model,
                     sec_decoder,
                     batch,
@@ -265,6 +283,7 @@ def train(
                     device,
                     lambda_nsec,
                     lambda_s2,
+                    lambda_balance,
                 )
                 optimizer.zero_grad()
                 loss.backward()
@@ -277,6 +296,7 @@ def train(
                 train_s1_sum += l_s1.item() * B
                 train_nsec_sum += l_nsec.item() * B
                 train_s2_sum += l_s2.item() * B
+                train_balance_sum += l_balance.item() * B
                 train_n += B
                 ema_loss = (
                     batch_loss if train_n == B else 0.95 * ema_loss + 0.05 * batch_loss
@@ -299,10 +319,11 @@ def train(
             val_s1_sum = 0.0
             val_nsec_sum = 0.0
             val_s2_sum = 0.0
+            val_balance_sum = 0.0
             val_n = 0
             with torch.no_grad():
                 for batch in val_loader:
-                    loss, l_s1, l_nsec, l_s2 = _compute_losses(
+                    loss, l_s1, l_nsec, l_s2, l_balance = _compute_losses(
                         stage1_model,
                         sec_decoder,
                         batch,
@@ -311,12 +332,14 @@ def train(
                         device,
                         lambda_nsec,
                         lambda_s2,
+                        lambda_balance,
                     )
                     B = batch[0].size(0)
                     val_loss_sum += loss.item() * B
                     val_s1_sum += l_s1.item() * B
                     val_nsec_sum += l_nsec.item() * B
                     val_s2_sum += l_s2.item() * B
+                    val_balance_sum += l_balance.item() * B
                     val_n += B
             val_loss = val_loss_sum / max(val_n, 1)
             epoch_time = time.monotonic() - epoch_start
@@ -328,7 +351,8 @@ def train(
                 f"  train {train_loss:.4f}"
                 f"  (s1={train_s1_sum / max(train_n, 1):.3f}"
                 f"  nsec={train_nsec_sum / max(train_n, 1):.3f}"
-                f"  s2={train_s2_sum / max(train_n, 1):.3f})"
+                f"  s2={train_s2_sum / max(train_n, 1):.3f}"
+                f"  bal={train_balance_sum / max(train_n, 1):.3f})"
                 f"  val {val_loss:.4f}"
                 f"  lr {current_lr:.2e}  {epoch_time:.1f}s{marker}"
             )
@@ -339,10 +363,12 @@ def train(
                     "train_loss_s1": train_s1_sum / max(train_n, 1),
                     "train_loss_nsec": train_nsec_sum / max(train_n, 1),
                     "train_loss_s2": train_s2_sum / max(train_n, 1),
+                    "train_loss_balance": train_balance_sum / max(train_n, 1),
                     "val_loss": val_loss,
                     "val_loss_s1": val_s1_sum / max(val_n, 1),
                     "val_loss_nsec": val_nsec_sum / max(val_n, 1),
                     "val_loss_s2": val_s2_sum / max(val_n, 1),
+                    "val_loss_balance": val_balance_sum / max(val_n, 1),
                     "lr": current_lr,
                     "epoch_time_s": epoch_time,
                 }

tests/test_router.py

@@ -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)