Deduplicate n_sec_head/type_head MLPs into build_mlp_head (gitea #36)

The same two-layer classifier head (Linear(cond_out_dim, hidden_dim // 2)
-> SiLU -> Linear(hidden_dim // 2, out_dim)) was hand-rolled five times in
giant/model/models.py: Stage1Model.n_sec_head, Stage2OneShot.n_sec_head/
.type_head, and Stage2Autoregressive.n_sec_head/.type_head. The `// 2`
ratio and fixed 2-layer depth were undocumented magic numbers, and both
n_sec accuracy and secondary-species accuracy are known weak spots that
were untunable independently of the trunk they hang off.

Adds `build_mlp_head(in_dim, out_dim, hidden, depth, act)` to
giant/model/layers.py (depth=1 is a bare Linear; depth>=2 matches the old
hardcoded shape exactly), and a new `HeadConfig` (hidden_ratio, depth)
dataclass in giant/config.py, wired in as `stage1_model.heads.n_sec` and
`stage2_model.heads.{n_sec,type}` — split per head type (not one shared
block per stage) since n_sec and species prediction are called out as
separate weak spots that may want independent capacity. Defaults
(hidden_ratio=0.5, depth=2) reproduce the old hardcoded architecture
bit-for-bit, so every existing config.toml and migrated v0.2 checkpoint
is unaffected; no changes were needed to migrate_config or the legacy
migration surfaces. No new CLI flags, matching how other nested
sub-config (router.*, trunk.*) is set via config.toml rather than
per-field flags.

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

giant/config.py

@@ -516,6 +516,64 @@ class AutoregressiveConfig:
         }
 
 
+@dataclass(frozen=True)
+class HeadConfig:
+    """A single classifier head's shape — `n_sec_head`/`type_head` (gitea
+    #36 deduplicated their five identical hand-rolled
+    `Linear -> SiLU -> Linear` definitions into
+    `giant.model.layers.build_mlp_head`, which this config drives).
+    `hidden_ratio=0.5`/`depth=2` are the exact pre-#36 hardcoded values
+    (hidden width = `hidden_dim // 2`, one hidden layer), so omitting a
+    `heads` block — including every migrated v0.2 config — reproduces the
+    old architecture bit-for-bit."""
+
+    hidden_ratio: float = 0.5  # hidden width = round(hidden_dim * hidden_ratio)
+    depth: int = 2  # matches build_mlp_head's depth
+
+    @classmethod
+    def from_dict(cls, d: dict | None) -> "HeadConfig":
+        d = d or {}
+        return cls(hidden_ratio=d.get("hidden_ratio", 0.5), depth=d.get("depth", 2))
+
+    def to_dict(self) -> dict:
+        return {"hidden_ratio": self.hidden_ratio, "depth": self.depth}
+
+
+@dataclass(frozen=True)
+class Stage1HeadsConfig:
+    """Stage 1 only ever owns `n_sec_head`, and only for a migrated v0.2
+    checkpoint (`stage2_model.n_sec.owner = "stage1"`) — see
+    `Stage1Model`'s docstring."""
+
+    n_sec: HeadConfig = field(default_factory=HeadConfig)
+
+    @classmethod
+    def from_dict(cls, d: dict | None) -> "Stage1HeadsConfig":
+        d = d or {}
+        return cls(n_sec=HeadConfig.from_dict(d.get("n_sec")))
+
+    def to_dict(self) -> dict:
+        return {"n_sec": self.n_sec.to_dict()}
+
+
+@dataclass(frozen=True)
+class Stage2HeadsConfig:
+    """`n_sec` and `type` are independently configurable — n_sec accuracy
+    and secondary-species accuracy are separately known weak spots (gitea
+    #36)."""
+
+    n_sec: HeadConfig = field(default_factory=HeadConfig)
+    type: HeadConfig = field(default_factory=HeadConfig)
+
+    @classmethod
+    def from_dict(cls, d: dict | None) -> "Stage2HeadsConfig":
+        d = d or {}
+        return cls(n_sec=HeadConfig.from_dict(d.get("n_sec")), type=HeadConfig.from_dict(d.get("type")))
+
+    def to_dict(self) -> dict:
+        return {"n_sec": self.n_sec.to_dict(), "type": self.type.to_dict()}
+
+
 @dataclass(frozen=True)
 class Stage1ModelConfig:
     # false skips building/training stage 1 entirely. The resulting
@@ -540,6 +598,7 @@ class Stage1ModelConfig:
     wgan: Stage1WganConfig = field(default_factory=Stage1WganConfig)
     router: RouterConfig = field(default_factory=RouterConfig)
     trunk: TrunkConfig = field(default_factory=TrunkConfig)
+    heads: Stage1HeadsConfig = field(default_factory=Stage1HeadsConfig)
 
     @classmethod
     def from_dict(cls, d: dict | None) -> "Stage1ModelConfig":
@@ -556,6 +615,7 @@ class Stage1ModelConfig:
             wgan=Stage1WganConfig.from_dict(d.get("wgan")),
             router=RouterConfig.from_dict(d.get("router")),
             trunk=TrunkConfig.from_dict(d.get("trunk")),
+            heads=Stage1HeadsConfig.from_dict(d.get("heads")),
         )
 
     def to_dict(self) -> dict:
@@ -571,6 +631,7 @@ class Stage1ModelConfig:
             "wgan": self.wgan.to_dict(),
             "router": self.router.to_dict(),
             "trunk": self.trunk.to_dict(),
+            "heads": self.heads.to_dict(),
         }
 
 
@@ -609,6 +670,7 @@ class Stage2ModelConfig:
     wgan: Stage2WganConfig = field(default_factory=Stage2WganConfig)
     router: Stage2RouterConfig = field(default_factory=Stage2RouterConfig)
     trunk: TrunkConfig = field(default_factory=TrunkConfig)
+    heads: Stage2HeadsConfig = field(default_factory=Stage2HeadsConfig)
 
     @classmethod
     def from_dict(cls, d: dict | None) -> "Stage2ModelConfig":
@@ -632,6 +694,7 @@ class Stage2ModelConfig:
             wgan=Stage2WganConfig.from_dict(d.get("wgan")),
             router=Stage2RouterConfig.from_dict(d.get("router")),
             trunk=TrunkConfig.from_dict(d.get("trunk")),
+            heads=Stage2HeadsConfig.from_dict(d.get("heads")),
         )
 
     def to_dict(self) -> dict:
