Make ResBlock's conditioning-injection mechanism selectable (gitea #34)

ResBlock injected conditioning exactly one way — h = linear1(h) +
cond_proj(cond), a conditional bias, the weakest standard option for a
model whose entire job is to be conditional. Adds BLOCK_REGISTRY
(giant/model/layers.py), mirroring the TRUNK_REGISTRY/ROUTER_REGISTRY
registry+factory idiom (gitea #33), with two new drop-in alternatives:
FilmResBlock (per-channel scale+shift modulating the norm output,
zero-init so conditioning has no effect at construction) and
AdaLNResBlock (DiT-style AdaLN-Zero — the norm's own affine is replaced
by a conditioning-derived scale/shift, plus a zero-init gate on the
residual branch, making the block the exact identity function at init).

Selected per stage via a new stage{1,2}_model.trunk.block_conditioning
config leaf ("add" | "film" | "adaln", default "add"), threaded through
build_trunk/build_expert_body/RoutedTrunk and the three stage model
constructors. Default stays "add" and ResBlock's body is unchanged, so
existing configs/checkpoints are bit-identical to before this change.

Decided during planning: the new field lives on the existing TrunkConfig
rather than a new top-level block/blocks config section; the WGAN
CriticModel (which builds its own ResBlock stack outside TRUNK_REGISTRY)
and the issue's mentioned blocks.norm/blocks.activation axes are both
left out of scope.

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

giant/config.py

@@ -346,17 +346,23 @@ class TrunkConfig:
     stage, unaffected by this block. A future body's own hyperparameters
     (e.g. a transformer's `n_heads`/`n_layers`) would get their own sibling
     field here, matching how `flow`/`ddpm`/`wgan` already coexist selected by
-    `generator`."""
+    `generator`.
+
+    `block_conditioning` selects each body's conditioning-injection mechanism
+    from `giant.model.layers.BLOCK_REGISTRY` — `"add"` (default, today's
+    conditional-bias `ResBlock`, bit-identical to pre-gitea-#34 behaviour),
+    `"film"`, or `"adaln"`."""
 
     type: str = "resmlp"
+    block_conditioning: str = "add"
 
     @classmethod
     def from_dict(cls, d: dict | None) -> "TrunkConfig":
         d = d or {}
-        return cls(type=d.get("type", "resmlp"))
+        return cls(type=d.get("type", "resmlp"), block_conditioning=d.get("block_conditioning", "add"))
 
     def to_dict(self) -> dict:
-        return {"type": self.type}
+        return {"type": self.type, "block_conditioning": self.block_conditioning}
 
 
 @dataclass(frozen=True)
@@ -1431,6 +1437,8 @@ _OUT_DIR_NAME_CANDIDATES = [
     ),
     ("stage1_trunk_type", _path_candidate("stage1_model.trunk.type", "s1t-")),
     ("stage2_trunk_type", _path_candidate("stage2_model.trunk.type", "s2t-")),
+    ("stage1_block_cond", _path_candidate("stage1_model.trunk.block_conditioning", "s1bc-")),
+    ("stage2_block_cond", _path_candidate("stage2_model.trunk.block_conditioning", "s2bc-")),
     ("stage1_router", _router_candidate("stage1_model", "s1")),
     ("stage2_router", _router_candidate("stage2_model", "s2")),
     (

giant/model/builders.py

@@ -84,6 +84,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
             noise_dim=s1_spec.wgan.noise_dim,
             router=stage1_router,
             trunk_type=s1_spec.trunk.type,
+            block_conditioning=s1_spec.trunk.block_conditioning,
             n_sec_head_k_max=n_sec_head_k_max,
             cond_enc=shared_cond_enc,
         )
@@ -123,6 +124,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
                 k_max=k_max,
                 router=stage2_router,
                 trunk_type=s2_spec.trunk.type,
+                block_conditioning=s2_spec.trunk.block_conditioning,
                 build_n_sec_head=n_sec_owner != "stage1",
                 particle_type_cfg=particle_type_cfg,
                 history=ar_cfg.history,
@@ -151,6 +153,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
                 k_max=k_max,
                 router=stage2_router,
                 trunk_type=s2_spec.trunk.type,
+                block_conditioning=s2_spec.trunk.block_conditioning,
                 build_n_sec_head=n_sec_owner != "stage1",
                 particle_type_cfg=particle_type_cfg,
                 cond_enc=shared_cond_enc,

giant/model/layers.py

@@ -57,6 +57,26 @@ class ContextAdapter(nn.Module):
         return torch.tanh(self.proj(x))
 
 
+BLOCK_REGISTRY: dict[str, type[nn.Module]] = {}
+
+
+def register_block(name: str):
+    def decorator(cls: type[nn.Module]) -> type[nn.Module]:
+        BLOCK_REGISTRY[name] = cls
+        return cls
+
+    return decorator
+
+
+def build_block(name: str, dim: int, cond_dim: int, dropout: float = 0.0) -> nn.Module:
+    """Factory: look up a registered conditioning-injection block by name and
+    construct one instance — `trunk.block_conditioning` (gitea #34)."""
+    if name not in BLOCK_REGISTRY:
+        raise ValueError(f"unknown block conditioning type {name!r}; available: {sorted(BLOCK_REGISTRY)}")
+    return BLOCK_REGISTRY[name](dim, cond_dim, dropout)
+
+
+@register_block("add")
 class ResBlock(nn.Module):
     def __init__(self, dim: int, cond_dim: int, dropout: float = 0.0) -> None:
         super().__init__()
@@ -74,3 +94,66 @@ class ResBlock(nn.Module):
         h = self.dropout(h)
         h = self.linear2(h)
         return x + h
+
+
+@register_block("film")
+class FilmResBlock(nn.Module):
+    """FiLM conditioning (Perez et al. 2018): a per-channel scale+shift
+    modulates the normalized features, on top of the norm's own affine —
+    an *additional* modulation, unlike `AdaLNResBlock` below, which replaces
+    the norm's affine outright. `film_proj` is zero-initialized so
+    `gamma=beta=0` at construction — conditioning has no effect on the
+    output until training moves it, a stable starting point (though not a
+    literal identity block, since `linear1`/`linear2` aren't zero-init)."""
+
+    def __init__(self, dim: int, cond_dim: int, dropout: float = 0.0) -> None:
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.linear1 = nn.Linear(dim, dim)
+        self.film_proj = nn.Linear(cond_dim, 2 * dim)
+        nn.init.zeros_(self.film_proj.weight)
+        nn.init.zeros_(self.film_proj.bias)
+        self.act = nn.SiLU()
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim, dim)
+
+    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+        h = self.norm(x)
+        gamma, beta = self.film_proj(cond).chunk(2, dim=-1)
+        h = h * (1 + gamma) + beta
+        h = self.linear1(h)
+        h = self.act(h)
+        h = self.dropout(h)
+        h = self.linear2(h)
+        return x + h
+
+
+@register_block("adaln")
+class AdaLNResBlock(nn.Module):
+    """AdaLN-Zero conditioning (DiT, Peebles & Xie 2022): the norm's own
+    affine is replaced by a conditioning-derived scale/shift, and the
+    residual branch is scaled by a conditioning-derived gate. `adaln_proj`
+    is zero-initialized, so `scale=shift=gate=0` at construction — the block
+    is the exact identity function at init (`x + 0 * h' == x`), regardless
+    of `x`/`cond`."""
+
+    def __init__(self, dim: int, cond_dim: int, dropout: float = 0.0) -> None:
+        super().__init__()
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False)
+        self.linear1 = nn.Linear(dim, dim)
+        self.adaln_proj = nn.Linear(cond_dim, 3 * dim)
+        nn.init.zeros_(self.adaln_proj.weight)
+        nn.init.zeros_(self.adaln_proj.bias)
+        self.act = nn.SiLU()
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim, dim)
+
+    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+        h = self.norm(x)
+        scale, shift, gate = self.adaln_proj(cond).chunk(3, dim=-1)
+        h = h * (1 + scale) + shift
+        h = self.linear1(h)
+        h = self.act(h)
+        h = self.dropout(h)
+        h = self.linear2(h)
+        return x + gate * h

