Merge pull request 'Add learnable per-expert width and shared temperature to EnergyRouter' (#25) from feat/router-learnable-width into master

Reviewed-on: #25

giant/config.py

@@ -77,7 +77,28 @@ DEFAULT_CONFIG: dict = {
             "expert_n_blocks": 0,
             "temperature": 0.5,  # energy/pdg-router kwarg
             "learn_centers": True,  # energy/pdg-router kwarg
+            # energy-router kwargs: mutually exclusive optional learnable
+            # gate-sharpness modes (see giant.model.network.EnergyRouter).
+            # learn_width generalizes the shared `temperature` to one
+            # learnable width per expert; learn_temperature instead makes
+            # the single shared `temperature` itself learnable. Both are
+            # bounded to [width_min_ratio, width_max_ratio] * temperature
+            # (sigmoid-parameterized, warm-started to reproduce `temperature`
+            # exactly at init) so gate sharpness can't run away to a
+            # collapse-inducing extreme during training.
+            "learn_width": False,
+            "learn_temperature": False,
+            "width_min_ratio": 0.1,
+            "width_max_ratio": 10.0,
             "lambda_balance": 0.0,  # optional load-balance aux loss weight
+            # optional entropy-regularization aux loss weight (generic
+            # Router.entropy_loss, penalizes uniform/collapsed gating) — a
+            # secondary guard against all experts' widths/temperature
+            # co-inflating together, which lambda_balance alone can't see
+            # since per-expert usage shares stay even throughout that
+            # failure mode. Off by default; bounding above is the primary
+            # defense. See giant.model.network.Router.entropy_loss.
+            "lambda_entropy": 0.0,
             "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

@@ -566,6 +566,25 @@ class Router(nn.Module):
         """
         return torch.zeros((), device=cond_cont.device)
 
+    def entropy_loss(
+        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
+    ) -> torch.Tensor:
+        """Optional auxiliary loss rewarding sharper (lower-entropy) routing.
+
+        Reuses `gate_stats`'s `norm_entropy` (already in [0, 1], 1.0 =
+        uniform/collapsed) directly as the loss, so minimizing it pushes
+        every router's gate toward decisiveness. A generic base-class
+        default — works for any Router via gate_stats, no per-subclass
+        override needed. Off by default (see `lambda_entropy` in
+        giant.train): bounded width/temperature (EnergyRouter's
+        `learn_width`/`learn_temperature`) is the primary defense against
+        gate collapse; this is a secondary, use-with-caution lever, since
+        indiscriminately penalizing entropy can also suppress legitimate
+        soft ambiguity near a router's own decision boundary.
+        """
+        norm_entropy, _ = self.gate_stats(cond_cont, cond_cat)
+        return norm_entropy
+
     def gate_stats(
         self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
     ) -> tuple[torch.Tensor, torch.Tensor]:
@@ -620,6 +639,22 @@ def build_router(name: str, n_experts: int, **kwargs) -> Router:
     return cls(n_experts=n_experts, **filtered)
 
 
+def _bounded_interp(raw: torch.Tensor, lo: float, hi: float) -> torch.Tensor:
+    """Sigmoid interpolation into `[lo, hi]` — smooth, always-positive-gradient
+    bound (unlike `clamp`, which zeroes gradient past the boundary) used for
+    EnergyRouter's `learn_width`/`learn_temperature` modes."""
+    return lo + (hi - lo) * torch.sigmoid(raw)
+
+
+def _inverse_bounded_interp(value: float, lo: float, hi: float) -> float:
+    """Inverse of `_bounded_interp`, used once at construction to warm-start
+    `raw` so `_bounded_interp(raw, lo, hi) == value` — lets `learn_width`/
+    `learn_temperature` start out exactly reproducing the fixed-`temperature`
+    gate before any training moves them."""
+    p = min(max((value - lo) / (hi - lo), 1e-6), 1 - 1e-6)
+    return math.log(p / (1 - p))
+
+
 @register_router("energy")
 class EnergyRouter(Router):
     """Soft turn-on gate over normalized pre-step log-energy.
@@ -633,6 +668,28 @@ class EnergyRouter(Router):
     `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.
+
+    `temperature` is normally a single fixed scalar shared by every expert.
+    Two mutually exclusive optional modes generalize it:
+      - `learn_width`: each expert gets its own learnable width, so
+        `gate(e) = softmax_i(-(e - c_i)^2 / width_i)` — experts can learn
+        independently how much of the energy axis they cover.
+      - `learn_temperature`: the single shared `temperature` itself becomes
+        learnable (still one scalar for every expert).
+    Both parameterize their raw learnable value through a sigmoid bounded
+    into `[width_min_ratio, width_max_ratio] * temperature` (see
+    `_bounded_interp`), warm-started so the initial effective width/
+    temperature exactly equals `temperature` — enabling either mode is a
+    no-op at init. The bound is deliberately not raw `softplus`/`exp`
+    (unbounded above): an unbounded width lets one expert's width run away
+    to infinity, making its logit `-d2/width -> 0` almost everywhere so it
+    wins nearly every row regardless of true distance to its center — the
+    same "experts overlap instead of partitioning" failure this whole
+    router design is trying to avoid, just via a new mechanism. See
+    `Router.entropy_loss`/`giant.train`'s `lambda_entropy` for a secondary,
+    optional guard against all experts' widths co-inflating together
+    (which bounding caps but doesn't forbid, and which the load-balance
+    loss alone can't see since usage shares stay even throughout).
     """
 
     def __init__(
@@ -642,10 +699,31 @@ class EnergyRouter(Router):
         learn_centers: bool = True,
         energy_idx: int = 3,
         centers_init: Sequence[float] | None = None,
+        learn_width: bool = False,
+        learn_temperature: bool = False,
+        width_min_ratio: float = 0.1,
+        width_max_ratio: float = 10.0,
     ) -> None:
         super().__init__(n_experts)
+        if learn_width and learn_temperature:
+            raise ValueError("learn_width and learn_temperature are mutually exclusive")
         self.temperature = temperature
         self.energy_idx = energy_idx
+        self.learn_width = learn_width
+        self.learn_temperature = learn_temperature
+        if learn_width or learn_temperature:
+            if not (width_min_ratio < 1.0 < width_max_ratio):
+                raise ValueError(
+                    f"width_min_ratio ({width_min_ratio}) and width_max_ratio "
+                    f"({width_max_ratio}) must bracket 1.0"
+                )
+            self._width_lo = width_min_ratio * temperature
+            self._width_hi = width_max_ratio * temperature
+            raw0 = _inverse_bounded_interp(temperature, self._width_lo, self._width_hi)
+            if learn_width:
+                self.raw_width = nn.Parameter(torch.full((n_experts,), raw0))
+            else:
+                self.raw_temperature = nn.Parameter(torch.tensor(raw0))
         if centers_init is None:
             centers = torch.linspace(-2.0, 2.0, n_experts)
         else:
@@ -660,10 +738,20 @@ class EnergyRouter(Router):
         else:
             self.register_buffer("centers", centers)
 
+    def effective_width(self) -> torch.Tensor | float:
+        """Softmax denominator used by `gate()`: a fixed scalar `temperature`
+        (default), a per-expert `(n_experts,)` bounded width (`learn_width`),
+        or a single bounded learnable scalar (`learn_temperature`)."""
+        if self.learn_width:
+            return _bounded_interp(self.raw_width, self._width_lo, self._width_hi)
+        if self.learn_temperature:
+            return _bounded_interp(self.raw_temperature, self._width_lo, self._width_hi)
+        return self.temperature
+
     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)