@@ -654,6 +717,7 @@ class Stage2ModelConfig:
             "wgan": self.wgan.to_dict(),
             "router": self.router.to_dict(),
             "trunk": self.trunk.to_dict(),
+            "heads": self.heads.to_dict(),
         }
 
 

giant/model/builders.py

@@ -89,6 +89,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
             block_conditioning=s1_spec.trunk.block_conditioning,
             n_sec_head_k_max=n_sec_head_k_max,
             cond_enc=shared_cond_enc,
+            n_sec_head_cfg=s1_spec.heads.n_sec.to_dict(),
         )
 
     if s2_spec.active:
@@ -134,6 +135,8 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
                 attn_n_heads=ar_cfg.attn_n_heads,
                 attn_n_layers=ar_cfg.attn_n_layers,
                 cond_enc=shared_cond_enc,
+                n_sec_head_cfg=s2_spec.heads.n_sec.to_dict(),
+                type_head_cfg=s2_spec.heads.type.to_dict(),
             )
         else:
             sec_dim = stage2_trunk_sec_dim(
@@ -160,6 +163,8 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
                 build_n_sec_head=n_sec_owner != "stage1",
                 particle_type_cfg=particle_type_cfg,
                 cond_enc=shared_cond_enc,
+                n_sec_head_cfg=s2_spec.heads.n_sec.to_dict(),
+                type_head_cfg=s2_spec.heads.type.to_dict(),
             )
 
     return result

giant/model/layers.py

@@ -42,6 +42,31 @@ def _make_axis_mlp(in_dim: int, emb_dim: int, n_layers: int) -> nn.Sequential:
     return nn.Sequential(*layers)
 
 