giant/model/models.py

@@ -93,6 +93,7 @@ class Stage1Model(nn.Module):
         noise_dim: int = 64,
         router: Router | None = None,
         trunk_type: str = "resmlp",
+        block_conditioning: str = "add",
         n_sec_head_k_max: int | None = None,
         cond_enc: ConditionEncoder | None = None,
     ) -> None:
@@ -108,7 +109,9 @@ class Stage1Model(nn.Module):
         self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None
         merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim
         in_dim = noise_dim if generator == "wgan" else x_dim
-        self.trunk = build_trunk(router, trunk_type, in_dim, x_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout)
+        self.trunk = build_trunk(
+            router, trunk_type, in_dim, x_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout, block_conditioning
+        )
         self.n_sec_head = None
         if n_sec_head_k_max is not None:
             self.n_sec_head = nn.Sequential(
@@ -187,6 +190,7 @@ class Stage2OneShot(nn.Module):
         k_max: int = K_MAX,
         router: Router | None = None,
         trunk_type: str = "resmlp",
+        block_conditioning: str = "add",
         build_n_sec_head: bool = True,
         particle_type_cfg: dict | None = None,
         cond_enc: ConditionEncoder | None = None,
@@ -214,7 +218,15 @@ class Stage2OneShot(nn.Module):
         merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim
         in_dim = noise_dim if generator == "wgan" else sec_dim
         self.trunk = build_trunk(
-            router, trunk_type, in_dim, sec_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout
+            router,
+            trunk_type,
+            in_dim,
+            sec_dim,
+            hidden_dim,
+            n_res_blocks,
+            merged_cond_dim,
+            dropout,
+            block_conditioning,
         )
         self.n_sec_head = None
         if build_n_sec_head:
@@ -333,6 +345,7 @@ class Stage2Autoregressive(nn.Module):
         k_max: int = K_MAX,
         router: Router | None = None,
         trunk_type: str = "resmlp",
+        block_conditioning: str = "add",
         build_n_sec_head: bool = True,
         particle_type_cfg: dict | None = None,
         history: str = "markov",
@@ -392,6 +405,7 @@ class Stage2Autoregressive(nn.Module):
             n_res_blocks,
             merged_cond_dim,
             dropout,
+            block_conditioning,
         )
 
         self.n_sec_head = None

