From c9d255b1c59dac136a9e6a893409d3bdb817d2a8 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Fri, 7 Aug 2026 09:36:49 +0200 Subject: [PATCH] =?UTF-8?q?v0.3.0=20step=205:=20Stage2Autoregressive=20(hi?= =?UTF-8?q?story=3Dmarkov)=20+=20=C2=A711.4=20grad=20instrumentation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces the Stage2Autoregressive stub with a real per-token secondary decoder: MarkovHistory summarizes the previous secondary, remaining-energy fraction and slot index round out the per-token conditioning, and the existing Trunk/MonolithicTrunk/RoutedTrunk machinery is reused unchanged by batching all K_MAX tokens together under teacher forcing (one parallel pass, no new trunk code). build_models/build_critics wire it in; the WGAN critic stays whole-sequence, so build_critics needs no AR-specific path. train.py's FlowDDPMStageTrainer/WGANStageTrainer gain a decoder branch, sharing optimizer/EMA/checkpoint machinery with the one-shot path. _assemble_stage2_real is now defined in terms of the new unflattened _assemble_stage2_ar_target helper, removing a near-duplicate branch. Also lands the §11.4 differentiability validation-obligation instrumentation (trunk-gradient norm from the particle-type slice vs. the continuous slices, for generator=wgan + particle_type.target=onehot) via backward hooks in _relax_onehot_type_slice, decoder-agnostic and surfaced as two new metrics.csv columns. This also fixes the standing regression where any config not explicitly overriding decoder="one_shot" crashed at build_models, since stage2_model.decoder defaults to "autoregressive" — confirmed by removing tests/test_pipeline.py's now-stale override so the default config runs end-to-end against real synthetic data. Co-Authored-By: Claude Sonnet 5 --- giant/config.py | 16 ++ giant/model/network.py | 325 +++++++++++++++++++++++++++++++---- giant/model/schedule.py | 67 ++++++++ giant/train.py | 366 ++++++++++++++++++++++++++++++++++------ tests/test_config.py | 49 ++++++ tests/test_network.py | 228 +++++++++++++++++++++++++ tests/test_phase2.py | 118 ++++++++++++- tests/test_pipeline.py | 8 +- tests/test_train.py | 252 ++++++++++++++++++++++++++- 9 files changed, 1328 insertions(+), 101 deletions(-) diff --git a/giant/config.py b/giant/config.py index 1cb44f3..93972e6 100644 --- a/giant/config.py +++ b/giant/config.py @@ -733,6 +733,22 @@ def validate_config(cfg: dict) -> None: "(standalone stage-2 evaluation only, never for rollout)" ) + if _get_path(cfg, "stage2_model.decoder") == "autoregressive": + history = _get_path(cfg, "stage2_model.autoregressive.history") + if history != "markov": + raise ValueError( + f"stage2_model.autoregressive.history = {history!r} is " + "accepted by the schema but not implemented until v0.3.0 " + "step 7 — use 'markov'" + ) + teacher_forcing = _get_path(cfg, "stage2_model.autoregressive.teacher_forcing") + if teacher_forcing != "always": + raise ValueError( + "stage2_model.autoregressive.teacher_forcing = " + f"{teacher_forcing!r} is accepted by the schema but not " + "implemented until v0.3.0 step 7 — use 'always'" + ) + _CONDITIONING_CODE = {"physical": "phys", "embedding": "emb", "onehot": "oh"} diff --git a/giant/model/network.py b/giant/model/network.py index 3a52b32..581e1dd 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -782,6 +782,43 @@ def build_trunk( return MonolithicTrunk(in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout) +# --------------------------------------------------------------------------- +# History encoders — stage-2 autoregressive only (docs/v0.3.0-design.md §6.2) +# --------------------------------------------------------------------------- + + +class HistoryEncoder(nn.Module): + """Interface for stage-2 autoregressive per-token history summaries: + `forward(feat, has_prev) -> (B, K, out_dim)`. `MarkovHistory` is the only + implementation until v0.3.0 step 7 (`AttentionHistory`, + `history = "attention"`).""" + + def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor: + raise NotImplementedError + + +class MarkovHistory(HistoryEncoder): + """Summarizes the previous secondary's own `(energy_fraction, direction, + type_representation)` through one small MLP — the "markov" history + (docs/v0.3.0-design.md §6.2): token i+1 only ever sees token i plus the + running scalars (`remaining_frac`/`slot_idx`, fused in separately by + `Stage2Autoregressive._token_cond`), not the full prefix. + + At slot 0 (`has_prev` False) substitutes a learned start vector rather + than zeros — a reasonable default, not specified by the design doc. + """ + + def __init__(self, in_dim: int, out_dim: int) -> None: + super().__init__() + self.start = nn.Parameter(torch.zeros(in_dim)) + self.mlp = nn.Sequential(nn.Linear(in_dim, out_dim), nn.SiLU()) + + def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor: + start = self.start.view(1, 1, -1).expand_as(feat) + x = torch.where(has_prev.unsqueeze(-1), feat, start) + return self.mlp(x) + + # --------------------------------------------------------------------------- # Stage models (docs/v0.3.0-design.md §5.3) # --------------------------------------------------------------------------- @@ -1041,16 +1078,214 @@ class Stage2OneShot(nn.Module): class Stage2Autoregressive(nn.Module): - """Not implemented until v0.3.0 steps 4-7 (docs/v0.3.0-design.md §6, §12) - — stub so `stage2_model.decoder = "autoregressive"` (the v0.3.0 default) - fails loudly instead of silently no-op-ing.""" + """Emits secondaries one at a time in descending-energy order + (docs/v0.3.0-design.md §6), instead of `Stage2OneShot`'s simultaneous + k_max-slot prediction. Only `history = "markov"` is implemented (v0.3.0 + step 5) — `history = "attention"` raises immediately at construction. + `teacher_forcing` handling lives entirely in the trainer + (`giant/train.py`), since it only affects how training inputs are + assembled, not this module's architecture. - def __init__(self, *args, **kwargs) -> None: + Under teacher forcing every token's conditioning is built from ground + truth, so a whole K-token sequence trains in one parallel batched pass: + `forward` accepts `(B, K, ...)` tensors for an arbitrary K (not hardcoded + to `k_max`) — this also means a future one-token-at-a-time inference loop + (`K=1` per call, step 6) needs no interface change here. + + Two independent conditioning paths, mirroring `Stage2OneShot`'s + `_cond_embed` but split in two: `_base_cond` (`cond_enc` + + `context_adapter` only) feeds `predict_n_sec`, since n_sec doesn't depend + on token position; `_token_cond` additionally fuses in the history + encoding and two running scalars (remaining energy-budget fraction, + normalized slot index), and feeds `forward`/`predict_type`/the trunk. + """ + + def __init__( + self, + pdg_vocab: int, + mat_vocab: int, + particle_cfg: dict, + material_cfg: dict, + hidden_dim: int = 256, + n_res_blocks: int = 6, + cond_out_dim: int = 128, + context_dim: int = 64, + x_dim: int = X_DIM, + dropout: float = 0.0, + generator: str = "wgan", + time_dim: int = 64, + noise_dim: int = 64, + k_max: int = K_MAX, + router: Router | None = None, + build_n_sec_head: bool = True, + particle_type_cfg: dict | None = None, + history: str = "markov", + ) -> None: super().__init__() - raise NotImplementedError( - "stage2_model.decoder = 'autoregressive' is not implemented yet " - "(design doc v0.3.0 steps 4-7) — use decoder = 'one_shot' for now" + if history != "markov": + raise NotImplementedError( + f"stage2_model.autoregressive.history={history!r} is not " + "implemented until v0.3.0 step 7 — use 'markov'" + ) + self.generator_kind = generator + self.noise_dim = noise_dim + self.k_max = k_max + self.particle_type_cfg = dict(particle_type_cfg or {"target": "physical"}) + emb_dim = particle_cfg["emb_dim"] + self.type_dim = stage2_type_dim(self.particle_type_cfg, emb_dim) + + self.cond_enc = ConditionEncoder( + pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim ) + self.context_adapter = ContextAdapter(x_dim, context_dim) + self.base_fuse = nn.Sequential( + nn.Linear(cond_out_dim + context_dim, cond_out_dim), + nn.SiLU(), + ) + + # Reuses conditioning.out_dim for the history encoder's own output + # width — there's no dedicated stage2_model.autoregressive key for + # this, a reasonable default rather than a design-doc-specified value. + history_dim = cond_out_dim + self.history_encoder = MarkovHistory(CONT_SLOT_DIM + self.type_dim, history_dim) + token_fuse_in = ( + cond_out_dim + context_dim + history_dim + 2 + ) # +2: remaining_frac, slot_idx + self.token_fuse = nn.Sequential( + nn.Linear(token_fuse_in, cond_out_dim), + nn.SiLU(), + ) + + has_time = generator in ("flow", "ddpm") + 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 + self.trunk = build_trunk( + router, + in_dim, + token_dim, + hidden_dim, + n_res_blocks, + merged_cond_dim, + dropout, + ) + + 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), + ) + self.type_head = None + target = self.particle_type_cfg.get("target", "physical") + if target != "physical" and generator in ("flow", "ddpm"): + self.type_head = nn.Sequential( + nn.Linear(cond_out_dim, hidden_dim // 2), + nn.SiLU(), + nn.Linear(hidden_dim // 2, self.type_dim), + ) + + 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) + ctx = self.context_adapter(stage1_out) + return self.base_fuse(torch.cat([base, ctx], dim=-1)) + + def _token_cond( + self, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + stage1_out: torch.Tensor, + history_feat: torch.Tensor, + has_prev: torch.Tensor, + remaining_frac: torch.Tensor, + slot_idx: torch.Tensor, + ) -> torch.Tensor: + K = history_feat.size(1) + base = self.cond_enc(cond_cont, cond_cat).unsqueeze(1).expand(-1, K, -1) + ctx = self.context_adapter(stage1_out).unsqueeze(1).expand(-1, K, -1) + hist = self.history_encoder(history_feat, has_prev) + scalars = torch.stack([remaining_frac, slot_idx], dim=-1) + return self.token_fuse(torch.cat([base, ctx, hist, scalars], dim=-1)) + + def forward( + self, + x_t: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + stage1_out: torch.Tensor, + history_feat: torch.Tensor, + has_prev: torch.Tensor, + remaining_frac: torch.Tensor, + slot_idx: torch.Tensor, + t: torch.Tensor | None = None, + ) -> torch.Tensor: + B, K = x_t.shape[0], x_t.shape[1] + c_emb = self._token_cond( + cond_cont, + cond_cat, + stage1_out, + history_feat, + has_prev, + remaining_frac, + slot_idx, + ) + if self.time_emb is not None: + assert t is not None + t_emb = self.time_emb(t.reshape(-1)).view(B, K, -1) + cond = torch.cat([t_emb, c_emb], dim=-1) + else: + cond = c_emb + x_flat = x_t.reshape(B * K, -1) + cond_flat = cond.reshape(B * K, -1) + cond_cont_flat = cond_cont.unsqueeze(1).expand(-1, K, -1).reshape(B * K, -1) + cond_cat_flat = cond_cat.unsqueeze(1).expand(-1, K, -1).reshape(B * K, -1) + out = self.trunk(x_flat, cond_flat, cond_cont_flat, cond_cat_flat) + return out.view(B, K, -1) + + def predict_n_sec( + self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor + ) -> torch.Tensor: + if self.n_sec_head is None: + raise RuntimeError( + "this Stage2Autoregressive has no n_sec_head — it belongs to " + "a migrated v0.2 checkpoint (legacy_owner='stage1'); call " + "stage1.predict_n_sec(cond_cont, cond_cat) instead" + ) + return self.n_sec_head(self._base_cond(cond_cont, cond_cat, stage1_out)) + + def predict_type( + self, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + stage1_out: torch.Tensor, + history_feat: torch.Tensor, + has_prev: torch.Tensor, + remaining_frac: torch.Tensor, + slot_idx: torch.Tensor, + ) -> torch.Tensor: + if self.type_head is None: + raise RuntimeError( + "this Stage2Autoregressive has no type_head — either " + "particle_type.target='physical' (the type slice is part of " + "forward()'s own output) or generator='wgan' (the WGAN " + "trainer reads the type slice out of forward()'s output " + "directly instead)" + ) + c_emb = self._token_cond( + cond_cont, + cond_cat, + stage1_out, + history_feat, + has_prev, + remaining_frac, + slot_idx, + ) + B, K, _ = c_emb.shape + return self.type_head(c_emb.reshape(B * K, -1)).view(B, K, self.type_dim) class CriticModel(nn.Module): @@ -1299,9 +1534,6 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]: if s2cfg.get("active", True): decoder = s2cfg.get("decoder", "one_shot") - if decoder == "autoregressive": - result["stage2"] = Stage2Autoregressive() - return result router_cfg = s2cfg.get("router") or {} stage2_router: Router | None = None if router_cfg.get("enabled"): @@ -1316,28 +1548,51 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]: legacy_owner = (s2cfg.get("n_sec") or {}).get("legacy_owner") k_max = s2cfg.get("k_max", K_MAX) particle_type_cfg = s2cfg.get("particle_type") or {"target": "physical"} - sec_dim = stage2_trunk_sec_dim( - particle_type_cfg, generator, k_max, particle_cfg["emb_dim"] - ) - result["stage2"] = Stage2OneShot( - pdg_vocab=pdg_vocab, - mat_vocab=mat_vocab, - particle_cfg=particle_cfg, - material_cfg=material_cfg, - hidden_dim=s2cfg.get("hidden_dim", 256), - n_res_blocks=s2cfg.get("n_res_blocks", 6), - cond_out_dim=cond_out_dim, - context_dim=s2cfg.get("context_dim", 64), - sec_dim=sec_dim, - dropout=s2cfg.get("dropout", 0.0), - generator=generator, - time_dim=gen_sub.get("time_dim", 64), - noise_dim=(s2cfg.get("wgan") or {}).get("noise_dim", 64), - k_max=k_max, - router=stage2_router, - build_n_sec_head=legacy_owner != "stage1", - particle_type_cfg=particle_type_cfg, - ) + + if decoder == "autoregressive": + ar_cfg = s2cfg.get("autoregressive") or {} + result["stage2"] = Stage2Autoregressive( + pdg_vocab=pdg_vocab, + mat_vocab=mat_vocab, + particle_cfg=particle_cfg, + material_cfg=material_cfg, + hidden_dim=s2cfg.get("hidden_dim", 256), + n_res_blocks=s2cfg.get("n_res_blocks", 6), + cond_out_dim=cond_out_dim, + context_dim=s2cfg.get("context_dim", 64), + dropout=s2cfg.get("dropout", 0.0), + generator=generator, + time_dim=gen_sub.get("time_dim", 64), + noise_dim=(s2cfg.get("wgan") or {}).get("noise_dim", 64), + k_max=k_max, + router=stage2_router, + build_n_sec_head=legacy_owner != "stage1", + particle_type_cfg=particle_type_cfg, + history=ar_cfg.get("history", "markov"), + ) + else: + sec_dim = stage2_trunk_sec_dim( + particle_type_cfg, generator, k_max, particle_cfg["emb_dim"] + ) + result["stage2"] = Stage2OneShot( + pdg_vocab=pdg_vocab, + mat_vocab=mat_vocab, + particle_cfg=particle_cfg, + material_cfg=material_cfg, + hidden_dim=s2cfg.get("hidden_dim", 256), + n_res_blocks=s2cfg.get("n_res_blocks", 6), + cond_out_dim=cond_out_dim, + context_dim=s2cfg.get("context_dim", 64), + sec_dim=sec_dim, + dropout=s2cfg.get("dropout", 0.0), + generator=generator, + time_dim=gen_sub.get("time_dim", 64), + noise_dim=(s2cfg.get("wgan") or {}).get("noise_dim", 64), + k_max=k_max, + router=stage2_router, + build_n_sec_head=legacy_owner != "stage1", + particle_type_cfg=particle_type_cfg, + ) return result @@ -1377,11 +1632,7 @@ def build_critics(model_config: dict) -> dict[str, nn.Module | None]: stage="stage1", ) - if ( - s2cfg.get("active", True) - and s2cfg.get("generator") == "wgan" - and s2cfg.get("decoder", "one_shot") != "autoregressive" - ): + if s2cfg.get("active", True) and s2cfg.get("generator") == "wgan": k_max = s2cfg.get("k_max", K_MAX) particle_type_cfg = s2cfg.get("particle_type") or {"target": "physical"} in_dim = stage2_trunk_sec_dim( diff --git a/giant/model/schedule.py b/giant/model/schedule.py index 0a5bf3d..5c1e019 100644 --- a/giant/model/schedule.py +++ b/giant/model/schedule.py @@ -132,3 +132,70 @@ def flow_matching_loss_secondary( phys_err = err[..., CONT_SLOT_DIM : CONT_SLOT_DIM + type_dim].mean(dim=-1) phys_loss = (phys_err * mask).sum() / denom return cont_loss + phys_loss + + +def flow_matching_loss_secondary_ar( + model: torch.nn.Module, + x1: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + stage1_out: torch.Tensor, + history_feat: torch.Tensor, + has_prev: torch.Tensor, + remaining_frac: torch.Tensor, + slot_idx: torch.Tensor, + sec_mask: torch.Tensor, + type_dim: int | None = None, +) -> torch.Tensor: + """`Stage2Autoregressive` analogue of `flow_matching_loss_secondary`, same + masked, per-block (continuous vs. type) loss recipe — but native to + `Stage2Autoregressive`'s `(B, K_MAX, token_dim)` I/O and its extra + per-token conditioning args, rather than a flattened `(B, K_MAX*token_dim)` + vector. Kept as a sibling rather than unified with the flat version: the + model call signature differs enough (four extra per-token conditioning + tensors) that merging would need an awkward shape-flag + closure. + + Under teacher forcing (docs/v0.3.0-design.md §6.2 point 3) this is still a + single parallel pass over all K_MAX tokens — `x1`/`history_feat`/etc. are + already built from ground truth for every slot by the caller + (`giant.train._assemble_stage2_ar_inputs`/`_assemble_stage2_ar_target`). + + x1: (B, K_MAX, CONT_SLOT_DIM + type_dim) — per-token flattened target + (stick_logit, dir, then a `type_dim`-wide type slice) + sec_mask: (B, K_MAX) bool — True for valid secondary slots + type_dim: as `flow_matching_loss_secondary` — defaults to + `PARTICLE_PHYS_DIM`, `0` means no type slice is in `x1` at all. + """ + from giant.constants import CONT_SLOT_DIM, PARTICLE_PHYS_DIM + + if type_dim is None: + type_dim = PARTICLE_PHYS_DIM + + B, K, _ = x1.shape + t = torch.rand(B, K, device=x1.device) + x0 = torch.randn_like(x1) + x_t = (1.0 - t.unsqueeze(-1)) * x0 + t.unsqueeze(-1) * x1 + u_t = x1 - x0 + v_t = model( + x_t, + cond_cont, + cond_cat, + stage1_out, + history_feat, + has_prev, + remaining_frac, + slot_idx, + t=t, + ) + + err = (v_t - u_t) ** 2 + cont_err = err[..., :CONT_SLOT_DIM].mean(dim=-1) # (B, K_MAX) + + mask = sec_mask.float() + denom = mask.sum().clamp(min=1) + cont_loss = (cont_err * mask).sum() / denom + if type_dim == 0: + return cont_loss + phys_err = err[..., CONT_SLOT_DIM : CONT_SLOT_DIM + type_dim].mean(dim=-1) + phys_loss = (phys_err * mask).sum() / denom + return cont_loss + phys_loss diff --git a/giant/train.py b/giant/train.py index 53403c9..a117d2f 100644 --- a/giant/train.py +++ b/giant/train.py @@ -16,7 +16,7 @@ import torch.optim as optim from torch.utils.data import DataLoader from tqdm import tqdm -from giant.constants import CONT_SLOT_DIM, K_MAX +from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM from giant.data.loader import TopNMap from giant.data.setup_cache import topnmap_to_json from giant.model.network import Router, stage2_type_dim @@ -24,6 +24,7 @@ from giant.model.schedule import ( CosineSchedule, flow_matching_loss, flow_matching_loss_secondary, + flow_matching_loss_secondary_ar, ) from giant.model.wgan import gradient_penalty, generator_loss from giant.validate import validate_marginals @@ -102,6 +103,67 @@ def _batch_to_device(batch: tuple, device: torch.device) -> tuple: return tuple(t.to(device) for t in batch) +def _type_repr( + sec_type_idx: torch.Tensor, + sec_cont: torch.Tensor, + particle_type_cfg: dict, + cond_enc: torch.nn.Module, + emb_dim: int, +) -> torch.Tensor: + """(B, K_MAX, type_dim) ground-truth type representation, generator- + independent (unlike `_assemble_stage2_ar_target`'s training *target*, + which varies by generator/objective — see its docstring): `"physical"` -> + `(log_mass, charge)`; `"onehot"` -> one-hot of the true class; + `"embedding"` -> the conditioning's own detached embedding-table row. + + Used both to build `_assemble_stage2_ar_target`'s wgan+onehot/embedding + branch and as the AR history features' previous-secondary identity — the + latter must always reflect the true physical secondary that came before, + regardless of what the *current* token's own training objective is. + """ + target = particle_type_cfg.get("target", "physical") + if target == "physical": + return sec_cont[..., CONT_SLOT_DIM : CONT_SLOT_DIM + PARTICLE_PHYS_DIM] + if target == "onehot": + return F.one_hot(sec_type_idx, num_classes=emb_dim).float() + return cond_enc.pdg_emb(sec_type_idx).detach() + + +def _assemble_stage2_ar_target( + sec_cont: torch.Tensor, + sec_type_idx: torch.Tensor, + particle_type_cfg: dict, + generator: str, + cond_enc: torch.nn.Module, + emb_dim: int, +) -> torch.Tensor: + """(B, K_MAX, token_dim) ground-truth per-token target — the unflattened + analogue of `_assemble_stage2_real` (defined below in terms of this), + matching whatever width `Stage2Autoregressive`'s (or `Stage2OneShot`'s) + own trunk produces for this (target, generator) combination + (`giant.model.network.stage2_trunk_sec_dim`; docs/v0.3.0-design.md + decision 2/3): + + - `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 = particle_type_cfg.get("target", "physical") + if target == "physical": + return sec_cont + cont = sec_cont[..., :CONT_SLOT_DIM] + if generator != "wgan": + 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) + + def _assemble_stage2_real( sec_cont: torch.Tensor, sec_type_idx: torch.Tensor, @@ -110,38 +172,87 @@ def _assemble_stage2_real( cond_enc: torch.nn.Module, emb_dim: int, ) -> torch.Tensor: - """Ground-truth flattened stage-2 vector, matching whatever width - `Stage2OneShot`'s own trunk produces for this (target, generator) - combination (`giant.model.network.stage2_trunk_sec_dim`; - docs/v0.3.0-design.md decision 2/3): + """Ground-truth flattened stage-2 vector for `Stage2OneShot` — the + flattened form of `_assemble_stage2_ar_target`, which + `Stage2Autoregressive`'s per-token target also uses; the two must stay in + lockstep. See `_assemble_stage2_ar_target`'s docstring for the + (target, generator) width rules.""" + return _assemble_stage2_ar_target( + sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim + ).flatten(1) - - `target = "physical"`: unchanged from v0.2 — `sec_cont` (stick_logit, - dir, log_mass, charge) flattened as-is. - - `target` in `("onehot", "embedding")` + `generator in ("flow", "ddpm")`: - just the continuous stick/dir slots — the type slice isn't part of - this vector at all (`Stage2OneShot.type_head` handles it separately). - - `target` in `("onehot", "embedding")` + `generator == "wgan"`: stick/dir - slots concatenated with the per-slot type vector — a one-hot of the - true class (`"onehot"`, relaxed on the *generated* side only, by the - caller) or the conditioning's own detached embedding-table row - (`"embedding"`, already continuous — no relaxation needed either side, - §2.1). - """ - target = particle_type_cfg.get("target", "physical") - if target == "physical": - return sec_cont.flatten(1) - cont = sec_cont[..., :CONT_SLOT_DIM] - if generator != "wgan": - return cont.flatten(1) - if target == "onehot": - type_vec = F.one_hot(sec_type_idx, num_classes=emb_dim).float() - else: - type_vec = cond_enc.pdg_emb(sec_type_idx).detach() - return torch.cat([cont, type_vec], dim=-1).flatten(1) + +def _stick_fraction(sec_cont: torch.Tensor) -> torch.Tensor: + """(B, K_MAX) — sigmoid of each slot's own stick-breaking logit + (`sec_cont[...,0]`); scale-free (see `giant.data.transforms. + encode_secondaries`), so this needs no absolute `e_sec`.""" + return torch.sigmoid(sec_cont[..., 0]) + + +def _remaining_energy_fraction(fraction: torch.Tensor) -> torch.Tensor: + """(B, K_MAX) — fraction of the original e_sec budget unclaimed entering + slot i: `1.0` at `i=0`, `prod_{j=1` + (docs/v0.3.0-design.md §6.3 — "no re-derivation needed": the existing + stick-breaking encoding is already scale-free, so this is derivable from + the batch's ground-truth stick logits alone, no `e_sec` required).""" + cumprod = torch.cumprod(1.0 - fraction, dim=1) + return torch.cat([torch.ones_like(cumprod[:, :1]), cumprod[:, :-1]], dim=1) + + +def _shift_prev(x: torch.Tensor) -> torch.Tensor: + """`(B, K, ...)` -> same shape, slot i holds slot i-1's value; slot 0 gets + an arbitrary zero placeholder (never read as-is — see `_ar_has_prev`; + `MarkovHistory` substitutes its own learned start vector there instead).""" + return torch.cat([torch.zeros_like(x[:, :1]), x[:, :-1]], dim=1) + + +def _ar_has_prev(k_max: int, device: torch.device) -> torch.Tensor: + """`(1, K_MAX)` bool: True for slot index `>= 1`. Correct without + `n_sec`: `sec_mask` is a prefix mask, so any *valid* token at `k>=1` + always has a valid predecessor at `k-1`; the only wrong cases are tokens + that are themselves padding, already masked out of every loss.""" + return (torch.arange(k_max, device=device) >= 1).unsqueeze(0) + + +def _assemble_stage2_ar_inputs( + sec_cont: torch.Tensor, + sec_type_idx: torch.Tensor, + particle_type_cfg: dict, + cond_enc: torch.nn.Module, + emb_dim: int, +) -> dict[str, torch.Tensor]: + """Ground-truth per-token AR conditioning tensors — all `(B, K_MAX, ...)` + or `(B, K_MAX)`, built in one vectorized pass (teacher forcing means + every token's input is ground truth, docs/v0.3.0-design.md §6.2 point 3). + Keys match `Stage2Autoregressive.forward`'s trailing kwargs.""" + device = sec_cont.device + B, K = sec_cont.shape[0], sec_cont.shape[1] + fraction = _stick_fraction(sec_cont) + type_repr = _type_repr(sec_type_idx, sec_cont, particle_type_cfg, cond_enc, emb_dim) + history_feat = torch.cat( + [ + _shift_prev(fraction).unsqueeze(-1), + _shift_prev(sec_cont[..., 1:CONT_SLOT_DIM]), + _shift_prev(type_repr), + ], + dim=-1, + ) + slot_idx = (torch.arange(K, device=device).float() / max(K - 1, 1)).unsqueeze(0) + return { + "history_feat": history_feat, + "has_prev": _ar_has_prev(K, device).expand(B, -1), + "remaining_frac": _remaining_energy_fraction(fraction), + "slot_idx": slot_idx.expand(B, -1), + } def _relax_onehot_type_slice( - x_flat: torch.Tensor, k_max: int, cont_dim: int, type_dim: int, tau: float + x_flat: torch.Tensor, + k_max: int, + cont_dim: int, + type_dim: int, + tau: float, + grad_probe: dict[str, float] | None = None, ) -> torch.Tensor: """Straight-through Gumbel-softmax relaxation of the per-slot type slice inside a flattened `(B, k_max * (cont_dim + type_dim))` WGAN generator @@ -149,10 +260,26 @@ def _relax_onehot_type_slice( hard one-hot (matching what the critic sees from real data), the backward pass flows smooth gradient. Continuous slots (stick/dir, and the type slice itself under `target = "embedding"`, which never calls - this) pass through unchanged.""" + this) pass through unchanged. + + `grad_probe`, if given, gets `["cont"]`/`["type"]` populated with the L2 + norm of the gradient reaching this split point during the next + `.backward()` call that touches it — a backward hook, not a second + backward pass. This is the §11.4 differentiability validation-obligation + instrumentation (docs/v0.3.0-design.md): the trunk-gradient contribution + from the type slice vs. the continuous slices, for + `particle_type.target="onehot"` + `generator="wgan"`. Only ever populated + on a `did_g_step` batch — the critic step backprops through + `fake.detach()`, which never reaches these hooks — so it stays empty + (callers default to `0.0`) otherwise.""" B = x_flat.size(0) x = x_flat.view(B, k_max, cont_dim + type_dim) cont, type_logits = x[..., :cont_dim], x[..., cont_dim:] + if grad_probe is not None: + cont.register_hook(lambda g: grad_probe.__setitem__("cont", g.norm().item())) + type_logits.register_hook( + lambda g: grad_probe.__setitem__("type", g.norm().item()) + ) type_soft = F.gumbel_softmax(type_logits, tau=tau, hard=True, dim=-1) return torch.cat([cont, type_soft], dim=-1).reshape(B, -1) @@ -234,6 +361,7 @@ class FlowDDPMStageTrainer(StageTrainer): device: torch.device, particle_type_cfg: dict | None = None, particle_type_emb_dim: int = 16, + decoder: str = "one_shot", ) -> None: if is_stage2 and generator not in ("flow",): raise NotImplementedError( @@ -244,6 +372,7 @@ class FlowDDPMStageTrainer(StageTrainer): self.name = name self.is_stage2 = is_stage2 self.generator = generator + self.decoder = decoder self.device = device self.model = model.to(device) self.lambda_weight = lambda_weight @@ -302,12 +431,29 @@ class FlowDDPMStageTrainer(StageTrainer): return self.model.predict_n_sec(cond_cont, cond_cat, stage1_ctx) return self.model.predict_n_sec(cond_cont, cond_cat) - def _generator_loss(self, cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx): + 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( self.model, x1_s2, @@ -319,10 +465,17 @@ class FlowDDPMStageTrainer(StageTrainer): ) def _type_loss( - self, cond_cont, cond_cat, stage1_ctx, sec_type_idx, sec_mask, device + self, + cond_cont, + cond_cat, + stage1_ctx, + sec_type_idx, + sec_mask, + device, + ar_inputs=None, ): - """CE (`target="onehot"`) or MSE (`target="embedding"`) loss for - `Stage2OneShot.type_head` — the non-adversarial counterpart to + """CE (`target="onehot"`) or MSE (`target="embedding"`) loss for the + stage-2 model's `type_head` — the non-adversarial counterpart to WGANStageTrainer's ST-Gumbel-into-the-critic path (decision 2/5). Zero when this stage has no `type_head` (stage 1, or `particle_type.target = "physical"`).""" @@ -331,7 +484,19 @@ class FlowDDPMStageTrainer(StageTrainer): type_head = getattr(self.model, "type_head", None) if not self.is_stage2 or type_head is None: return l_type, type_acc - type_out = self.model.predict_type(cond_cont, cond_cat, stage1_ctx) + if self.decoder == "autoregressive": + assert ar_inputs is not None + type_out = self.model.predict_type( + cond_cont, + cond_cat, + stage1_ctx, + ar_inputs["history_feat"], + ar_inputs["has_prev"], + ar_inputs["remaining_frac"], + ar_inputs["slot_idx"], + ) + else: + type_out = self.model.predict_type(cond_cont, cond_cat, stage1_ctx) mask = sec_mask.float() denom = mask.sum().clamp(min=1) if self.particle_type_cfg.get("target") == "onehot": @@ -360,8 +525,27 @@ class FlowDDPMStageTrainer(StageTrainer): ) = _batch_to_device(batch, device) sec_mask = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec.unsqueeze(1) stage1_ctx = x1_s1.detach() - x1_s2 = ( - _assemble_stage2_real( + + x1_s2 = None + ar_inputs = None + if self.is_stage2 and self.decoder == "autoregressive": + ar_inputs = _assemble_stage2_ar_inputs( + sec_cont, + sec_type_idx, + self.particle_type_cfg, + self.model.cond_enc, + self.particle_type_emb_dim, + ) + x1_s2 = _assemble_stage2_ar_target( + sec_cont, + sec_type_idx, + self.particle_type_cfg, + self.generator, + self.model.cond_enc, + self.particle_type_emb_dim, + ) + elif self.is_stage2: + x1_s2 = _assemble_stage2_real( sec_cont, sec_type_idx, self.particle_type_cfg, @@ -369,12 +553,9 @@ class FlowDDPMStageTrainer(StageTrainer): self.model.cond_enc, self.particle_type_emb_dim, ) - if self.is_stage2 - else None - ) l_gen = self._generator_loss( - cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx + cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx, ar_inputs=ar_inputs ) n_sec_logits = self._predict_n_sec(cond_cont, cond_cat, stage1_ctx) l_nsec = torch.zeros((), device=device) @@ -384,7 +565,13 @@ class FlowDDPMStageTrainer(StageTrainer): nsec_acc = (n_sec_logits.argmax(dim=-1) == n_sec).float().mean() l_type, type_acc = self._type_loss( - cond_cont, cond_cat, stage1_ctx, sec_type_idx, sec_mask, device + cond_cont, + cond_cat, + stage1_ctx, + sec_type_idx, + sec_mask, + device, + ar_inputs=ar_inputs, ) l_balance = l_proc = l_entropy = torch.zeros((), device=device) @@ -531,9 +718,11 @@ class WGANStageTrainer(StageTrainer): particle_type_emb_dim: int = 16, type_gumbel_tau_start: float = 1.0, type_gumbel_tau_end: float = 0.1, + decoder: str = "one_shot", ) -> None: self.name = name self.is_stage2 = is_stage2 + self.decoder = decoder self.device = device self.model = model.to(device) self.critic = critic.to(device) @@ -592,6 +781,7 @@ class WGANStageTrainer(StageTrainer): ) = _batch_to_device(batch, device) B = cond_cont.size(0) stage1_ctx = x1_s1.detach() + grad_probe: dict[str, float] = {} if not self.is_stage2: real = x1_s1 @@ -615,28 +805,62 @@ class WGANStageTrainer(StageTrainer): mask = ( sec_mask.unsqueeze(-1).expand(-1, -1, slot_width).reshape(B, -1).float() ) - real = ( - _assemble_stage2_real( - sec_cont, - sec_type_idx, - self.particle_type_cfg, - "wgan", - self.model.cond_enc, - self.particle_type_emb_dim, - ) - * mask - ) def critic_fn(x): return self.critic(x, cond_cont, cond_cat, stage1_ctx) - z = torch.randn(B, self.model.noise_dim, device=device) - fake_raw = self.model(z, cond_cont, cond_cat, stage1_ctx) + if self.decoder == "autoregressive": + ar = _assemble_stage2_ar_inputs( + sec_cont, + sec_type_idx, + self.particle_type_cfg, + self.model.cond_enc, + self.particle_type_emb_dim, + ) + real = ( + _assemble_stage2_ar_target( + sec_cont, + sec_type_idx, + self.particle_type_cfg, + "wgan", + self.model.cond_enc, + self.particle_type_emb_dim, + ).reshape(B, -1) + * mask + ) + z = torch.randn(B, K_MAX, self.model.noise_dim, device=device) + fake_raw = self.model( + z, + cond_cont, + cond_cat, + stage1_ctx, + ar["history_feat"], + ar["has_prev"], + ar["remaining_frac"], + ar["slot_idx"], + ).reshape(B, -1) + else: + real = ( + _assemble_stage2_real( + sec_cont, + sec_type_idx, + self.particle_type_cfg, + "wgan", + self.model.cond_enc, + self.particle_type_emb_dim, + ) + * mask + ) + z = torch.randn(B, self.model.noise_dim, device=device) + fake_raw = self.model(z, cond_cont, cond_cat, stage1_ctx) + if target == "onehot": # Straight-through Gumbel-softmax relaxation of the type # slice only (decision 5) — the critic must see a hard # one-hot forward (matching what "real" data looks like) # while gradient still flows smoothly to the generator. + # grad_probe captures the §11.4 gradient-magnitude + # instrumentation — see _relax_onehot_type_slice's docstring. tau = _gumbel_tau( global_step, self.total_steps, @@ -644,7 +868,7 @@ class WGANStageTrainer(StageTrainer): self.type_gumbel_tau_end, ) fake_raw = _relax_onehot_type_slice( - fake_raw, K_MAX, CONT_SLOT_DIM, type_dim, tau + fake_raw, K_MAX, CONT_SLOT_DIM, type_dim, tau, grad_probe=grad_probe ) fake = fake_raw * mask @@ -711,6 +935,8 @@ class WGANStageTrainer(StageTrainer): "grad_norm": grad_norm_d.item() + grad_norm_g.item(), "grad_norm_d": grad_norm_d.item(), "grad_norm_g": grad_norm_g.item(), + "grad_norm_type_slice": grad_probe.get("type", 0.0), + "grad_norm_cont_slice": grad_probe.get("cont", 0.0), "lr": self.optimizer.param_groups[0]["lr"], "critic_lr": self.optimizer_d.param_groups[0]["lr"], "batch_size": B, @@ -820,6 +1046,17 @@ def _build_stage_trainers( "target": "physical" } particle_type_emb_dim = cfg["conditioning"]["particle"]["emb_dim"] + decoder = stage_cfg.get("decoder", "one_shot") if is_stage2 else "one_shot" + if is_stage2 and decoder == "autoregressive": + teacher_forcing = (stage_cfg.get("autoregressive") or {}).get( + "teacher_forcing", "always" + ) + if teacher_forcing != "always": + raise NotImplementedError( + "stage2_model.autoregressive.teacher_forcing=" + f"{teacher_forcing!r} is not implemented until v0.3.0 " + "step 7 — use 'always'" + ) if generator == "wgan": critic = critics.get(name) @@ -848,6 +1085,7 @@ def _build_stage_trainers( particle_type_emb_dim=particle_type_emb_dim, type_gumbel_tau_start=wgan_cfg.get("gumbel_tau_start", 1.0), type_gumbel_tau_end=wgan_cfg.get("gumbel_tau_end", 0.1), + decoder=decoder, ) else: ddpm_n_steps = stage_cfg.get("ddpm", {}).get("n_steps", 1000) @@ -873,6 +1111,7 @@ def _build_stage_trainers( device=device, particle_type_cfg=particle_type_cfg, particle_type_emb_dim=particle_type_emb_dim, + decoder=decoder, ) return trainers @@ -892,6 +1131,14 @@ def _metrics_fields(trainers: dict[str, StageTrainer]) -> list[str]: f"{name}_train_grad_norm_d", f"{name}_train_grad_norm_g", ] + if ( + trainer.is_stage2 + and trainer.particle_type_cfg.get("target") == "onehot" + ): + fields += [ + f"{name}_train_grad_norm_type_slice", + f"{name}_train_grad_norm_cont_slice", + ] else: fields += [ f"{name}_train_loss", @@ -1358,6 +1605,13 @@ def train( metrics_row[f"{name}_train_grad_norm_g"] = ( sums.get("grad_norm_g", 0.0) / n_train ) + if tr.is_stage2 and tr.particle_type_cfg.get("target") == "onehot": + metrics_row[f"{name}_train_grad_norm_type_slice"] = ( + sums.get("grad_norm_type_slice", 0.0) / n_train + ) + metrics_row[f"{name}_train_grad_norm_cont_slice"] = ( + sums.get("grad_norm_cont_slice", 0.0) / n_train + ) grad_norm_total += sums.get("grad_norm", 0.0) / n_train metrics_row[f"{name}_critic_lr"] = tr.optimizer_d.param_groups[0][ "lr" diff --git a/tests/test_config.py b/tests/test_config.py index 653fca1..4fb9b00 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -637,6 +637,55 @@ def test_validate_config_stop_token_not_implemented(): assert "stop_token" in str(e) +def test_validate_config_ar_default_markov_always_passes(): + """DEFAULT_CONFIG already has decoder='autoregressive', + history='markov', teacher_forcing='always' — must not raise (v0.3.0 + step 5; see also test_validate_config_default_config_passes).""" + cfg = _cfg_with(**{"stage2_model.decoder": "autoregressive"}) + gconfig.validate_config(cfg) # must not raise + + +def test_validate_config_ar_history_attention_not_implemented(): + cfg = _cfg_with( + **{ + "stage2_model.decoder": "autoregressive", + "stage2_model.autoregressive.history": "attention", + } + ) + try: + gconfig.validate_config(cfg) + assert False, "expected ValueError" + except ValueError as e: + assert "history" in str(e) + + +def test_validate_config_ar_teacher_forcing_scheduled_not_implemented(): + cfg = _cfg_with( + **{ + "stage2_model.decoder": "autoregressive", + "stage2_model.autoregressive.teacher_forcing": "scheduled", + } + ) + try: + gconfig.validate_config(cfg) + assert False, "expected ValueError" + except ValueError as e: + assert "teacher_forcing" in str(e) + + +def test_validate_config_ar_checks_skipped_under_one_shot(): + """history/teacher_forcing values that would fail under AR are irrelevant + (and unchecked) when decoder='one_shot'.""" + cfg = _cfg_with( + **{ + "stage2_model.decoder": "one_shot", + "stage2_model.autoregressive.history": "attention", + "stage2_model.autoregressive.teacher_forcing": "scheduled", + } + ) + gconfig.validate_config(cfg) # must not raise + + # --------------------------------------------------------------------------- # checkpoint config-mismatch warnings (unchanged surface, still exercised) # --------------------------------------------------------------------------- diff --git a/tests/test_network.py b/tests/test_network.py index 2117898..dde4bf9 100644 --- a/tests/test_network.py +++ b/tests/test_network.py @@ -1,9 +1,12 @@ +import pytest import torch from giant.constants import CONT_SLOT_DIM, COND_DIM, PARTICLE_PHYS_DIM, SEC_SLOT_DIM from giant.model.network import ( ConditionEncoder, + MarkovHistory, SinusoidalEmbedding, Stage1Model, + Stage2Autoregressive, Stage2OneShot, cat_col_layout, stage2_trunk_sec_dim, @@ -290,3 +293,228 @@ def test_stage2_oneshot_forward_shape_onehot_flow_excludes_type(): t = torch.rand(B) out = model(x_t, cond_cont, cond_cat, stage1_out, t=t) assert out.shape == (B, k_max * CONT_SLOT_DIM) + + +# --- MarkovHistory (docs/v0.3.0-design.md §6.2) ----------------------------- + + +def test_markov_history_shape(): + hist = MarkovHistory(in_dim=7, out_dim=12) + B, K = 3, 5 + feat = torch.randn(B, K, 7) + has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1) + out = hist(feat, has_prev) + assert out.shape == (B, K, 12) + + +def test_markov_history_uses_start_vector_when_no_prev(): + """Slot 0's own raw feature must be ignored — a learned start vector is + substituted there instead (a reasonable default not specified by the + design doc, see Stage2Autoregressive's docstring).""" + hist = MarkovHistory(in_dim=4, out_dim=6) + B, K = 2, 3 + has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1) + feat_a = torch.randn(B, K, 4) + feat_b = feat_a.clone() + feat_b[:, 0] = torch.randn(B, 4) * 100 + out_a = hist(feat_a, has_prev) + out_b = hist(feat_b, has_prev) + assert torch.allclose(out_a[:, 0], out_b[:, 0]) + assert torch.allclose(out_a[:, 1:], out_b[:, 1:]) + + +# --- Stage2Autoregressive (docs/v0.3.0-design.md §6, v0.3.0 step 5) --------- + + +def _build_stage2_ar( + target: str, + generator: str, + emb_dim: int = 6, + k_max: int = 5, + history: str = "markov", +) -> Stage2Autoregressive: + particle_cfg = {"type": "physical", "emb_dim": emb_dim, "n_layers": 1} + if target == "embedding": + particle_cfg = dict(particle_cfg) + particle_cfg["type"] = "embedding" + return Stage2Autoregressive( + pdg_vocab=5, + mat_vocab=3, + particle_cfg=particle_cfg, + material_cfg=MATERIAL_CFG, + hidden_dim=16, + n_res_blocks=1, + cond_out_dim=16, + context_dim=8, + generator=generator, + k_max=k_max, + particle_type_cfg={"target": target, "lambda": 1.0}, + history=history, + ) + + +def _ar_inputs(B: int, K: int, hist_dim: int): + history_feat = torch.randn(B, K, hist_dim) + has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1) + remaining_frac = torch.rand(B, K) + slot_idx = torch.linspace(0, 1, K).unsqueeze(0).expand(B, -1) + return history_feat, has_prev, remaining_frac, slot_idx + + +def test_stage2_autoregressive_history_attention_raises(): + with pytest.raises(NotImplementedError): + _build_stage2_ar("onehot", "wgan", history="attention") + + +@pytest.mark.parametrize("target", ["physical", "onehot", "embedding"]) +@pytest.mark.parametrize("generator", ["wgan", "flow"]) +def test_stage2_autoregressive_forward_shape(target, generator): + B, K, emb_dim = 4, 5, 6 + model = _build_stage2_ar(target, generator, emb_dim=emb_dim, k_max=K) + cond_cont = torch.randn(B, COND_DIM) + cond_cat = torch.zeros(B, 2, dtype=torch.long) + stage1_out = torch.randn(B, 9) + type_dim = stage2_type_dim({"target": target}, emb_dim) + history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs( + B, K, CONT_SLOT_DIM + type_dim + ) + token_dim = stage2_trunk_sec_dim({"target": target}, generator, 1, emb_dim) + if generator == "wgan": + x_t = torch.randn(B, K, model.noise_dim) + t = None + else: + x_t = torch.randn(B, K, token_dim) + t = torch.rand(B, K) + out = model( + x_t, + cond_cont, + cond_cat, + stage1_out, + history_feat, + has_prev, + remaining_frac, + slot_idx, + t=t, + ) + assert out.shape == (B, K, token_dim) + + +def test_stage2_autoregressive_predict_n_sec_shape(): + B, k_max, emb_dim = 4, 5, 6 + model = _build_stage2_ar("physical", "wgan", emb_dim=emb_dim, k_max=k_max) + cond_cont = torch.randn(B, COND_DIM) + cond_cat = torch.zeros(B, 2, dtype=torch.long) + stage1_out = torch.randn(B, 9) + logits = model.predict_n_sec(cond_cont, cond_cat, stage1_out) + assert logits.shape == (B, k_max + 1) + + +def test_stage2_autoregressive_predict_type_shape(): + B, K, emb_dim = 4, 5, 6 + model = _build_stage2_ar("onehot", "flow", emb_dim=emb_dim, k_max=K) + cond_cont = torch.randn(B, COND_DIM) + cond_cat = torch.zeros(B, 2, dtype=torch.long) + stage1_out = torch.randn(B, 9) + type_dim = stage2_type_dim({"target": "onehot"}, emb_dim) + history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs( + B, K, CONT_SLOT_DIM + type_dim + ) + out = model.predict_type( + cond_cont, + cond_cat, + stage1_out, + history_feat, + has_prev, + remaining_frac, + slot_idx, + ) + assert out.shape == (B, K, emb_dim) + + +@pytest.mark.parametrize("target,generator", [("physical", "flow"), ("onehot", "wgan")]) +def test_stage2_autoregressive_predict_type_raises_when_no_type_head(target, generator): + B, K, emb_dim = 2, 5, 6 + model = _build_stage2_ar(target, generator, emb_dim=emb_dim, k_max=K) + cond_cont = torch.randn(B, COND_DIM) + cond_cat = torch.zeros(B, 2, dtype=torch.long) + stage1_out = torch.randn(B, 9) + type_dim = stage2_type_dim({"target": target}, emb_dim) + history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs( + B, K, CONT_SLOT_DIM + type_dim + ) + with pytest.raises(RuntimeError): + model.predict_type( + cond_cont, + cond_cat, + stage1_out, + history_feat, + has_prev, + remaining_frac, + slot_idx, + ) + + +def test_stage2_autoregressive_gradients_flow_wgan_onehot(): + B, K, emb_dim = 4, 5, 6 + model = _build_stage2_ar("onehot", "wgan", emb_dim=emb_dim, k_max=K) + cond_cont = torch.randn(B, COND_DIM) + cond_cat = torch.zeros(B, 2, dtype=torch.long) + stage1_out = torch.randn(B, 9) + type_dim = stage2_type_dim({"target": "onehot"}, emb_dim) + history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs( + B, K, CONT_SLOT_DIM + type_dim + ) + z = torch.randn(B, K, model.noise_dim) + gen_out = model( + z, + cond_cont, + cond_cat, + stage1_out, + history_feat, + has_prev, + remaining_frac, + slot_idx, + ).sum() + nsec_out = model.predict_n_sec(cond_cont, cond_cat, stage1_out).sum() + (gen_out + nsec_out).backward() + for name, p in model.named_parameters(): + assert p.grad is not None, f"no grad for {name}" + + +def test_stage2_autoregressive_gradients_flow_onehot(): + B, K, emb_dim = 4, 5, 6 + model = _build_stage2_ar("onehot", "flow", emb_dim=emb_dim, k_max=K) + cond_cont = torch.randn(B, COND_DIM) + cond_cat = torch.zeros(B, 2, dtype=torch.long) + stage1_out = torch.randn(B, 9) + type_dim = stage2_type_dim({"target": "onehot"}, emb_dim) + history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs( + B, K, CONT_SLOT_DIM + type_dim + ) + token_dim = stage2_trunk_sec_dim({"target": "onehot"}, "flow", 1, emb_dim) + x_t = torch.randn(B, K, token_dim) + t = torch.rand(B, K) + flow_out = model( + x_t, + cond_cont, + cond_cat, + stage1_out, + history_feat, + has_prev, + remaining_frac, + slot_idx, + t=t, + ).sum() + nsec_out = model.predict_n_sec(cond_cont, cond_cat, stage1_out).sum() + type_out = model.predict_type( + cond_cont, + cond_cat, + stage1_out, + history_feat, + has_prev, + remaining_frac, + slot_idx, + ).sum() + (flow_out + nsec_out + type_out).backward() + for name, p in model.named_parameters(): + assert p.grad is not None, f"no grad for {name}" diff --git a/tests/test_phase2.py b/tests/test_phase2.py index deeb9c7..f548f7a 100644 --- a/tests/test_phase2.py +++ b/tests/test_phase2.py @@ -4,9 +4,19 @@ import numpy as np import pytest import torch -from giant.constants import COND_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_DIM, X_DIM -from giant.model.network import Stage1Model, Stage2OneShot -from giant.model.schedule import flow_matching_loss_secondary +from giant.constants import ( + COND_DIM, + CONT_SLOT_DIM, + K_MAX, + PARTICLE_PHYS_DIM, + SEC_DIM, + X_DIM, +) +from giant.model.network import Stage1Model, Stage2Autoregressive, Stage2OneShot +from giant.model.schedule import ( + flow_matching_loss_secondary, + flow_matching_loss_secondary_ar, +) from giant.sample import sample_secondaries _SAMPLE_SECONDARIES_XFAIL_REASON = ( @@ -185,6 +195,108 @@ def test_flow_matching_loss_secondary_has_grad(): assert any(p.grad is not None for p in decoder.parameters()) +# ── masked flow matching loss — autoregressive (v0.3.0 step 5) ───────────── + + +def _sec_decoder_ar(pdg=3, mat=2, k_max=K_MAX): + particle_cfg, material_cfg = _particle_material_cfg("embedding") + return Stage2Autoregressive( + pdg_vocab=pdg, + mat_vocab=mat, + particle_cfg=particle_cfg, + material_cfg=material_cfg, + hidden_dim=32, + n_res_blocks=2, + generator="flow", + time_dim=16, + k_max=k_max, + ) + + +def _ar_history_inputs(B, K, hist_dim): + history_feat = torch.randn(B, K, hist_dim) + has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1) + remaining_frac = torch.rand(B, K) + slot_idx = torch.linspace(0, 1, K).unsqueeze(0).expand(B, -1) + return history_feat, has_prev, remaining_frac, slot_idx + + +def test_flow_matching_loss_secondary_ar_scalar(): + B, K, pdg, mat = 8, K_MAX, 3, 2 + decoder = _sec_decoder_ar(pdg, mat, k_max=K) + x1 = torch.randn(B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM) + cond_cont, cond_cat = _cond(B, pdg, mat) + stage1_out = torch.randn(B, X_DIM) + history_feat, has_prev, remaining_frac, slot_idx = _ar_history_inputs( + B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM + ) + sec_mask = torch.ones(B, K, dtype=torch.bool) + loss = flow_matching_loss_secondary_ar( + decoder, + x1, + cond_cont, + cond_cat, + stage1_out, + history_feat, + has_prev, + remaining_frac, + slot_idx, + sec_mask, + ) + assert loss.shape == () + assert loss.item() >= 0.0 + + +def test_flow_matching_loss_secondary_ar_mask_zeros_padding(): + B, K, pdg, mat = 4, K_MAX, 3, 2 + decoder = _sec_decoder_ar(pdg, mat, k_max=K) + x1 = torch.randn(B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM) + cond_cont, cond_cat = _cond(B, pdg, mat) + stage1_out = torch.randn(B, X_DIM) + history_feat, has_prev, remaining_frac, slot_idx = _ar_history_inputs( + B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM + ) + sec_mask = torch.zeros(B, K, dtype=torch.bool) + loss = flow_matching_loss_secondary_ar( + decoder, + x1, + cond_cont, + cond_cat, + stage1_out, + history_feat, + has_prev, + remaining_frac, + slot_idx, + sec_mask, + ) + assert loss.item() == pytest.approx(0.0, abs=1e-6) + + +def test_flow_matching_loss_secondary_ar_has_grad(): + B, K, pdg, mat = 4, K_MAX, 3, 2 + decoder = _sec_decoder_ar(pdg, mat, k_max=K) + x1 = torch.randn(B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM) + cond_cont, cond_cat = _cond(B, pdg, mat) + stage1_out = torch.randn(B, X_DIM) + history_feat, has_prev, remaining_frac, slot_idx = _ar_history_inputs( + B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM + ) + sec_mask = torch.ones(B, K, dtype=torch.bool) + flow_matching_loss_secondary_ar( + decoder, + x1, + cond_cont, + cond_cat, + stage1_out, + history_feat, + has_prev, + remaining_frac, + slot_idx, + sec_mask, + ).backward() + assert any(p.grad is not None for p in decoder.parameters()) + + # ── sampling ────────────────────────────────────────────────────────────────── diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index f49c38d..f4c093c 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -105,10 +105,10 @@ def _tiny_cfg(**train_overrides): cfg["train"].update(train_overrides) cfg["stage1_model"].update({"hidden_dim": 8, "n_res_blocks": 1, "dropout": 0.0}) cfg["stage2_model"].update( - # decoder="autoregressive" is DEFAULT_CONFIG's default (the finished - # v0.3.0 target) but Stage2Autoregressive isn't implemented until - # design doc step 4/5 — every run must override to "one_shot" for now. - {"decoder": "one_shot", "hidden_dim": 8, "n_res_blocks": 1, "dropout": 0.0} + # decoder="autoregressive" is DEFAULT_CONFIG's default (v0.3.0 step 5) + # and left as-is here on purpose, so this pipeline-level fixture + # exercises the real default end-to-end against actual data. + {"hidden_dim": 8, "n_res_blocks": 1, "dropout": 0.0} ) cfg["conditioning"]["particle"]["emb_dim"] = 4 cfg["conditioning"]["material"]["emb_dim"] = 4 diff --git a/tests/test_train.py b/tests/test_train.py index e20223d..8318add 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -1,18 +1,36 @@ """Tests for giant/train.py.""" import copy +import csv import tempfile from pathlib import Path import pytest import torch -from giant.constants import COND_DIM, K_MAX, SEC_SLOT_DIM, X_DIM +from giant.constants import ( + COND_DIM, + CONT_SLOT_DIM, + K_MAX, + PARTICLE_PHYS_DIM, + SEC_SLOT_DIM, + X_DIM, +) from giant.model.network import build_critics, build_models from giant.train import ( FlowDDPMStageTrainer, WGANStageTrainer, + _ar_has_prev, + _assemble_stage2_ar_inputs, + _assemble_stage2_ar_target, + _assemble_stage2_real, + _build_stage_trainers, _gumbel_tau, + _relax_onehot_type_slice, + _remaining_energy_fraction, + _shift_prev, + _stick_fraction, + _type_repr, _wandb_run_config, train, ) @@ -67,6 +85,122 @@ def test_wandb_run_config_handles_missing_model_config(): assert wcfg["model_config"] == {} +# --- AR helper functions (v0.3.0 step 5, docs/v0.3.0-design.md §6) --------- + + +def test_stick_fraction_matches_sigmoid_of_logit(): + sec_cont = torch.zeros(2, 3, SEC_SLOT_DIM) + sec_cont[..., 0] = torch.tensor([[0.0, 2.0, -2.0], [1.0, -1.0, 0.0]]) + frac = _stick_fraction(sec_cont) + assert torch.allclose(frac, torch.sigmoid(sec_cont[..., 0])) + + +def test_remaining_energy_fraction_hand_computed(): + fraction = torch.tensor([[0.5, 0.5, 1.0]]) + remaining = _remaining_energy_fraction(fraction) + assert torch.allclose(remaining, torch.tensor([[1.0, 0.5, 0.25]])) + + +def test_shift_prev_shifts_and_zero_pads_slot0(): + x = torch.arange(2 * 4 * 3).reshape(2, 4, 3).float() + shifted = _shift_prev(x) + assert torch.all(shifted[:, 0] == 0) + assert torch.equal(shifted[:, 1:], x[:, :-1]) + + +def test_ar_has_prev_false_only_at_slot_zero(): + has_prev = _ar_has_prev(5, torch.device("cpu")) + assert has_prev.shape == (1, 5) + assert has_prev.tolist() == [[False, True, True, True, True]] + + +@pytest.mark.parametrize("target", ["physical", "onehot", "embedding"]) +def test_type_repr_shapes_and_values(target): + B, K, emb_dim = 3, 4, 6 + sec_cont = torch.randn(B, K, SEC_SLOT_DIM) + sec_type_idx = torch.randint(0, emb_dim, (B, K)) + cond_enc = torch.nn.Module() + if target == "embedding": + cond_enc.pdg_emb = torch.nn.Embedding(emb_dim, emb_dim) + repr_ = _type_repr(sec_type_idx, sec_cont, {"target": target}, cond_enc, emb_dim) + expected_width = PARTICLE_PHYS_DIM if target == "physical" else emb_dim + assert repr_.shape == (B, K, expected_width) + if target == "physical": + assert torch.equal( + repr_, sec_cont[..., CONT_SLOT_DIM : CONT_SLOT_DIM + PARTICLE_PHYS_DIM] + ) + if target == "onehot": + assert torch.all(repr_.sum(-1) == 1.0) + + +@pytest.mark.parametrize( + "target,generator", + [ + ("physical", "flow"), + ("physical", "wgan"), + ("onehot", "flow"), + ("onehot", "wgan"), + ("embedding", "flow"), + ("embedding", "wgan"), + ], +) +def test_assemble_stage2_ar_target_matches_assemble_stage2_real_flattened( + target, generator +): + """Regression test tying the refactor together: _assemble_stage2_real is + now defined as _assemble_stage2_ar_target(...).flatten(1).""" + B, emb_dim = 4, 6 + sec_cont = torch.randn(B, K_MAX, SEC_SLOT_DIM) + sec_type_idx = torch.randint(0, emb_dim, (B, K_MAX)) + cond_enc = torch.nn.Module() + if target == "embedding": + cond_enc.pdg_emb = torch.nn.Embedding(emb_dim, emb_dim) + particle_type_cfg = {"target": target} + flat = _assemble_stage2_real( + sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim + ) + unflat = _assemble_stage2_ar_target( + sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim + ) + assert torch.equal(unflat.flatten(1), flat) + + +def test_assemble_stage2_ar_inputs_shapes_and_history_feat_width(): + B, emb_dim = 3, 6 + sec_cont = torch.randn(B, K_MAX, SEC_SLOT_DIM) + sec_type_idx = torch.randint(0, emb_dim, (B, K_MAX)) + cond_enc = torch.nn.Module() + out = _assemble_stage2_ar_inputs( + sec_cont, sec_type_idx, {"target": "physical"}, cond_enc, emb_dim + ) + assert out["history_feat"].shape == (B, K_MAX, CONT_SLOT_DIM + PARTICLE_PHYS_DIM) + assert out["has_prev"].shape == (B, K_MAX) + assert out["remaining_frac"].shape == (B, K_MAX) + assert out["slot_idx"].shape == (B, K_MAX) + assert torch.all(out["slot_idx"][:, 0] == 0.0) + assert torch.all(out["slot_idx"][:, -1] == 1.0) + + +def test_relax_onehot_type_slice_grad_probe_populates_both_norms(): + B, k_max, cont_dim, type_dim = 4, K_MAX, CONT_SLOT_DIM, 6 + x_flat = torch.randn(B, k_max * (cont_dim + type_dim), requires_grad=True) + grad_probe: dict[str, float] = {} + out = _relax_onehot_type_slice( + x_flat, k_max, cont_dim, type_dim, tau=0.5, grad_probe=grad_probe + ) + out.sum().backward() + assert grad_probe["cont"] >= 0.0 + assert grad_probe["type"] >= 0.0 + + +def test_relax_onehot_type_slice_grad_probe_none_is_backward_compatible(): + B, k_max, cont_dim, type_dim = 4, K_MAX, CONT_SLOT_DIM, 6 + x_flat = torch.randn(B, k_max * (cont_dim + type_dim), requires_grad=True) + out = _relax_onehot_type_slice(x_flat, k_max, cont_dim, type_dim, tau=0.5) + out.sum().backward() + assert x_flat.grad is not None + + # --- end-to-end train() integration tests ----------------------------------- PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1} @@ -273,6 +407,59 @@ def _run_train(cfg, out_dir, resume_path=None): ), ), ), + ( + "ar_wgan_onehot", + lambda cfg: ( + cfg["stage2_model"].__setitem__("decoder", "autoregressive"), + cfg["stage2_model"].__setitem__( + "particle_type", {"target": "onehot", "lambda": 1.0} + ), + ), + ), + ( + "ar_wgan_physical", + lambda cfg: cfg["stage2_model"].__setitem__("decoder", "autoregressive"), + ), + ( + "ar_flow_onehot", + lambda cfg: ( + cfg["stage2_model"].__setitem__("decoder", "autoregressive"), + cfg["stage2_model"].__setitem__("generator", "flow"), + cfg["stage2_model"].__setitem__( + "particle_type", {"target": "onehot", "lambda": 1.0} + ), + ), + ), + ( + "ar_flow_embedding", + lambda cfg: ( + cfg["conditioning"]["particle"].__setitem__("type", "embedding"), + cfg["conditioning"]["material"].__setitem__("type", "embedding"), + cfg["stage2_model"].__setitem__("decoder", "autoregressive"), + cfg["stage2_model"].__setitem__("generator", "flow"), + cfg["stage2_model"].__setitem__( + "particle_type", {"target": "embedding", "lambda": 1.0} + ), + ), + ), + ( + "ar_stage2_only", + lambda cfg: ( + cfg["stage1_model"].__setitem__("active", False), + cfg["stage2_model"].__setitem__("decoder", "autoregressive"), + ), + ), + ( + "ar_mixed_stage1_wgan_stage2_flow_onehot", + lambda cfg: ( + cfg["stage1_model"].__setitem__("generator", "wgan"), + cfg["stage2_model"].__setitem__("generator", "flow"), + cfg["stage2_model"].__setitem__("decoder", "autoregressive"), + cfg["stage2_model"].__setitem__( + "particle_type", {"target": "onehot", "lambda": 1.0} + ), + ), + ), ], ) def test_train_end_to_end(label, mutate): @@ -400,3 +587,66 @@ def test_flow_stage_trainer_ddpm_not_implemented_for_stage2(): ddpm_n_steps=50, device=torch.device("cpu"), ) + + +# --- AR trainer wiring (v0.3.0 step 5) -------------------------------------- + + +def test_build_stage_trainers_rejects_scheduled_teacher_forcing(): + cfg = _base_cfg() + cfg["stage2_model"]["decoder"] = "autoregressive" + cfg["stage2_model"]["autoregressive"] = { + "history": "markov", + "teacher_forcing": "scheduled", + } + model_config = _model_config(cfg) + models = build_models(model_config) + critics = build_critics(model_config) + with pytest.raises(NotImplementedError): + _build_stage_trainers( + cfg, models, critics, torch.device("cpu"), total_train_batches=4 + ) + + +def test_ar_wgan_onehot_grad_norm_instrumentation_populates_metrics(): + """§11.4 differentiability validation-obligation instrumentation: the + trunk-gradient-norm-by-slice columns must appear and actually fire for + generator='wgan' + particle_type.target='onehot' under decoder= + 'autoregressive' (added at v0.3.0 step 5 per the design doc's + instruction to accrue evidence during the architecture comparison).""" + cfg = _base_cfg() + cfg["stage2_model"]["decoder"] = "autoregressive" + cfg["stage2_model"]["particle_type"] = {"target": "onehot", "lambda": 1.0} + with tempfile.TemporaryDirectory() as tmp: + out_dir = Path(tmp) / "run" + _run_train(cfg, out_dir) + with open(out_dir / "metrics.csv", newline="") as f: + rows = list(csv.DictReader(f)) + assert "stage2_train_grad_norm_type_slice" in rows[0] + assert "stage2_train_grad_norm_cont_slice" in rows[0] + assert any(float(r["stage2_train_grad_norm_type_slice"]) > 0 for r in rows) + assert any(float(r["stage2_train_grad_norm_cont_slice"]) > 0 for r in rows) + + +def test_wgan_onehot_one_shot_also_gets_grad_norm_instrumentation(): + """The instrumentation is decoder-agnostic — one_shot + wgan + onehot + must populate the same columns.""" + cfg = _base_cfg() + cfg["stage2_model"]["particle_type"] = {"target": "onehot", "lambda": 1.0} + with tempfile.TemporaryDirectory() as tmp: + out_dir = Path(tmp) / "run" + _run_train(cfg, out_dir) + with open(out_dir / "metrics.csv", newline="") as f: + rows = list(csv.DictReader(f)) + assert any(float(r["stage2_train_grad_norm_type_slice"]) > 0 for r in rows) + assert any(float(r["stage2_train_grad_norm_cont_slice"]) > 0 for r in rows) + + +def test_wgan_physical_omits_grad_norm_slice_columns(): + cfg = _base_cfg() # default stage2_model has no particle_type -> "physical" + with tempfile.TemporaryDirectory() as tmp: + out_dir = Path(tmp) / "run" + _run_train(cfg, out_dir) + header = (out_dir / "metrics.csv").read_text().splitlines()[0].split(",") + assert "stage2_train_grad_norm_type_slice" not in header + assert "stage2_train_grad_norm_cont_slice" not in header