+def build_mlp_head(
+    in_dim: int, out_dim: int, hidden: int, depth: int = 2, act: type[nn.Module] = nn.SiLU
+) -> nn.Sequential:
+    """`depth`-layer MLP head (gitea #36) — factors out the n_sec_head/
+    type_head pattern duplicated five times across `giant.model.models`.
+
+    `depth=1` is a bare `Linear(in_dim, out_dim)` (no hidden layer/
+    activation); `depth>=2` is `Linear(in_dim, hidden) -> act -> [Linear
+    (hidden, hidden) -> act] * (depth-2) -> Linear(hidden, out_dim)` —
+    `depth=2` reproduces every pre-#36 n_sec_head/type_head exactly when
+    `hidden == hidden_dim // 2`. Mirrors `_make_axis_mlp`'s depth
+    convention above, but takes `hidden` and `out_dim` as independent
+    widths (n_sec_head/type_head's hidden width is not their output width,
+    unlike the particle/material axis MLPs)."""
+    if depth < 1:
+        raise ValueError(f"depth must be >= 1, got {depth}")
+    if depth == 1:
+        return nn.Sequential(nn.Linear(in_dim, out_dim))
+    layers: list[nn.Module] = [nn.Linear(in_dim, hidden), act()]
+    for _ in range(depth - 2):
+        layers += [nn.Linear(hidden, hidden), act()]
+    layers.append(nn.Linear(hidden, out_dim))
+    return nn.Sequential(*layers)
+
+
 class ContextAdapter(nn.Module):
     """Projects a stage's outcome (e.g. Stage 1's 9D target) down to a
     fixed-width context vector for a downstream stage's conditioning —

giant/model/models.py

@@ -4,10 +4,11 @@
 import torch
 import torch.nn as nn
 
+from giant.config import HeadConfig
 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
+from giant.model.layers import ContextAdapter, ResBlock, SinusoidalEmbedding, build_mlp_head
 from giant.model.objectives import build_objective
 from giant.model.routers import Router
 from giant.model.trunks import build_trunk
@@ -97,6 +98,7 @@ class Stage1Model(nn.Module):
         block_conditioning: str = "add",
         n_sec_head_k_max: int | None = None,
         cond_enc: ConditionEncoder | None = None,
+        n_sec_head_cfg: dict | None = None,
     ) -> None:
         super().__init__()
         self.generator_kind = generator
@@ -116,11 +118,9 @@ class Stage1Model(nn.Module):
         )
         self.n_sec_head = None
         if n_sec_head_k_max is not None:
-            self.n_sec_head = nn.Sequential(
-                nn.Linear(cond_out_dim, hidden_dim // 2),
-                nn.SiLU(),
-                nn.Linear(hidden_dim // 2, n_sec_head_k_max + 1),
-            )
+            head_cfg = HeadConfig.from_dict(n_sec_head_cfg)
+            hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio))
+            self.n_sec_head = build_mlp_head(cond_out_dim, n_sec_head_k_max + 1, hidden, head_cfg.depth)
 
     def forward(
         self,
@@ -197,6 +197,8 @@ class Stage2OneShot(nn.Module):
         build_n_sec_head: bool = True,
         particle_type_cfg: dict | None = None,
         cond_enc: ConditionEncoder | None = None,
+        n_sec_head_cfg: dict | None = None,
+        type_head_cfg: dict | None = None,
     ) -> None:
         super().__init__()
         self.generator_kind = generator
@@ -234,20 +236,16 @@ class Stage2OneShot(nn.Module):
         )
         self.n_sec_head = None
         if build_n_sec_head:
-            self.n_sec_head = nn.Sequential(
-                nn.Linear(cond_out_dim, hidden_dim // 2),
-                nn.SiLU(),
-                nn.Linear(hidden_dim // 2, k_max + 1),
-            )
+            head_cfg = HeadConfig.from_dict(n_sec_head_cfg)
+            hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio))
+            self.n_sec_head = build_mlp_head(cond_out_dim, k_max + 1, hidden, head_cfg.depth)
         self.type_head = None
         target = self.particle_type_cfg.get("target", "physical")
         if target != "physical" and not objective.folds_type_slice:
             emb_dim = resolve_type_n_classes(self.particle_type_cfg, particle_cfg["emb_dim"])
-            self.type_head = nn.Sequential(
-                nn.Linear(cond_out_dim, hidden_dim // 2),
-                nn.SiLU(),
-                nn.Linear(hidden_dim // 2, k_max * emb_dim),
-            )
+            head_cfg = HeadConfig.from_dict(type_head_cfg)
+            hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio))
+            self.type_head = build_mlp_head(cond_out_dim, k_max * emb_dim, hidden, head_cfg.depth)
             self._type_k_max = k_max
             self._type_emb_dim = emb_dim
 
@@ -356,6 +354,8 @@ class Stage2Autoregressive(nn.Module):
         attn_n_heads: int = 4,
         attn_n_layers: int = 2,
         cond_enc: ConditionEncoder | None = None,
+        n_sec_head_cfg: dict | None = None,
+        type_head_cfg: dict | None = None,
     ) -> None:
         super().__init__()
         self.history_kind = history
@@ -411,19 +411,15 @@ class Stage2Autoregressive(nn.Module):
 
         self.n_sec_head = None
         if build_n_sec_head:
-            self.n_sec_head = nn.Sequential(
-                nn.Linear(cond_out_dim, hidden_dim // 2),
-                nn.SiLU(),
-                nn.Linear(hidden_dim // 2, k_max + 1),
-            )
+            head_cfg = HeadConfig.from_dict(n_sec_head_cfg)
+            hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio))
+            self.n_sec_head = build_mlp_head(cond_out_dim, k_max + 1, hidden, head_cfg.depth)
         self.type_head = None
         target = self.particle_type_cfg.get("target", "physical")
         if target != "physical" and not objective.folds_type_slice:
-            self.type_head = nn.Sequential(
-                nn.Linear(cond_out_dim, hidden_dim // 2),
-                nn.SiLU(),
-                nn.Linear(hidden_dim // 2, self.type_dim),
-            )
+            head_cfg = HeadConfig.from_dict(type_head_cfg)
+            hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio))
+            self.type_head = build_mlp_head(cond_out_dim, self.type_dim, hidden, head_cfg.depth)
 
     def _base_cond(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor:
         base = self.cond_enc(cond_cont, cond_cat)

giant/model/network.py

@@ -28,6 +28,7 @@ from giant.model.layers import (
     SinusoidalEmbedding,
     _make_axis_mlp,
     build_block,
+    build_mlp_head,
     register_block,
 )
 from giant.model.models import (
@@ -116,6 +117,7 @@ __all__ = [
     "build_critics",
     "build_expert_body",
     "build_history",
+    "build_mlp_head",
     "build_models",
     "build_objective",
     "build_router",

tests/test_config.py

@@ -55,6 +55,9 @@ def test_giant_config_to_dict_matches_default_config():
         gconfig.NSecConfig,
         gconfig.ParticleTypeConfig,
         gconfig.AutoregressiveConfig,
+        gconfig.HeadConfig,
+        gconfig.Stage1HeadsConfig,
+        gconfig.Stage2HeadsConfig,
         gconfig.Stage1ModelConfig,
         gconfig.Stage2ModelConfig,
         gconfig.TrainConfig,
@@ -94,6 +97,21 @@ def test_trunk_config_defaults_block_conditioning_to_add_for_both_stages():
     assert gconfig.DEFAULT_CONFIG["stage2_model"]["trunk"]["block_conditioning"] == "add"
 
 
+def test_heads_config_defaults_reproduce_pre_gitea_36_hardcoded_shape():
+    """gitea #36: a pre-existing config with no `heads` key must reproduce
+    today's hardcoded `hidden_dim // 2`, one-hidden-layer architecture
+    exactly."""
+    assert gconfig.Stage1ModelConfig().heads.n_sec.hidden_ratio == 0.5
+    assert gconfig.Stage1ModelConfig().heads.n_sec.depth == 2
+    assert gconfig.Stage2ModelConfig().heads.n_sec.hidden_ratio == 0.5
+    assert gconfig.Stage2ModelConfig().heads.n_sec.depth == 2
+    assert gconfig.Stage2ModelConfig().heads.type.hidden_ratio == 0.5
+    assert gconfig.Stage2ModelConfig().heads.type.depth == 2
+    assert gconfig.DEFAULT_CONFIG["stage1_model"]["heads"]["n_sec"] == {"hidden_ratio": 0.5, "depth": 2}
+    assert gconfig.DEFAULT_CONFIG["stage2_model"]["heads"]["n_sec"] == {"hidden_ratio": 0.5, "depth": 2}
+    assert gconfig.DEFAULT_CONFIG["stage2_model"]["heads"]["type"] == {"hidden_ratio": 0.5, "depth": 2}
+
+
 def test_particle_type_config_n_classes_defaults_to_zero_and_round_trips():
     """gitea #29: n_classes=0 means "inherit conditioning.particle.emb_dim"
     — the default must stay 0 so an existing config.toml with no

