Implement Phase 2: secondary particle prediction

Two-stage factorisation: Stage 1 predicts 9D primary kinematics + n_sec
classification head (COND_DIM reduced to 8, dropping n_sec/e_sec inputs);
Stage 2 (SecondaryDecoder) generates K_MAX=15 secondary slots via masked
flow matching over (stick_logit, local_dir, type_emb) conditioned on Stage 1
output. Joint training with combined loss L_s1 + λ_nsec*L_nsec + λ_s2*L_s2.

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

giant/cli.py

@@ -409,7 +409,8 @@ def predict(
 
         cc = torch.from_numpy(cond_cont).float().to(_device)
         ck = torch.from_numpy(cond_cat).long().to(_device)
-        pred = sample_flow(model, cc, ck, steps=steps).cpu().numpy()  # normalised
+        pred, _n_sec = sample_flow(model, cc, ck, steps=steps)
+        pred = pred.cpu().numpy()  # normalised
 
         # Inverse-normalise → local frame, log-scaled scalars
         raw = tgt_norm.inverse_transform(pred)

giant/config.py

@@ -20,6 +20,8 @@ DEFAULT_CONFIG: dict = {
         "validate_every": 10,
         "validate_steps": 10,
         "warmup_epochs": 5,
+        "lambda_nsec": 0.1,
+        "lambda_s2": 1.0,
     },
     "model": {
         "hidden_dim": 256,

giant/constants.py

@@ -1,16 +1,24 @@
 X_DIM = 9
 
-# Conditioning continuous-feature width: pre_pos(3), log(pre_E)(1), pre_dir(3),
-# layer_id(1), n_sec(1), log(e_sec)(1). One wider than X_DIM because e_sec
-# (secondary energy) is a conditioning input in the energy-conservation PoC.
-COND_DIM = 10
+# Conditioning continuous-feature width (Phase 2): pre_pos(3), log(pre_E)(1),
+# pre_dir(3), layer_id(1). n_sec and log(e_sec) are removed — they are now
+# *outputs* predicted by Stage 1, not conditioning inputs.
+COND_DIM = 8
 
-# The two energy columns are additive-log-ratio (ALR) coordinates of the
-# deposit/secondary/post energy simplex (fractions of pre_E that sum to 1),
-# referenced to the post-energy fraction — see giant.data.transforms
-# .energy_simplex_encode/.energy_simplex_decode. They replace the former
-# independent log_delta_e / log_edep targets so energy conservation holds by
-# construction after decoding.
+# Maximum number of secondary slots. From data: max(n_sec)=14 in PbWO4 dataset;
+# K_MAX=15 covers it with one spare slot.
+K_MAX = 15
+
+# Per-slot secondary target dimension: 1 (stick-breaking logit) + 3 (local dir) +
+# EMB_DIM (continuous type embedding). EMB_DIM must match DenoisingMLP.emb_dim.
+# Default emb_dim=16 → SEC_SLOT_DIM=20.
+SEC_SLOT_DIM = 20  # 1 + 3 + 16
+EMB_DIM = 16       # must match model emb_dim default
+
+# Flattened Stage-2 target dimension
+SEC_DIM = K_MAX * SEC_SLOT_DIM  # 15 * 20 = 300
+
+# Stage-1 9D target names (unchanged from energy-conservation PoC)
 LOCAL_TARGET_NAMES = [
     "log_step_length",
     "edep_logit",

giant/data/dataset.py

@@ -4,53 +4,12 @@ from pathlib import Path
 
 import numpy as np
 import torch
-from torch.utils.data import Dataset, IterableDataset
+from torch.utils.data import IterableDataset
 
 from giant.data.loader import iter_file_chunks
 from giant.data.transforms import Normalizer, build_features
 
 
-class StepsDataset(Dataset):
-    def __init__(
-        self,
-        cond_cont: np.ndarray,
-        cond_cat: np.ndarray,
-        target: np.ndarray,
-    ) -> None:
-        self.cond_cont = torch.from_numpy(cond_cont).float()
-        self.cond_cat = torch.from_numpy(cond_cat).long()
-        self.target = torch.from_numpy(target).float()
-
-    def __len__(self) -> int:
-        return len(self.target)
-
-    def __getitem__(self, index):
-        return self.cond_cont[index], self.cond_cat[index], self.target[index]
-
-
-def train_val_split(
-    data: dict,
-    cond_cont: np.ndarray,
-    cond_cat: np.ndarray,
-    target: np.ndarray,
-    val_fraction: float = 0.1,
-    seed: int = 42,
-) -> tuple[StepsDataset, StepsDataset]:
-    rng = np.random.default_rng(seed)
-    unique_events = np.unique(data["event_id"])
-    rng.shuffle(unique_events)
-    n_val = max(1, int(len(unique_events) * val_fraction))
-    val_events = set(unique_events[:n_val].tolist())
-
-    val_mask = np.array([e in val_events for e in data["event_id"]])
-    train_mask = ~val_mask
-
-    return (
-        StepsDataset(cond_cont[train_mask], cond_cat[train_mask], target[train_mask]),
-        StepsDataset(cond_cont[val_mask], cond_cat[val_mask], target[val_mask]),
-    )
-
-
 def make_event_split(
     all_event_ids: np.ndarray,
     val_fraction: float = 0.1,
@@ -75,6 +34,16 @@ class StreamingStepsDataset(IterableDataset):
     Yields whole batches (use with `DataLoader(..., batch_size=None)`)
     rather than single rows, so the batch is assembled with vectorized
     numpy slicing instead of a per-row Python loop in the default collate.
+
+    Each batch is a tuple:
+        (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx)
+    where:
+        cond_cont:   (B, COND_DIM) float32
+        cond_cat:    (B, 2) int64
+        target_s1:   (B, 9) float32 — normalised Stage-1 primary target
+        n_sec:       (B,) int64 — true secondary count per step
+        sec_cont:    (B, K_MAX, 4) float32 — [stick_logit, local_dir] per slot
+        sec_pdg_idx: (B, K_MAX) int64 — PDG model-index per secondary slot
     """
 
     def __init__(
@@ -91,13 +60,12 @@ class StreamingStepsDataset(IterableDataset):
     ) -> None:
         self.files = list(files)
         self.split_events = split_events
-        self._events_arr = np.array(sorted(split_events))  # for np.isin
+        self._events_arr = np.array(sorted(split_events))
         self.pdg_map = pdg_map
         self.mat_map = mat_map
         self.cond_normalizer = cond_normalizer
         self.target_normalizer = target_normalizer
         self.batch_size = batch_size
-        # Buffer must hold at least one batch or we could never emit one.
         self.shuffle_buffer = max(shuffle_buffer, batch_size)
         self.shuffle = shuffle
 
@@ -114,6 +82,9 @@ class StreamingStepsDataset(IterableDataset):
         buf_cont: list[np.ndarray] = []
         buf_cat: list[np.ndarray] = []
         buf_tgt: list[np.ndarray] = []
+        buf_nsec: list[np.ndarray] = []
+        buf_sec: list[np.ndarray] = []
+        buf_spdg: list[np.ndarray] = []
         buf_n = 0
 
         for path in files:
@@ -123,43 +94,57 @@ class StreamingStepsDataset(IterableDataset):
                     continue
                 chunk = {k: v[mask] for k, v in chunk.items()}
 
-                cond_cont, cond_cat, target, _, _ = build_features(
-                    chunk,
-                    self.pdg_map,
-                    self.mat_map,
-                    cond_normalizer=self.cond_normalizer,
-                    target_normalizer=self.target_normalizer,
+                cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx, _, _ = (
+                    build_features(
+                        chunk,
+                        self.pdg_map,
+                        self.mat_map,
+                        cond_normalizer=self.cond_normalizer,
+                        target_normalizer=self.target_normalizer,
+                    )
                 )
                 buf_cont.append(cond_cont)
                 buf_cat.append(cond_cat)
-                buf_tgt.append(target)
+                buf_tgt.append(target_s1)
+                buf_nsec.append(n_sec)
+                buf_sec.append(sec_cont)
+                buf_spdg.append(sec_pdg_idx)
                 buf_n += len(cond_cont)
 
                 if buf_n >= self.shuffle_buffer:
-                    buf_cont, buf_cat, buf_tgt, buf_n = yield from self._flush(
-                        buf_cont, buf_cat, buf_tgt, final=False
+                    buf_cont, buf_cat, buf_tgt, buf_nsec, buf_sec, buf_spdg, buf_n = (
+                        yield from self._flush(
+                            buf_cont, buf_cat, buf_tgt, buf_nsec, buf_sec, buf_spdg,
+                            final=False,
+                        )
                     )
 
         if buf_n > 0:
-            yield from self._flush(buf_cont, buf_cat, buf_tgt, final=True)
+            yield from self._flush(
+                buf_cont, buf_cat, buf_tgt, buf_nsec, buf_sec, buf_spdg, final=True
+            )
 
     def _flush(
         self,
         buf_cont: list[np.ndarray],
         buf_cat: list[np.ndarray],
         buf_tgt: list[np.ndarray],
+        buf_nsec: list[np.ndarray],
+        buf_sec: list[np.ndarray],
+        buf_spdg: list[np.ndarray],
         final: bool,
     ):
-        """Yield full batches of `batch_size`; carry any remainder back to the caller.
-
-        All batching is done via vectorized numpy slicing (no per-row Python loop).
-        """
         cont = np.concatenate(buf_cont)
         cat = np.concatenate(buf_cat)
         tgt = np.concatenate(buf_tgt)
+        nsec = np.concatenate(buf_nsec)
+        sec = np.concatenate(buf_sec)
+        spdg = np.concatenate(buf_spdg)
+
         if self.shuffle:
             idx = np.random.permutation(len(cont))
             cont, cat, tgt = cont[idx], cat[idx], tgt[idx]
+            nsec, sec, spdg = nsec[idx], sec[idx], spdg[idx]
 
         bs = self.batch_size
         n = len(cont)
@@ -170,9 +155,16 @@ class StreamingStepsDataset(IterableDataset):
                 torch.from_numpy(cont[start:end]).float(),
                 torch.from_numpy(cat[start:end]).long(),
                 torch.from_numpy(tgt[start:end]).float(),
+                torch.from_numpy(nsec[start:end]).long(),
+                torch.from_numpy(sec[start:end]).float(),
+                torch.from_numpy(spdg[start:end]).long(),
             )
 
         if final:
-            return [], [], [], 0
+            return [], [], [], [], [], [], 0
         rem = n_full * bs
-        return [cont[rem:]], [cat[rem:]], [tgt[rem:]], n - rem
+        return (
+            [cont[rem:]], [cat[rem:]], [tgt[rem:]],
+            [nsec[rem:]], [sec[rem:]], [spdg[rem:]],
+            n - rem,
+        )

giant/data/loader.py

@@ -40,8 +40,52 @@ def find_parquet_files(path: str | Path) -> list[Path]:
     return [p]
 
 
+def _pad_list_col(series: pd.Series, K: int, fill: float = 0.0) -> np.ndarray:
+    """Pad / truncate a list-valued Series to fixed width K → (N, K) float32."""
+    out = np.full((len(series), K), fill, dtype=np.float32)
+    for i, lst in enumerate(series):
+        if lst is not None and len(lst) > 0:
+            n = min(len(lst), K)
+            out[i, :n] = lst[:n]
+    return out
+
+
+def _pad_list_col_int(series: pd.Series, K: int, fill: int = 0) -> np.ndarray:
+    """Pad / truncate a list-valued integer Series to fixed width K → (N, K) int64."""
+    out = np.full((len(series), K), fill, dtype=np.int64)
+    for i, lst in enumerate(series):
+        if lst is not None and len(lst) > 0:
+            n = min(len(lst), K)
+            out[i, :n] = lst[:n]
+    return out
+
+
+def _pad_dir_col(
+    dx: pd.Series, dy: pd.Series, dz: pd.Series, K: int
+) -> np.ndarray:
+    """Pad three list-valued direction columns → (N, K, 3) float32.
+
+    Padding direction defaults to (0,0,1) (forward) so it is a valid unit vector.
+    """
+    N = len(dx)
+    out = np.zeros((N, K, 3), dtype=np.float32)
+    out[:, :, 2] = 1.0
+    for i in range(N):
+        lx, ly, lz = dx.iloc[i], dy.iloc[i], dz.iloc[i]
+        if lx is not None and len(lx) > 0:
+            n = min(len(lx), K)
+            out[i, :n, 0] = lx[:n]
+            out[i, :n, 1] = ly[:n]
+            out[i, :n, 2] = lz[:n]
+    return out
+
+
 def _df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
-    return {
+    from giant.constants import K_MAX
+
+    has_sec_lists = "sec_E_list" in df.columns
+
+    d: dict[str, np.ndarray] = {
         "event_id": df["event_id"].to_numpy(),
         "pdg": df["pdg"].to_numpy(dtype=np.int32),
         "pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
@@ -59,6 +103,15 @@ def _df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
         "post_pos": df[["post_x", "post_y", "post_z"]].to_numpy(dtype=np.float32),
     }
 
+    if has_sec_lists:
+        d["sec_E_list"] = _pad_list_col(df["sec_E_list"], K_MAX)
+        d["sec_pdg_list"] = _pad_list_col_int(df["sec_pdg_list"], K_MAX)
+        d["sec_dir_list"] = _pad_dir_col(
+            df["sec_dx_list"], df["sec_dy_list"], df["sec_dz_list"], K_MAX
+        )
+
+    return d
+
 
 def load_steps(path: str | Path) -> dict[str, np.ndarray]:
     return _df_to_dict(pd.read_parquet(path))

giant/data/transforms.py

@@ -227,6 +227,114 @@ def inv_local_frame_rotation(
     )
 
 
+_STICK_LOGIT_CLIP = 10.0  # logit value used for the last valid secondary slot
+
+
+def encode_secondaries(
+    sec_E_list: np.ndarray,
+    sec_dir_list: np.ndarray,
+    sec_valid: np.ndarray,
+    e_sec: np.ndarray,
+    pre_dir: np.ndarray,
+) -> np.ndarray:
+    """Encode per-secondary attributes into continuous per-slot targets.
+
+    Secondaries must already be sorted descending by energy (as stored in the
+    parquet). Returns sec_cont of shape (N, K_MAX, 4):
+      slot[i] = [stick_break_logit, local_dir_x, local_dir_y, local_dir_z]
+
+    Stick-breaking logit: for slot i, f_i = E_i / remaining_budget, where
+    remaining_budget = e_sec - sum(E_0..E_{i-1}). The logit is log(f/(1-f)),
+    clipped to ±_STICK_LOGIT_CLIP. The last valid slot gets +_STICK_LOGIT_CLIP
+    (takes the full remaining budget). Padding slots get 0.
+
+    sec_pdg_idx (integer) is not processed here — kept separate so the loss
+    function can look up the embedding table at training time.
+    """
+    N, K = sec_E_list.shape
+    e_sec = np.asarray(e_sec, dtype=np.float64)
+
+    stick_logits = np.zeros((N, K), dtype=np.float32)
+    for i in range(K):
+        if i == 0:
+            remaining = e_sec
+        else:
+            remaining = np.maximum(e_sec - sec_E_list[:, :i].sum(axis=1), _EPS)
+        f = np.clip(sec_E_list[:, i].astype(np.float64) / remaining, _EPS, 1.0 - _EPS)
+        logit = np.log(f / (1.0 - f)).astype(np.float32)
+        # Last valid slot: give it the full remaining budget
+        is_last = sec_valid[:, i] & ~(sec_valid[:, i + 1] if i + 1 < K else np.zeros(N, dtype=bool))
+        logit = np.where(is_last, _STICK_LOGIT_CLIP, logit)
+        logit = np.where(sec_valid[:, i], np.clip(logit, -_STICK_LOGIT_CLIP, _STICK_LOGIT_CLIP), 0.0)
+        stick_logits[:, i] = logit.astype(np.float32)
+
+    # Rotate each slot's direction into the local frame of the primary.
+    # pre_dir is broadcast across all K slots.
+    dir_local = np.zeros((N, K, 3), dtype=np.float32)
+    for i in range(K):
+        # Only rotate valid slots; leave padded slots as (0,0,1) or whatever.
+        valid_mask = sec_valid[:, i]
+        if valid_mask.any():
+            dir_local[valid_mask, i] = local_frame_rotation(
+                pre_dir[valid_mask], sec_dir_list[valid_mask, i]
+            )
+
+    sec_cont = np.concatenate(
+        [stick_logits[:, :, None], dir_local], axis=-1
+    )  # (N, K, 4)
+    return sec_cont.astype(np.float32)
+
+
+def decode_secondaries(
+    sec_cont: np.ndarray,
+    sec_pdg_pred: np.ndarray,
+    n_sec: np.ndarray,
+    e_sec: np.ndarray,
+    pre_dir: np.ndarray,
+    pdg_map_inv: dict[int, int],
+) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    """Inverse of encode_secondaries: continuous targets → physical secondary attrs.
+
+    sec_cont: (N, K_MAX, 4) — [stick_logit, local_dir_x, local_dir_y, local_dir_z]
+    sec_pdg_pred: (N, K_MAX) integer PDG indices (from nearest-neighbor snap)
+    n_sec: (N,) integer secondary counts
+    e_sec: (N,) total secondary energy budget [MeV]
+    pre_dir: (N, 3) pre-step world-frame direction
+    pdg_map_inv: maps model index → PDG code
+
+    Returns (sec_E, sec_dir_world, sec_pdg_code, sec_valid) each shape (N, K_MAX).
+    """
+    N, K, _ = sec_cont.shape
+    stick_logits = sec_cont[:, :, 0]       # (N, K)
+    dir_local = sec_cont[:, :, 1:]         # (N, K, 3)
+
+    fractions = 1.0 / (1.0 + np.exp(-stick_logits.astype(np.float64)))
+
+    sec_E = np.zeros((N, K), dtype=np.float32)
+    e_sec = np.asarray(e_sec, dtype=np.float64)
+    remaining = e_sec.copy()
+    for i in range(K):
+        sec_E[:, i] = (fractions[:, i] * remaining).astype(np.float32)
+        remaining = np.maximum(remaining - sec_E[:, i].astype(np.float64), 0.0)
+
+    sec_valid = np.arange(K)[None, :] < n_sec[:, None]  # (N, K)
+
+    sec_dir_world = np.zeros((N, K, 3), dtype=np.float32)
+    for i in range(K):
+        valid = sec_valid[:, i]
+        if valid.any():
+            sec_dir_world[valid, i] = inv_local_frame_rotation(
+                pre_dir[valid], dir_local[valid, i]
+            )
+
+    sec_pdg_code = np.array(
+        [[pdg_map_inv.get(int(sec_pdg_pred[n, i]), 0) for i in range(K)] for n in range(N)],
+        dtype=np.int32,
+    )
+
+    return sec_E, sec_dir_world, sec_pdg_code, sec_valid
+
+
 def build_cond_features(
     data: dict[str, np.ndarray],
     pdg_map: dict[int, int],
@@ -240,8 +348,6 @@ def build_cond_features(
             log_transform(data["pre_E"]),
             data["pre_dir"],
             data["layer_id"].astype(np.float32),
-            data["n_sec"].astype(np.float32),
-            log_transform(data["e_sec"]),
         ]
     ).astype(np.float32)
 
@@ -262,11 +368,17 @@ def build_features(
     cond_normalizer: Normalizer | None = None,
     target_normalizer: Normalizer | None = None,
     fit: bool = False,
-) -> tuple[np.ndarray, np.ndarray, np.ndarray, Normalizer | None, Normalizer | None]:
-    """Assemble (cond_cont, cond_cat, target) arrays ready for StepsDataset.
+) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, Normalizer | None, Normalizer | None]:
+    """Assemble (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx) arrays.
 
-    When fit=True, new Normalizers are fitted on the supplied arrays.
+    target_s1: (N, 9) Stage-1 primary post-step target (unchanged from Phase 1)
+    n_sec: (N,) integer secondary counts (target for n_sec head)
+    sec_cont: (N, K_MAX, 4) continuous secondary targets [stick_logit, dir_local]
+    sec_pdg_idx: (N, K_MAX) integer PDG model-indices; used to look up embedding
+                 targets in the training loop
     """
+    from giant.constants import K_MAX
+
     post_dir_local = local_frame_rotation(data["pre_dir"], data["post_dir"])
     travel_dir_local = local_frame_rotation(
         data["pre_dir"], travel_direction(data["pre_pos"], data["post_pos"])
@@ -274,9 +386,9 @@ def build_features(
 
     energy_z = energy_simplex_encode(
         data["edep"], data["e_sec"], data["post_E"], data["pre_E"]
-    )  # (N, 2): ALR coords of the deposit/secondary/post energy simplex
+    )  # (N, 2)
 
-    target = np.column_stack(
+    target_s1 = np.column_stack(
         [
             log_transform(data["step_length"]),
             energy_z,
@@ -285,28 +397,44 @@ def build_features(
         ]
     ).astype(np.float32)  # (N, 9)
 
+    # Phase 2: conditioning drops n_sec and log(e_sec)
     cond_cont = np.column_stack(
         [
             data["pre_pos"],
             log_transform(data["pre_E"]),
             data["pre_dir"],
             data["layer_id"].astype(np.float32),
-            data["n_sec"].astype(np.float32),
-            log_transform(data["e_sec"]),
         ]
-    ).astype(np.float32)  # (N, COND_DIM)
+    ).astype(np.float32)  # (N, COND_DIM=8)
 
     pdg_idx = np.array([pdg_map[int(p)] for p in data["pdg"]], dtype=np.int64)
     mat_idx = np.array([mat_map[str(m)] for m in data["material"]], dtype=np.int64)
     cond_cat = np.column_stack([pdg_idx, mat_idx])  # (N, 2)
 
+    n_sec = data["n_sec"].astype(np.int64)  # (N,)
+
+    # Secondary continuous targets
+    sec_E_list = data.get("sec_E_list")
+    sec_dir_list = data.get("sec_dir_list")
+    sec_pdg_idx = data.get("sec_pdg_idx")
+
+    if sec_E_list is not None and sec_dir_list is not None and sec_pdg_idx is not None:
+        sec_valid = np.arange(K_MAX)[None, :] < n_sec[:, None]  # (N, K_MAX)
+        sec_cont = encode_secondaries(
+            sec_E_list, sec_dir_list, sec_valid, data["e_sec"], data["pre_dir"]
+        )  # (N, K_MAX, 4)
+    else:
+        N = len(n_sec)
+        sec_cont = np.zeros((N, K_MAX, 4), dtype=np.float32)
+        sec_pdg_idx = np.zeros((N, K_MAX), dtype=np.int64)
+
     if fit:
         cond_normalizer = Normalizer().fit(cond_cont)
-        target_normalizer = Normalizer().fit(target)
+        target_normalizer = Normalizer().fit(target_s1)
 
     if cond_normalizer is not None:
         cond_cont = cond_normalizer.transform(cond_cont)
     if target_normalizer is not None:
-        target = target_normalizer.transform(target)
+        target_s1 = target_normalizer.transform(target_s1)
 
-    return cond_cont, cond_cat, target, cond_normalizer, target_normalizer
+    return cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx, cond_normalizer, target_normalizer

giant/model/network.py

@@ -3,7 +3,7 @@ import math
 import torch
 import torch.nn as nn
 
-from giant.constants import COND_DIM, X_DIM
+from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM
 
 
 class SinusoidalEmbedding(nn.Module):
@@ -70,6 +70,12 @@ class ResBlock(nn.Module):
 
 
 class DenoisingMLP(nn.Module):
+    """Stage-1 model: predicts the 9D primary post-step vector field + n_sec logits.
+
+    The n_sec head runs on the condition encoding only (no diffusion noise),
+    so it can be called at inference time independently via `predict_n_sec`.
+    """
+
     def __init__(
         self,
         pdg_vocab: int,
@@ -81,6 +87,7 @@ class DenoisingMLP(nn.Module):
         cond_out_dim: int = 128,
         x_dim: int = X_DIM,
         dropout: float = 0.1,
+        k_max: int = K_MAX,
     ) -> None:
         super().__init__()
         self.time_emb = SinusoidalEmbedding(time_dim)
@@ -99,6 +106,13 @@ class DenoisingMLP(nn.Module):
             ]
         )
         self.out_proj = nn.Linear(hidden_dim, x_dim)
+        # Predicts n_sec as classification over {0, 1, ..., k_max}.
+        # Applied to the condition encoding (not the diffused latent).
+        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,
@@ -107,9 +121,117 @@ class DenoisingMLP(nn.Module):
         cond_cont: torch.Tensor,
         cond_cat: torch.Tensor,
     ) -> torch.Tensor:
-        t_emb = self.time_emb(t)  # (B, time_dim)
+        t_emb = self.time_emb(t)               # (B, time_dim)
         c_emb = self.cond_enc(cond_cont, cond_cat)  # (B, cond_out_dim)
-        cond = torch.cat([t_emb, c_emb], dim=-1)  # (B, time_dim+cond_out_dim)
+        cond = torch.cat([t_emb, c_emb], dim=-1)
+        x = self.input_proj(x_t)
+        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)
+
+    def pdg_embedding_weight(self) -> torch.Tensor:
+        """Return the PDG embedding table weights for secondary type targets."""
+        return self.cond_enc.pdg_emb.weight
+
+
+class SecondaryConditionEncoder(nn.Module):
+    """Encodes pre-step conditioning + Stage-1 output for the secondary decoder."""
+
+    def __init__(
+        self,
+        pdg_vocab: int,
+        mat_vocab: int,
+        emb_dim: int = 16,
+        cond_out_dim: int = 128,
+        stage1_dim: int = X_DIM,
+        stage1_proj_dim: int = 64,
+        out_dim: int = 128,
+    ) -> None:
+        super().__init__()
+        self.base = ConditionEncoder(
+            pdg_vocab=pdg_vocab,
+            mat_vocab=mat_vocab,
+            emb_dim=emb_dim,
+            out_dim=cond_out_dim,
+        )
+        self.stage1_proj = nn.Linear(stage1_dim, stage1_proj_dim)
+        fused_dim = cond_out_dim + stage1_proj_dim
+        self.fuse = nn.Sequential(
+            nn.Linear(fused_dim, out_dim),
+            nn.SiLU(),
+        )
+
+    def forward(
+        self,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_out: torch.Tensor,
+    ) -> torch.Tensor:
+        base = self.base(cond_cont, cond_cat)                    # (B, cond_out_dim)
+        s1 = self.stage1_proj(stage1_out).tanh()                 # (B, stage1_proj_dim)
+        return self.fuse(torch.cat([base, s1], dim=-1))           # (B, out_dim)
+
+
+class SecondaryDecoder(nn.Module):
+    """Stage-2 model: predicts vector field over K_MAX secondary slots simultaneously.
+
+    Each slot encodes (stick_break_logit, local_dir_3D, type_emb) for one
+    secondary ordered by descending energy. Padded slots are masked from loss.
+    """
+
+    def __init__(
+        self,
+        pdg_vocab: int,
+        mat_vocab: int,
+        hidden_dim: int = 256,
+        n_blocks: int = 6,
+        emb_dim: int = 16,
+        time_dim: int = 64,
+        cond_out_dim: int = 128,
+        stage1_proj_dim: int = 64,
+        sec_dim: int = SEC_DIM,
+        dropout: float = 0.1,
+    ) -> None:
+        super().__init__()
+        self.time_emb = SinusoidalEmbedding(time_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,
+        )
+        merged_cond_dim = time_dim + cond_out_dim
+        self.input_proj = nn.Linear(sec_dim, hidden_dim)
+        self.blocks = nn.ModuleList(
+            [
+                ResBlock(hidden_dim, merged_cond_dim, dropout=dropout)
+                for _ in range(n_blocks)
+            ]
+        )
+        self.out_proj = nn.Linear(hidden_dim, sec_dim)
+
+    def forward(
+        self,
+        x_t: torch.Tensor,
+        t: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_out: torch.Tensor,
+    ) -> torch.Tensor:
+        t_emb = self.time_emb(t)
+        c_emb = self.cond_enc(cond_cont, cond_cat, stage1_out)
+        cond = torch.cat([t_emb, c_emb], dim=-1)
         x = self.input_proj(x_t)
         for block in self.blocks:
             x = block(x, cond)

