Give CriticModel a registry-built trunk and StageModel base (gitea #57)

CriticModel was the one stage-shaped class left out of the trunk-registry
(gitea #33), block-conditioning-registry (gitea #34), and StageModel-base
(gitea #39) refactors: it hand-rolled a plain ResBlock stack, so a
routed/FiLM/AdaLN trunk was available to every generative stage model except
the critic competing against them under WGAN-GP.

CriticModel now subclasses StageModel (reusing its cond_enc construction, and
a stage-2 context-fusion helper factored out of Stage2OneShot onto the base)
and builds its body via build_trunk (output width 1) instead of a bespoke
ResBlock loop, so trunk.type/trunk.block_conditioning now affect the critic
too. Each stage's critic inherits its own generator's trunk config rather
than a new critic_trunk config key, mirroring the existing
critic_hidden_dim/critic_n_res_blocks "0 = inherit from generator" pattern.
Router mixing (MoE) for the critic stays out of scope. Since CriticModel is
training-only and never persisted for inference, and WGAN-GP is still
unbenchmarked, its state_dict shape has no back-compat burden.

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

giant/model/builders.py

@@ -202,6 +202,8 @@ def build_critics(model_config: dict) -> dict[str, nn.Module | None]:
             cond_out_dim=cond_out_dim,
             dropout=s1_spec.dropout,
             stage="stage1",
+            trunk_type=s1_spec.trunk.type,
+            block_conditioning=s1_spec.trunk.block_conditioning,
         )
 
     if s2_spec.active and build_objective(s2_spec.generator).is_adversarial:
@@ -225,6 +227,8 @@ def build_critics(model_config: dict) -> dict[str, nn.Module | None]:
             dropout=s2_spec.dropout,
             stage="stage2",
             context_dim=s2_spec.context_dim,
+            trunk_type=s2_spec.trunk.type,
+            block_conditioning=s2_spec.trunk.block_conditioning,
         )
 
     return result

giant/model/models.py

@@ -8,7 +8,7 @@ from giant.config import ConditioningAxisConfig, HeadConfig, ParticleTypeConfig
 from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_DIM, SEC_SLOT_DIM, X_DIM
 from giant.model.encoders import ConditionEncoder
 from giant.model.history import HistoryEncoder, build_history
-from giant.model.layers import ContextAdapter, ResBlock, SinusoidalEmbedding, build_mlp_head
+from giant.model.layers import ContextAdapter, SinusoidalEmbedding, build_mlp_head
 from giant.model.objectives import build_objective
 from giant.model.routers import Router
 from giant.model.trunks import build_trunk
@@ -188,6 +188,26 @@ class StageModel(nn.Module):
             hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio))
             self.stop_head = build_mlp_head(cond_out_dim, 1, hidden, head_cfg.depth)
 
+    def _build_context_fusion(self, x_dim: int, context_dim: int, cond_out_dim: int) -> None:
+        """Builds `self.context_adapter`/`self.fuse` — the stage-2-style
+        context-fusion pattern (project the previous stage's outcome down to
+        `context_dim` via `ContextAdapter`, concat onto the base conditioning,
+        project back to `cond_out_dim`) shared by `Stage2OneShot` and a
+        `stage="stage2"` `CriticModel` (gitea #57). Call from a subclass's
+        `__init__` before using `_cond_embed`."""
+        self.context_adapter = ContextAdapter(x_dim, context_dim)
+        self.fuse = nn.Sequential(
+            nn.Linear(cond_out_dim + context_dim, cond_out_dim),
+            nn.SiLU(),
+        )
+
+    def _cond_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor:
+        """Fuses base conditioning with the previous stage's outcome — pairs
+        with `_build_context_fusion`."""
+        base = self.cond_enc(cond_cont, cond_cat)
+        ctx = self.context_adapter(stage1_out)
+        return self.fuse(torch.cat([base, ctx], dim=-1))
+
     def _require_n_sec_head(self) -> None:
         if self.n_sec_head is None:
             raise RuntimeError(
@@ -363,11 +383,7 @@ class Stage2OneShot(StageModel):
             particle_type_cfg=particle_type_cfg,
             cond_enc=cond_enc,
         )