tests/test_layers.py

@@ -0,0 +1,42 @@
+import pytest
+import torch
+
+from giant.model.layers import build_mlp_head
+
+
+def test_build_mlp_head_depth_1_is_bare_linear():
+    head = build_mlp_head(8, 4, hidden=16, depth=1)
+    assert len(head) == 1
+    assert isinstance(head[0], torch.nn.Linear)
+    assert head[0].in_features == 8
+    assert head[0].out_features == 4
+    out = head(torch.randn(3, 8))
+    assert out.shape == (3, 4)
+
+
+def test_build_mlp_head_depth_2_matches_pre_gitea_36_shape():
+    head = build_mlp_head(8, 4, hidden=16, depth=2)
+    assert len(head) == 3
+    assert isinstance(head[0], torch.nn.Linear)
+    assert head[0].in_features == 8
+    assert head[0].out_features == 16
+    assert isinstance(head[1], torch.nn.SiLU)
+    assert isinstance(head[2], torch.nn.Linear)
+    assert head[2].in_features == 16
+    assert head[2].out_features == 4
+    out = head(torch.randn(5, 8))
+    assert out.shape == (5, 4)
+
+
+def test_build_mlp_head_depth_3_has_extra_hidden_layer():
+    head = build_mlp_head(8, 4, hidden=16, depth=3)
+    assert len(head) == 5
+    widths = [(m.in_features, m.out_features) for m in head if isinstance(m, torch.nn.Linear)]
+    assert widths == [(8, 16), (16, 16), (16, 4)]
+    out = head(torch.randn(2, 8))
+    assert out.shape == (2, 4)
+
+
+def test_build_mlp_head_depth_0_raises():
+    with pytest.raises(ValueError, match="depth"):
+        build_mlp_head(8, 4, hidden=16, depth=0)

