Add opt-in straight-through Gumbel-softmax combine weights to MoE router

Trains the routed trunk's forward combination as a hard one-hot sample
(matching eval-time top-1 dispatch exactly) while keeping a smooth gradient
on the backward pass, targeting the train/eval mismatch identified as a
likely contributor to experts overlapping instead of partitioning in the
first energy-router rollout benchmark. Off by default (model.router.gumbel);
existing routed configs/checkpoints are unaffected.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

giant/config.py

@@ -99,6 +99,19 @@ DEFAULT_CONFIG: dict = {
             # failure mode. Off by default; bounding above is the primary
             # defense. See giant.model.network.Router.entropy_loss.
             "lambda_entropy": 0.0,
+            # Opt-in straight-through Gumbel-softmax train-time combine weights
+            # (see giant.model.network.Router.combine_weights): the training
+            # forward pass samples a hard one-hot combination — matching
+            # eval-time top-1 dispatch exactly — while the backward pass still
+            # flows a smooth gradient to every expert. Targets the train/eval
+            # mismatch identified as a likely contributor to experts
+            # overlapping instead of partitioning (see CLAUDE.md roadmap).
+            # gumbel_tau_start/_end are annealed linearly over training
+            # (giant.train._gumbel_tau); off by default, no effect unless
+            # gumbel = true.
+            "gumbel": False,
+            "gumbel_tau_start": 1.0,
+            "gumbel_tau_end": 0.1,
             "emb_dim": 8,  # process/pdg-router kwarg: own pdg(/mat) embedding width
             "hidden_dim": 64,  # process-router kwarg: its classifier's hidden width
             "lambda_proc": 0.0,  # process-router kwarg: supervised process-CE weight

giant/model/network.py

@@ -537,11 +537,49 @@ class Router(nn.Module):
     def __init__(self, n_experts: int) -> None:
         super().__init__()
         self.n_experts = n_experts
+        # Opt-in straight-through Gumbel-softmax combine weights (see
+        # combine_weights below) — off by default, set from model.router.gumbel
+        # by _build_router_from_cfg. gumbel_tau is annealed per training step
+        # by giant.train (model.router.gumbel_tau_start/_end); neither is an
+        # nn.Parameter/buffer since neither is learned or needs checkpointing —
+        # the tau schedule is deterministic in global_step, so it recomputes
+        # correctly on resume.
+        self.gumbel = False
+        self.gumbel_tau = 1.0
 
     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 combine_weights(
+        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
+    ) -> torch.Tensor:
+        """(B, n_experts) train-time expert-combination weights.
+
+        Default (`gumbel=False`): identical to `gate()` — the original dense
+        soft-mixture combination. Opt-in straight-through Gumbel-softmax
+        (`gumbel=True`, train mode only): samples a Gumbel-perturbed
+        categorical draw from the same distribution `gate()` defines
+        (`log(gate())` is a valid unnormalized-logit input to
+        `F.gumbel_softmax` since softmax is shift-invariant, so no subclass
+        needs to expose separate pre-softmax logits), then hardens it to a
+        one-hot vector on the forward pass while keeping the soft sample's
+        gradient on the backward pass. This makes the training-time forward
+        combination match eval-time top-1 dispatch exactly (one expert's
+        output, unweighted) instead of the smooth blend `gate()` gives —
+        intended to close the train/eval mismatch identified as a likely
+        cause of experts overlapping instead of partitioning (see the
+        router_gating write-up referenced in CLAUDE.md's roadmap).
+        `gate()` itself is untouched and still backs `balance_loss`/
+        `entropy_loss`/`gate_stats`, so those diagnostics keep reading the
+        smooth distribution rather than a noisy sample.
+        """
+        probs = self.gate(cond_cont, cond_cat)
+        if not (self.gumbel and self.training):
+            return probs
+        log_probs = torch.log(probs.clamp_min(1e-8))
+        return F.gumbel_softmax(log_probs, tau=self.gumbel_tau, hard=True, dim=-1)
+
     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)
