From b51eafcfa5fa58326a53a488f7a282f5f424a8c5 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Thu, 30 Jul 2026 16:29:20 +0200 Subject: [PATCH] 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 --- giant/config.py | 13 ++++++ giant/model/network.py | 67 +++++++++++++++++++++++--- giant/pipeline.py | 2 + giant/train.py | 24 ++++++++++ tests/test_config.py | 10 ++++ tests/test_router.py | 104 +++++++++++++++++++++++++++++++++++++++++ tests/test_train.py | 26 +++++++++++ 7 files changed, 239 insertions(+), 7 deletions(-) create mode 100644 tests/test_train.py diff --git a/giant/config.py b/giant/config.py index d21be99..37d91b1 100644 --- a/giant/config.py +++ b/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 diff --git a/giant/model/network.py b/giant/model/network.py index f517d02..d04e4c4 100644 --- a/giant/model/network.py +++ b/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]: diff --git a/giant/pipeline.py b/giant/pipeline.py index 53bdc61..ed12e55 100644 --- a/giant/pipeline.py +++ b/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(), diff --git a/giant/train.py b/giant/train.py index b0d8654..bcc8a38 100644 --- a/giant/train.py +++ b/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: diff --git a/tests/test_config.py b/tests/test_config.py index 76f1da0..3795367 100644 --- a/tests/test_config.py +++ b/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) diff --git a/tests/test_router.py b/tests/test_router.py index 05105dc..34c669e 100644 --- a/tests/test_router.py +++ b/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 ──────────────────────────────────────────────────────────────── diff --git a/tests/test_train.py b/tests/test_train.py new file mode 100644 index 0000000..0b0a25d --- /dev/null +++ b/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