Add WGAN-GP mode as a throwaway fast-eval experiment

Adds --mode wgan alongside flow/ddpm: both stages get a WGAN-GP
generator/critic pair (giant.model.wgan) instead of flow matching, so
inference is a single forward pass per stage rather than a 10-step ODE
integration — the fast-eval architecture noted in the roadmap.
predict/rollout auto-detect the mode from the checkpoint's model_config.
Best-checkpoint selection for wgan uses marginal-KL against the EMA
generators every epoch, since a critic loss isn't a monotone quality
signal. --router is not supported together with --mode wgan.

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

giant/cli.py

@@ -45,7 +45,12 @@ from giant.model.network import build_models
 from giant.particles import nearest_known_pdg
 from giant.pipeline import run_train_job
 from giant.rollout import rollout as run_rollout
-from giant.sample import sample_flow, sample_secondaries
+from giant.sample import (
+    sample_flow,
+    sample_secondaries,
+    sample_wgan,
+    sample_secondaries_wgan,
+)
 
 app = typer.Typer(no_args_is_help=True)
 
@@ -177,6 +182,7 @@ def _main() -> None:
 class Mode(str, Enum):
     flow = "flow"
     ddpm = "ddpm"
+    wgan = "wgan"
 
 
 class Conditioning(str, Enum):
@@ -308,6 +314,38 @@ def train(
             "--router-axis 'pdg:n_experts=3,emb_dim=8'",
         ),
     ] = None,
+    n_critic: Annotated[
+        Optional[int],
+        typer.Option(
+            "--n-critic",
+            help="WGAN-GP (--mode wgan only): critic updates per generator "
+            "update (default: 5)",
+        ),
+    ] = None,
+    gp_weight: Annotated[
+        Optional[float],
+        typer.Option(
+            "--gp-weight",
+            help="WGAN-GP (--mode wgan only): gradient-penalty coefficient "
+            "(default: 10.0)",
+        ),
+    ] = None,
+    noise_dim: Annotated[
+        Optional[int],
+        typer.Option(
+            "--noise-dim",
+            help="WGAN (--mode wgan only): generator input noise-vector "
+            "width (default: 64)",
+        ),
+    ] = None,
+    critic_lr: Annotated[
+        Optional[float],
+        typer.Option(
+            "--critic-lr",
+            help="WGAN-GP (--mode wgan only): critic learning rate "
+            "(default: same as --lr)",
+        ),
+    ] = None,
     val_fraction: Annotated[
         Optional[float], typer.Option("--val-fraction", "-f")
     ] = None,
@@ -395,6 +433,9 @@ def train(
             "validate_every": validate_every,
             "validate_steps": validate_steps,
             "max_val_batches": max_val_batches,
+            "n_critic": n_critic,
+            "gp_weight": gp_weight,
+            "critic_lr": critic_lr,
         }.items()
         if v is not None
     }
@@ -406,6 +447,7 @@ def train(
             "emb_dim": emb_dim,
             "dropout": dropout,
             "conditioning": conditioning.value if conditioning is not None else None,
+            "noise_dim": noise_dim,
         }.items()
         if v is not None
     }
