CriticModel doesn't use the trunk/block registries or the StageModel base #57

New Issue

2026-08-14T14:39:43+02:00

lars commented

2026-08-14 14:39:43 +02:00

CriticModel (giant/model/models.py) is the one stage-shaped class left out of two refactors:

It hand-rolls its own body — self.blocks = nn.ModuleList([ResBlock(...) for _ in range(n_res_blocks)]) — instead of going through build_trunk/TRUNK_REGISTRY (gitea #33), so a routed/transformer/FiLM-conditioned trunk is available to every generative stage model except the critic that's supposed to compete against them under WGAN-GP.
It always uses plain ResBlock (additive conditioning) directly rather than build_block/BLOCK_REGISTRY (gitea #34), so trunk.block_conditioning = "film" | "adaln" has no effect on the critic side of a WGAN pair.
It wasn't folded into the StageModel base added in #39: it duplicates cond_enc construction and (for stage="stage2") the same context_adapter + fuse pattern as Stage2OneShot/Stage2Autoregressive's _cond_embed, by hand.

This was deliberately left out of #39 rather than folded in, since CriticModel would need reshaping to a registry-trunk body first — its forward takes a raw x/stage1_out pair and returns a scalar, not the (x_t, cond, cond_cont, cond_cat) -> vector shape Trunk implementations expect, so it isn't a drop-in _build_trunk_and_heads caller today.

Proposal: give CriticModel an output-width-1 trunk via build_trunk/build_block (probably wants a TrunkConfig-shaped stage{1,2}_model.wgan.critic_trunk or reuse of the generator's own trunk/block_conditioning config, tbd), and subclass StageModel for the cond_enc/context-fusion half — its stage2 context_adapter+fuse pair is already identical to Stage2OneShot's. Config-key naming and back-compat for existing WGAN checkpoints need deciding before implementation, per this tracker's usual process.

Related: #33 (trunk registry), #34 (block-conditioning registry), #39 (StageModel base, this issue's sibling for the three generator-side stage classes).

`CriticModel` (`giant/model/models.py`) is the one stage-shaped class left out of two refactors: - It hand-rolls its own body — `self.blocks = nn.ModuleList([ResBlock(...) for _ in range(n_res_blocks)])` — instead of going through `build_trunk`/`TRUNK_REGISTRY` (gitea #33), so a routed/transformer/FiLM-conditioned trunk is available to every generative stage model except the critic that's supposed to compete against them under WGAN-GP. - It always uses plain `ResBlock` (additive conditioning) directly rather than `build_block`/`BLOCK_REGISTRY` (gitea #34), so `trunk.block_conditioning = "film" | "adaln"` has no effect on the critic side of a WGAN pair. - It wasn't folded into the `StageModel` base added in #39: it duplicates `cond_enc` construction and (for `stage="stage2"`) the same `context_adapter` + `fuse` pattern as `Stage2OneShot`/`Stage2Autoregressive`'s `_cond_embed`, by hand. This was deliberately left out of #39 rather than folded in, since `CriticModel` would need reshaping to a registry-trunk body first — its `forward` takes a raw `x`/`stage1_out` pair and returns a scalar, not the `(x_t, cond, cond_cont, cond_cat) -> vector` shape `Trunk` implementations expect, so it isn't a drop-in `_build_trunk_and_heads` caller today. **Proposal:** give `CriticModel` an output-width-1 trunk via `build_trunk`/`build_block` (probably wants a `TrunkConfig`-shaped `stage{1,2}_model.wgan.critic_trunk` or reuse of the generator's own `trunk`/`block_conditioning` config, tbd), and subclass `StageModel` for the `cond_enc`/context-fusion half — its stage2 `context_adapter`+`fuse` pair is already identical to `Stage2OneShot`'s. Config-key naming and back-compat for existing WGAN checkpoints need deciding before implementation, per this tracker's usual process. Related: #33 (trunk registry), #34 (block-conditioning registry), #39 (StageModel base, this issue's sibling for the three generator-side stage classes).

lars added the architecture modularity labels 2026-08-14 14:39:43 +02:00

lars referenced this issue from a commit

2026-08-24 09:48:35 +02:00

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

lars referenced this issue

2026-08-24 09:51:57 +02:00

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

lars referenced this issue from a commit

2026-08-24 09:58:08 +02:00

Merge pull request 'Give CriticModel a registry-built trunk and StageModel base (gitea #57)' (#74) from fix/issue-57 into master

lars closed this issue

2026-08-24 09:58:16 +02:00