-        self.context_adapter = ContextAdapter(x_dim, context_dim)
-        self.fuse = nn.Sequential(
-            nn.Linear(cond_out_dim + context_dim, cond_out_dim),
-            nn.SiLU(),
-        )
+        self._build_context_fusion(x_dim, context_dim, cond_out_dim)
         target = self.particle_type_cfg.target
         type_head_out_dim = None if target == "physical" else k_max * self.type_dim
         self._build_trunk_and_heads(
@@ -386,11 +402,6 @@ class Stage2OneShot(StageModel):
             type_head_cfg=type_head_cfg,
         )
 
-    def _cond_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor:
-        base = self.cond_enc(cond_cont, cond_cat)
-        ctx = self.context_adapter(stage1_out)
-        return self.fuse(torch.cat([base, ctx], dim=-1))
-
     def forward(
         self,
         x_t: torch.Tensor,
@@ -702,11 +713,25 @@ class Stage2Autoregressive(StageModel):
         return self.stop_head(c_emb.reshape(B * K, -1)).view(B, K)
 
 
-class CriticModel(nn.Module):
+class CriticModel(StageModel):
     """Generator-agnostic WGAN-GP critic body: a scalar realism score, for
     either stage (`stage="stage1"` mirrors v0.2 `Critic`; `stage="stage2"`
     mirrors v0.2 `SecondaryCritic`, adding the same context-fusion path as
-    `Stage2OneShot`). Used only when that stage's `generator == "wgan"`."""
+    `Stage2OneShot`, via `StageModel._build_context_fusion`/`_cond_embed`).
+    Used only when that stage's `generator == "wgan"`.
+
+    Subclasses `StageModel` for the `cond_enc` construction and (stage 2)
+    context-fusion scaffolding only — its trunk is built directly via
+    `build_trunk` (output width 1) rather than through
+    `_build_trunk_and_heads`, since that helper is shaped around a
+    generator's `Objective`/time-embedding/flow-matching concerns
+    (`forward`'s `(x_t, cond) -> vector` shape) that don't apply to a critic's
+    `(x, cond) -> scalar` (gitea #57). `generator="wgan"` is passed to the
+    base purely because that's factually when a critic exists; nothing here
+    ever calls `_build_trunk_and_heads`, so no head/time-embedding machinery
+    is built from it. Never routed (MoE) — that's a separate, unrequested
+    axis of scope; see gitea #57's proposal, which covers only the trunk/
+    block registries."""
 
     def __init__(
         self,
@@ -722,22 +747,26 @@ class CriticModel(nn.Module):
         stage: str = "stage1",
         context_dim: int = 64,
         context_in_dim: int = X_DIM,
+        trunk_type: str = "resmlp",
+        block_conditioning: str = "add",
     ) -> None:
-        super().__init__()
+        super().__init__(
+            pdg_vocab,
+            mat_vocab,
+            particle_cfg,
+            material_cfg,
+            cond_out_dim=cond_out_dim,
+            generator="wgan",
+            noise_dim=0,
+        )
         if stage not in ("stage1", "stage2"):
             raise ValueError(f"stage must be 'stage1' or 'stage2', got {stage!r}")
         self.stage = stage
-        self.cond_enc = ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
         if stage == "stage2":
-            self.context_adapter = ContextAdapter(context_in_dim, context_dim)
-            self.fuse = nn.Sequential(
-                nn.Linear(cond_out_dim + context_dim, cond_out_dim),
-                nn.SiLU(),
-            )
-        self.input_proj = nn.Linear(in_dim, hidden_dim)
-        self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_res_blocks)])
-        self.out_norm = nn.LayerNorm(hidden_dim)
-        self.out_proj = nn.Linear(hidden_dim, 1)
+            self._build_context_fusion(context_in_dim, context_dim, cond_out_dim)
+        self.trunk = build_trunk(
+            None, trunk_type, in_dim, 1, hidden_dim, n_res_blocks, cond_out_dim, dropout, block_conditioning
+        )
 
     def forward(
         self,
@@ -746,13 +775,9 @@ class CriticModel(nn.Module):
         cond_cat: torch.Tensor,
         stage1_out: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        base = self.cond_enc(cond_cont, cond_cat)
         if self.stage == "stage2":
-            ctx = self.context_adapter(stage1_out)
-            cond = self.fuse(torch.cat([base, ctx], dim=-1))
+            assert stage1_out is not None, "stage='stage2' CriticModel requires stage1_out"
+            cond = self._cond_embed(cond_cont, cond_cat, stage1_out)
         else:
-            cond = base
-        h = self.input_proj(x)
-        for block in self.blocks:
-            h = block(h, cond)
-        return self.out_proj(self.out_norm(h)).squeeze(-1)
+            cond = self.cond_enc(cond_cont, cond_cat)
+        return self.trunk(x, cond, cond_cont, cond_cat).squeeze(-1)

tests/test_network.py

@@ -8,7 +8,10 @@ from giant.model.network import (
     HISTORY_REGISTRY,
     AttentionHistory,
     ConditionEncoder,
+    CriticModel,
+    FilmResBlock,
     HistoryEncoder,
+    LinearTrunk,
     MarkovHistory,
     NoHistory,
     SinusoidalEmbedding,
@@ -944,7 +947,7 @@ def test_build_critics_particle_type_n_classes_overrides_conditioning_emb_dim():
     assert wider_critic is not None
     # k_max=3 slots, each CONT_SLOT_DIM + n_classes wide under wgan folding —
     # widening n_classes alone (emb_dim stays 4) must widen the critic input.
-    assert wider_critic.input_proj.in_features > default_n_classes_critic.input_proj.in_features
+    assert wider_critic.trunk.input_proj.in_features > default_n_classes_critic.trunk.input_proj.in_features
 
 
 # ── build_models/build_critics: DEFAULT_CONFIG fallback drift (issues.md #1) ─
@@ -1021,12 +1024,12 @@ def test_build_critics_omitted_particle_type_matches_default_config():
     cfg["stage2_model"]["generator"] = "wgan"
     onehot_critic = build_critics(cfg)["stage2"]
     assert onehot_critic is not None
-    onehot_in_dim = onehot_critic.input_proj.in_features
+    onehot_in_dim = onehot_critic.trunk.input_proj.in_features
 
     cfg["stage2_model"]["particle_type"] = {"target": "physical"}
     physical_critic = build_critics(cfg)["stage2"]
     assert physical_critic is not None
-    physical_in_dim = physical_critic.input_proj.in_features
+    physical_in_dim = physical_critic.trunk.input_proj.in_features
 
     # onehot's per-slot type width is emb_dim classes vs. physical's fixed
     # (log-mass, charge) pair — different unless emb_dim happens to be 2, so
@@ -1046,15 +1049,15 @@ def test_build_critics_stage1_critic_hidden_dim_and_n_res_blocks_override_genera
 
     inherited = build_critics(cfg)["stage1"]
     assert inherited is not None
-    assert inherited.input_proj.out_features == 8
-    assert len(inherited.blocks) == 1
+    assert inherited.trunk.input_proj.out_features == 8
+    assert len(inherited.trunk.blocks) == 1
 
     cfg["stage1_model"]["wgan"]["critic_hidden_dim"] = 16
     cfg["stage1_model"]["wgan"]["critic_n_res_blocks"] = 3
     overridden = build_critics(cfg)["stage1"]
     assert overridden is not None
-    assert overridden.input_proj.out_features == 16
-    assert len(overridden.blocks) == 3
+    assert overridden.trunk.input_proj.out_features == 16
+    assert len(overridden.trunk.blocks) == 3
 
 
 def test_build_critics_stage2_critic_hidden_dim_and_n_res_blocks_override_generator_size():
@@ -1065,15 +1068,15 @@ def test_build_critics_stage2_critic_hidden_dim_and_n_res_blocks_override_genera
 
     inherited = build_critics(cfg)["stage2"]
     assert inherited is not None
-    assert inherited.input_proj.out_features == 8
-    assert len(inherited.blocks) == 1
+    assert inherited.trunk.input_proj.out_features == 8
+    assert len(inherited.trunk.blocks) == 1
 
     cfg["stage2_model"]["wgan"]["critic_hidden_dim"] = 16
     cfg["stage2_model"]["wgan"]["critic_n_res_blocks"] = 3
     overridden = build_critics(cfg)["stage2"]
     assert overridden is not None
-    assert overridden.input_proj.out_features == 16
-    assert len(overridden.blocks) == 3
+    assert overridden.trunk.input_proj.out_features == 16
+    assert len(overridden.trunk.blocks) == 3
 
 
 # ── StageModel base (gitea #39): Stage1Model/Stage2OneShot/Stage2Autoregressive
@@ -1234,3 +1237,34 @@ def test_stagemodel_time_emb_matches_objective_needs_time(cls, generator):
     assert model.generator_kind == generator
     assert model.noise_dim == 8
     assert (model.time_emb is not None) == build_objective(generator).needs_time
+
+
+# ── CriticModel uses the trunk/block registries + StageModel base (gitea #57) ─
+
+
+def test_critic_model_is_stagemodel_subclass():
+    assert issubclass(CriticModel, StageModel)
+
+
+@pytest.mark.parametrize("stage", ["stage1", "stage2"])
+def test_build_critics_threads_trunk_type_from_generator_config(stage):
+    cfg = _minimal_model_config(share_stages=False)
+    cfg["stage1_model"]["generator"] = "wgan"
+    cfg["stage2_model"]["generator"] = "wgan"
+    cfg[f"{stage}_model"]["trunk"] = {"type": "linear"}
+
+    critic = build_critics(cfg)[stage]
+    assert critic is not None
+    assert isinstance(critic.trunk, LinearTrunk)
+
+
+@pytest.mark.parametrize("stage", ["stage1", "stage2"])
+def test_build_critics_threads_block_conditioning_from_generator_config(stage):
+    cfg = _minimal_model_config(share_stages=False)
+    cfg["stage1_model"]["generator"] = "wgan"
+    cfg["stage2_model"]["generator"] = "wgan"
+    cfg[f"{stage}_model"]["trunk"] = {"block_conditioning": "film"}
+
+    critic = build_critics(cfg)[stage]
+    assert critic is not None
+    assert all(isinstance(block, FilmResBlock) for block in critic.trunk.blocks)

tests/test_wgan.py

@@ -2,7 +2,7 @@ import torch
 
 from giant.config import ConditioningAxisConfig
 from giant.constants import COND_DIM, K_MAX, SEC_DIM, SEC_SLOT_DIM, X_DIM
-from giant.model.network import CriticModel, Stage1Model, Stage2OneShot
+from giant.model.network import CriticModel, LinearTrunk, Stage1Model, Stage2OneShot
 from giant.model.wgan import critic_loss, generator_loss, gradient_penalty
 from giant.sample import sample_secondaries_wgan, sample_wgan
 
@@ -115,6 +115,33 @@ def test_critic_output_shape():
     assert out.shape == (B,)
 
 
+def test_critic_model_honours_trunk_type_and_block_conditioning():
+    """gitea #57: CriticModel routes its body through build_trunk/build_block
+    like every generator stage model, instead of hand-rolling a plain
+    ResBlock stack."""
+    B = 8
+    critic = CriticModel(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=PARTICLE_CFG,
+        material_cfg=MATERIAL_CFG,
+        in_dim=X_DIM,
+        hidden_dim=32,
+        n_res_blocks=2,
+        stage="stage1",
+        trunk_type="linear",
+        block_conditioning="adaln",
+    )
+    assert isinstance(critic.trunk, LinearTrunk)
+    cond_cont, cond_cat = _cond(B)
+    real = torch.randn(B, X_DIM)
+    fake = torch.randn(B, X_DIM)
+    loss = critic_loss(lambda x: critic(x, cond_cont, cond_cat), real, fake.detach(), gp_weight=10.0)
+    loss.backward()
+    for name, p in critic.named_parameters():
+        assert p.grad is not None, f"no grad for {name}"
+
+
 def test_sample_wgan_shape():
     B = 6
     model = _small_generator()