@@ -970,13 +1008,17 @@ def _route_forward(
 ) -> 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.
+    Train mode: full mixture `sum_i weight_i * expert_i(x)` — always
+    N-expert dense compute, fully differentiable. `weight` is
+    `router.combine_weights(...)`: the plain soft `gate()` by default, or (see
+    `Router.combine_weights`) a straight-through Gumbel-softmax one-hot sample
+    when `router.gumbel` is enabled — either way, no change to the compute
+    cost of this branch. 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)
+        weights = router.combine_weights(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)
@@ -1191,10 +1233,19 @@ def _build_router_from_cfg(router_cfg: dict, pdg_vocab: int, mat_vocab: int) ->
 
     `router_cfg["type"] == "composed"` reads `axis{i}_{field}` flat keys
     (see `_parse_composed_axes`) instead of a single `type`/`n_experts` pair.
+
+    `gumbel` is set as a post-construction attribute here rather than a
+    per-subclass constructor kwarg, same reasoning as `lambda_balance`/
+    `lambda_proc`/`lambda_entropy` living in `router_cfg` without being a
+    `Router` subclass constructor param: it's a training-time toggle shared by
+    every router type, not a per-type hyperparameter (`build_router`'s
+    kwarg-filtering would otherwise just silently drop it).
     """
     shared_vocab = dict(pdg_vocab=pdg_vocab, mat_vocab=mat_vocab)
     if router_cfg["type"] == "composed":
-        return build_composed_router(_parse_composed_axes(router_cfg), **shared_vocab)
+        router = build_composed_router(_parse_composed_axes(router_cfg), **shared_vocab)
+        router.gumbel = bool(router_cfg.get("gumbel", False))
+        return router
     router_kwargs = {
         k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts")
     }
@@ -1204,7 +1255,9 @@ def _build_router_from_cfg(router_cfg: dict, pdg_vocab: int, mat_vocab: int) ->
     # vocab, same as the trunk's ConditionEncoder.
     router_kwargs.setdefault("pdg_vocab", pdg_vocab)
     router_kwargs.setdefault("mat_vocab", mat_vocab)
-    return build_router(router_cfg["type"], router_cfg["n_experts"], **router_kwargs)
+    router = build_router(router_cfg["type"], router_cfg["n_experts"], **router_kwargs)
+    router.gumbel = bool(router_cfg.get("gumbel", False))
+    return router
 
 
 def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]:

giant/pipeline.py

@@ -417,6 +417,8 @@ def run_train_job(
         lambda_balance=router_cfg.get("lambda_balance", 0.0),
         lambda_proc=router_cfg.get("lambda_proc", 0.0),
         lambda_entropy=router_cfg.get("lambda_entropy", 0.0),
+        gumbel_tau_start=router_cfg.get("gumbel_tau_start", 1.0),
+        gumbel_tau_end=router_cfg.get("gumbel_tau_end", 0.1),
         normalizer_dict={
             "cond": cond_norm.to_dict(),
             "target": tgt_norm.to_dict(),

giant/train.py

@@ -114,6 +114,18 @@ def _update_ema(
         ema_p.mul_(decay).add_(p, alpha=1 - decay)
 
 
+def _gumbel_tau(step: int, total_steps: int, tau_start: float, tau_end: float) -> float:
+    """Linear anneal of the straight-through Gumbel-softmax temperature.
+
+    Deterministic in `step`/`total_steps` alone (no extra state), so it
+    recomputes correctly on `--resume` from a checkpoint's saved `global_step`
+    without needing to persist anything new (see
+    giant.model.network.Router.combine_weights).
+    """
+    progress = min(step / max(total_steps, 1), 1.0)
+    return tau_start + (tau_end - tau_start) * progress
+
+
 def _compute_losses(
     stage1_model: torch.nn.Module,
     sec_decoder: torch.nn.Module,
@@ -347,6 +359,8 @@ def train(
     lambda_balance: float = 0.0,
     lambda_proc: float = 0.0,
     lambda_entropy: float = 0.0,
+    gumbel_tau_start: float = 1.0,
+    gumbel_tau_end: float = 0.1,
     normalizer_dict: dict | None = None,
     pdg_map: dict | None = None,
     mat_map: dict | None = None,
@@ -410,6 +424,8 @@ def train(
                 "lambda_balance": lambda_balance,
                 "lambda_proc": lambda_proc,
                 "lambda_entropy": lambda_entropy,
+                "gumbel_tau_start": gumbel_tau_start,
+                "gumbel_tau_end": gumbel_tau_end,
                 "n_critic": n_critic,
                 "gp_weight": gp_weight,
                 "model": model_config or {},
@@ -596,6 +612,13 @@ def train(
                 dynamic_ncols=True,
             )
             for batch in bar:
+                if has_router:
+                    gumbel_tau = _gumbel_tau(
+                        global_step, total_steps, gumbel_tau_start, gumbel_tau_end
+                    )
+                    stage1_model.router.gumbel_tau = gumbel_tau
+                    sec_decoder.router.gumbel_tau = gumbel_tau
+
                 if mode == "wgan":
                     assert (
                         critic is not None
@@ -737,6 +760,7 @@ def train(
                             )
                         log_payload["batch/router_s1_entropy"] = s1_entropy.item()
                         log_payload["batch/router_s2_entropy"] = s2_entropy.item()
+                        log_payload["batch/gumbel_tau"] = gumbel_tau
                     wandb_run.log(log_payload, step=global_step)
 
                 if shutdown.requested:

tests/test_config.py

@@ -137,6 +137,16 @@ def test_resolve_expert_dims_missing_keys_also_inherit():
     assert (hidden_dim, n_blocks) == (512, 6)
 
 
+def test_default_config_gumbel_router_defaults_off():
+    # Straight-through Gumbel-softmax combine weights (giant.model.network.
+    # Router.combine_weights) must be opt-in — existing routed configs and
+    # checkpoints should be unaffected unless gumbel is explicitly enabled.
+    router_cfg = gconfig.DEFAULT_CONFIG["model"]["router"]
+    assert router_cfg["gumbel"] is False
+    assert router_cfg["gumbel_tau_start"] == 1.0
+    assert router_cfg["gumbel_tau_end"] == 0.1
+
+
 def test_resolve_expert_dims_explicit_override_wins():
     router_cfg = {"expert_hidden_dim": 128, "expert_n_blocks": 3}
     hidden_dim, n_blocks = gconfig.resolve_expert_dims(router_cfg, 512, 6)

tests/test_router.py

@@ -286,6 +286,110 @@ def test_router_entropy_loss_is_nonnegative_bounded_scalar():
     assert 0.0 <= loss.item() <= 1.0
 
 
+# ── Router.combine_weights (straight-through Gumbel-softmax) ───────────────
+
+
+def test_combine_weights_defaults_to_gate():
+    """gumbel=False (the default) must be a pure pass-through to gate()."""
+    router = EnergyRouter(n_experts=4)
+    cond_cont, cond_cat = _cond(16)
+    torch.testing.assert_close(
+        router.combine_weights(cond_cont, cond_cat),
+        router.gate(cond_cont, cond_cat),
+    )
+
+
+def test_combine_weights_gumbel_train_mode_is_hard_one_hot():
+    router = EnergyRouter(n_experts=4)
+    router.gumbel = True
+    router.gumbel_tau = 0.5
+    router.train()
+    cond_cont, cond_cat = _cond(16)
+    weights = router.combine_weights(cond_cont, cond_cat)
+    assert weights.shape == (16, 4)
+    torch.testing.assert_close(weights.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0)
+    assert torch.all((weights.max(dim=-1).values - 1.0).abs() < 1e-5)
+
+
+def test_combine_weights_gumbel_eval_mode_falls_back_to_gate():
+    """No Gumbel noise at eval — combine_weights must match gate() exactly,
+    same as the gumbel=False path, once the router is in eval mode."""
+    router = EnergyRouter(n_experts=4)
+    router.gumbel = True
+    router.eval()
+    cond_cont, cond_cat = _cond(16)
+    torch.testing.assert_close(
+        router.combine_weights(cond_cont, cond_cat),
+        router.gate(cond_cont, cond_cat),
+    )
+
+
+def test_combine_weights_gumbel_straight_through_gradient_reaches_centers():
+    router = EnergyRouter(n_experts=4, learn_centers=True)
+    router.gumbel = True
+    router.gumbel_tau = 0.5
+    router.train()
+    cond_cont, cond_cat = _cond(16)
+    weights = router.combine_weights(cond_cont, cond_cat)
+    weights.sum().backward()
+    assert router.centers.grad is not None
+    assert torch.any(router.centers.grad != 0.0)
+
+
+def test_build_router_from_cfg_sets_gumbel_from_config():
+    from giant.model.network import _build_router_from_cfg
+
+    router = _build_router_from_cfg(
+        {"enabled": True, "type": "energy", "n_experts": 4, "gumbel": True},
+        pdg_vocab=3,
+        mat_vocab=2,
+    )
+    assert router.gumbel is True
+
+    router_off = _build_router_from_cfg(
+        {"enabled": True, "type": "energy", "n_experts": 4},
+        pdg_vocab=3,
+        mat_vocab=2,
+    )
+    assert router_off.gumbel is False
+
+
+def test_build_router_from_cfg_sets_gumbel_for_composed_router():
+    from giant.model.network import _build_router_from_cfg
+
+    router = _build_router_from_cfg(
+        {
+            "enabled": True,
+            "type": "composed",
+            "gumbel": True,
+            "axis0_type": "energy",
+            "axis0_n_experts": 4,
+            "axis1_type": "pdg",
+            "axis1_n_experts": 3,
+        },
+        pdg_vocab=5,
+        mat_vocab=2,
+    )
+    assert isinstance(router, ComposedRouter)
+    assert router.gumbel is True
+
+
+def test_routed_denoising_mlp_forward_runs_with_gumbel_enabled():
+    """End-to-end forward through _route_forward's train branch with
+    straight-through Gumbel-softmax combine weights enabled."""
+    B = 8
+    model = _routed_stage1(n_experts=3)
+    model.router.gumbel = True
+    model.router.gumbel_tau = 0.5
+    model.train()
+    x_t = torch.randn(B, X_DIM)
+    t = torch.rand(B)
+    cond_cont, cond_cat = _cond(B)
+    out = model(x_t, t, cond_cont, cond_cat)
+    assert out.shape == (B, X_DIM)
+    assert torch.isfinite(out).all()
+
+
 # ── PdgRouter ────────────────────────────────────────────────────────────────
 
 

tests/test_train.py

@@ -0,0 +1,26 @@
+"""Tests for giant/train.py helpers."""
+
+from giant.train import _gumbel_tau
+
+
+def test_gumbel_tau_at_step_zero_is_start():
+    assert _gumbel_tau(0, 1000, 1.0, 0.1) == 1.0
+
+
+def test_gumbel_tau_at_total_steps_is_end():
+    assert abs(_gumbel_tau(1000, 1000, 1.0, 0.1) - 0.1) < 1e-9
+
+
+def test_gumbel_tau_interpolates_linearly_midway():
+    assert abs(_gumbel_tau(500, 1000, 1.0, 0.1) - 0.55) < 1e-9
+
+
+def test_gumbel_tau_clamps_beyond_total_steps():
+    assert _gumbel_tau(5000, 1000, 1.0, 0.1) == _gumbel_tau(1000, 1000, 1.0, 0.1)
+
+
+def test_gumbel_tau_handles_zero_total_steps():
+    # total_steps=0 is guarded to 1 internally: step=0 gives zero progress
+    # (still tau_start), any step>=1 immediately clamps to full progress.
+    assert _gumbel_tau(0, 0, 1.0, 0.1) == 1.0
+    assert abs(_gumbel_tau(1, 0, 1.0, 0.1) - 0.1) < 1e-9