tests/test_network.py

@@ -93,6 +93,42 @@ def test_stage1_model_no_n_sec_head_by_default():
     assert model.n_sec_head is None
 
 
+def test_stage1_model_n_sec_head_default_cfg_matches_pre_gitea_36_shape():
+    """No n_sec_head_cfg given must reproduce the old hardcoded
+    hidden_dim // 2, one-hidden-layer architecture exactly (gitea #36)."""
+    model = Stage1Model(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=PARTICLE_CFG,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=40,
+        cond_out_dim=12,
+        n_sec_head_k_max=15,
+    )
+    assert model.n_sec_head is not None
+    assert len(model.n_sec_head) == 3
+    assert model.n_sec_head[0].in_features == 12
+    assert model.n_sec_head[0].out_features == 20  # hidden_dim // 2
+    assert model.n_sec_head[2].out_features == 16  # k_max + 1
+
+
+def test_stage1_model_n_sec_head_cfg_controls_hidden_width_and_depth():
+    model = Stage1Model(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=PARTICLE_CFG,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=32,
+        cond_out_dim=16,
+        n_sec_head_k_max=15,
+        n_sec_head_cfg={"hidden_ratio": 0.25, "depth": 1},
+    )
+    assert model.n_sec_head is not None
+    assert len(model.n_sec_head) == 1
+    assert model.n_sec_head[0].in_features == 16
+    assert model.n_sec_head[0].out_features == 16
+
+
 # --- cat_col_layout / stage2_type_dim / stage2_trunk_sec_dim ---------------
 
 
@@ -268,6 +304,38 @@ def test_stage2_oneshot_predict_type_raises_when_no_type_head():
         pass
 
 
+def test_stage2_oneshot_n_sec_head_and_type_head_cfg_control_hidden_width_and_depth():
+    """gitea #36: n_sec_head_cfg/type_head_cfg are independently tunable."""
+    k_max, emb_dim = 5, 6
+    particle_cfg = {"type": "onehot", "emb_dim": emb_dim, "n_layers": 1}
+    sec_dim = stage2_trunk_sec_dim({"target": "onehot"}, "flow", k_max, emb_dim)
+    model = Stage2OneShot(
+        pdg_vocab=5,
+        mat_vocab=3,
+        particle_cfg=particle_cfg,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=40,
+        n_res_blocks=1,
+        cond_out_dim=12,
+        context_dim=8,
+        sec_dim=sec_dim,
+        generator="flow",
+        k_max=k_max,
+        particle_type_cfg={"target": "onehot", "lambda": 1.0},
+        n_sec_head_cfg={"hidden_ratio": 0.25, "depth": 1},
+        type_head_cfg={"hidden_ratio": 0.75, "depth": 2},
+    )
+    assert model.n_sec_head is not None
+    assert len(model.n_sec_head) == 1
+    assert model.n_sec_head[0].in_features == 12
+    assert model.n_sec_head[0].out_features == k_max + 1
+
+    assert model.type_head is not None
+    assert len(model.type_head) == 3
+    assert model.type_head[0].out_features == 30  # round(40 * 0.75)
+    assert model.type_head[2].out_features == k_max * emb_dim
+
+
 def test_stage2_oneshot_forward_shape_onehot_wgan():
     B, k_max, emb_dim = 4, 5, 6
     model = _build_stage2("onehot", "wgan", emb_dim=emb_dim)
