diff --git a/giant/config.py b/giant/config.py index d9e4265..d21be99 100644 --- a/giant/config.py +++ b/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 diff --git a/giant/model/network.py b/giant/model/network.py index 35ec968..f517d02 100644 --- a/giant/model/network.py +++ b/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") diff --git a/giant/pipeline.py b/giant/pipeline.py index d8c022d..53bdc61 100644 --- a/giant/pipeline.py +++ b/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(), diff --git a/giant/train.py b/giant/train.py index 473652c..b0d8654 100644 --- a/giant/train.py +++ b/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, diff --git a/tests/test_router.py b/tests/test_router.py index 62a8c86..05105dc 100644 --- a/tests/test_router.py +++ b/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 ────────────────────────────────────────────────────────────────