+        return torch.softmax(-d2 / self.effective_width(), dim=-1)
 
 
 @register_router("pdg")

giant/pipeline.py

@@ -416,6 +416,7 @@ def run_train_job(
         lambda_s2=t.get("lambda_s2", 1.0),
         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),
         normalizer_dict={
             "cond": cond_norm.to_dict(),
             "target": tgt_norm.to_dict(),

giant/train.py

@@ -32,6 +32,7 @@ _METRICS_FIELDS = [
     "train_loss_s2",
     "train_loss_balance",
     "train_loss_proc",
+    "train_loss_entropy",
     "train_nsec_acc",
     "d_loss",
     "g_loss",
@@ -43,6 +44,7 @@ _METRICS_FIELDS = [
     "val_loss_s2",
     "val_loss_balance",
     "val_loss_proc",
+    "val_loss_entropy",
     "val_nsec_acc",
     "val_marginal_kl",
     "router_s1_entropy",
@@ -123,6 +125,7 @@ def _compute_losses(
     lambda_s2: float,
     lambda_balance: float = 0.0,
     lambda_proc: float = 0.0,
+    lambda_entropy: float = 0.0,
 ) -> tuple[
     torch.Tensor,
     torch.Tensor,
@@ -131,8 +134,9 @@ def _compute_losses(
     torch.Tensor,
     torch.Tensor,
     torch.Tensor,
+    torch.Tensor,
 ]:
-    """Compute (total_loss, L_s1, L_nsec, L_s2, L_balance, L_proc, nsec_acc) for one batch."""
+    """Compute (total_loss, L_s1, L_nsec, L_s2, L_balance, L_proc, L_entropy, nsec_acc) for one batch."""
     cond_cont, cond_cat, x1_s1, n_sec, sec_cont, proc_idx = batch
     cond_cont = cond_cont.to(device)
     cond_cat = cond_cat.to(device)
@@ -185,16 +189,25 @@ def _compute_losses(
         l_proc = stage1_model.router.classify_loss(
             cond_cont, cond_cat, proc_idx
         ) + sec_decoder.router.classify_loss(cond_cont, cond_cat, proc_idx)
+        # Optional entropy-regularization aux loss (see Router.entropy_loss):
+        # penalizes uniform/collapsed gating, a secondary guard against
+        # gate-sharpness collapse that lambda_balance alone can't see.
+        l_entropy = stage1_model.router.entropy_loss(
+            cond_cont, cond_cat
+        ) + sec_decoder.router.entropy_loss(cond_cont, cond_cat)
     else:
         l_balance = torch.zeros((), device=device)
         l_proc = torch.zeros((), device=device)
+        l_entropy = torch.zeros((), device=device)
 
     total = l_s1 + lambda_nsec * l_nsec + lambda_s2 * l_s2
     if lambda_balance > 0:
         total = total + lambda_balance * l_balance
     if lambda_proc > 0:
         total = total + lambda_proc * l_proc
-    return total, l_s1, l_nsec, l_s2, l_balance, l_proc, nsec_acc
+    if lambda_entropy > 0:
+        total = total + lambda_entropy * l_entropy
+    return total, l_s1, l_nsec, l_s2, l_balance, l_proc, l_entropy, nsec_acc
 
 
 def _wgan_train_step(
@@ -333,6 +346,7 @@ def train(
     lambda_s2: float = 1.0,
     lambda_balance: float = 0.0,
     lambda_proc: float = 0.0,
+    lambda_entropy: float = 0.0,
     normalizer_dict: dict | None = None,
     pdg_map: dict | None = None,
     mat_map: dict | None = None,
@@ -395,6 +409,7 @@ def train(
                 "lambda_s2": lambda_s2,
                 "lambda_balance": lambda_balance,
                 "lambda_proc": lambda_proc,
+                "lambda_entropy": lambda_entropy,
                 "n_critic": n_critic,
                 "gp_weight": gp_weight,
                 "model": model_config or {},
@@ -559,6 +574,7 @@ def train(
             train_s2_sum = 0.0
             train_balance_sum = 0.0
             train_proc_sum = 0.0
+            train_entropy_sum = 0.0
             train_d_sum = 0.0
             train_g_sum = 0.0
             train_wasserstein_sum = 0.0
@@ -626,7 +642,7 @@ def train(
                     train_grad_norm_d_sum += stats["grad_norm_d"] * B
                     train_grad_norm_g_sum += stats["grad_norm_g"] * B
                 else:
-                    loss, l_s1, l_nsec, l_s2, l_balance, l_proc, nsec_acc = (
+                    loss, l_s1, l_nsec, l_s2, l_balance, l_proc, l_entropy, nsec_acc = (
                         _compute_losses(
                             stage1_model,
                             sec_decoder,
@@ -638,6 +654,7 @@ def train(
                             lambda_s2,
                             lambda_balance,
                             lambda_proc,
+                            lambda_entropy,
                         )
                     )
                     optimizer.zero_grad()
@@ -661,6 +678,7 @@ def train(
                     train_s2_sum += l_s2.item() * B
                     train_balance_sum += l_balance.item() * B
                     train_proc_sum += l_proc.item() * B
+                    train_entropy_sum += l_entropy.item() * B
                     train_nsec_acc_sum += nsec_acc.item() * B
 
                 train_n += B
@@ -772,7 +790,7 @@ def train(
                 val_loss = val_marginal_kl
                 val_s1_sum = val_nsec_sum = val_s2_sum = val_balance_sum = (
                     val_proc_sum
-                ) = val_nsec_acc_sum = 0.0
+                ) = val_entropy_sum = val_nsec_acc_sum = 0.0
                 val_n = 1
                 val_nsec_acc = 0.0
                 router_s1_entropy = router_s2_entropy = 0.0
@@ -785,6 +803,7 @@ def train(
                 val_s2_sum = 0.0
                 val_balance_sum = 0.0
                 val_proc_sum = 0.0
+                val_entropy_sum = 0.0
                 val_nsec_acc_sum = 0.0
                 val_n = 0
                 if has_router:
@@ -802,19 +821,27 @@ def train(
                     for val_batch_idx, batch in enumerate(val_loader):
                         if max_val_batches > 0 and val_batch_idx >= max_val_batches:
                             break
-                        loss, l_s1, l_nsec, l_s2, l_balance, l_proc, nsec_acc = (
-                            _compute_losses(
-                                stage1_model,
-                                sec_decoder,
-                                batch,
-                                mode,
-                                ddpm_schedule,
-                                device,
-                                lambda_nsec,
-                                lambda_s2,
-                                lambda_balance,
-                                lambda_proc,
-                            )
+                        (
+                            loss,
+                            l_s1,
+                            l_nsec,
+                            l_s2,
+                            l_balance,
+                            l_proc,
+                            l_entropy,
+                            nsec_acc,
+                        ) = _compute_losses(
+                            stage1_model,
+                            sec_decoder,
+                            batch,
+                            mode,
+                            ddpm_schedule,
+                            device,
+                            lambda_nsec,
+                            lambda_s2,
+                            lambda_balance,
+                            lambda_proc,
+                            lambda_entropy,
                         )
                         B = batch[0].size(0)
                         val_loss_sum += loss.item() * B
@@ -823,6 +850,7 @@ def train(
                         val_s2_sum += l_s2.item() * B
                         val_balance_sum += l_balance.item() * B
                         val_proc_sum += l_proc.item() * B
+                        val_entropy_sum += l_entropy.item() * B
                         val_nsec_acc_sum += nsec_acc.item() * B
                         if has_router:
                             cond_cont = batch[0].to(device)
@@ -895,6 +923,7 @@ def train(
                 f"  s2={train_s2_sum / max(train_n, 1):.3f}"
                 f"  bal={train_balance_sum / max(train_n, 1):.3f}"
                 f"  proc={train_proc_sum / max(train_n, 1):.3f}"
+                f"  entropy={train_entropy_sum / max(train_n, 1):.3f}"
                 f"  d={train_d_sum / max(train_n, 1):.3f}"
                 f"  g={train_g_sum / max(train_n, 1):.3f})"
                 f"  val {val_loss:.4f}"
@@ -909,6 +938,7 @@ def train(
                 "train_loss_s2": train_s2_sum / max(train_n, 1),
                 "train_loss_balance": train_balance_sum / max(train_n, 1),
                 "train_loss_proc": train_proc_sum / max(train_n, 1),
+                "train_loss_entropy": train_entropy_sum / max(train_n, 1),
                 "train_nsec_acc": train_nsec_acc,
                 "d_loss": train_d_sum / max(train_n, 1),
                 "g_loss": train_g_sum / max(train_n, 1),
@@ -920,6 +950,7 @@ def train(
                 "val_loss_s2": val_s2_sum / max(val_n, 1),
                 "val_loss_balance": val_balance_sum / max(val_n, 1),
                 "val_loss_proc": val_proc_sum / max(val_n, 1),
+                "val_loss_entropy": val_entropy_sum / max(val_n, 1),
                 "val_nsec_acc": val_nsec_acc,
                 "val_marginal_kl": val_marginal_kl,
                 "router_s1_entropy": router_s1_entropy,

tests/test_router.py

@@ -141,6 +141,151 @@ def test_build_router_unknown_type_raises():
     raise AssertionError("expected ValueError for unknown router type")
 
 
+# ── EnergyRouter learn_width / learn_temperature ────────────────────────────
+
+
+def test_energy_router_learn_width_matches_fixed_temperature_at_init():
+    """Enabling learn_width should be a no-op at init — the warm-started
+    per-expert width must reproduce the fixed-temperature gate exactly."""
+    centers_init = [-1.0, 0.0, 0.5, 1.5]
+    fixed = EnergyRouter(n_experts=4, temperature=0.3, centers_init=centers_init)
+    learned = EnergyRouter(
+        n_experts=4, temperature=0.3, centers_init=centers_init, learn_width=True
+    )
+    cond_cont, cond_cat = _cond(16)
+    torch.testing.assert_close(
+        learned.gate(cond_cont, cond_cat),
+        fixed.gate(cond_cont, cond_cat),
+        atol=1e-5,
+        rtol=0,
+    )
+
+
+def test_energy_router_learn_temperature_matches_fixed_temperature_at_init():
+    centers_init = [-1.0, 0.0, 0.5, 1.5]
+    fixed = EnergyRouter(n_experts=4, temperature=0.3, centers_init=centers_init)
+    learned = EnergyRouter(
+        n_experts=4,
+        temperature=0.3,
+        centers_init=centers_init,
+        learn_temperature=True,
+    )
+    cond_cont, cond_cat = _cond(16)
+    torch.testing.assert_close(
+        learned.gate(cond_cont, cond_cat),
+        fixed.gate(cond_cont, cond_cat),
+        atol=1e-5,
+        rtol=0,
+    )
+
+
+def test_energy_router_learn_width_and_temperature_mutually_exclusive_raises():
+    try:
+        EnergyRouter(n_experts=4, learn_width=True, learn_temperature=True)
+    except ValueError:
+        return
+    raise AssertionError(
+        "expected ValueError for learn_width and learn_temperature both set"
+    )
+
+
+def test_energy_router_width_ratio_bounds_must_bracket_one_raises():
+    try:
+        EnergyRouter(
+            n_experts=4, learn_width=True, width_min_ratio=1.0, width_max_ratio=2.0
+        )
+    except ValueError:
+        return
+    raise AssertionError(
+        "expected ValueError for width_min_ratio/width_max_ratio not bracketing 1.0"
+    )
+
+
+def test_energy_router_effective_width_stays_within_bounds():
+    router = EnergyRouter(
+        n_experts=4,
+        temperature=0.5,
+        learn_width=True,
+        width_min_ratio=0.1,
+        width_max_ratio=10.0,
+    )
+    lo, hi = 0.1 * 0.5, 10.0 * 0.5
+    with torch.no_grad():
+        router.raw_width.fill_(1e6)
+    width = router.effective_width()
+    assert isinstance(width, torch.Tensor)
+    assert torch.all(width <= hi + 1e-4)
+    with torch.no_grad():
+        router.raw_width.fill_(-1e6)
+    width = router.effective_width()
+    assert isinstance(width, torch.Tensor)
+    assert torch.all(width >= lo - 1e-4)
+
+
+def test_energy_router_learn_width_gate_still_partition_of_unity():
+    router = EnergyRouter(n_experts=4, learn_width=True)
+    with torch.no_grad():
+        router.raw_width.copy_(torch.randn(4) * 3)
+    cond_cont, cond_cat = _cond(16)
+    g = router.gate(cond_cont, cond_cat)
+    torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0)
+
+
+def test_energy_router_learn_width_hardens_when_pushed_to_floor():
+    """Pushing every expert's width toward the (tiny) floor should harden the
+    gate to a one-hot at the nearest center, generalizing the fixed-
+    temperature->0 hardening test to the per-expert path."""
+    router = EnergyRouter(
+        n_experts=4, learn_width=True, width_min_ratio=1e-4, width_max_ratio=10.0
+    )
+    with torch.no_grad():
+        router.raw_width.fill_(-1e6)
+    cond_cont, cond_cat = _cond(16)
+    g = router.gate(cond_cont, cond_cat)
+    e = cond_cont[:, router.energy_idx].unsqueeze(-1)
+    d2 = (e - router.centers.unsqueeze(0)) ** 2
+    onehot = torch.nn.functional.one_hot(d2.argmin(dim=-1), num_classes=4).float()
+    torch.testing.assert_close(g, onehot, atol=1e-3, rtol=0)
+
+
+def test_energy_router_own_width_controls_own_coverage_independent_of_others():
+    """Widening one expert's width should monotonically grow only that
+    expert's own gate share, without needing to touch any other expert's
+    width — the "each expert learns its own coverage independently" property
+    this feature is meant to add."""
+    router = EnergyRouter(
+        n_experts=2, temperature=1.0, learn_width=True, centers_init=[0.0, 10.0]
+    )
+    cond_cont, cond_cat = _cond(4)
+    cond_cont[:, 3] = 3.0  # fixed energy, unequal distance to each center
+
+    shares = []
+    with torch.no_grad():
+        for raw in torch.linspace(-8.0, 8.0, 9):
+            router.raw_width[0] = raw
+            shares.append(router.gate(cond_cont, cond_cat)[0, 0].item())
+
+    assert all(a <= b + 1e-6 for a, b in zip(shares, shares[1:]))
+
+
+def test_build_router_threads_learn_width_kwargs_through():
+    router = build_router(
+        "energy", 4, learn_width=True, width_min_ratio=0.2, width_max_ratio=8.0
+    )
+    assert isinstance(router, EnergyRouter)
+    assert router.learn_width is True
+    assert isinstance(router.raw_width, torch.nn.Parameter)
+    assert router.raw_width.shape == (4,)
+
+
+def test_router_entropy_loss_is_nonnegative_bounded_scalar():
+    router = EnergyRouter(n_experts=4)
+    cond_cont, cond_cat = _cond(16)
+    loss = router.entropy_loss(cond_cont, cond_cat)
+    assert loss.shape == ()
+    assert 0.0 <= loss.item() <= 1.0
+
+
 # ── PdgRouter ────────────────────────────────────────────────────────────────