Add learnable per-expert width and shared temperature to EnergyRouter

EnergyRouter's gate sharpness was a single fixed temperature shared by
every expert, with no way for an expert to independently learn how much
of the energy axis it covers. Adds two mutually exclusive, default-off
modes: learn_width (per-expert learnable width) and learn_temperature
(single learnable shared scalar), both bounded via a sigmoid
interpolation warm-started to reproduce today's fixed-temperature gate
exactly at init, to compare against each other without risking the
unbounded-width collapse failure mode. Also promotes gate_stats's
entropy into a generic, optional Router.entropy_loss (lambda_entropy) as
a secondary guard against all experts' widths co-inflating together.

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

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