@@ -510,7 +552,12 @@ def predict(
         ),
     ] = "4096",
     steps: Annotated[
-        int, typer.Option("--steps", "-s", help="Flow matching ODE steps")
+        int,
+        typer.Option(
+            "--steps",
+            "-s",
+            help="Flow matching ODE steps (ignored for a wgan checkpoint)",
+        ),
     ] = 10,
     weights: Annotated[
         Weights,
@@ -649,12 +696,20 @@ def predict(
 
         cc = torch.from_numpy(cond_cont).float().to(_device)
         ck = torch.from_numpy(cond_cat).long().to(_device)
-        stage1_norm, n_sec_pred = sample_flow(model, cc, ck, steps=steps)
+        if model_cfg.get("mode") == "wgan":
+            stage1_norm, n_sec_pred = sample_wgan(model, cc, ck)
+        else:
+            stage1_norm, n_sec_pred = sample_flow(model, cc, ck, steps=steps)
 
         if coord == Coord.global_:
-            sec_cont, sec_phys, _sec_valid_pred = sample_secondaries(
-                sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps=steps
-            )
+            if model_cfg.get("mode") == "wgan":
+                sec_cont, sec_phys, _sec_valid_pred = sample_secondaries_wgan(
+                    sec_decoder, cc, ck, stage1_norm, n_sec_pred
+                )
+            else:
+                sec_cont, sec_phys, _sec_valid_pred = sample_secondaries(
+                    sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps=steps
+                )
             sec_full_np = torch.cat([sec_cont, sec_phys], dim=-1).cpu().numpy()
 
         n_sec_pred_np = n_sec_pred.cpu().numpy()
@@ -898,7 +953,11 @@ def rollout(
     ] = 1000,
     steps: Annotated[
         int,
-        typer.Option("--steps", "-s", help="Flow matching ODE steps per model call"),
+        typer.Option(
+            "--steps",
+            "-s",
+            help="Flow matching ODE steps per model call (ignored for a wgan checkpoint)",
+        ),
     ] = 10,
     weights: Annotated[
         Weights,
@@ -1030,6 +1089,7 @@ def rollout(
         escape_threshold=escape_threshold,
         on_chunk=_write_chunk,
         conditioning=conditioning,
+        mode=model_cfg.get("mode", "flow"),
     )
     if writer is not None:
         writer.close()

giant/config.py

@@ -27,12 +27,22 @@ DEFAULT_CONFIG: dict = {
         "warmup_epochs": 5,
         "lambda_nsec": 0.1,
         "lambda_s2": 1.0,
+        # WGAN-GP-only knobs (mode == "wgan"; ignored by flow/ddpm). n_critic:
+        # critic updates per generator update. gp_weight: gradient-penalty
+        # coefficient (Gulrajani et al. 2017). critic_lr: 0.0 means "use
+        # `lr`" — not None, since save_config's TOML writer has no null
+        # literal to round-trip.
+        "n_critic": 5,
+        "gp_weight": 10.0,
+        "critic_lr": 0.0,
     },
     "model": {
         "hidden_dim": 256,
         "n_blocks": 6,
         "emb_dim": 16,
         "dropout": 0.1,
+        # WGAN generator noise-vector width (mode == "wgan" only).
+        "noise_dim": 64,
         # "physical" conditions on material/particle physical properties via
         # a small MLP (giant.model.network.ConditionEncoder); "embedding"
         # keeps the original learned pdg/material embedding tables — kept

giant/model/network.py

@@ -297,6 +297,233 @@ class SecondaryDecoder(nn.Module):
         return self.out_proj(x)
 
 
+class WGANGenerator(nn.Module):
+    """Stage-1 WGAN-GP generator: single forward pass, no diffusion/flow time.
+
+    Same `ConditionEncoder` + `ResBlock` trunk as `DenoisingMLP`, but the
+    input is a noise vector `z` (not a diffused/interpolated `x_t`) and the
+    ResBlocks condition on the condition encoding alone (no time embedding to
+    concatenate) — see `giant/model/wgan.py` for the adversarial losses, and
+    `giant.sample.sample_wgan` for single-pass sampling.
+    """
+
+    def __init__(
+        self,
+        pdg_vocab: int,
+        mat_vocab: int,
+        hidden_dim: int = 256,
+        n_blocks: int = 6,
+        emb_dim: int = 16,
+        cond_out_dim: int = 128,
+        x_dim: int = X_DIM,
+        noise_dim: int = 64,
+        dropout: float = 0.1,
+        k_max: int = K_MAX,
+        conditioning: str = "embedding",
+    ) -> None:
+        super().__init__()
+        self.noise_dim = noise_dim
+        self.cond_enc = ConditionEncoder(
+            pdg_vocab=pdg_vocab,
+            mat_vocab=mat_vocab,
+            emb_dim=emb_dim,
+            out_dim=cond_out_dim,
+            conditioning=conditioning,
+        )
+        self.input_proj = nn.Linear(noise_dim, hidden_dim)
+        self.blocks = nn.ModuleList(
+            [
+                ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
+                for _ in range(n_blocks)
+            ]
+        )
+        self.out_proj = nn.Linear(hidden_dim, x_dim)
+        self.n_sec_head = nn.Sequential(
+            nn.Linear(cond_out_dim, hidden_dim // 2),
+            nn.SiLU(),
+            nn.Linear(hidden_dim // 2, k_max + 1),
+        )
+
+    def forward(
+        self,
+        z: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+    ) -> torch.Tensor:
+        cond = self.cond_enc(cond_cont, cond_cat)
+        x = self.input_proj(z)
+        for block in self.blocks:
+            x = block(x, cond)
+        return self.out_proj(x)
+
+    def predict_n_sec(
+        self,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+    ) -> torch.Tensor:
+        """Return n_sec logits (B, K_MAX+1) from conditioning alone."""
+        c_emb = self.cond_enc(cond_cont, cond_cat)
+        return self.n_sec_head(c_emb)
+
+
+class Critic(nn.Module):
+    """Stage-1 WGAN-GP critic: scalar realism score, own `ConditionEncoder`.
+
+    Kept structurally parallel to `WGANGenerator` (own condition encoder —
+    separate weights from the generator's, standard GAN practice) but has no
+    n_sec head: n_sec is never adversarial, it stays a plain classifier on
+    the generator side.
+    """
+
+    def __init__(
+        self,
+        pdg_vocab: int,
+        mat_vocab: int,
+        hidden_dim: int = 256,
+        n_blocks: int = 6,
+        emb_dim: int = 16,
+        cond_out_dim: int = 128,
+        x_dim: int = X_DIM,
+        dropout: float = 0.1,
+        conditioning: str = "embedding",
+    ) -> None:
+        super().__init__()
+        self.cond_enc = ConditionEncoder(
+            pdg_vocab=pdg_vocab,
+            mat_vocab=mat_vocab,
+            emb_dim=emb_dim,
+            out_dim=cond_out_dim,
+            conditioning=conditioning,
+        )
+        self.input_proj = nn.Linear(x_dim, hidden_dim)
+        self.blocks = nn.ModuleList(
+            [
+                ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
+                for _ in range(n_blocks)
+            ]
+        )
+        self.out_norm = nn.LayerNorm(hidden_dim)
+        self.out_proj = nn.Linear(hidden_dim, 1)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+    ) -> torch.Tensor:
+        cond = self.cond_enc(cond_cont, cond_cat)
+        h = self.input_proj(x)
+        for block in self.blocks:
+            h = block(h, cond)
+        return self.out_proj(self.out_norm(h)).squeeze(-1)
+
+
+class WGANSecondaryGenerator(nn.Module):
+    """Stage-2 WGAN-GP generator: single forward pass over all K_MAX slots.
+
+    Mirrors `SecondaryDecoder` minus the time embedding, the same way
+    `WGANGenerator` mirrors `DenoisingMLP` — takes noise `z` instead of `x_t`,
+    conditions on `SecondaryConditionEncoder`'s output alone.
+    """
+
+    def __init__(
+        self,
+        pdg_vocab: int,
+        mat_vocab: int,
+        hidden_dim: int = 256,
+        n_blocks: int = 6,
+        emb_dim: int = 16,
+        cond_out_dim: int = 128,
+        stage1_proj_dim: int = 64,
+        sec_dim: int = SEC_DIM,
+        noise_dim: int = 64,
+        dropout: float = 0.1,
+        conditioning: str = "embedding",
+    ) -> None:
+        super().__init__()
+        self.noise_dim = noise_dim
+        self.cond_enc = SecondaryConditionEncoder(
+            pdg_vocab=pdg_vocab,
+            mat_vocab=mat_vocab,
+            emb_dim=emb_dim,
+            cond_out_dim=cond_out_dim,
+            stage1_proj_dim=stage1_proj_dim,
+            out_dim=cond_out_dim,
+            conditioning=conditioning,
+        )
+        self.input_proj = nn.Linear(noise_dim, hidden_dim)
+        self.blocks = nn.ModuleList(
+            [
+                ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
+                for _ in range(n_blocks)
+            ]
+        )
+        self.out_proj = nn.Linear(hidden_dim, sec_dim)
+
+    def forward(
+        self,
+        z: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_out: torch.Tensor,
+    ) -> torch.Tensor:
+        cond = self.cond_enc(cond_cont, cond_cat, stage1_out)
+        x = self.input_proj(z)
+        for block in self.blocks:
+            x = block(x, cond)
+        return self.out_proj(x)
+
+
+class SecondaryCritic(nn.Module):
+    """Stage-2 WGAN-GP critic: scalar realism score over the flattened 90D slots."""
+
+    def __init__(
+        self,
+        pdg_vocab: int,
+        mat_vocab: int,
+        hidden_dim: int = 256,
+        n_blocks: int = 6,
+        emb_dim: int = 16,
+        cond_out_dim: int = 128,
+        stage1_proj_dim: int = 64,
+        sec_dim: int = SEC_DIM,
+        dropout: float = 0.1,
+        conditioning: str = "embedding",
+    ) -> None:
+        super().__init__()
+        self.cond_enc = SecondaryConditionEncoder(
+            pdg_vocab=pdg_vocab,
+            mat_vocab=mat_vocab,
+            emb_dim=emb_dim,
+            cond_out_dim=cond_out_dim,
+            stage1_proj_dim=stage1_proj_dim,
+            out_dim=cond_out_dim,
+            conditioning=conditioning,
+        )
+        self.input_proj = nn.Linear(sec_dim, hidden_dim)
+        self.blocks = nn.ModuleList(
+            [
+                ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
+                for _ in range(n_blocks)
+            ]
+        )
+        self.out_norm = nn.LayerNorm(hidden_dim)
+        self.out_proj = nn.Linear(hidden_dim, 1)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_out: torch.Tensor,
+    ) -> torch.Tensor:
+        cond = self.cond_enc(cond_cont, cond_cat, stage1_out)
+        h = self.input_proj(x)
+        for block in self.blocks:
+            h = block(h, cond)
+        return self.out_proj(self.out_norm(h)).squeeze(-1)
+
+
 class Router(nn.Module):
     """Contract for a pluggable mixture-of-experts routing axis.
 
@@ -799,6 +1026,11 @@ _SEC_DECODER_MODEL_KEYS = {
     "dropout",
     "conditioning",
 }
+_WGAN_GENERATOR_MODEL_KEYS = _STAGE1_MODEL_KEYS | {"noise_dim"}
+_WGAN_SEC_GENERATOR_MODEL_KEYS = _SEC_DECODER_MODEL_KEYS | {"noise_dim"}
+# Critic has no n_sec head (n_sec is never adversarial), so it doesn't accept
+# k_max the way DenoisingMLP/WGANGenerator do.
+_CRITIC_MODEL_KEYS = _STAGE1_MODEL_KEYS - {"k_max"}
 
 
 _AXIS_KEY_RE = re.compile(r"^axis(\d+)_(.+)$")
@@ -864,6 +1096,19 @@ def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]:
     once and passed to both stage1/sec_decoder, so they structurally always
     share one mode.
     """
+    if model_config.get("mode") == "wgan":
+        stage1 = WGANGenerator(
+            **{k: v for k, v in model_config.items() if k in _WGAN_GENERATOR_MODEL_KEYS}
+        )
+        sec_decoder = WGANSecondaryGenerator(
+            **{
+                k: v
+                for k, v in model_config.items()
+                if k in _WGAN_SEC_GENERATOR_MODEL_KEYS
+            }
+        )
+        return stage1, sec_decoder
+
     router_cfg = model_config.get("router")
     if router_cfg and router_cfg.get("enabled"):
         pdg_vocab = model_config["pdg_vocab"]
@@ -895,3 +1140,18 @@ def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]:
         **{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS}
     )
     return stage1, sec_decoder
+
+
+def build_critics(model_config: dict) -> tuple[nn.Module, nn.Module]:
+    """Construct (critic, sec_critic) for `--mode wgan` training.
+
+    Training-only — never persisted for inference the way `build_models`'s
+    pair is, since `predict`/`rollout` only ever run the generators.
+    """
+    critic = Critic(
+        **{k: v for k, v in model_config.items() if k in _CRITIC_MODEL_KEYS}
+    )
+    sec_critic = SecondaryCritic(
+        **{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS}
+    )
+    return critic, sec_critic

giant/model/wgan.py

@@ -0,0 +1,47 @@
+from typing import Callable
+
+import torch
+
+CriticFn = Callable[[torch.Tensor], torch.Tensor]
+
+
+def gradient_penalty(
+    critic_fn: CriticFn,
+    real: torch.Tensor,
+    fake: torch.Tensor,
+    mask: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """WGAN-GP penalty (Gulrajani et al. 2017): (||grad||_2 - 1)^2 at a random interpolate.
+
+    `mask` (same shape as `real`/`fake`, 1 for real content / 0 for padding)
+    is for Stage 2's variable-length slot vector: both the interpolate and the
+    critic's gradient are zeroed on padded dims first, so the norm target of 1
+    is only ever asked of genuine content, not the padding convention shared
+    by both `real` and `fake`.
+    """
+    eps = torch.rand(real.size(0), 1, device=real.device)
+    x_hat = eps * real + (1 - eps) * fake
+    if mask is not None:
+        x_hat = x_hat * mask
+    x_hat = x_hat.requires_grad_(True)
+    scores = critic_fn(x_hat)
+    grad = torch.autograd.grad(outputs=scores.sum(), inputs=x_hat, create_graph=True)[0]
+    if mask is not None:
+        grad = grad * mask
+    return ((grad.norm(2, dim=1) - 1) ** 2).mean()
+
+
+def critic_loss(
+    critic_fn: CriticFn,
+    real: torch.Tensor,
+    fake: torch.Tensor,
+    gp_weight: float,
+    mask: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """WGAN-GP critic loss. `fake` must already be `.detach()`'d by the caller."""
+    gp = gradient_penalty(critic_fn, real, fake, mask=mask)
+    return critic_fn(fake).mean() - critic_fn(real).mean() + gp_weight * gp
+
+
+def generator_loss(critic_fn: CriticFn, fake: torch.Tensor) -> torch.Tensor:
+    return -critic_fn(fake).mean()

giant/pipeline.py

@@ -22,7 +22,7 @@ from giant.data.loader import (
 )
 from giant.data.transforms import build_features, _WelfordAccumulator
 from giant.data.dataset import make_event_split, StreamingStepsDataset
-from giant.model.network import build_models
+from giant.model.network import build_models, build_critics
 from giant.train import train as run_training
 
 
@@ -63,6 +63,11 @@ def run_train_job(
     echo(f"  {len(pdg_map)} PDG codes  |  {len(mat_map)} materials")
 
     router_cfg = m["router"]
+    if t["mode"] == "wgan" and router_cfg.get("enabled"):
+        raise ValueError(
+            "--mode wgan does not support --router (no routed WGAN generator/"
+            "critic exists) — disable one or the other"
+        )
     proc_map: dict[str, int] | None = None
     if router_cfg.get("enabled") and router_cfg.get("type") == "process":
         echo("building process vocabulary …")
@@ -159,6 +164,10 @@ def run_train_job(
         "router": dict(router_cfg),
         "expert_hidden_dim": router_cfg["expert_hidden_dim"],
         "expert_n_blocks": router_cfg["expert_n_blocks"],
+        # Read by `predict`/`rollout` (which never receive their own --mode
+        # flag) to auto-detect which sampler a checkpoint needs.
+        "mode": t["mode"],
+        "noise_dim": m.get("noise_dim", 64),
     }
 
     stage1_model, sec_decoder = build_models(model_config)
@@ -167,6 +176,15 @@ def run_train_job(
         f"sec_decoder: {sum(p.numel() for p in sec_decoder.parameters()):,} parameters"
     )
 
+    critic = None
+    sec_critic = None
+    if t["mode"] == "wgan":
+        critic, sec_critic = build_critics(model_config)
+        echo(
+            f"critic: {sum(p.numel() for p in critic.parameters()):,} parameters  |  "
+            f"sec_critic: {sum(p.numel() for p in sec_critic.parameters()):,} parameters"
+        )
+
     out_dir.mkdir(parents=True, exist_ok=True)
     meta = config.build_run_meta(
         data=data,
@@ -210,4 +228,9 @@ def run_train_job(
         validate_steps=t["validate_steps"],
         max_val_batches=t["max_val_batches"],
         total_train_batches=total_train_batches,
+        critic=critic,
+        sec_critic=sec_critic,
+        n_critic=t.get("n_critic", 5),
+        gp_weight=t.get("gp_weight", 10.0),
+        critic_lr=t.get("critic_lr") or None,
     )

giant/rollout.py

@@ -39,7 +39,12 @@ from giant.data.transforms import (
     reconstruct_post_pos,
 )
 from giant.particles import nearest_known_pdg, particle_phys_array
-from giant.sample import sample_flow, sample_secondaries
+from giant.sample import (
+    sample_flow,
+    sample_secondaries,
+    sample_wgan,
+    sample_secondaries_wgan,
+)
 
 # Record columns produced per step / per terminal marker.
 _RECORD_KEYS = [
@@ -306,6 +311,7 @@ def rollout(
     escape_threshold: float | None = None,
     on_chunk: Callable[[dict[str, np.ndarray]], None] | None = None,
     conditioning: str = "embedding",
+    mode: str = "flow",
 ) -> dict[str, np.ndarray] | RolloutSummary:
     """Run showers to completion.
 
@@ -359,6 +365,7 @@ def rollout(
                     device,
                     max_tracks_per_event,
                     conditioning,
+                    mode,
                 )
             )
         frontier = _concat_frontiers(next_parts)
@@ -389,6 +396,7 @@ def _step_chunk(
     device,
     max_tracks_per_event,
     conditioning,
+    mode="flow",
 ) -> dict[str, np.ndarray]:
     """Advance one chunk of tracks by a single step; return the next frontier."""
     n = len(tr["event_id"])
@@ -463,7 +471,10 @@ def _step_chunk(
     cc = torch.from_numpy(cond_cont).float().to(device)
     ck = torch.from_numpy(cond_cat).long().to(device)
 
-    stage1_norm, n_sec_pred = sample_flow(stage1_model, cc, ck, steps=steps)
+    if mode == "wgan":
+        stage1_norm, n_sec_pred = sample_wgan(stage1_model, cc, ck)
+    else:
+        stage1_norm, n_sec_pred = sample_flow(stage1_model, cc, ck, steps=steps)
     raw = tgt_norm.inverse_transform(stage1_norm.cpu().numpy())
 
     step_length = inv_log_transform(raw[:, 0])
@@ -490,9 +501,14 @@ def _step_chunk(
     # identity, used as-is for the spawned track's own future conditioning.
     # sec_pdg_code below is a *separate*, reporting-only nearest-known-PDG
     # label (never fed back into the model) — see giant/particles.py.
-    sec_cont, sec_phys, _valid = sample_secondaries(
-        sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps=steps
-    )
+    if mode == "wgan":
+        sec_cont, sec_phys, _valid = sample_secondaries_wgan(
+            sec_decoder, cc, ck, stage1_norm, n_sec_pred
+        )
+    else:
+        sec_cont, sec_phys, _valid = sample_secondaries(
+            sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps=steps
+        )
     sec_full = torch.cat([sec_cont, sec_phys], dim=-1).cpu().numpy()
     sec_E, sec_dir_world, sec_mass, sec_charge, sec_valid = decode_secondaries(
         sec_full,

giant/sample.py

@@ -3,6 +3,26 @@ import torch
 from giant.constants import K_MAX, SEC_DIM, SEC_SLOT_DIM, X_DIM
 
 
+def _slots_from_flat(
+    x: torch.Tensor, n_sec_pred: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Reshape a flat (B, SEC_DIM) decoder output into per-slot tensors.
+
+    Returns (sec_cont, sec_phys, sec_valid) — see `sample_secondaries`'s
+    docstring for their shapes/meaning. Shared by both the flow-matching and
+    WGAN Stage-2 samplers, which differ only in how `x` was produced.
+    """
+    B = x.size(0)
+    device = x.device
+    x_slots = x.view(B, K_MAX, SEC_SLOT_DIM)
+    sec_cont = x_slots[:, :, :4]
+    sec_phys = x_slots[:, :, 4:]
+    sec_valid = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(
+        1
+    )
+    return sec_cont, sec_phys, sec_valid
+
+
 @torch.no_grad()
 def sample_flow(
     model: torch.nn.Module,
@@ -64,14 +84,7 @@ def sample_secondaries(
         v = sec_decoder(x, t, cond_cont, cond_cat, stage1_out)
         x = x + v * dt
 
-    x_slots = x.view(B, K_MAX, SEC_SLOT_DIM)
-    sec_cont = x_slots[:, :, :4]
-    sec_phys = x_slots[:, :, 4:]
-
-    sec_valid = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(
-        1
-    )
-    return sec_cont, sec_phys, sec_valid
+    return _slots_from_flat(x, n_sec_pred)
 
 
 @torch.no_grad()
@@ -102,6 +115,39 @@ def sample_ddpm(
     return x, n_sec_pred
 
 
+@torch.no_grad()
+def sample_wgan(
+    generator: torch.nn.Module,
+    cond_cont: torch.Tensor,
+    cond_cat: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Single-pass Stage-1 WGAN generator sample. Returns (sample, n_sec_pred)."""
+    generator.eval()
+    B = cond_cont.size(0)
+    z = torch.randn(B, generator.noise_dim, device=cond_cont.device)
+    x = generator(z, cond_cont, cond_cat)
+    n_sec_logits = generator.predict_n_sec(cond_cont, cond_cat)
+    n_sec_pred = n_sec_logits.argmax(dim=-1)
+    return x, n_sec_pred
+
+
+@torch.no_grad()
+def sample_secondaries_wgan(
+    sec_decoder: torch.nn.Module,
+    cond_cont: torch.Tensor,
+    cond_cat: torch.Tensor,
+    stage1_out: torch.Tensor,
+    n_sec_pred: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Single-pass Stage-2 WGAN generator sample; see `sample_secondaries`'s
+    docstring for the returned (sec_cont, sec_phys, sec_valid) shapes."""
+    sec_decoder.eval()
+    B = cond_cont.size(0)
+    z = torch.randn(B, sec_decoder.noise_dim, device=cond_cont.device)
+    x = sec_decoder(z, cond_cont, cond_cat, stage1_out)
+    return _slots_from_flat(x, n_sec_pred)
+
+
 @torch.no_grad()
 def sample_ddim(
     model: torch.nn.Module,

giant/train.py

@@ -8,17 +8,20 @@ from pathlib import Path
 from types import FrameType
 from typing import Callable
 
+import numpy as np
 import torch
 import torch.nn.functional as F
 import torch.optim as optim
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
+from giant.constants import K_MAX, SEC_SLOT_DIM
 from giant.model.schedule import (
     CosineSchedule,
     flow_matching_loss,
     flow_matching_loss_secondary,
 )
+from giant.model.wgan import gradient_penalty, generator_loss
 from giant.validate import validate_marginals
 
 _METRICS_FIELDS = [
@@ -29,12 +32,17 @@ _METRICS_FIELDS = [
     "train_loss_s2",
     "train_loss_balance",
     "train_loss_proc",
+    "d_loss",
+    "g_loss",
+    "wasserstein_estimate",
+    "gp_loss",
     "val_loss",
     "val_loss_s1",
     "val_loss_nsec",
     "val_loss_s2",
     "val_loss_balance",
     "val_loss_proc",
+    "val_marginal_kl",
     "lr",
     "grad_norm",
     "epoch_time_s",
@@ -129,8 +137,6 @@ def _compute_losses(
     # giant.data.transforms.encode_secondaries) rather than a learned/moving
     # one, so — unlike the embedding-table target this replaced — no
     # detaching is needed to keep the target from chasing the decoder.
-    from giant.constants import K_MAX
-
     x1_s2 = sec_cont.flatten(1)  # (B, SEC_DIM)
 
     sec_mask = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec.unsqueeze(1)
@@ -168,6 +174,120 @@ def _compute_losses(
     return total, l_s1, l_nsec, l_s2, l_balance, l_proc
 
 
+def _wgan_train_step(
+    generator: torch.nn.Module,
+    sec_generator: torch.nn.Module,
+    critic: torch.nn.Module,
+    sec_critic: torch.nn.Module,
+    batch: tuple,
+    device: torch.device,
+    optimizer_g: optim.Optimizer,
+    optimizer_d: optim.Optimizer,
+    g_params: list,
+    d_params: list,
+    step_count: int,
+    n_critic: int,
+    gp_weight: float,
+    lambda_nsec: float,
+    lambda_s2: float,
+) -> dict:
+    """One WGAN-GP training step, both stages (see giant/model/wgan.py for the losses).
+
+    Both critics update every batch. Every `n_critic`-th batch additionally
+    updates both generators. The (non-adversarial) n_sec classifier updates
+    every batch regardless — folded into whichever `optimizer_g` step happens
+    this batch (full adversarial g_loss on generator batches, n_sec-only in
+    between) rather than throttled to the generator's cadence, since n_sec
+    accuracy is a headline flow-vs-wgan comparison metric and shares the
+    generator's ConditionEncoder.
+
+    Stage 2's real/fake target is a flattened (B, SEC_DIM) vector with
+    `K_MAX - n_sec` padded slots per row; both critic's input and its
+    gradient-penalty gradient are masked to the valid slots (see
+    `giant.model.wgan.gradient_penalty`) so the critic can't key on padding
+    instead of genuine content. Stage 2 is conditioned on the *real*
+    ground-truth Stage-1 target (`x1_s1`, detached) rather than the
+    generator's own fake Stage-1 output — same precedent as the flow-matching
+    path's `flow_matching_loss_secondary` call, avoiding compounding errors
+    during training.
+    """
+    cond_cont, cond_cat, x1_s1, n_sec, sec_cont, _proc_idx = batch
+    cond_cont = cond_cont.to(device)
+    cond_cat = cond_cat.to(device)
+    x1_s1 = x1_s1.to(device)
+    n_sec = n_sec.to(device)
+    sec_cont = sec_cont.to(device)
+
+    B = x1_s1.size(0)
+    x1_s2 = sec_cont.flatten(1)  # (B, SEC_DIM)
+    sec_mask = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec.unsqueeze(1)
+    mask_flat = (
+        sec_mask.unsqueeze(-1).expand(-1, -1, SEC_SLOT_DIM).reshape(B, -1).float()
+    )
+    stage1_ctx = x1_s1.detach()
+
+    def critic_fn1(x: torch.Tensor) -> torch.Tensor:
+        return critic(x, cond_cont, cond_cat)
+
+    def critic_fn2(x: torch.Tensor) -> torch.Tensor:
+        return sec_critic(x, cond_cont, cond_cat, stage1_ctx)
+
+    z1 = torch.randn(B, generator.noise_dim, device=device)
+    fake1 = generator(z1, cond_cont, cond_cat)
+    z2 = torch.randn(B, sec_generator.noise_dim, device=device)
+    fake2 = sec_generator(z2, cond_cont, cond_cat, stage1_ctx)
+    fake2_masked = fake2 * mask_flat
+    real2_masked = x1_s2 * mask_flat
+
+    # --- Critic step (every batch) ---
+    fake1_detached = fake1.detach()
+    real1_score = critic_fn1(x1_s1)
+    fake1_score = critic_fn1(fake1_detached)
+    gp1 = gradient_penalty(critic_fn1, x1_s1, fake1_detached)
+    d1 = fake1_score.mean() - real1_score.mean() + gp_weight * gp1
+    wasserstein_estimate = (real1_score.mean() - fake1_score.mean()).detach()
+
+    fake2_detached_masked = fake2_masked.detach()
+    real2_score = critic_fn2(real2_masked)
+    fake2_score = critic_fn2(fake2_detached_masked)
+    gp2 = gradient_penalty(
+        critic_fn2, real2_masked, fake2_detached_masked, mask=mask_flat
+    )
+    d2 = fake2_score.mean() - real2_score.mean() + gp_weight * gp2
+
+    d_loss = d1 + lambda_s2 * d2
+    optimizer_d.zero_grad()
+    d_loss.backward()
+    grad_norm_d = torch.nn.utils.clip_grad_norm_(d_params, 1.0)
+    optimizer_d.step()
+
+    # --- Generator (+ n_sec) step ---
+    did_g_step = step_count % n_critic == 0
+    l_nsec = F.cross_entropy(generator.predict_n_sec(cond_cont, cond_cat), n_sec)
+    optimizer_g.zero_grad()
+    if did_g_step:
+        g1 = generator_loss(critic_fn1, fake1)
+        g2 = generator_loss(critic_fn2, fake2_masked)
+        g_loss = g1 + lambda_nsec * l_nsec + lambda_s2 * g2
+    else:
+        g1 = torch.zeros((), device=device)
+        g2 = torch.zeros((), device=device)
+        g_loss = lambda_nsec * l_nsec
+    g_loss.backward()
+    grad_norm_g = torch.nn.utils.clip_grad_norm_(g_params, 1.0)
+    optimizer_g.step()
+
+    return {
+        "d_loss": d_loss.detach(),
+        "g_loss": (g1 + lambda_s2 * g2).detach(),
+        "wasserstein_estimate": wasserstein_estimate,
+        "gp_loss": (gp1 + lambda_s2 * gp2).detach(),
+        "l_nsec": l_nsec.detach(),
+        "did_g_step": did_g_step,
+        "grad_norm": grad_norm_d.item() + grad_norm_g.item(),
+    }
+
+
 def train(
     stage1_model: torch.nn.Module,
     sec_decoder: torch.nn.Module,
@@ -195,12 +315,23 @@ def train(
     validate_steps: int = 10,
     max_val_batches: int = 0,
     total_train_batches: int = 0,
+    critic: torch.nn.Module | None = None,
+    sec_critic: torch.nn.Module | None = None,
+    n_critic: int = 5,
+    gp_weight: float = 10.0,
+    critic_lr: float | None = None,
 ) -> None:
     out_dir = Path(out_dir)
     out_dir.mkdir(parents=True, exist_ok=True)
 
     stage1_model = stage1_model.to(device)
     sec_decoder = sec_decoder.to(device)
+    if mode == "wgan":
+        assert critic is not None and sec_critic is not None, (
+            "mode='wgan' requires critic/sec_critic (see giant.model.network.build_critics)"
+        )
+        critic = critic.to(device)
+        sec_critic = sec_critic.to(device)
 
     # Flow-matching/diffusion models sample noticeably better from an EMA of
     # the weights than from the raw SGD-noisy ones — buffers (e.g. the fixed
@@ -217,14 +348,36 @@ def train(
             p.requires_grad_(False)
 
     all_params = list(stage1_model.parameters()) + list(sec_decoder.parameters())
-    optimizer = optim.AdamW(all_params, lr=lr, weight_decay=weight_decay)
+    all_params_d: list = []
+    optimizer_d: optim.Optimizer | None = None
+    if mode == "wgan":
+        assert critic is not None and sec_critic is not None
+        # Standard WGAN-GP recipe (Gulrajani et al. 2017): Adam with
+        # beta1=0 (momentum destabilizes critic training) and no weight
+        # decay, rather than the AdamW(weight_decay=...) used for flow/ddpm.
+        optimizer = optim.Adam(all_params, lr=lr, betas=(0.0, 0.9))
+        all_params_d = list(critic.parameters()) + list(sec_critic.parameters())
+        optimizer_d = optim.Adam(
+            all_params_d,
+            lr=critic_lr if critic_lr is not None else lr,
+            betas=(0.0, 0.9),
+        )
+    else:
+        optimizer = optim.AdamW(all_params, lr=lr, weight_decay=weight_decay)
 
     # Warmup/decay in units of optimizer steps rather than epochs: at large
     # dataset sizes a single epoch can be tens of thousands of steps, and an
     # epoch-granularity schedule would leave warmup/cosine decay unable to
     # move within it. Requires an accurate `total_train_batches` (steps per
     # epoch); the only caller, run_train_job, always supplies one.
+    #
+    # In wgan mode, `lr_sched.step()`/EMA only fire on generator steps (see
+    # the per-batch loop below) — 1 in every `n_critic` batches — so the
+    # schedule's own step-counting must be in those same units, or warmup
+    # would never finish and cosine decay would barely move.
     steps_per_epoch = max(total_train_batches, 1)
+    if mode == "wgan":
+        steps_per_epoch = max(total_train_batches // (n_critic + 1), 1)
     warmup_steps = warmup_epochs * steps_per_epoch
     total_steps = max(epochs * steps_per_epoch, 1)
 
@@ -251,6 +404,15 @@ def train(
             ema_sec_decoder.load_state_dict(
                 ckpt.get("sec_decoder_ema", ckpt["sec_decoder"])
             )
+        if mode == "wgan":
+            assert (
+                critic is not None
+                and sec_critic is not None
+                and optimizer_d is not None
+            )
+            critic.load_state_dict(ckpt["critic"])
+            sec_critic.load_state_dict(ckpt["sec_critic"])
+            optimizer_d.load_state_dict(ckpt["optimizer_d"])
         optimizer.load_state_dict(ckpt["optimizer"])
         lr_sched.load_state_dict(ckpt["lr_sched"])
         start_epoch = ckpt.get("epoch", 0) + 1
@@ -284,17 +446,26 @@ def train(
 
     epoch_w = len(str(epochs))
     last_completed_epoch = start_epoch - 1
+    global_step = 0
     with _GracefulShutdown() as shutdown:
         for epoch in range(start_epoch, epochs + 1):
             epoch_start = time.monotonic()
             stage1_model.train()
             sec_decoder.train()
+            if mode == "wgan":
+                assert critic is not None and sec_critic is not None
+                critic.train()
+                sec_critic.train()
             train_loss_sum = 0.0
             train_s1_sum = 0.0
             train_nsec_sum = 0.0
             train_s2_sum = 0.0
             train_balance_sum = 0.0
             train_proc_sum = 0.0
+            train_d_sum = 0.0
+            train_g_sum = 0.0
+            train_wasserstein_sum = 0.0
+            train_gp_sum = 0.0
             train_n = 0
             train_batches = 0
             grad_norm_sum = 0.0
@@ -309,37 +480,84 @@ def train(
                 dynamic_ncols=True,
             )
             for batch in bar:
-                loss, l_s1, l_nsec, l_s2, l_balance, l_proc = _compute_losses(
-                    stage1_model,
-                    sec_decoder,
-                    batch,
-                    mode,
-                    ddpm_schedule,
-                    device,
-                    lambda_nsec,
-                    lambda_s2,
-                    lambda_balance,
-                    lambda_proc,
-                )
-                optimizer.zero_grad()
-                loss.backward()
-                grad_norm = torch.nn.utils.clip_grad_norm_(all_params, 1.0)
-                optimizer.step()
-                lr_sched.step()
-                if ema_decay > 0:
-                    assert ema_stage1_model is not None and ema_sec_decoder is not None
-                    _update_ema(ema_stage1_model, stage1_model, ema_decay)
-                    _update_ema(ema_sec_decoder, sec_decoder, ema_decay)
+                if mode == "wgan":
+                    assert (
+                        critic is not None
+                        and sec_critic is not None
+                        and optimizer_d is not None
+                    )
+                    stats = _wgan_train_step(
+                        stage1_model,
+                        sec_decoder,
+                        critic,
+                        sec_critic,
+                        batch,
+                        device,
+                        optimizer,
+                        optimizer_d,
+                        all_params,
+                        all_params_d,
+                        global_step,
+                        n_critic,
+                        gp_weight,
+                        lambda_nsec,
+                        lambda_s2,
+                    )
+                    global_step += 1
+                    if stats["did_g_step"]:
+                        lr_sched.step()
+                        if ema_decay > 0:
+                            assert (
+                                ema_stage1_model is not None
+                                and ema_sec_decoder is not None
+                            )
+                            _update_ema(ema_stage1_model, stage1_model, ema_decay)
+                            _update_ema(ema_sec_decoder, sec_decoder, ema_decay)
+
+                    B = batch[0].size(0)
+                    batch_loss = stats["d_loss"].item() + stats["g_loss"].item()
+                    batch_grad_norm = stats["grad_norm"]
+                    train_loss_sum += batch_loss * B
+                    train_nsec_sum += stats["l_nsec"].item() * B
+                    train_d_sum += stats["d_loss"].item() * B
+                    train_g_sum += stats["g_loss"].item() * B
+                    train_wasserstein_sum += stats["wasserstein_estimate"].item() * B
+                    train_gp_sum += stats["gp_loss"].item() * B
+                else:
+                    loss, l_s1, l_nsec, l_s2, l_balance, l_proc = _compute_losses(
+                        stage1_model,
+                        sec_decoder,
+                        batch,
+                        mode,
+                        ddpm_schedule,
+                        device,
+                        lambda_nsec,
+                        lambda_s2,
+                        lambda_balance,
+                        lambda_proc,
+                    )
+                    optimizer.zero_grad()
+                    loss.backward()
+                    grad_norm = torch.nn.utils.clip_grad_norm_(all_params, 1.0)
+                    optimizer.step()
+                    lr_sched.step()
+                    if ema_decay > 0:
+                        assert (
+                            ema_stage1_model is not None and ema_sec_decoder is not None
+                        )
+                        _update_ema(ema_stage1_model, stage1_model, ema_decay)
+                        _update_ema(ema_sec_decoder, sec_decoder, ema_decay)
+
+                    B = batch[0].size(0)
+                    batch_loss = loss.item()
+                    batch_grad_norm = grad_norm.item()
+                    train_loss_sum += batch_loss * B
+                    train_s1_sum += l_s1.item() * B
+                    train_nsec_sum += l_nsec.item() * B
+                    train_s2_sum += l_s2.item() * B
+                    train_balance_sum += l_balance.item() * B
+                    train_proc_sum += l_proc.item() * B
 
-                B = batch[0].size(0)
-                batch_loss = loss.item()
-                batch_grad_norm = grad_norm.item()
-                train_loss_sum += batch_loss * B
-                train_s1_sum += l_s1.item() * B
-                train_nsec_sum += l_nsec.item() * B
-                train_s2_sum += l_s2.item() * B
-                train_balance_sum += l_balance.item() * B
-                train_proc_sum += l_proc.item() * B
                 train_n += B
                 train_batches += 1
                 grad_norm_sum += batch_grad_norm
@@ -368,38 +586,87 @@ def train(
 
             stage1_model.eval()
             sec_decoder.eval()
-            val_loss_sum = 0.0
-            val_s1_sum = 0.0
-            val_nsec_sum = 0.0
-            val_s2_sum = 0.0
-            val_balance_sum = 0.0
-            val_proc_sum = 0.0
-            val_n = 0
-            with torch.no_grad():
-                for val_batch_idx, batch in enumerate(val_loader):
-                    if max_val_batches > 0 and val_batch_idx >= max_val_batches:
-                        break
-                    loss, l_s1, l_nsec, l_s2, l_balance, l_proc = _compute_losses(
+            if mode == "wgan":
+                assert critic is not None and sec_critic is not None
+                critic.eval()
+                sec_critic.eval()
+
+            val_marginal_kl = float("nan")
+            if mode == "wgan":
+                # WGANGenerator.forward(z, cond_cont, cond_cat) has no
+                # diffusion/flow `t` argument, so the usual _compute_losses
+                # val loop below (which calls flow_matching_loss ->
+                # stage1_model(x_t, t, ...)) doesn't apply — and a WGAN
+                # critic loss isn't a monotone quality signal fit for
+                # best-checkpoint selection anyway. Select on marginal KL
+                # against the EMA generators instead (matches what
+                # predict/rollout sample from by default, --weights ema).
+                eval_stage1 = (
+                    ema_stage1_model if ema_stage1_model is not None else stage1_model
+                )
+                eval_sec_decoder = (
+                    ema_sec_decoder if ema_sec_decoder is not None else sec_decoder
+                )
+                marginal_result = validate_marginals(
+                    eval_stage1,
+                    val_loader,
+                    mode=mode,
+                    device=device,
+                    sec_decoder=eval_sec_decoder,
+                )
+                val_marginal_kl = float(np.mean(marginal_result["kl_divergence"]))
+                val_loss = val_marginal_kl
+                val_s1_sum = val_nsec_sum = val_s2_sum = val_balance_sum = (
+                    val_proc_sum
+                ) = 0.0
+                val_n = 1
+            else:
+                val_loss_sum = 0.0
+                val_s1_sum = 0.0
+                val_nsec_sum = 0.0
+                val_s2_sum = 0.0
+                val_balance_sum = 0.0
+                val_proc_sum = 0.0
+                val_n = 0
+                with torch.no_grad():
+                    for val_batch_idx, batch in enumerate(val_loader):
+                        if max_val_batches > 0 and val_batch_idx >= max_val_batches:
+                            break
+                        loss, l_s1, l_nsec, l_s2, l_balance, l_proc = _compute_losses(
+                            stage1_model,
+                            sec_decoder,
+                            batch,
+                            mode,
+                            ddpm_schedule,
+                            device,
+                            lambda_nsec,
+                            lambda_s2,
+                            lambda_balance,
+                            lambda_proc,
+                        )
+                        B = batch[0].size(0)
+                        val_loss_sum += loss.item() * B
+                        val_s1_sum += l_s1.item() * B
+                        val_nsec_sum += l_nsec.item() * B
+                        val_s2_sum += l_s2.item() * B
+                        val_balance_sum += l_balance.item() * B
+                        val_proc_sum += l_proc.item() * B
+                        val_n += B
+                val_loss = val_loss_sum / max(val_n, 1)
+
+                if validate_every > 0 and epoch % validate_every == 0:
+                    print(f"[epoch {epoch}] marginal validation:")
+                    marginal_result = validate_marginals(
                         stage1_model,
-                        sec_decoder,
-                        batch,
-                        mode,
-                        ddpm_schedule,
-                        device,
-                        lambda_nsec,
-                        lambda_s2,
-                        lambda_balance,
-                        lambda_proc,
+                        val_loader,
+                        mode=mode,
+                        schedule=ddpm_schedule,
+                        device=device,
+                        steps=validate_steps,
+                        sec_decoder=sec_decoder,
                     )
-                    B = batch[0].size(0)
-                    val_loss_sum += loss.item() * B
-                    val_s1_sum += l_s1.item() * B
-                    val_nsec_sum += l_nsec.item() * B
-                    val_s2_sum += l_s2.item() * B
-                    val_balance_sum += l_balance.item() * B
-                    val_proc_sum += l_proc.item() * B
-                    val_n += B
-            val_loss = val_loss_sum / max(val_n, 1)
+                    val_marginal_kl = float(np.mean(marginal_result["kl_divergence"]))
+
             epoch_time = time.monotonic() - epoch_start
 
             is_best = val_loss < best_val_loss
@@ -411,7 +678,9 @@ def train(
                 f"  nsec={train_nsec_sum / max(train_n, 1):.3f}"
                 f"  s2={train_s2_sum / max(train_n, 1):.3f}"
                 f"  bal={train_balance_sum / max(train_n, 1):.3f}"
-                f"  proc={train_proc_sum / max(train_n, 1):.3f})"
+                f"  proc={train_proc_sum / max(train_n, 1):.3f}"
+                f"  d={train_d_sum / max(train_n, 1):.3f}"
+                f"  g={train_g_sum / max(train_n, 1):.3f})"
                 f"  val {val_loss:.4f}"
                 f"  lr {current_lr:.2e}  gnorm {train_grad_norm:.3f}"
                 f"  {epoch_time:.1f}s{marker}"
@@ -425,12 +694,17 @@ def train(
                     "train_loss_s2": train_s2_sum / max(train_n, 1),
                     "train_loss_balance": train_balance_sum / max(train_n, 1),
                     "train_loss_proc": train_proc_sum / max(train_n, 1),
+                    "d_loss": train_d_sum / max(train_n, 1),
+                    "g_loss": train_g_sum / max(train_n, 1),
+                    "wasserstein_estimate": train_wasserstein_sum / max(train_n, 1),
+                    "gp_loss": train_gp_sum / max(train_n, 1),
                     "val_loss": val_loss,
                     "val_loss_s1": val_s1_sum / max(val_n, 1),
                     "val_loss_nsec": val_nsec_sum / max(val_n, 1),
                     "val_loss_s2": val_s2_sum / max(val_n, 1),
                     "val_loss_balance": val_balance_sum / max(val_n, 1),
                     "val_loss_proc": val_proc_sum / max(val_n, 1),
+                    "val_marginal_kl": val_marginal_kl,
                     "lr": current_lr,
                     "grad_norm": train_grad_norm,
                     "epoch_time_s": epoch_time,
@@ -438,18 +712,6 @@ def train(
             )
             metrics_file.flush()
 
-            if validate_every > 0 and epoch % validate_every == 0:
-                print(f"[epoch {epoch}] marginal validation:")
-                validate_marginals(
-                    stage1_model,
-                    val_loader,
-                    mode=mode,
-                    schedule=ddpm_schedule,
-                    device=device,
-                    steps=validate_steps,
-                    sec_decoder=sec_decoder,
-                )
-
             ckpt: dict = {
                 "model": stage1_model.state_dict(),
                 "sec_decoder": sec_decoder.state_dict(),
@@ -458,6 +720,15 @@ def train(
                 "epoch": epoch,
                 "best_val_loss": best_val_loss,
             }
+            if mode == "wgan":
+                assert (
+                    critic is not None
+                    and sec_critic is not None
+                    and optimizer_d is not None
+                )
+                ckpt["critic"] = critic.state_dict()
+                ckpt["sec_critic"] = sec_critic.state_dict()
+                ckpt["optimizer_d"] = optimizer_d.state_dict()
             if ema_decay > 0:
                 assert ema_stage1_model is not None and ema_sec_decoder is not None
                 ckpt["model_ema"] = ema_stage1_model.state_dict()

giant/validate.py

@@ -8,6 +8,8 @@ from giant.sample import (
     sample_ddpm,
     sample_ddim,
     sample_secondaries,
+    sample_wgan,
+    sample_secondaries_wgan,
 )
 
 _SEC_PHYS_NAMES = ["log_mass", "charge"]
@@ -98,6 +100,8 @@ def validate_marginals(
             gen, n_sec_pred = sample_flow(model, cond_cont, cond_cat, **_kw(steps))
         elif mode == "ddpm":
             gen, n_sec_pred = sample_ddpm(model, cond_cont, cond_cat, schedule)
+        elif mode == "wgan":
+            gen, n_sec_pred = sample_wgan(model, cond_cont, cond_cat)
         else:
             gen, n_sec_pred = sample_ddim(
                 model, cond_cont, cond_cat, schedule, **_kw(steps)
@@ -118,14 +122,19 @@ def validate_marginals(
         real_frac = 1.0 / (1.0 + np.exp(-sec_cont[:, :, 0].numpy().astype(np.float64)))
         real_phys = sec_cont[:, :, 4:6].numpy()  # (B, K_MAX, 2) [log_mass, charge]
 
-        sec_cont_pred, sec_phys_pred, sec_valid_pred = sample_secondaries(
-            sec_decoder,
-            cond_cont,
-            cond_cat,
-            gen,
-            n_sec_pred,
-            steps=steps if steps is not None else 10,
-        )
+        if mode == "wgan":
+            sec_cont_pred, sec_phys_pred, sec_valid_pred = sample_secondaries_wgan(
+                sec_decoder, cond_cont, cond_cat, gen, n_sec_pred
+            )
+        else:
+            sec_cont_pred, sec_phys_pred, sec_valid_pred = sample_secondaries(
+                sec_decoder,
+                cond_cont,
+                cond_cat,
+                gen,
+                n_sec_pred,
+                steps=steps if steps is not None else 10,
+            )
         gen_frac = 1.0 / (
             1.0 + np.exp(-sec_cont_pred[:, :, 0].cpu().numpy().astype(np.float64))
         )

tests/test_wgan.py

@@ -0,0 +1,185 @@
+import torch
+
+from giant.constants import COND_DIM, K_MAX, SEC_DIM, SEC_SLOT_DIM, X_DIM
+from giant.model.network import (
+    Critic,
+    SecondaryCritic,
+    WGANGenerator,
+    WGANSecondaryGenerator,
+)
+from giant.model.wgan import critic_loss, generator_loss, gradient_penalty
+from giant.sample import sample_secondaries_wgan, sample_wgan
+
+
+def _cond(B=8):
+    cond_cont = torch.randn(B, COND_DIM)
+    cond_cat = torch.zeros(B, 2, dtype=torch.long)
+    return cond_cont, cond_cat
+
+
+def _small_generator():
+    return WGANGenerator(
+        pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2, noise_dim=8
+    )
+
+
+def _small_critic():
+    return Critic(pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2)
+
+
+def _small_sec_generator():
+    return WGANSecondaryGenerator(
+        pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2, noise_dim=8
+    )
+
+
+def _small_sec_critic():
+    return SecondaryCritic(pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2)
+
+
+def _mask(B, n_sec):
+    sec_mask = torch.arange(K_MAX).unsqueeze(0) < n_sec.unsqueeze(1)
+    return sec_mask.unsqueeze(-1).expand(-1, -1, SEC_SLOT_DIM).reshape(B, -1).float()
+
+
+# --- Stage-1 generator/critic ---
+
+
+def test_wgan_generator_output_shape():
+    B = 8
+    model = _small_generator()
+    cond_cont, cond_cat = _cond(B)
+    z = torch.randn(B, model.noise_dim)
+    out = model(z, cond_cont, cond_cat)
+    assert out.shape == (B, X_DIM)
+
+
+def test_wgan_generator_predict_n_sec_shape():
+    B = 6
+    model = _small_generator()
+    cond_cont, cond_cat = _cond(B)
+    logits = model.predict_n_sec(cond_cont, cond_cat)
+    assert logits.shape == (B, K_MAX + 1)
+
+
+def test_wgan_generator_gradients_flow():
+    B = 4
+    model = _small_generator()
+    cond_cont, cond_cat = _cond(B)
+    z = torch.randn(B, model.noise_dim)
+    gen_loss = model(z, cond_cont, cond_cat).sum()
+    nsec_loss = model.predict_n_sec(cond_cont, cond_cat).sum()
+    (gen_loss + nsec_loss).backward()
+    for name, p in model.named_parameters():
+        assert p.grad is not None, f"no grad for {name}"
+
+
+def test_critic_output_shape():
+    B = 8
+    critic = _small_critic()
+    cond_cont, cond_cat = _cond(B)
+    x = torch.randn(B, X_DIM)
+    out = critic(x, cond_cont, cond_cat)
+    assert out.shape == (B,)
+
+
+def test_sample_wgan_shape():
+    B = 6
+    model = _small_generator()
+    cond_cont, cond_cat = _cond(B)
+    sample, n_sec = sample_wgan(model, cond_cont, cond_cat)
+    assert sample.shape == (B, X_DIM)
+    assert n_sec.shape == (B,)
+
+
+# --- Stage-2 generator/critic ---
+
+
+def test_wgan_secondary_generator_output_shape():
+    B = 8
+    model = _small_sec_generator()
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    z = torch.randn(B, model.noise_dim)
+    out = model(z, cond_cont, cond_cat, stage1_out)
+    assert out.shape == (B, SEC_DIM)
+
+
+def test_secondary_critic_output_shape():
+    B = 8
+    critic = _small_sec_critic()
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    x = torch.randn(B, SEC_DIM)
+    out = critic(x, cond_cont, cond_cat, stage1_out)
+    assert out.shape == (B,)
+
+
+def test_sample_secondaries_wgan_shape():
+    B = 5
+    model = _small_sec_generator()
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    n_sec_pred = torch.randint(0, K_MAX, (B,))
+    sec_cont, sec_phys, sec_valid = sample_secondaries_wgan(
+        model, cond_cont, cond_cat, stage1_out, n_sec_pred
+    )
+    assert sec_cont.shape == (B, K_MAX, 4)
+    assert sec_phys.shape == (B, K_MAX, 2)
+    assert sec_valid.shape == (B, K_MAX)
+
+
+# --- Losses ---
+
+
+def test_gradient_penalty_nonneg():
+    B = 8
+    critic = _small_critic()
+    cond_cont, cond_cat = _cond(B)
+    real = torch.randn(B, X_DIM)
+    fake = torch.randn(B, X_DIM)
+    gp = gradient_penalty(lambda x: critic(x, cond_cont, cond_cat), real, fake)
+    assert gp.item() >= 0.0
+    assert gp.shape == ()
+
+
+def test_gradient_penalty_masked():
+    B = 8
+    sec_critic = _small_sec_critic()
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    n_sec = torch.randint(0, K_MAX, (B,))
+    mask = _mask(B, n_sec)
+    real = torch.randn(B, SEC_DIM) * mask
+    fake = torch.randn(B, SEC_DIM) * mask
+    gp = gradient_penalty(
+        lambda x: sec_critic(x, cond_cont, cond_cat, stage1_out), real, fake, mask=mask
+    )
+    assert gp.item() >= 0.0
+
+
+def test_critic_loss_scalar_and_grad():
+    B = 8
+    critic = _small_critic()
+    cond_cont, cond_cat = _cond(B)
+    real = torch.randn(B, X_DIM)
+    fake = torch.randn(B, X_DIM)
+    loss = critic_loss(
+        lambda x: critic(x, cond_cont, cond_cat), real, fake.detach(), gp_weight=10.0
+    )
+    assert loss.shape == ()
+    loss.backward()
+    assert any(p.grad is not None for p in critic.parameters())
+
+
+def test_generator_loss_scalar_and_grad():
+    B = 4
+    generator = _small_generator()
+    critic = _small_critic()
+    cond_cont, cond_cat = _cond(B)
+    z = torch.randn(B, generator.noise_dim)
+    fake = generator(z, cond_cont, cond_cat)
+    loss = generator_loss(lambda x: critic(x, cond_cont, cond_cat), fake)
+    assert loss.shape == ()
+    loss.backward()
+    assert any(p.grad is not None for p in generator.parameters())