giant/model/network.py

@@ -11,7 +11,17 @@ from giant.model._legacy import _migrate_legacy_model_config, migrate_legacy_sta
 from giant.model.builders import build_critics, build_models
 from giant.model.encoders import ConditionEncoder, cat_col_layout
 from giant.model.history import AttentionHistory, HistoryEncoder, MarkovHistory, _CausalAttnBlock
-from giant.model.layers import ContextAdapter, ResBlock, SinusoidalEmbedding, _make_axis_mlp
+from giant.model.layers import (
+    BLOCK_REGISTRY,
+    AdaLNResBlock,
+    ContextAdapter,
+    FilmResBlock,
+    ResBlock,
+    SinusoidalEmbedding,
+    _make_axis_mlp,
+    build_block,
+    register_block,
+)
 from giant.model.models import (
     CriticModel,
     Stage1Model,
@@ -47,13 +57,16 @@ from giant.model.trunks import (
 )
 
 __all__ = [
+    "AdaLNResBlock",
     "AttentionHistory",
+    "BLOCK_REGISTRY",
     "ComposedRouter",
     "ConditionEncoder",
     "ContextAdapter",
     "CriticModel",
     "EnergyRouter",
     "ExpertTrunk",
+    "FilmResBlock",
     "HistoryEncoder",
     "MarkovHistory",
     "PdgRouter",
@@ -75,6 +88,7 @@ __all__ = [
     "_migrate_legacy_model_config",
     "_parse_composed_axes",
     "_route_forward",
+    "build_block",
     "build_composed_router",
     "build_critics",
     "build_expert_body",
@@ -83,6 +97,7 @@ __all__ = [
     "build_trunk",
     "cat_col_layout",
     "migrate_legacy_state_dict",
+    "register_block",
     "register_router",
     "register_trunk",
     "resolve_type_n_classes",

giant/model/trunks.py

@@ -12,7 +12,7 @@ free — no separate "routed transformer trunk" class needed.
 import torch
 import torch.nn as nn
 
-from giant.model.layers import ResBlock
+from giant.model.layers import build_block
 from giant.model.routers import Router
 
 TRUNK_REGISTRY: dict[str, type[nn.Module]] = {}
@@ -34,14 +34,19 @@ def build_expert_body(
     n_blocks: int,
     cond_dim: int,
     dropout: float = 0.0,
+    block_conditioning: str = "add",
 ) -> nn.Module:
     """Factory: look up a registered trunk body by name and construct one
     instance of it — used both for a standalone (unrouted) trunk and for each
-    expert inside a `RoutedTrunk`."""
+    expert inside a `RoutedTrunk`. `block_conditioning` selects the
+    `BLOCK_REGISTRY` entry each body's internal `ResBlock`-family blocks use
+    (`trunk.block_conditioning`, gitea #34) — an optional trailing kwarg a
+    future non-`ResBlock`-based body can simply ignore, same idiom as
+    `Trunk.forward`'s accept-and-ignore `cond_cont`/`cond_cat`."""
     if name not in TRUNK_REGISTRY:
         raise ValueError(f"unknown trunk type {name!r}; available: {sorted(TRUNK_REGISTRY)}")
     cls = TRUNK_REGISTRY[name]
-    return cls(in_dim, out_dim, hidden_dim, n_blocks, cond_dim, dropout)
+    return cls(in_dim, out_dim, hidden_dim, n_blocks, cond_dim, dropout, block_conditioning=block_conditioning)
 
 
 @register_trunk("resmlp")
@@ -65,11 +70,14 @@ class ExpertTrunk(nn.Module):
         n_blocks: int,
         cond_dim: int,
         dropout: float = 0.0,
+        block_conditioning: str = "add",
     ) -> None:
         super().__init__()
         self.out_dim = out_dim
         self.input_proj = nn.Linear(in_dim, hidden_dim)
-        self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_dim, dropout=dropout) for _ in range(n_blocks)])
+        self.blocks = nn.ModuleList(
+            [build_block(block_conditioning, hidden_dim, cond_dim, dropout) for _ in range(n_blocks)]
+        )
         self.out_proj = nn.Linear(hidden_dim, out_dim)
 
     def forward(
@@ -145,12 +153,22 @@ class RoutedTrunk(Trunk):
         n_res_blocks: int,
         cond_dim: int,
         dropout: float = 0.0,
+        block_conditioning: str = "add",
     ) -> None:
         super().__init__()
         self.router = router
         self.experts = nn.ModuleList(
             [
-                build_expert_body(trunk_type, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout)
+                build_expert_body(
+                    trunk_type,
+                    in_dim,
+                    out_dim,
+                    hidden_dim,
+                    n_res_blocks,
+                    cond_dim,
+                    dropout,
+                    block_conditioning=block_conditioning,
+                )
                 for _ in range(router.n_experts)
             ]
         )
@@ -174,6 +192,7 @@ def build_trunk(
     n_res_blocks: int,
     cond_dim: int,
     dropout: float = 0.0,
+    block_conditioning: str = "add",
 ) -> nn.Module:
     """Build a stage's trunk: `trunk_type` (a `TRUNK_REGISTRY` key, e.g.
     `"resmlp"`) selects the expert body architecture; `router`, if given,
@@ -182,8 +201,14 @@ def build_trunk(
     class), which is what makes an unrouted trunk's state-dict keys land
     directly under `trunk.*` instead of `trunk.experts.0.*` (see
     `giant.model._legacy.migrate_legacy_state_dict`, which assumes exactly
-    this flat layout for a v0.2 monolithic checkpoint).
+    this flat layout for a v0.2 monolithic checkpoint). `block_conditioning`
+    (a `BLOCK_REGISTRY` key, e.g. `"add"`/`"film"`/`"adaln"`) selects each
+    body's conditioning-injection mechanism (gitea #34).
     """
     if router is not None:
-        return RoutedTrunk(router, trunk_type, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout)
-    return build_expert_body(trunk_type, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout)
+        return RoutedTrunk(
+            router, trunk_type, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout, block_conditioning
+        )
+    return build_expert_body(
+        trunk_type, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout, block_conditioning
+    )