diff --git a/giant/config.py b/giant/config.py index d0ba445..7e4f8e7 100644 --- a/giant/config.py +++ b/giant/config.py @@ -34,10 +34,10 @@ DEFAULT_CONFIG: dict = { "n_experts": 4, "expert_hidden_dim": 128, "expert_n_blocks": 3, - "temperature": 0.5, # energy-router kwarg - "learn_centers": True, # energy-router kwarg + "temperature": 0.5, # energy/pdg-router kwarg + "learn_centers": True, # energy/pdg-router kwarg "lambda_balance": 0.0, # optional load-balance aux loss weight - "emb_dim": 8, # process-router kwarg: its own pdg/mat embedding width + "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 # (0.0 still trains a working router — the gate gets gradient diff --git a/giant/model/network.py b/giant/model/network.py index 05487d5..7a6e8ca 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -344,6 +344,48 @@ class EnergyRouter(Router): return torch.softmax(-d2 / self.temperature, dim=-1) +@register_router("pdg") +class PdgRouter(Router): + """Soft turn-on gate over a learned PDG embedding. + + Unlike ProcessRouter's process label, PDG code is already known at + pre-step time (it's a conditioning input, `cond_cat[:, 0]`), so no + supervision is needed — `classify_loss` falls back to the Router base + class's zero-loss default, same as EnergyRouter. Because PDG is + categorical rather than a scalar, this generalizes EnergyRouter's + soft-turn-on-then-Voronoi trick from a 1-D distance to a distance in a + small embedding space: its own embedding table (kept separate from the + trunk's ConditionEncoder, same reasoning as ProcessRouter's own + pdg/mat embeddings) maps each PDG code to a point, and `n_experts` + learnable (or fixed) centers partition that space. + `gate(pdg) = softmax_i(-||emb(pdg) - c_i||^2 / tau)`. + """ + + def __init__( + self, + n_experts: int, + pdg_vocab: int, + emb_dim: int = 8, + temperature: float = 0.5, + learn_centers: bool = True, + ) -> None: + super().__init__(n_experts) + self.temperature = temperature + self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim) + centers = torch.randn(n_experts, emb_dim) * 0.1 + if learn_centers: + self.centers = nn.Parameter(centers) + else: + self.register_buffer("centers", centers) + + def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + e = self.pdg_emb(cond_cat[:, 0]) # (B, emb_dim) + d2 = ((e.unsqueeze(1) - self.centers.unsqueeze(0)) ** 2).sum( + -1 + ) # (B, n_experts) + return torch.softmax(-d2 / self.temperature, dim=-1) + + @register_router("process") class ProcessRouter(Router): """Routes on the physics process expected to end the step. diff --git a/tests/test_router.py b/tests/test_router.py index 0755ebb..eef1c3f 100644 --- a/tests/test_router.py +++ b/tests/test_router.py @@ -6,6 +6,7 @@ from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM from giant.model.network import ( DenoisingMLP, EnergyRouter, + PdgRouter, ProcessRouter, ROUTER_REGISTRY, RoutedDenoisingMLP, @@ -102,6 +103,100 @@ def test_build_router_unknown_type_raises(): raise AssertionError("expected ValueError for unknown router type") +# ── PdgRouter ──────────────────────────────────────────────────────────────── + + +def test_pdg_router_registered(): + assert ROUTER_REGISTRY["pdg"] is PdgRouter + + +def test_pdg_router_gate_partition_of_unity(): + router = PdgRouter(n_experts=4, pdg_vocab=3) + cond_cont, cond_cat = _cond(16) + g = router.gate(cond_cont, cond_cat) + assert g.shape == (16, 4) + torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0) + + +def test_pdg_router_top1_matches_gate_argmax(): + router = PdgRouter(n_experts=4, pdg_vocab=3) + cond_cont, cond_cat = _cond(16) + assert torch.equal( + router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1) + ) + + +def test_pdg_router_hardens_as_temperature_shrinks(): + """As tau -> 0 the soft gate should converge to a one-hot at the argmax.""" + router = PdgRouter(n_experts=4, pdg_vocab=3, temperature=1e-4) + cond_cont, cond_cat = _cond(16) + g = router.gate(cond_cont, cond_cat) + top1 = router.top1(cond_cont, cond_cat) + onehot = torch.nn.functional.one_hot(top1, num_classes=4).float() + torch.testing.assert_close(g, onehot, atol=1e-3, rtol=0) + + +def test_pdg_router_balance_loss_is_nonnegative_scalar(): + router = PdgRouter(n_experts=4, pdg_vocab=3) + cond_cont, cond_cat = _cond(16) + loss = router.balance_loss(cond_cont, cond_cat) + assert loss.shape == () + assert loss.item() >= 0.0 + + +def test_pdg_router_classify_loss_defaults_to_zero(): + """PDG is already known at gate time (unlike ProcessRouter's process + label), so no supervision is needed — falls back to Router's default.""" + router = PdgRouter(n_experts=4, pdg_vocab=3) + cond_cont, cond_cat = _cond(16) + labels = torch.randint(0, 4, (16,)) + loss = router.classify_loss(cond_cont, cond_cat, labels) + assert loss.shape == () + assert loss.item() == 0.0 + + +def test_pdg_router_only_reads_pdg_column(): + """Gate must depend on cond_cat[:, 0] (pdg) only, not cond_cont or material.""" + router = PdgRouter(n_experts=4, pdg_vocab=3) + cond_cont, cond_cat = _cond(16) + g_before = router.gate(cond_cont, cond_cat) + + cond_cont_perturbed = torch.randn_like(cond_cont) + cond_cat_diff_mat = cond_cat.clone() + cond_cat_diff_mat[:, 1] = (cond_cat_diff_mat[:, 1] + 1) % 2 + g_after = router.gate(cond_cont_perturbed, cond_cat_diff_mat) + + torch.testing.assert_close(g_before, g_after, atol=1e-6, rtol=0) + + +def test_build_router_pdg_type_uses_pdg_vocab(): + router = build_router("pdg", 4, pdg_vocab=5, mat_vocab=3, emb_dim=8) + assert isinstance(router, PdgRouter) + assert router.pdg_emb.num_embeddings == 5 + + +def test_build_models_routed_with_pdg_router(): + model_config = dict( + pdg_vocab=4, + mat_vocab=2, + emb_dim=16, + dropout=0.1, + k_max=K_MAX, + expert_hidden_dim=16, + expert_n_blocks=2, + router={ + "enabled": True, + "type": "pdg", + "n_experts": 3, + }, + ) + stage1, sec_decoder = build_models(model_config) + assert isinstance(stage1, RoutedDenoisingMLP) + assert isinstance(stage1.router, PdgRouter) + assert len(stage1.experts) == 3 + assert stage1.router.pdg_emb.num_embeddings == 4 + + # ── ProcessRouter ────────────────────────────────────────────────────────────