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>

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