@@ -518,6 +586,37 @@ def test_stage2_autoregressive_particle_type_n_classes_overrides_conditioning_em
     assert model.type_head[-1].out_features == 20
 
 
+def test_stage2_autoregressive_n_sec_head_and_type_head_cfg_control_hidden_width_and_depth():
+    """gitea #36, Stage2Autoregressive side — see the Stage2OneShot version
+    of this test for the full rationale."""
+    particle_cfg = {"type": "physical", "emb_dim": 6, "n_layers": 1}
+    particle_type_cfg = {"target": "onehot", "lambda": 1.0}
+    model = Stage2Autoregressive(
+        pdg_vocab=5,
+        mat_vocab=3,
+        particle_cfg=particle_cfg,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=40,
+        n_res_blocks=1,
+        cond_out_dim=12,
+        context_dim=8,
+        generator="flow",
+        k_max=5,
+        particle_type_cfg=particle_type_cfg,
+        n_sec_head_cfg={"hidden_ratio": 0.25, "depth": 1},
+        type_head_cfg={"hidden_ratio": 0.75, "depth": 2},
+    )
+    assert model.n_sec_head is not None
+    assert len(model.n_sec_head) == 1
+    assert model.n_sec_head[0].in_features == 12
+    assert model.n_sec_head[0].out_features == 6  # k_max + 1
+
+    assert model.type_head is not None
+    assert len(model.type_head) == 3
+    assert model.type_head[0].out_features == 30  # round(40 * 0.75)
+    assert model.type_head[2].out_features == model.type_dim
+
+
 @pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
 @pytest.mark.parametrize("generator", ["wgan", "flow"])
 @pytest.mark.parametrize("history", ["markov", "attention"])
@@ -813,6 +912,39 @@ def test_build_models_omitted_decoder_and_particle_type_match_default_config():
     assert built["stage2"].particle_type_cfg["target"] == "onehot"
 
 
+def test_build_models_custom_heads_block_controls_head_shapes():
+    """gitea #36: stage{1,2}_model.heads flows all the way from config dict
+    through build_models to the actual constructed head shapes."""
+    cfg = _partial_model_config()
+    cfg["stage1_model"] = {
+        "active": True,
+        "hidden_dim": 40,
+        "n_res_blocks": 1,
+        "heads": {"n_sec": {"hidden_ratio": 0.25, "depth": 1}},
+    }
+    cfg["stage2_model"]["decoder"] = "one_shot"
+    cfg["stage2_model"]["generator"] = "flow"  # wgan folds the type slice; no separate type_head
+    cfg["stage2_model"]["n_sec"] = {"owner": "stage1"}
+    cfg["stage2_model"]["heads"] = {
+        "n_sec": {"hidden_ratio": 0.25, "depth": 1},
+        "type": {"hidden_ratio": 0.75, "depth": 2},
+    }
+
+    built = build_models(cfg)
+    stage1, stage2 = built["stage1"], built["stage2"]
+    assert stage1 is not None
+    assert stage2 is not None
+
+    assert stage1.n_sec_head is not None
+    assert len(stage1.n_sec_head) == 1  # owner=stage1, so stage1 builds it
+    assert stage2.n_sec_head is None  # owner=stage1, so stage2 doesn't
+
+    assert isinstance(stage2, Stage2OneShot)
+    assert stage2.type_head is not None
+    assert len(stage2.type_head) == 3
+    assert stage2.type_head[0].out_features == 6  # round(8 * 0.75)
+
+
 def test_build_critics_omitted_particle_type_matches_default_config():
     cfg = _partial_model_config()
     cfg["stage2_model"]["generator"] = "wgan"