giant/model/schedule.py

@@ -69,3 +69,37 @@ def flow_matching_loss(
     u_t = x1 - x0
     v_t = model(x_t, t, cond_cont, cond_cat)
     return F.mse_loss(v_t, u_t)
+
+
+def flow_matching_loss_secondary(
+    model: torch.nn.Module,
+    x1: torch.Tensor,
+    cond_cont: torch.Tensor,
+    cond_cat: torch.Tensor,
+    stage1_out: torch.Tensor,
+    sec_mask: torch.Tensor,
+) -> torch.Tensor:
+    """Flow matching loss for the secondary decoder with per-slot masking.
+
+    x1: (B, SEC_DIM) — flattened secondary target (stick_logit, dir, type_emb)
+    sec_mask: (B, K_MAX) bool — True for valid secondary slots
+
+    Only valid-slot dimensions contribute to the loss; padded slots are zeroed
+    before averaging, so the loss is not diluted by empty slots.
+    """
+    from giant.constants import SEC_SLOT_DIM
+
+    B = x1.size(0)
+    t = torch.rand(B, device=x1.device)
+    x0 = torch.randn_like(x1)
+    x_t = (1.0 - t.view(-1, 1)) * x0 + t.view(-1, 1) * x1
+    u_t = x1 - x0
+    v_t = model(x_t, t, cond_cont, cond_cat, stage1_out)
+
+    # Expand mask: (B, K_MAX) → (B, K_MAX * SEC_SLOT_DIM)
+    mask_expanded = (
+        sec_mask.float().unsqueeze(-1).expand(-1, -1, SEC_SLOT_DIM).reshape(B, -1)
+    )
+    err = (v_t - u_t) ** 2
+    denom = mask_expanded.sum().clamp(min=1)
+    return (err * mask_expanded).sum() / denom

giant/pipeline.py

@@ -5,7 +5,7 @@ import torch
 from torch.utils.data import DataLoader
 
 from giant import config
-from giant.constants import COND_DIM, X_DIM
+from giant.constants import COND_DIM, EMB_DIM, K_MAX, SEC_SLOT_DIM, X_DIM
 from giant.data.loader import (
     find_parquet_files,
     load_event_ids,
@@ -14,7 +14,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 DenoisingMLP
+from giant.model.network import DenoisingMLP, SecondaryDecoder
 from giant.train import train as run_training
 
 
@@ -63,9 +63,11 @@ def run_train_job(
             if not mask.any():
                 continue
             chunk_tr = {k: v[mask] for k, v in chunk.items()}
-            cond_cont, _, target, _, _ = build_features(chunk_tr, pdg_map, mat_map)
+            cond_cont, _, target_s1, _n_sec, _sec_cont, _sec_pdg, _, _ = build_features(
+                chunk_tr, pdg_map, mat_map
+            )
             cond_acc.update(cond_cont)
-            tgt_acc.update(target)
+            tgt_acc.update(target_s1)
     cond_norm = cond_acc.to_normalizer()
     tgt_norm = tgt_acc.to_normalizer()
 
@@ -91,8 +93,6 @@ def run_train_job(
         shuffle=False,
     )
 
-    # Dataset yields whole batches already, so batch_size=None tells DataLoader
-    # to pass them through instead of re-collating row-by-row in Python.
     pin = device.type == "cuda"
     train_loader = DataLoader(
         train_ds,
@@ -107,15 +107,34 @@ def run_train_job(
         pin_memory=pin,
     )
 
-    model = DenoisingMLP(
+    emb_dim = m.get("emb_dim", EMB_DIM)
+    # SEC_SLOT_DIM must match constants (1 stick + 3 dir + emb_dim)
+    assert SEC_SLOT_DIM == 1 + 3 + emb_dim, (
+        f"SEC_SLOT_DIM={SEC_SLOT_DIM} must equal 1+3+emb_dim={1+3+emb_dim}; "
+        "update giant/constants.py if emb_dim changed"
+    )
+
+    stage1_model = DenoisingMLP(
         pdg_vocab=len(pdg_map),
         mat_vocab=len(mat_map),
         hidden_dim=m["hidden_dim"],
         n_blocks=m["n_blocks"],
-        emb_dim=m["emb_dim"],
+        emb_dim=emb_dim,
+        dropout=m["dropout"],
+        k_max=K_MAX,
+    )
+    sec_decoder = SecondaryDecoder(
+        pdg_vocab=len(pdg_map),
+        mat_vocab=len(mat_map),
+        hidden_dim=m["hidden_dim"],
+        n_blocks=m["n_blocks"],
+        emb_dim=emb_dim,
         dropout=m["dropout"],
     )
-    echo(f"model: {sum(p.numel() for p in model.parameters()):,} parameters")
+    echo(
+        f"stage1: {sum(p.numel() for p in stage1_model.parameters()):,} parameters  |  "
+        f"sec_decoder: {sum(p.numel() for p in sec_decoder.parameters()):,} parameters"
+    )
 
     out_dir.mkdir(parents=True, exist_ok=True)
     meta = config.build_run_meta(
@@ -134,12 +153,15 @@ def run_train_job(
         "mat_vocab": len(mat_map),
         "hidden_dim": m["hidden_dim"],
         "n_blocks": m["n_blocks"],
-        "emb_dim": m["emb_dim"],
+        "emb_dim": emb_dim,
         "dropout": m["dropout"],
+        "k_max": K_MAX,
+        "sec_slot_dim": SEC_SLOT_DIM,
     }
 
     run_training(
-        model=model,
+        stage1_model=stage1_model,
+        sec_decoder=sec_decoder,
         train_loader=train_loader,
         val_loader=val_loader,
         mode=t["mode"],
@@ -148,6 +170,8 @@ def run_train_job(
         warmup_epochs=t["warmup_epochs"],
         device=device,
         out_dir=out_dir,
+        lambda_nsec=t.get("lambda_nsec", 0.1),
+        lambda_s2=t.get("lambda_s2", 1.0),
         normalizer_dict={"cond": cond_norm.to_dict(), "target": tgt_norm.to_dict()},
         pdg_map={str(k): v for k, v in pdg_map.items()},
         mat_map={str(k): v for k, v in mat_map.items()},

giant/sample.py

@@ -1,6 +1,6 @@
 import torch
 
-from giant.constants import X_DIM
+from giant.constants import K_MAX, SEC_DIM, SEC_SLOT_DIM, X_DIM
 
 
 @torch.no_grad()
@@ -9,8 +9,13 @@ def sample_flow(
     cond_cont: torch.Tensor,
     cond_cat: torch.Tensor,
     steps: int = 10,
-) -> torch.Tensor:
-    """Euler integration of the learned vector field from t=0 to t=1."""
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Euler integration of the Stage-1 vector field from t=0 to t=1.
+
+    Returns (primary_sample, n_sec_pred):
+        primary_sample: (B, X_DIM) — normalised 9D primary post-step output
+        n_sec_pred: (B,) int64 — predicted secondary count
+    """
     model.eval()
     B = cond_cont.size(0)
     device = cond_cont.device
@@ -20,7 +25,63 @@ def sample_flow(
         t = torch.full((B,), i * dt, device=device)
         v = model(x, t, cond_cont, cond_cat)
         x = x + v * dt
-    return x
+    n_sec_logits = model.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(
+    sec_decoder: torch.nn.Module,
+    cond_cont: torch.Tensor,
+    cond_cat: torch.Tensor,
+    stage1_out: torch.Tensor,
+    n_sec_pred: torch.Tensor,
+    steps: int = 10,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Euler integration of the Stage-2 vector field; return raw slot outputs.
+
+    n_sec_pred: (B,) int64 — number of valid secondaries per step
+
+    Returns (sec_cont, sec_type_emb, sec_valid):
+        sec_cont:     (B, K_MAX, 4) — [stick_logit, local_dir_x, local_dir_y, local_dir_z]
+        sec_type_emb: (B, K_MAX, emb_dim) — predicted type embedding per slot
+        sec_valid:    (B, K_MAX) bool — True for slots i < n_sec_pred
+    """
+    sec_decoder.eval()
+    B = cond_cont.size(0)
+    device = cond_cont.device
+
+    x = torch.randn(B, SEC_DIM, device=device)
+    dt = 1.0 / steps
+    for i in range(steps):
+        t = torch.full((B,), i * dt, device=device)
+        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_type_emb = x_slots[:, :, 4:]
+
+    sec_valid = (
+        torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(1)
+    )
+    return sec_cont, sec_type_emb, sec_valid
+
+
+def snap_type_to_pdg_idx(
+    sec_type_emb: torch.Tensor,
+    pdg_emb_weight: torch.Tensor,
+) -> torch.Tensor:
+    """Nearest-neighbour snap: predicted type embedding → PDG model-index.
+
+    sec_type_emb: (B, K_MAX, emb_dim)
+    Returns (B, K_MAX) int64 with model-indices.
+    """
+    B, K, D = sec_type_emb.shape
+    flat = sec_type_emb.reshape(-1, D)
+    dists = torch.cdist(flat.float(), pdg_emb_weight.float())
+    return dists.argmin(dim=-1).reshape(B, K)
 
 
 @torch.no_grad()
@@ -29,8 +90,8 @@ def sample_ddpm(
     cond_cont: torch.Tensor,
     cond_cat: torch.Tensor,
     schedule,
-) -> torch.Tensor:
-    """Full DDPM ancestral sampling (T reverse steps)."""
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Full DDPM ancestral sampling (T reverse steps). Returns (sample, n_sec_pred)."""
     model.eval()
     B = cond_cont.size(0)
     device = cond_cont.device
@@ -46,7 +107,9 @@ def sample_ddpm(
         x = (1.0 / alpha.sqrt()) * (
             x - (1.0 - alpha) / (1.0 - alpha_bar).sqrt() * eps_pred
         ) + beta.sqrt() * z
-    return x
+    n_sec_logits = model.predict_n_sec(cond_cont, cond_cat)
+    n_sec_pred = n_sec_logits.argmax(dim=-1)
+    return x, n_sec_pred
 
 
 @torch.no_grad()
@@ -56,8 +119,8 @@ def sample_ddim(
     cond_cat: torch.Tensor,
     schedule,
     steps: int = 50,
-) -> torch.Tensor:
-    """DDIM deterministic sampling (Song et al. 2020) with `steps` substeps."""
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """DDIM deterministic sampling (Song et al. 2020). Returns (sample, n_sec_pred)."""
     model.eval()
     B = cond_cont.size(0)
     device = cond_cont.device
@@ -75,4 +138,6 @@ def sample_ddim(
             ab_prev = torch.ones(1, device=device)
         x0_pred = (x - (1.0 - ab_t).sqrt() * eps_pred) / ab_t.sqrt()
         x = ab_prev.sqrt() * x0_pred + (1.0 - ab_prev).sqrt() * eps_pred
-    return x
+    n_sec_logits = model.predict_n_sec(cond_cont, cond_cat)
+    n_sec_pred = n_sec_logits.argmax(dim=-1)
+    return x, n_sec_pred

giant/train.py

@@ -8,14 +8,31 @@ from types import FrameType
 from typing import Callable
 
 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.model.schedule import CosineSchedule, flow_matching_loss
+from giant.model.schedule import (
+    CosineSchedule,
+    flow_matching_loss,
+    flow_matching_loss_secondary,
+)
 from giant.validate import validate_marginals
 
-_METRICS_FIELDS = ["epoch", "train_loss", "val_loss", "lr", "epoch_time_s"]
+_METRICS_FIELDS = [
+    "epoch",
+    "train_loss",
+    "train_loss_s1",
+    "train_loss_nsec",
+    "train_loss_s2",
+    "val_loss",
+    "val_loss_s1",
+    "val_loss_nsec",
+    "val_loss_s2",
+    "lr",
+    "epoch_time_s",
+]
 
 _CATCHABLE_SIGNALS = (signal.SIGINT, signal.SIGTERM)
 
@@ -57,8 +74,80 @@ class _GracefulShutdown:
         )
 
 
+def _build_sec_x1(
+    sec_cont: torch.Tensor,
+    sec_pdg_idx: torch.Tensor,
+    pdg_emb_weight: torch.Tensor,
+) -> torch.Tensor:
+    """Assemble the Stage-2 flow target by appending type embeddings.
+
+    sec_cont: (B, K_MAX, 4) — [stick_logit, dir_local]
+    sec_pdg_idx: (B, K_MAX) — integer PDG model-indices
+    pdg_emb_weight: (pdg_vocab, emb_dim) — live embedding table weights
+
+    Returns (B, SEC_DIM) = (B, K_MAX * (4 + emb_dim)).
+    """
+    type_emb = pdg_emb_weight[sec_pdg_idx]           # (B, K_MAX, emb_dim)
+    x1_s2 = torch.cat([sec_cont, type_emb], dim=-1)  # (B, K_MAX, 4+emb_dim)
+    return x1_s2.flatten(1)                           # (B, SEC_DIM)
+
+
+def _compute_losses(
+    stage1_model: torch.nn.Module,
+    sec_decoder: torch.nn.Module,
+    batch: tuple,
+    mode: str,
+    ddpm_schedule,
+    device: torch.device,
+    lambda_nsec: float,
+    lambda_s2: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Compute (total_loss, L_s1, L_nsec, L_s2) for one batch."""
+    cond_cont, cond_cat, x1_s1, n_sec, sec_cont, sec_pdg_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)
+    sec_pdg_idx = sec_pdg_idx.to(device)
+
+    # Stage-1 flow loss
+    if mode == "flow":
+        l_s1 = flow_matching_loss(stage1_model, x1_s1, cond_cont, cond_cat)
+    else:
+        assert ddpm_schedule is not None
+        l_s1 = ddpm_schedule.loss(stage1_model, x1_s1, cond_cont, cond_cat)
+
+    # n_sec classification loss
+    n_sec_logits = stage1_model.predict_n_sec(cond_cont, cond_cat)
+    l_nsec = F.cross_entropy(n_sec_logits, n_sec)
+
+    # Stage-2 secondary flow loss
+    # Use a noiseless Stage-1 target as context (detach to avoid back-prop
+    # coupling between the two flow paths through the same embedding table).
+    # The embedding table still receives gradients from the type-embedding loss.
+    from giant.constants import K_MAX
+
+    pdg_emb_weight = stage1_model.pdg_embedding_weight()
+    x1_s2 = _build_sec_x1(sec_cont, sec_pdg_idx, pdg_emb_weight)
+
+    sec_mask = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec.unsqueeze(1)
+    l_s2 = flow_matching_loss_secondary(
+        sec_decoder,
+        x1_s2,
+        cond_cont,
+        cond_cat,
+        x1_s1.detach(),
+        sec_mask,
+    )
+
+    total = l_s1 + lambda_nsec * l_nsec + lambda_s2 * l_s2
+    return total, l_s1, l_nsec, l_s2
+
+
 def train(
-    model: torch.nn.Module,
+    stage1_model: torch.nn.Module,
+    sec_decoder: torch.nn.Module,
     train_loader: DataLoader,
     val_loader: DataLoader,
     mode: str,
@@ -67,6 +156,8 @@ def train(
     warmup_epochs: int,
     device: torch.device,
     out_dir: str | Path,
+    lambda_nsec: float = 0.1,
+    lambda_s2: float = 1.0,
     normalizer_dict: dict | None = None,
     pdg_map: dict | None = None,
     mat_map: dict | None = None,
@@ -79,8 +170,11 @@ def train(
     out_dir = Path(out_dir)
     out_dir.mkdir(parents=True, exist_ok=True)
 
-    model = model.to(device)
-    optimizer = optim.AdamW(model.parameters(), lr=lr)
+    stage1_model = stage1_model.to(device)
+    sec_decoder = sec_decoder.to(device)
+
+    all_params = list(stage1_model.parameters()) + list(sec_decoder.parameters())
+    optimizer = optim.AdamW(all_params, lr=lr)
 
     def _lr_lambda(epoch: int) -> float:
         if warmup_epochs > 0 and epoch < warmup_epochs:
@@ -97,7 +191,8 @@ def train(
     best_val_loss = float("inf")
     if resume_path is not None:
         ckpt = torch.load(resume_path, map_location=device, weights_only=False)
-        model.load_state_dict(ckpt["model"])
+        stage1_model.load_state_dict(ckpt["model"])
+        sec_decoder.load_state_dict(ckpt["sec_decoder"])
         optimizer.load_state_dict(ckpt["optimizer"])
         lr_sched.load_state_dict(ckpt["lr_sched"])
         start_epoch = ckpt.get("epoch", 0) + 1
@@ -116,8 +211,12 @@ def train(
         for epoch in range(start_epoch, epochs + 1):
             epoch_start = time.monotonic()
             current_lr = optimizer.param_groups[0]["lr"]
-            model.train()
+            stage1_model.train()
+            sec_decoder.train()
             train_loss_sum = 0.0
+            train_s1_sum = 0.0
+            train_nsec_sum = 0.0
+            train_s2_sum = 0.0
             train_n = 0
             ema_loss = 0.0
             bar = tqdm(
@@ -128,27 +227,26 @@ def train(
                 unit="batch",
                 dynamic_ncols=True,
             )
-            for cond_cont, cond_cat, x1 in bar:
-                cond_cont = cond_cont.to(device)
-                cond_cat = cond_cat.to(device)
-                x1 = x1.to(device)
-
-                if mode == "flow":
-                    loss = flow_matching_loss(model, x1, cond_cont, cond_cat)
-                else:
-                    assert ddpm_schedule is not None
-                    loss = ddpm_schedule.loss(model, x1, cond_cont, cond_cat)
-
+            for batch in bar:
+                loss, l_s1, l_nsec, l_s2 = _compute_losses(
+                    stage1_model, sec_decoder, batch, mode, ddpm_schedule, device,
+                    lambda_nsec, lambda_s2,
+                )
                 optimizer.zero_grad()
                 loss.backward()
-                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+                torch.nn.utils.clip_grad_norm_(all_params, 1.0)
                 optimizer.step()
+
+                B = batch[0].size(0)
                 batch_loss = loss.item()
-                train_loss_sum += batch_loss * x1.size(0)
-                train_n += x1.size(0)
+                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_n += B
                 ema_loss = (
                     batch_loss
-                    if train_n == x1.size(0)
+                    if train_n == B
                     else 0.95 * ema_loss + 0.05 * batch_loss
                 )
                 bar.set_postfix_str(f"loss={ema_loss:.4f}", refresh=False)
@@ -158,28 +256,30 @@ def train(
             bar.close()
 
             if shutdown.requested:
-                # Mid-epoch: discard the partial epoch rather than persist an
-                # inconsistent (lr_sched not stepped, no validation) checkpoint.
                 break
 
             train_loss = train_loss_sum / max(train_n, 1)
             lr_sched.step()
 
-            model.eval()
+            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_n = 0
             with torch.no_grad():
-                for cond_cont, cond_cat, x1 in val_loader:
-                    cond_cont = cond_cont.to(device)
-                    cond_cat = cond_cat.to(device)
-                    x1 = x1.to(device)
-                    if mode == "flow":
-                        loss = flow_matching_loss(model, x1, cond_cont, cond_cat)
-                    else:
-                        assert ddpm_schedule is not None
-                        loss = ddpm_schedule.loss(model, x1, cond_cont, cond_cat)
-                    val_loss_sum += loss.item() * x1.size(0)
-                    val_n += x1.size(0)
+                for batch in val_loader:
+                    loss, l_s1, l_nsec, l_s2 = _compute_losses(
+                        stage1_model, sec_decoder, batch, mode, ddpm_schedule, device,
+                        lambda_nsec, lambda_s2,
+                    )
+                    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_n += B
             val_loss = val_loss_sum / max(val_n, 1)
             epoch_time = time.monotonic() - epoch_start
 
@@ -187,14 +287,24 @@ def train(
             marker = "  [best]" if is_best else ""
             print(
                 f"epoch {epoch:{epoch_w}d}/{epochs}"
-                f"  train {train_loss:.4f}  val {val_loss:.4f}"
+                f"  train {train_loss:.4f}"
+                f"  (s1={train_s1_sum/max(train_n,1):.3f}"
+                f"  nsec={train_nsec_sum/max(train_n,1):.3f}"
+                f"  s2={train_s2_sum/max(train_n,1):.3f})"
+                f"  val {val_loss:.4f}"
                 f"  lr {current_lr:.2e}  {epoch_time:.1f}s{marker}"
             )
             metrics_writer.writerow(
                 {
                     "epoch": epoch,
                     "train_loss": train_loss,
+                    "train_loss_s1": train_s1_sum / max(train_n, 1),
+                    "train_loss_nsec": train_nsec_sum / max(train_n, 1),
+                    "train_loss_s2": train_s2_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),
                     "lr": current_lr,
                     "epoch_time_s": epoch_time,
                 }
@@ -204,7 +314,7 @@ def train(
             if validate_every > 0 and epoch % validate_every == 0:
                 print(f"[epoch {epoch}] marginal validation:")
                 validate_marginals(
-                    model,
+                    stage1_model,
                     val_loader,
                     mode=mode,
                     schedule=ddpm_schedule,
@@ -213,7 +323,8 @@ def train(
                 )
 
             ckpt: dict = {
-                "model": model.state_dict(),
+                "model": stage1_model.state_dict(),
+                "sec_decoder": sec_decoder.state_dict(),
                 "optimizer": optimizer.state_dict(),
                 "lr_sched": lr_sched.state_dict(),
                 "epoch": epoch,

giant/validate.py

@@ -53,18 +53,21 @@ def validate_marginals(
     model.eval()
 
     all_real, all_gen = [], []
-    for i, (cond_cont, cond_cat, x1) in enumerate(val_loader):
+    for i, batch in enumerate(val_loader):
         if n_batches is not None and i >= n_batches:
             break
+        # Batch is (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx);
+        # validate_marginals only needs the Stage-1 primary target.
+        cond_cont, cond_cat, x1 = batch[0], batch[1], batch[2]
         cond_cont = cond_cont.to(device)
         cond_cat = cond_cat.to(device)
 
         if mode == "flow":
-            gen = sample_flow(model, cond_cont, cond_cat, **_kw(steps))
+            gen, _n_sec = sample_flow(model, cond_cont, cond_cat, **_kw(steps))
         elif mode == "ddpm":
-            gen = sample_ddpm(model, cond_cont, cond_cat, schedule)
+            gen, _n_sec = sample_ddpm(model, cond_cont, cond_cat, schedule)
         else:
-            gen = sample_ddim(model, cond_cont, cond_cat, schedule, **_kw(steps))
+            gen, _n_sec = sample_ddim(model, cond_cont, cond_cat, schedule, **_kw(steps))
 
         all_real.append(x1.numpy())
         all_gen.append(gen.cpu().numpy())