Generative objective (flow/ddpm/wgan) is a bare string, not a plugin #32

New Issue

2026-08-13T15:07:02+02:00

lars commented

2026-08-13 15:07:02 +02:00

This is the big one.

generator ∈ {"flow", "ddpm", "wgan"} is tested as a bare string in ~72 places
across 15 files. Excluding config.py/cli.py (where a string is correct),
the hot spots are models.py (22), sample.py (13), builders.py (6),
trainers.py (6), stage2_inputs.py (4).

Each of those sites independently re-derives some consequence of the choice:

does this stage have a time embedding? — has_time = generator in ("flow", "ddpm"), written out three times in models.py (lines 89, 191, 357)
what does the trunk take as input? — in_dim = noise_dim if generator == "wgan" else x_dim, three more times
is the type slice folded into the trunk output? — stage2_trunk_sec_dim in models.py, _type_folded in sample.py, _assemble_stage2_ar_target in stage2_inputs.py: three implementations of one rule
which sampler? — sample_stage1/sample_stage2 dispatch chains
which loss / which optimizer topology? — FlowDDPMStageTrainer vs WGANStageTrainer, ~500 lines of near-parallel structure

Adding a fourth objective — rectified flow, consistency distillation, a shortcut
model, anything in the "hit the 10× native-Geant4 eval budget" family that the
roadmap is explicitly chasing — means finding and correctly editing all five
groups. That is exactly the kind of change the branch should make cheap and
currently makes expensive.

Proposal: giant/model/objectives.py, mirroring routers.py:

class Objective(nn.Module):
    """One generative objective. Registered in OBJECTIVE_REGISTRY."""
    needs_time: bool                    # replaces `generator in ("flow","ddpm")`
    is_adversarial: bool                # replaces `generator == "wgan"`
    folds_type_slice: bool              # replaces `_type_folded`

    def trunk_in_dim(self, out_dim: int, noise_dim: int) -> int: ...
    def loss(self, model, x1, cond, mask=None, **kw) -> tuple[Tensor, dict]: ...
    def sample(self, model, cond, steps: int, **kw) -> Tensor: ...
    def build_critic(self, spec) -> nn.Module | None: ...   # None for non-adversarial

FlowObjective / DdpmObjective / WganObjective are then three small classes,
and one StageTrainer parameterised by an objective replaces the two subclasses
(the optimizer topology difference becomes build_critic() is not None).

Payoff: a new objective is one file plus one registry line, testable against
the existing per-objective tests without touching the stage models at all.

Risk: the largest change in this document — touches the trainers, which are
the least-covered-by-fast-tests part of the codebase. Sequence it after the
trunk-registry and conditioning-injection issues (small, independent) so the
refactor lands on a codebase where the other plug points are already
registry-shaped. Keep model.generator_kind as the persisted string so no
checkpoint changes.

Migrated from issues.md (v0.3.0 branch review, 2026-08-13), Issue 7.
The Router contract (giant/model/routers.py:21-110) is already the model
of what good looks like here — an ABC with working defaults, a registry, a
name-keyed factory that filters kwargs by signature so per-type
hyperparameters coexist without special-casing. Every Part B proposal from
that review is "extend the Router pattern to the axis that doesn't have it
yet." None of them requires inventing a new idea, and none of them breaks a
checkpoint: the config string stays the registry key.

**This is the big one.** `generator ∈ {"flow", "ddpm", "wgan"}` is tested as a bare string in ~72 places across 15 files. Excluding `config.py`/`cli.py` (where a string is correct), the hot spots are `models.py` (22), `sample.py` (13), `builders.py` (6), `trainers.py` (6), `stage2_inputs.py` (4). Each of those sites independently re-derives some consequence of the choice: - *does this stage have a time embedding?* — `has_time = generator in ("flow", "ddpm")`, written out three times in `models.py` (lines 89, 191, 357) - *what does the trunk take as input?* — `in_dim = noise_dim if generator == "wgan" else x_dim`, three more times - *is the type slice folded into the trunk output?* — `stage2_trunk_sec_dim` in `models.py`, `_type_folded` in `sample.py`, `_assemble_stage2_ar_target` in `stage2_inputs.py`: three implementations of one rule - *which sampler?* — `sample_stage1`/`sample_stage2` dispatch chains - *which loss / which optimizer topology?* — `FlowDDPMStageTrainer` vs `WGANStageTrainer`, ~500 lines of near-parallel structure Adding a fourth objective — rectified flow, consistency distillation, a shortcut model, anything in the "hit the 10× native-Geant4 eval budget" family that the roadmap is explicitly chasing — means finding and correctly editing all five groups. That is exactly the kind of change the branch should make cheap and currently makes expensive. **Proposal:** `giant/model/objectives.py`, mirroring `routers.py`: ```python class Objective(nn.Module): """One generative objective. Registered in OBJECTIVE_REGISTRY.""" needs_time: bool # replaces `generator in ("flow","ddpm")` is_adversarial: bool # replaces `generator == "wgan"` folds_type_slice: bool # replaces `_type_folded` def trunk_in_dim(self, out_dim: int, noise_dim: int) -> int: ... def loss(self, model, x1, cond, mask=None, **kw) -> tuple[Tensor, dict]: ... def sample(self, model, cond, steps: int, **kw) -> Tensor: ... def build_critic(self, spec) -> nn.Module | None: ... # None for non-adversarial ``` `FlowObjective` / `DdpmObjective` / `WganObjective` are then three small classes, and one `StageTrainer` parameterised by an objective replaces the two subclasses (the optimizer topology difference becomes `build_critic() is not None`). **Payoff:** a new objective is one file plus one registry line, testable against the existing per-objective tests without touching the stage models at all. **Risk:** the largest change in this document — touches the trainers, which are the least-covered-by-fast-tests part of the codebase. Sequence it after the trunk-registry and conditioning-injection issues (small, independent) so the refactor lands on a codebase where the other plug points are already registry-shaped. Keep `model.generator_kind` as the persisted string so no checkpoint changes. --- Migrated from `issues.md` (v0.3.0 branch review, 2026-08-13), Issue 7. The `Router` contract (`giant/model/routers.py:21-110`) is already the model of what good looks like here — an ABC with working defaults, a registry, a name-keyed factory that filters kwargs by signature so per-type hyperparameters coexist without special-casing. Every Part B proposal from that review is "extend the `Router` pattern to the axis that doesn't have it yet." None of them requires inventing a new idea, and none of them breaks a checkpoint: the config string stays the registry key.

lars added the architecture modularity labels 2026-08-13 15:07:02 +02:00

lars added a new dependency 2026-08-13 15:23:28 +02:00

#33 Trunk architecture is not selectable (hardcoded ResMLP/MoE choice)

lars added a new dependency 2026-08-13 15:23:28 +02:00

#34 Conditioning is injected additively only, with no FiLM/AdaLN alternative

Sign in to join this conversation.