From f8722e347e86597b6da5122b91e6bfe26dc8ba9e Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Fri, 14 Aug 2026 10:10:58 +0200 Subject: [PATCH] Add an Objective registry for the flow/ddpm/wgan generator choice (gitea #32) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit generator ∈ {"flow", "ddpm", "wgan"} was tested as a bare string in ~45 sites across models.py, sample.py, builders.py, trainers.py, and stage2_inputs.py, each independently re-deriving one of five consequences of the choice (needs a time embedding? what does the trunk take as input? is the type slice folded into the trunk output? which sampler? which loss?). giant/model/objectives.py adds an Objective ABC + OBJECTIVE_REGISTRY + build_objective factory, mirroring routers.py's Router pattern, and every bare-string site now goes through it (needs_time, is_adversarial, folds_type_slice, trunk_in_dim, build_schedule, stage1_loss/stage2_loss). Per discussion: FlowDDPMStageTrainer and WGANStageTrainer stay separate classes rather than merging into one StageTrainer as the issue's sketch proposed — their training loops are genuinely different shapes (single loss vs. dual G/D step with gradient penalty/n_critic/ST-Gumbel), and trainers.py is the least-covered-by-fast-tests part of the codebase, so a full merge was judged out of proportion to this issue's risk budget. FlowDDPMStageTrainer's own loss dispatch (flow vs ddpm, one-shot vs AR) does move onto the objective, so a future non-adversarial objective (rectified flow, consistency distillation) is still a one-file, zero-trainer-edits addition. No config-schema change — stage{1,2}_model.generator stays the persisted string, just looked up in the registry instead of string-compared. An unrecognized generator value now fails fast with a clear ValueError instead of silently falling through some bare-string checks and not others (same behavior build_router/build_trunk already have for their own type keys). Co-Authored-By: Claude Opus 5 --- giant/model/builders.py | 16 +- giant/model/models.py | 41 +++--- giant/model/network.py | 16 ++ giant/model/objectives.py | 202 ++++++++++++++++++++++++++ giant/sample.py | 16 +- giant/training/stage2_inputs.py | 18 ++- giant/training/trainers.py | 53 +++---- tests/test_objectives.py | 250 ++++++++++++++++++++++++++++++++ 8 files changed, 537 insertions(+), 75 deletions(-) create mode 100644 giant/model/objectives.py create mode 100644 tests/test_objectives.py diff --git a/giant/model/builders.py b/giant/model/builders.py index 9fe2ca1..61c7318 100644 --- a/giant/model/builders.py +++ b/giant/model/builders.py @@ -15,6 +15,7 @@ from giant.model.models import ( resolve_type_n_classes, stage2_trunk_sec_dim, ) +from giant.model.objectives import build_objective from giant.model.routers import Router, _build_router_from_cfg # --------------------------------------------------------------------------- @@ -64,10 +65,11 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]: if s1_spec.router.enabled: stage1_router = _build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, particle_conditioning) generator = s1_spec.generator + objective = build_objective(generator) # wgan has no time_dim concept (no diffusion/flow time variable) — # matches the pre-dataclass .get("time_dim", 64) fallback, which # always hit its default for a wgan sub-block too. - time_dim = getattr(s1_spec, generator).time_dim if generator != "wgan" else 64 + time_dim = getattr(s1_spec, generator).time_dim if objective.needs_time else 64 n_sec_owner = s2_spec.n_sec.owner n_sec_head_k_max = s2_spec.k_max if n_sec_owner == "stage1" else None result["stage1"] = Stage1Model( @@ -99,9 +101,10 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]: else: stage2_router = _build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, particle_conditioning) generator = s2_spec.generator + objective = build_objective(generator) # wgan has no time_dim concept — see the matching comment in stage 1 # above. - time_dim = getattr(s2_spec, generator).time_dim if generator != "wgan" else 64 + time_dim = getattr(s2_spec, generator).time_dim if objective.needs_time else 64 n_sec_owner = s2_spec.n_sec.owner k_max = s2_spec.k_max particle_type_cfg = s2_spec.particle_type.to_dict() @@ -180,7 +183,7 @@ def build_critics(model_config: dict) -> dict[str, nn.Module | None]: result: dict[str, nn.Module | None] = {"stage1": None, "stage2": None} - if s1_spec.active and s1_spec.generator == "wgan": + if s1_spec.active and build_objective(s1_spec.generator).is_adversarial: result["stage1"] = CriticModel( pdg_vocab=pdg_vocab, mat_vocab=mat_vocab, @@ -194,11 +197,14 @@ def build_critics(model_config: dict) -> dict[str, nn.Module | None]: stage="stage1", ) - if s2_spec.active and s2_spec.generator == "wgan": + if s2_spec.active and build_objective(s2_spec.generator).is_adversarial: k_max = s2_spec.k_max particle_type_cfg = s2_spec.particle_type.to_dict() in_dim = stage2_trunk_sec_dim( - particle_type_cfg, "wgan", k_max, resolve_type_n_classes(particle_type_cfg, particle_cfg["emb_dim"]) + particle_type_cfg, + s2_spec.generator, + k_max, + resolve_type_n_classes(particle_type_cfg, particle_cfg["emb_dim"]), ) result["stage2"] = CriticModel( pdg_vocab=pdg_vocab, diff --git a/giant/model/models.py b/giant/model/models.py index 9dcceb7..3d5d3b2 100644 --- a/giant/model/models.py +++ b/giant/model/models.py @@ -8,6 +8,7 @@ from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_DIM, SE from giant.model.encoders import ConditionEncoder from giant.model.history import AttentionHistory, HistoryEncoder, MarkovHistory from giant.model.layers import ContextAdapter, ResBlock, SinusoidalEmbedding +from giant.model.objectives import build_objective from giant.model.routers import Router from giant.model.trunks import build_trunk @@ -49,18 +50,18 @@ def stage2_trunk_sec_dim(particle_type_cfg: dict, generator: str, k_max: int, em `k_max * SEC_SLOT_DIM`, the type slice folded into the same flow-matched/WGAN vector as the continuous stick/dir slots. - `target` in `("onehot", "embedding")`: under `generator == "wgan"` the - type slice is still folded in (adversarial for onehot via ST-Gumbel, - already-continuous for embedding), just `emb_dim` wide instead of - `PARTICLE_PHYS_DIM` wide: `k_max * (CONT_SLOT_DIM + emb_dim)`. Under - `generator in ("flow", "ddpm")` the type slice isn't part of this vector - at all — it's `Stage2OneShot.type_head`'s job instead — so the trunk - only covers `k_max * CONT_SLOT_DIM`. + `target` in `("onehot", "embedding")`: under an objective with + `folds_type_slice` (currently just wgan) the type slice is still folded + in (adversarial for onehot via ST-Gumbel, already-continuous for + embedding), just `emb_dim` wide instead of `PARTICLE_PHYS_DIM` wide: + `k_max * (CONT_SLOT_DIM + emb_dim)`. Otherwise (flow/ddpm) the type slice + isn't part of this vector at all — it's `Stage2OneShot.type_head`'s job + instead — so the trunk only covers `k_max * CONT_SLOT_DIM`. """ target = particle_type_cfg.get("target", "physical") if target == "physical": return k_max * SEC_SLOT_DIM - if generator == "wgan": + if build_objective(generator).folds_type_slice: return k_max * (CONT_SLOT_DIM + emb_dim) return k_max * CONT_SLOT_DIM @@ -105,10 +106,11 @@ class Stage1Model(nn.Module): if cond_enc is not None else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim) ) - has_time = generator in ("flow", "ddpm") + objective = build_objective(generator) + has_time = objective.needs_time 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 + in_dim = objective.trunk_in_dim(x_dim, noise_dim) self.trunk = build_trunk( router, trunk_type, in_dim, x_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout, block_conditioning ) @@ -159,11 +161,12 @@ class Stage2OneShot(nn.Module): the trunk's own flow-matched/WGAN output, unchanged from v0.2 (`sec_dim` — computed by the caller via `stage2_trunk_sec_dim` — already reflects this). Under - `"onehot"`/`"embedding"` with `generator in ("flow", "ddpm")`, the type + `"onehot"`/`"embedding"` with an objective (`giant.model.objectives`) that + doesn't fold the type slice (flow/ddpm), the type slice is predicted by a separate `type_head` instead (same shape pattern as `n_sec_head`) — `sec_dim` then covers only the continuous stick/dir slots, `type_head` covers `k_max * emb_dim` type logits/vectors. - Under `generator == "wgan"` the type slice stays folded into `sec_dim` + Under a folding objective (wgan) the type slice stays folded into `sec_dim` (just `emb_dim` instead of `PARTICLE_PHYS_DIM` wide) and `type_head` is unused (`None`) — the WGAN trainer handles the ST-Gumbel relaxation. @@ -213,10 +216,11 @@ class Stage2OneShot(nn.Module): nn.Linear(cond_out_dim + context_dim, cond_out_dim), nn.SiLU(), ) - has_time = generator in ("flow", "ddpm") + objective = build_objective(generator) + has_time = objective.needs_time 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 sec_dim + in_dim = objective.trunk_in_dim(sec_dim, noise_dim) self.trunk = build_trunk( router, trunk_type, @@ -237,7 +241,7 @@ class Stage2OneShot(nn.Module): ) self.type_head = None target = self.particle_type_cfg.get("target", "physical") - if target != "physical" and generator in ("flow", "ddpm"): + 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), @@ -391,11 +395,12 @@ class Stage2Autoregressive(nn.Module): nn.SiLU(), ) - has_time = generator in ("flow", "ddpm") + objective = build_objective(generator) + has_time = objective.needs_time self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim token_dim = stage2_trunk_sec_dim(self.particle_type_cfg, generator, 1, emb_dim) - in_dim = noise_dim if generator == "wgan" else token_dim + in_dim = objective.trunk_in_dim(token_dim, noise_dim) self.trunk = build_trunk( router, trunk_type, @@ -417,7 +422,7 @@ class Stage2Autoregressive(nn.Module): ) self.type_head = None target = self.particle_type_cfg.get("target", "physical") - if target != "physical" and generator in ("flow", "ddpm"): + if target != "physical" and not objective.folds_type_slice: self.type_head = nn.Sequential( nn.Linear(cond_out_dim, hidden_dim // 2), nn.SiLU(), diff --git a/giant/model/network.py b/giant/model/network.py index fc558ed..fbdcb44 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -31,6 +31,15 @@ from giant.model.models import ( stage2_trunk_sec_dim, stage2_type_dim, ) +from giant.model.objectives import ( + OBJECTIVE_REGISTRY, + DdpmObjective, + FlowObjective, + Objective, + WganObjective, + build_objective, + register_objective, +) from giant.model.routers import ( ROUTER_REGISTRY, ComposedRouter, @@ -64,11 +73,15 @@ __all__ = [ "ConditionEncoder", "ContextAdapter", "CriticModel", + "DdpmObjective", "EnergyRouter", "ExpertTrunk", "FilmResBlock", + "FlowObjective", "HistoryEncoder", "MarkovHistory", + "OBJECTIVE_REGISTRY", + "Objective", "PdgRouter", "ProcessRouter", "ROUTER_REGISTRY", @@ -81,6 +94,7 @@ __all__ = [ "Stage2OneShot", "TRUNK_REGISTRY", "Trunk", + "WganObjective", "_CausalAttnBlock", "_build_router_from_cfg", "_check_router_conditioning_compat", @@ -93,11 +107,13 @@ __all__ = [ "build_critics", "build_expert_body", "build_models", + "build_objective", "build_router", "build_trunk", "cat_col_layout", "migrate_legacy_state_dict", "register_block", + "register_objective", "register_router", "register_trunk", "resolve_type_n_classes", diff --git a/giant/model/objectives.py b/giant/model/objectives.py new file mode 100644 index 0000000..c11fa10 --- /dev/null +++ b/giant/model/objectives.py @@ -0,0 +1,202 @@ +"""Generative objectives (flow/ddpm/wgan): `Objective` base + registry, +mirroring `giant.model.routers`'s `Router` pattern (gitea #32). Each objective +answers, in one place, the handful of questions every stage model/sampler/ +trainer used to re-derive independently from a bare `generator` string: does +this stage need a time embedding, is it adversarial, does it fold the +secondary type slice into its own trunk output, what does the trunk take as +input, which stage-1/stage-2 loss does it train against. + +Self-contained (no dependency on `giant.model.models`, unlike `Router` which +`giant.model.trunks` depends on) — `Objective` never needs to construct a +stage model or critic itself, only describe one. This also sidesteps a +`models.py` <-> `objectives.py` import cycle, since `models.py` calls +`build_objective`. +""" + +import inspect + +import torch + +from giant.model.schedule import ( + CosineSchedule, + flow_matching_loss, + flow_matching_loss_secondary, + flow_matching_loss_secondary_ar, +) + +# --------------------------------------------------------------------------- +# Objective contract +# --------------------------------------------------------------------------- + + +class Objective: + """Contract for a pluggable generative objective. Not an `nn.Module` — + unlike `Router`, no objective owns learnable parameters, so a plain + strategy object is the honest fit. + + `needs_time`/`is_adversarial`/`folds_type_slice`/`supports_stage2_decoder` + are set by each concrete subclass (no defaults here — a new objective + should have to state all four, not silently inherit one that happens to + be wrong for it). See `FlowObjective`/`DdpmObjective`/`WganObjective`. + """ + + needs_time: bool + is_adversarial: bool + folds_type_slice: bool + supports_stage2_decoder: bool = True + + def trunk_in_dim(self, out_dim: int, noise_dim: int) -> int: + """Width of the trunk's own input — `out_dim` (denoising/flow-matching + a same-shape vector) for every non-adversarial objective; + `WganObjective` overrides to `noise_dim` (a single-pass noise-to-output + generator).""" + return out_dim + + def build_schedule(self, n_steps: int, device: torch.device) -> CosineSchedule | None: + """Objective-owned auxiliary state a stage trainer must build once + and hold onto (device-placed) across its training loop. `None` for + every objective except `DdpmObjective` (its noise schedule).""" + return None + + def stage1_loss( + self, + model: torch.nn.Module, + x1: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + *, + schedule: object | None = None, + ) -> torch.Tensor: + """Stage-1 training loss. Only implemented by non-adversarial + objectives — `WganObjective` is unused here, `WGANStageTrainer` has + its own G/D step instead.""" + raise NotImplementedError(f"{type(self).__name__} has no stage1_loss") + + def stage2_loss( + self, + model: torch.nn.Module, + x1_s2: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + stage1_ctx: torch.Tensor, + sec_mask: torch.Tensor, + *, + type_dim: int | None, + ar_inputs: dict[str, torch.Tensor] | None = None, + ) -> torch.Tensor: + """Stage-2 secondary-decoder training loss, one-shot or + autoregressive depending on whether `ar_inputs` is given. Same + adversarial caveat as `stage1_loss`.""" + raise NotImplementedError(f"{type(self).__name__} has no stage2_loss") + + +OBJECTIVE_REGISTRY: dict[str, type[Objective]] = {} + + +def register_objective(name: str): + def decorator(cls: type[Objective]) -> type[Objective]: + OBJECTIVE_REGISTRY[name] = cls + return cls + + return decorator + + +def build_objective(name: str, **kwargs) -> Objective: + """Factory: look up an `Objective` subclass by name (a `generator` + config value) from the registry. + + Every registered objective is fed the same kwargs; kwargs not declared by + that type's constructor are silently dropped, so per-type hyperparameters + (e.g. `DdpmObjective`'s `n_steps`) can coexist in one call without + special-casing — same convention as `giant.model.routers.build_router`. + """ + if name not in OBJECTIVE_REGISTRY: + raise ValueError(f"unknown generator/objective {name!r}; available: {sorted(OBJECTIVE_REGISTRY)}") + cls = OBJECTIVE_REGISTRY[name] + accepted = set(inspect.signature(cls.__init__).parameters) - {"self"} + filtered = {k: v for k, v in kwargs.items() if k in accepted} + return cls(**filtered) + + +# --------------------------------------------------------------------------- +# Concrete objectives +# --------------------------------------------------------------------------- + + +@register_objective("flow") +class FlowObjective(Objective): + """Conditional flow matching (Lipman et al. 2022) — the primary + objective. ~10 ODE steps at inference (`giant.sample.sample_flow`).""" + + needs_time = True + is_adversarial = False + folds_type_slice = False + + def stage1_loss(self, model, x1, cond_cont, cond_cat, *, schedule=None) -> torch.Tensor: + return flow_matching_loss(model, x1, cond_cont, cond_cat) + + def stage2_loss( + self, + model, + x1_s2, + cond_cont, + cond_cat, + stage1_ctx, + sec_mask, + *, + type_dim=None, + ar_inputs=None, + ) -> torch.Tensor: + if ar_inputs is not None: + return flow_matching_loss_secondary_ar( + model, + x1_s2, + cond_cont, + cond_cat, + stage1_ctx, + ar_inputs["history_feat"], + ar_inputs["has_prev"], + ar_inputs["remaining_frac"], + ar_inputs["slot_idx"], + sec_mask, + type_dim=type_dim, + ) + return flow_matching_loss_secondary(model, x1_s2, cond_cont, cond_cat, stage1_ctx, sec_mask, type_dim=type_dim) + + +@register_objective("ddpm") +class DdpmObjective(Objective): + """Full DDPM ancestral sampling (Nichol & Dhariwal 2021 cosine schedule) + — the throwaway baseline. Stage-1 only: no `Stage2*` class has ever been + trained with `generator="ddpm"` in practice, so there's no stage-2 ddpm + loss to dispatch to (matches `FlowDDPMStageTrainer`'s pre-existing + stage-2 guard).""" + + needs_time = True + is_adversarial = False + folds_type_slice = False + supports_stage2_decoder = False + + def __init__(self, n_steps: int = 1000) -> None: + self.n_steps = n_steps + + def build_schedule(self, n_steps: int, device: torch.device) -> CosineSchedule: + return CosineSchedule(T=n_steps).to(device) + + def stage1_loss(self, model, x1, cond_cont, cond_cat, *, schedule=None) -> torch.Tensor: + assert schedule is not None, "DdpmObjective.stage1_loss needs a schedule (see build_schedule)" + return schedule.loss(model, x1, cond_cont, cond_cat) + + +@register_objective("wgan") +class WganObjective(Objective): + """WGAN-GP (Gulrajani et al. 2017) — single forward pass instead of an + ODE loop. `stage1_loss`/`stage2_loss` are unused: `WGANStageTrainer` owns + its own dual generator/critic step instead of a single scalar loss.""" + + needs_time = False + is_adversarial = True + folds_type_slice = True + + def trunk_in_dim(self, out_dim: int, noise_dim: int) -> int: + return noise_dim diff --git a/giant/sample.py b/giant/sample.py index d4486b2..fede756 100644 --- a/giant/sample.py +++ b/giant/sample.py @@ -2,7 +2,7 @@ import torch import torch.nn.functional as F from giant.constants import CONT_SLOT_DIM, X_DIM -from giant.model.network import Stage2Autoregressive, stage2_trunk_sec_dim +from giant.model.network import DdpmObjective, Stage2Autoregressive, build_objective, stage2_trunk_sec_dim from giant.model.schedule import CosineSchedule @@ -132,7 +132,7 @@ def _stage2_flat_width(sec_decoder: torch.nn.Module) -> int: def _type_folded(sec_decoder: torch.nn.Module) -> bool: target = sec_decoder.particle_type_cfg.get("target", "physical") - return target == "physical" or sec_decoder.generator_kind == "wgan" + return target == "physical" or build_objective(sec_decoder.generator_kind).folds_type_slice def _decode_stage2_flat( @@ -281,7 +281,7 @@ def sample_secondaries_ar( device = cond_cont.device k_max = sec_decoder.k_max type_dim = sec_decoder.type_dim - generator = sec_decoder.generator_kind + objective = build_objective(sec_decoder.generator_kind) target = sec_decoder.particle_type_cfg.get("target", "physical") type_folded = _type_folded(sec_decoder) token_dim = CONT_SLOT_DIM + type_dim if type_folded else CONT_SLOT_DIM @@ -301,7 +301,7 @@ def sample_secondaries_ar( slot_idx = torch.full((B, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32) hist, history_cache = sec_decoder.history_step(history_feat, has_prev, history_cache) - if generator == "wgan": + if objective.is_adversarial: z = torch.randn(B, 1, sec_decoder.noise_dim, device=device) token = sec_decoder( z, @@ -383,10 +383,10 @@ def sample_stage1( ddpm_steps: int = 1000, ) -> tuple[torch.Tensor, torch.Tensor | None]: """Dispatches on `stage1_model.generator_kind`.""" - kind = stage1_model.generator_kind - if kind == "wgan": + objective = build_objective(stage1_model.generator_kind) + if objective.is_adversarial: return sample_wgan(stage1_model, cond_cont, cond_cat) - if kind == "ddpm": + if isinstance(objective, DdpmObjective): schedule = CosineSchedule(T=ddpm_steps).to(cond_cont.device) return sample_ddpm(stage1_model, cond_cont, cond_cat, schedule) return sample_flow(stage1_model, cond_cont, cond_cat, steps=steps) @@ -409,7 +409,7 @@ def sample_stage2( """ if isinstance(sec_decoder, Stage2Autoregressive): return sample_secondaries_ar(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps) - if sec_decoder.generator_kind == "wgan": + if build_objective(sec_decoder.generator_kind).is_adversarial: return sample_secondaries_wgan(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred) return sample_secondaries(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps) diff --git a/giant/training/stage2_inputs.py b/giant/training/stage2_inputs.py index e2f8367..c09011e 100644 --- a/giant/training/stage2_inputs.py +++ b/giant/training/stage2_inputs.py @@ -12,6 +12,7 @@ import torch import torch.nn.functional as F from giant.constants import CONT_SLOT_DIM, PARTICLE_PHYS_DIM +from giant.model.objectives import build_objective from giant.sample import sample_secondaries_ar @@ -69,19 +70,20 @@ def _assemble_stage2_ar_target( - `target = "physical"`: unchanged from v0.2 — `sec_cont` (stick_logit, dir, log_mass, charge) as-is. - - `target` in `("onehot", "embedding")` + `generator in ("flow", "ddpm")`: - just the continuous stick/dir slots — the type slice isn't part of - this tensor at all (`type_head` handles it separately). - - `target` in `("onehot", "embedding")` + `generator == "wgan"`: stick/dir - slots concatenated with the per-slot type representation (a one-hot of - the true class, relaxed on the *generated* side only, by the caller; - or the conditioning's own detached embedding-table row). + - `target` in `("onehot", "embedding")` + an objective that doesn't fold + the type slice (flow/ddpm): just the continuous stick/dir slots — the + type slice isn't part of this tensor at all (`type_head` handles it + separately). + - `target` in `("onehot", "embedding")` + a folding objective (wgan): + stick/dir slots concatenated with the per-slot type representation (a + one-hot of the true class, relaxed on the *generated* side only, by the + caller; or the conditioning's own detached embedding-table row). """ target = particle_type_cfg.get("target", "physical") if target == "physical": return sec_cont cont = sec_cont[..., :CONT_SLOT_DIM] - if generator != "wgan": + if not build_objective(generator).folds_type_slice: return cont type_repr = _type_repr(sec_type_idx, sec_cont, particle_type_cfg, cond_enc, emb_dim) return torch.cat([cont, type_repr], dim=-1) diff --git a/giant/training/trainers.py b/giant/training/trainers.py index a998e77..e91ff8e 100644 --- a/giant/training/trainers.py +++ b/giant/training/trainers.py @@ -25,13 +25,7 @@ import torch.optim as optim from giant.config import ParticleTypeConfig, Stage1ModelConfig, Stage2ModelConfig, TrainConfig from giant.constants import CONT_SLOT_DIM from giant.data.dataset import StepBatch -from giant.model.network import Router, resolve_type_n_classes, stage2_type_dim -from giant.model.schedule import ( - CosineSchedule, - flow_matching_loss, - flow_matching_loss_secondary, - flow_matching_loss_secondary_ar, -) +from giant.model.network import Router, build_objective, resolve_type_n_classes, stage2_type_dim from giant.model.wgan import generator_loss, gradient_penalty from giant.training.metrics import MetricSpec, stage_metric, train_metric, val_metric from giant.training.stage2_inputs import ( @@ -450,19 +444,22 @@ class FlowDDPMStageTrainer(StageTrainer): """flow or ddpm generator for a single stage.""" def __init__(self, spec: StageSpec, model: torch.nn.Module, device: torch.device) -> None: - if spec.is_stage2 and spec.generator not in ("flow",): + objective = build_objective(spec.generator, n_steps=spec.ddpm_n_steps) + if spec.is_stage2 and not objective.supports_stage2_decoder: raise NotImplementedError( f"stage2_model.generator={spec.generator!r} is accepted by the " "schema but not implemented in v0.3.0 for stage 2 (only " "'flow' and 'wgan' have a stage-2 secondary-decoder loss)" ) super().__init__(spec, model, device) + self.objective = objective self.particle_type_lambda = self.particle_type_cfg.get("lambda", 1.0) # Width of the type slice actually folded into x1_s2 by _sec_target, - # under this trainer's generator (flow/ddpm only — see the - # NotImplementedError above): "physical" keeps it folded in - # (PARTICLE_PHYS_DIM wide, unchanged from v0.2); "onehot"/"embedding" - # pull it out into model.type_head instead (0 here). + # under this trainer's objective (flow/ddpm only — see the + # NotImplementedError above, neither folds the type slice): "physical" + # keeps it folded in (PARTICLE_PHYS_DIM wide, unchanged from v0.2); + # "onehot"/"embedding" pull it out into model.type_head instead (0 + # here). self._flow_type_dim = None if self.particle_type_cfg.get("target", "physical") == "physical" else 0 self.params = list(self.model.parameters()) @@ -472,7 +469,7 @@ class FlowDDPMStageTrainer(StageTrainer): warmup_steps=spec.warmup_epochs * spec.steps_per_epoch, total_steps=max(spec.epochs * spec.steps_per_epoch, 1), ) - self.ddpm_schedule = CosineSchedule(T=spec.ddpm_n_steps).to(device) if spec.generator == "ddpm" else None + self.ddpm_schedule = self.objective.build_schedule(spec.ddpm_n_steps, device) self.train_metrics = [ train_metric(key) @@ -504,26 +501,9 @@ class FlowDDPMStageTrainer(StageTrainer): def _generator_loss(self, cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx, ar_inputs=None): if not self.is_stage2: - if self.generator == "flow": - return flow_matching_loss(self.model, x1_s1, cond_cont, cond_cat) - assert self.ddpm_schedule is not None - return self.ddpm_schedule.loss(self.model, x1_s1, cond_cont, cond_cat) - if self.decoder == "autoregressive": - assert ar_inputs is not None - return flow_matching_loss_secondary_ar( - self.model, - x1_s2, - cond_cont, - cond_cat, - stage1_ctx, - ar_inputs["history_feat"], - ar_inputs["has_prev"], - ar_inputs["remaining_frac"], - ar_inputs["slot_idx"], - sec_mask, - type_dim=self._flow_type_dim, - ) - return flow_matching_loss_secondary( + return self.objective.stage1_loss(self.model, x1_s1, cond_cont, cond_cat, schedule=self.ddpm_schedule) + assert self.decoder != "autoregressive" or ar_inputs is not None + return self.objective.stage2_loss( self.model, x1_s2, cond_cont, @@ -531,6 +511,7 @@ class FlowDDPMStageTrainer(StageTrainer): stage1_ctx, sec_mask, type_dim=self._flow_type_dim, + ar_inputs=ar_inputs, ) def _type_loss( @@ -774,7 +755,7 @@ class WGANStageTrainer(StageTrainer): if self.decoder == "autoregressive": epoch = global_step // self.spec.steps_per_epoch ar = self._ar_inputs(cond_cont, cond_cat, stage1_ctx, sec_cont, sec_type_idx, n_sec, epoch) - real = self._sec_target(sec_cont, sec_type_idx, "wgan", flatten=False).reshape(B, -1) * mask + real = self._sec_target(sec_cont, sec_type_idx, self.generator, flatten=False).reshape(B, -1) * mask z = torch.randn(B, k_max, self.model.noise_dim, device=device) fake_raw = self.model( z, @@ -787,7 +768,7 @@ class WGANStageTrainer(StageTrainer): ar["slot_idx"], ).reshape(B, -1) else: - real = self._sec_target(sec_cont, sec_type_idx, "wgan", flatten=True) * mask + real = self._sec_target(sec_cont, sec_type_idx, self.generator, flatten=True) * mask z = torch.randn(B, self.model.noise_dim, device=device) fake_raw = self.model(z, cond_cont, cond_cat, stage1_ctx) @@ -957,7 +938,7 @@ def build_stage_trainers( if model is None: continue spec = StageSpec.from_config(cfg, name, is_stage2, max(total_train_batches, 1)) - if spec.generator == "wgan": + if build_objective(spec.generator).is_adversarial: critic = critics.get(name) assert critic is not None, ( f"{name}_model.generator='wgan' requires a critic (see giant.model.network.build_critics)" diff --git a/tests/test_objectives.py b/tests/test_objectives.py new file mode 100644 index 0000000..aaf043a --- /dev/null +++ b/tests/test_objectives.py @@ -0,0 +1,250 @@ +"""Tests for `giant/model/objectives.py` — the generator/objective registry +(gitea #32) that replaced bare `generator in ("flow", "ddpm", "wgan")` +string checks scattered across models.py/sample.py/builders.py/ +stage2_inputs.py/trainers.py.""" + +import pytest +import torch + +from giant.constants import COND_DIM, CONT_SLOT_DIM, PARTICLE_PHYS_DIM, SEC_SLOT_DIM, X_DIM +from giant.model.network import ( + OBJECTIVE_REGISTRY, + DdpmObjective, + FlowObjective, + Stage1Model, + Stage2Autoregressive, + Stage2OneShot, + WganObjective, + build_objective, +) +from giant.model.schedule import CosineSchedule, flow_matching_loss, flow_matching_loss_secondary + +_PHYS_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1} + + +def _cond(B: int, pdg: int = 3, mat: int = 2) -> tuple[torch.Tensor, torch.Tensor]: + cond_cont = torch.randn(B, COND_DIM) + cond_cat = torch.stack([torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1) + return cond_cont, cond_cat + + +# ── registry ───────────────────────────────────────────────────────────── + + +def test_registry_has_exactly_the_three_known_objectives(): + assert set(OBJECTIVE_REGISTRY) == {"flow", "ddpm", "wgan"} + + +def test_build_objective_returns_correct_concrete_type(): + assert isinstance(build_objective("flow"), FlowObjective) + assert isinstance(build_objective("ddpm"), DdpmObjective) + assert isinstance(build_objective("wgan"), WganObjective) + + +def test_build_objective_unknown_name_raises(): + with pytest.raises(ValueError, match="unknown generator/objective"): + build_objective("bogus") + + +def test_build_objective_filters_kwargs_by_signature(): + # FlowObjective takes no constructor args — n_steps (a DdpmObjective-only + # kwarg) must be silently dropped, not raise a TypeError. + build_objective("flow", n_steps=500) + ddpm = build_objective("ddpm", n_steps=250) + assert isinstance(ddpm, DdpmObjective) + assert ddpm.n_steps == 250 + + +# ── flags ──────────────────────────────────────────────────────────────── + + +def test_flow_objective_flags(): + obj = build_objective("flow") + assert obj.needs_time is True + assert obj.is_adversarial is False + assert obj.folds_type_slice is False + assert obj.supports_stage2_decoder is True + + +def test_ddpm_objective_flags(): + obj = build_objective("ddpm") + assert obj.needs_time is True + assert obj.is_adversarial is False + assert obj.folds_type_slice is False + assert obj.supports_stage2_decoder is False + + +def test_wgan_objective_flags(): + obj = build_objective("wgan") + assert obj.needs_time is False + assert obj.is_adversarial is True + assert obj.folds_type_slice is True + assert obj.supports_stage2_decoder is True + + +# ── trunk_in_dim ───────────────────────────────────────────────────────── + + +def test_trunk_in_dim_flow_and_ddpm_pass_through_out_dim(): + assert build_objective("flow").trunk_in_dim(out_dim=9, noise_dim=8) == 9 + assert build_objective("ddpm").trunk_in_dim(out_dim=9, noise_dim=8) == 9 + + +def test_trunk_in_dim_wgan_uses_noise_dim(): + assert build_objective("wgan").trunk_in_dim(out_dim=9, noise_dim=8) == 8 + + +# ── ddpm schedule ──────────────────────────────────────────────────────── + + +def test_ddpm_build_schedule_has_requested_length(): + schedule = build_objective("ddpm").build_schedule(n_steps=17, device=torch.device("cpu")) + assert isinstance(schedule, CosineSchedule) + assert schedule.T == 17 + + +def test_flow_and_wgan_build_schedule_is_none(): + assert build_objective("flow").build_schedule(100, torch.device("cpu")) is None + assert build_objective("wgan").build_schedule(100, torch.device("cpu")) is None + + +# ── stage1_loss parity ────────────────────────────────────────────────── + + +def test_flow_objective_stage1_loss_matches_direct_call(): + torch.manual_seed(0) + model = Stage1Model( + pdg_vocab=3, mat_vocab=2, particle_cfg=_PHYS_CFG, material_cfg=_PHYS_CFG, hidden_dim=16, n_res_blocks=1 + ) + cond_cont, cond_cat = _cond(4) + x1 = torch.randn(4, X_DIM) + + torch.manual_seed(1) + expected = flow_matching_loss(model, x1, cond_cont, cond_cat) + torch.manual_seed(1) + actual = build_objective("flow").stage1_loss(model, x1, cond_cont, cond_cat) + assert torch.allclose(actual, expected) + + +def test_ddpm_objective_stage1_loss_matches_direct_call(): + torch.manual_seed(0) + model = Stage1Model( + pdg_vocab=3, + mat_vocab=2, + particle_cfg=_PHYS_CFG, + material_cfg=_PHYS_CFG, + hidden_dim=16, + n_res_blocks=1, + generator="ddpm", + ) + cond_cont, cond_cat = _cond(4) + x1 = torch.randn(4, X_DIM) + objective = build_objective("ddpm", n_steps=50) + schedule = objective.build_schedule(50, torch.device("cpu")) + assert isinstance(schedule, CosineSchedule) + + torch.manual_seed(1) + expected = schedule.loss(model, x1, cond_cont, cond_cat) + torch.manual_seed(1) + actual = objective.stage1_loss(model, x1, cond_cont, cond_cat, schedule=schedule) + assert torch.allclose(actual, expected) + + +def test_ddpm_objective_stage1_loss_requires_a_schedule(): + model = Stage1Model( + pdg_vocab=3, + mat_vocab=2, + particle_cfg=_PHYS_CFG, + material_cfg=_PHYS_CFG, + hidden_dim=16, + n_res_blocks=1, + generator="ddpm", + ) + cond_cont, cond_cat = _cond(4) + with pytest.raises(AssertionError): + build_objective("ddpm").stage1_loss(model, torch.randn(4, X_DIM), cond_cont, cond_cat, schedule=None) + + +# ── stage2_loss dispatch ───────────────────────────────────────────────── + + +def test_flow_objective_stage2_loss_one_shot_matches_direct_call(): + torch.manual_seed(0) + B, k_max = 4, 5 + sec_dim = k_max * SEC_SLOT_DIM + model = Stage2OneShot( + pdg_vocab=3, + mat_vocab=2, + particle_cfg=_PHYS_CFG, + material_cfg=_PHYS_CFG, + hidden_dim=16, + n_res_blocks=1, + generator="flow", + sec_dim=sec_dim, + k_max=k_max, + ) + cond_cont, cond_cat = _cond(B) + stage1_ctx = torch.randn(B, X_DIM) + x1_s2 = torch.randn(B, sec_dim) + sec_mask = torch.ones(B, k_max, dtype=torch.bool) + + torch.manual_seed(1) + expected = flow_matching_loss_secondary(model, x1_s2, cond_cont, cond_cat, stage1_ctx, sec_mask, type_dim=None) + torch.manual_seed(1) + actual = build_objective("flow").stage2_loss( + model, x1_s2, cond_cont, cond_cat, stage1_ctx, sec_mask, type_dim=None, ar_inputs=None + ) + assert torch.allclose(actual, expected) + + +def test_flow_objective_stage2_loss_dispatches_to_ar_when_ar_inputs_given(): + torch.manual_seed(0) + B, k_max = 4, 5 + model = Stage2Autoregressive( + pdg_vocab=3, + mat_vocab=2, + particle_cfg=_PHYS_CFG, + material_cfg=_PHYS_CFG, + hidden_dim=16, + n_res_blocks=1, + generator="flow", + k_max=k_max, + ) + cond_cont, cond_cat = _cond(B) + stage1_ctx = torch.randn(B, X_DIM) + token_dim = CONT_SLOT_DIM + PARTICLE_PHYS_DIM + x1_s2 = torch.randn(B, k_max, token_dim) + sec_mask = torch.ones(B, k_max, dtype=torch.bool) + ar_inputs = { + "history_feat": torch.randn(B, k_max, token_dim), + "has_prev": torch.ones(B, k_max, dtype=torch.bool), + "remaining_frac": torch.rand(B, k_max), + "slot_idx": torch.linspace(0, 1, k_max).unsqueeze(0).expand(B, -1), + } + + loss = build_objective("flow").stage2_loss( + model, x1_s2, cond_cont, cond_cat, stage1_ctx, sec_mask, type_dim=None, ar_inputs=ar_inputs + ) + assert loss.dim() == 0 + assert torch.isfinite(loss) + + +def test_ddpm_objective_stage2_loss_not_implemented(): + dummy_model = torch.nn.Module() + dummy = torch.zeros(1) + with pytest.raises(NotImplementedError): + build_objective("ddpm").stage2_loss( + dummy_model, dummy, dummy, dummy, dummy, torch.ones(1, 1, dtype=torch.bool), type_dim=None + ) + + +def test_wgan_objective_has_no_loss_methods(): + dummy_model = torch.nn.Module() + dummy = torch.zeros(1) + objective = build_objective("wgan") + with pytest.raises(NotImplementedError): + objective.stage1_loss(dummy_model, dummy, dummy, dummy) + with pytest.raises(NotImplementedError): + objective.stage2_loss( + dummy_model, dummy, dummy, dummy, dummy, torch.ones(1, 1, dtype=torch.bool), type_dim=None + )