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/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)