From 732d5f1cd227e777f5374b0ce872f1f8d45ca566 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Mon, 24 Aug 2026 09:22:36 +0200 Subject: [PATCH] Add "none" variants for router, history, and trunk (gitea #45) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Turns "is this component earning its parameters?" into a one-line config flip for each of the three pluggable network components: - router.type = "none" (NoneRouter, giant/model/routers.py): still builds n_experts expert trunks via RoutedTrunk, but replaces the learned gate with a uniform 1/n_experts weight for every row — no centers/embeddings/classifier. Distinct from router.enabled=false (which drops routing/mixing entirely): this isolates whether the *learned routing signal* specifically is earning its parameters, holding expert count fixed. - stage2_model.autoregressive.history = "none" (NoHistory, giant/model/history.py): ignores feat/has_prev entirely and always returns zeros, ablating whether the AR decoder's history conditioning earns its parameters. Already validated for free by gitea #35's generic HISTORY_REGISTRY membership check. - trunk.type = "linear" (LinearTrunk, giant/model/trunks.py): a bare nn.Linear(in_dim + cond_dim, out_dim) body, no ResBlock stack. Per gitea #33's design, this composes for free with router.enabled=true ("mixture of trivial linear experts"). Both blocking issues (#33 trunk registry, #35 pluggable history encoder) are closed, so this was unblocked. --- giant/model/history.py | 17 ++++++++++ giant/model/network.py | 6 ++++ giant/model/routers.py | 18 +++++++++++ giant/model/trunks.py | 43 +++++++++++++++++++++++++ tests/test_network.py | 43 +++++++++++++++++++++++-- tests/test_router.py | 71 ++++++++++++++++++++++++++++++++++++++++++ 6 files changed, 195 insertions(+), 3 deletions(-) diff --git a/giant/model/history.py b/giant/model/history.py index a3cc914..d000329 100644 --- a/giant/model/history.py +++ b/giant/model/history.py @@ -63,6 +63,23 @@ def build_history(name: str, in_dim: int, out_dim: int, **kwargs) -> HistoryEnco return cls(in_dim, out_dim, **filtered) +@register_history("none") +class NoHistory(HistoryEncoder): + """No history signal at all — ignores feat/has_prev entirely and always + returns zeros. Ablates whether the AR decoder's history conditioning is + earning its parameters. `init_cache`/`step` use the base class's O(1) + defaults unmodified (this encoder's own `forward` is already O(1) per + call regardless of prefix length).""" + + def __init__(self, in_dim: int, out_dim: int) -> None: + super().__init__() + self.out_dim = out_dim + + def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor: + B, K, _ = feat.shape + return torch.zeros(B, K, self.out_dim, device=feat.device, dtype=feat.dtype) + + @register_history("markov") class MarkovHistory(HistoryEncoder): """Summarizes the previous secondary's own `(energy_fraction, direction, diff --git a/giant/model/network.py b/giant/model/network.py index ffd037a..41be980 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -15,6 +15,7 @@ from giant.model.history import ( AttentionHistory, HistoryEncoder, MarkovHistory, + NoHistory, _CausalAttnBlock, build_history, register_history, @@ -54,6 +55,7 @@ from giant.model.routers import ( ROUTER_REGISTRY, ComposedRouter, EnergyRouter, + NoneRouter, PdgRouter, ProcessRouter, Router, @@ -67,6 +69,7 @@ from giant.model.routers import ( from giant.model.trunks import ( TRUNK_REGISTRY, ExpertTrunk, + LinearTrunk, RoutedTrunk, Trunk, _route_forward, @@ -90,7 +93,10 @@ __all__ = [ "FlowObjective", "HISTORY_REGISTRY", "HistoryEncoder", + "LinearTrunk", "MarkovHistory", + "NoHistory", + "NoneRouter", "OBJECTIVE_REGISTRY", "Objective", "PdgRouter", diff --git a/giant/model/routers.py b/giant/model/routers.py index 34df658..a9060d3 100644 --- a/giant/model/routers.py +++ b/giant/model/routers.py @@ -144,6 +144,24 @@ def _inverse_bounded_interp(value: float, lo: float, hi: float) -> float: return math.log(p / (1 - p)) +@register_router("none") +class NoneRouter(Router): + """Uniform 1/n_experts gate — no learned routing signal at all. + + Still builds n_experts expert trunks via RoutedTrunk (same parameter + budget as a real router), but every row gets an identical weight + regardless of conditioning. Ablates whether the *learned routing + signal* — as opposed to simply having multiple experts — is earning + its parameters. `top1()` (the base class default) always dispatches to + expert 0 (argmax of a uniform vector), which still exercises + RoutedTrunk's real per-expert grouped-dispatch code path at eval time. + """ + + def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + B = cond_cont.shape[0] + return torch.full((B, self.n_experts), 1.0 / self.n_experts, device=cond_cont.device) + + @register_router("energy") class EnergyRouter(Router): """Soft turn-on gate over normalized pre-step log-energy. diff --git a/giant/model/trunks.py b/giant/model/trunks.py index c8f546c..bdd5313 100644 --- a/giant/model/trunks.py +++ b/giant/model/trunks.py @@ -94,6 +94,49 @@ class ExpertTrunk(nn.Module): return self.out_proj(x) +@register_trunk("linear") +class LinearTrunk(nn.Module): + """`nn.Linear(in_dim + cond_dim, out_dim)` over `concat([x, cond])` — + the trivial trunk body: no hidden layer, no ResBlock stack, no + nonlinearity. Ablates whether trunk depth/nonlinearity is earning its + parameters, holding everything else (heads, ConditionEncoder, + generator, ...) fixed. Composes for free with `router.enabled = true` + (gitea #33): a RoutedTrunk of n_experts linear bodies is "mixture of + trivial linear experts". `hidden_dim`/`n_blocks`/`dropout`/ + `block_conditioning` are accepted and ignored, matching + `build_expert_body`'s shared factory signature. + + `x` — the trunk's own input (e.g. the noised primary vector for flow + matching) — does not already carry conditioning; that's fused in + per-body via `cond`. So this concatenates `x` and `cond` itself to + remain a valid, conditioning-dependent model. + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + hidden_dim: int, + n_blocks: int, + cond_dim: int, + dropout: float = 0.0, + block_conditioning: str = "add", + ) -> None: + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + self.linear = nn.Linear(in_dim + cond_dim, out_dim) + + def forward( + self, + x: torch.Tensor, + cond: torch.Tensor, + cond_cont: torch.Tensor | None = None, + cond_cat: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.linear(torch.cat([x, cond], dim=-1)) + + def _route_forward( experts: nn.ModuleList, router: Router, diff --git a/tests/test_network.py b/tests/test_network.py index c82f77a..bcd1877 100644 --- a/tests/test_network.py +++ b/tests/test_network.py @@ -10,6 +10,7 @@ from giant.model.network import ( ConditionEncoder, HistoryEncoder, MarkovHistory, + NoHistory, SinusoidalEmbedding, Stage1Model, Stage2Autoregressive, @@ -465,16 +466,52 @@ def test_attention_history_step_matches_forward(): assert torch.allclose(stepped, expected, atol=1e-5) +# --- NoHistory (gitea #45) ---------------------------------------------------- + + +def test_no_history_shape(): + hist = NoHistory(in_dim=7, out_dim=12) + B, K = 3, 5 + feat = torch.randn(B, K, 7) + has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1) + out = hist(feat, has_prev) + assert out.shape == (B, K, 12) + + +def test_no_history_ignores_feat_and_has_prev(): + hist = NoHistory(in_dim=4, out_dim=6) + B, K = 2, 3 + has_prev_a = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1) + has_prev_b = torch.zeros(B, K, dtype=torch.bool) + feat_a = torch.randn(B, K, 4) + feat_b = torch.randn(B, K, 4) * 100 + out_a = hist(feat_a, has_prev_a) + out_b = hist(feat_b, has_prev_b) + assert torch.equal(out_a, torch.zeros(B, K, 6)) + assert torch.equal(out_a, out_b) + + +def test_no_history_uses_base_class_o1_defaults(): + hist = NoHistory(in_dim=4, out_dim=6) + assert hist.init_cache() is None + feat = torch.randn(2, 1, 4) + has_prev = torch.ones(2, 1, dtype=torch.bool) + out, cache = hist.step(feat, has_prev, "unused-cache") + assert torch.equal(out, torch.zeros(2, 1, 6)) + assert cache == "unused-cache" + + # --- HISTORY_REGISTRY / build_history (gitea #35) ---------------------------- -def test_history_registry_has_exactly_the_two_known_histories(): - assert set(HISTORY_REGISTRY) == {"markov", "attention"} +def test_history_registry_has_exactly_the_known_histories(): + assert set(HISTORY_REGISTRY) == {"markov", "attention", "none"} def test_build_history_returns_correct_concrete_type(): assert isinstance(build_history("markov", 4, 6), MarkovHistory) assert isinstance(build_history("attention", 4, 8), AttentionHistory) + assert isinstance(build_history("none", 4, 6), NoHistory) def test_build_history_unknown_name_raises(): @@ -605,7 +642,7 @@ def test_stage2_autoregressive_n_sec_head_and_type_head_cfg_control_hidden_width @pytest.mark.parametrize("target", ["physical", "onehot", "embedding"]) @pytest.mark.parametrize("generator", ["wgan", "flow"]) -@pytest.mark.parametrize("history", ["markov", "attention"]) +@pytest.mark.parametrize("history", ["markov", "attention", "none"]) def test_stage2_autoregressive_forward_shape(target, generator, history): B, K, emb_dim = 4, 5, 6 model = _build_stage2_ar(target, generator, emb_dim=emb_dim, k_max=K, history=history) diff --git a/tests/test_router.py b/tests/test_router.py index 25284bd..b8057ce 100644 --- a/tests/test_router.py +++ b/tests/test_router.py @@ -13,6 +13,8 @@ from giant.model.network import ( EnergyRouter, ExpertTrunk, FilmResBlock, + LinearTrunk, + NoneRouter, PdgRouter, ProcessRouter, ROUTER_REGISTRY, @@ -73,6 +75,34 @@ def test_energy_router_registered(): assert ROUTER_REGISTRY["energy"] is EnergyRouter +# ── NoneRouter (gitea #45) ────────────────────────────────────────────────── + + +def test_none_router_registered(): + assert ROUTER_REGISTRY["none"] is NoneRouter + + +def test_none_router_gate_is_uniform(): + router = NoneRouter(n_experts=4) + cond_cont, cond_cat = _cond(16) + g = router.gate(cond_cont, cond_cat) + assert g.shape == (16, 4) + torch.testing.assert_close(g, torch.full((16, 4), 0.25)) + + +def test_none_router_gate_ignores_conditioning(): + router = NoneRouter(n_experts=3) + cond_cont_a, cond_cat_a = _cond(8) + cond_cont_b, cond_cat_b = _cond(8) + torch.testing.assert_close(router.gate(cond_cont_a, cond_cat_a), router.gate(cond_cont_b, cond_cat_b)) + + +def test_none_router_top1_always_expert_zero(): + router = NoneRouter(n_experts=4) + cond_cont, cond_cat = _cond(16) + assert torch.equal(router.top1(cond_cont, cond_cat), torch.zeros(16, dtype=torch.long)) + + def test_energy_router_gate_partition_of_unity(): router = EnergyRouter(n_experts=4) cond_cont, cond_cat = _cond(16) @@ -168,6 +198,47 @@ def test_build_expert_body_unknown_type_raises(): raise AssertionError("expected ValueError for unknown trunk type") +# ── LinearTrunk (gitea #45) ────────────────────────────────────────────────── + + +def test_trunk_registry_has_linear(): + assert "linear" in TRUNK_REGISTRY + assert TRUNK_REGISTRY["linear"] is LinearTrunk + + +def test_linear_trunk_forward_shape(): + trunk = build_expert_body("linear", in_dim=9, out_dim=9, hidden_dim=64, n_blocks=6, cond_dim=12) + assert trunk.in_dim == 9 + assert trunk.out_dim == 9 + x = torch.randn(5, 9) + cond = torch.randn(5, 12) + out = trunk(x, cond) + assert out.shape == (5, 9) + + +def test_linear_trunk_depends_on_x_and_cond(): + trunk = build_expert_body("linear", in_dim=9, out_dim=9, hidden_dim=64, n_blocks=6, cond_dim=12) + x = torch.randn(5, 9) + cond_a = torch.randn(5, 12) + cond_b = torch.randn(5, 12) + assert not torch.allclose(trunk(x, cond_a), trunk(x, cond_b)) + + +def test_routed_linear_trunk_is_mixture_of_trivial_experts(): + """trunk.type = 'linear' composes for free with router.enabled = true + (gitea #33's comment on this issue) — a RoutedTrunk of n_experts linear + bodies.""" + router = build_router("energy", n_experts=3) + trunk = RoutedTrunk(router, "linear", in_dim=9, out_dim=9, hidden_dim=64, n_res_blocks=6, cond_dim=12) + assert len(trunk.experts) == 3 + assert all(isinstance(e, LinearTrunk) for e in trunk.experts) + x = torch.randn(5, 9) + cond = torch.randn(5, 12) + cond_cont, cond_cat = _cond(5) + out = trunk(x, cond, cond_cont, cond_cat) + assert out.shape == (5, 9) + + # ── BLOCK_REGISTRY / build_block (gitea #34) ────────────────────────────────