Detach Stage-2 type-embedding target to stop self-referential collapse

The shared PDG embedding table was used, un-detached, as the regression
target for the Stage-2 flow-matching loss. Since that tensor becomes x1
in u_t = x1 - x0, gradients could pull the embedding table itself toward
the decoder's predictions instead of the decoder learning to match the
table, risking species-embedding collapse and degrading the
nearest-neighbor species decode at inference.

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

giant/train.py

@@ -86,8 +86,15 @@ def _build_sec_x1(
     pdg_emb_weight: (pdg_vocab, emb_dim) — live embedding table weights
 
     Returns (B, SEC_DIM) = (B, K_MAX * (4 + emb_dim)).
+
+    Detaches the looked-up rows: this tensor becomes x1 in the flow-matching
+    loss (u_t = x1 - x0), so without detaching, the Stage-2 loss could pull
+    the embedding table itself toward whatever the decoder already predicts
+    (a moving, self-referential regression target) instead of only pulling
+    the decoder toward the table. The table is still trained normally via
+    its Stage-1 conditioning role and `predict_n_sec`.
     """
-    type_emb = pdg_emb_weight[sec_pdg_idx]  # (B, K_MAX, emb_dim)
+    type_emb = pdg_emb_weight[sec_pdg_idx].detach()  # (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)
 
@@ -125,7 +132,9 @@ def _compute_losses(
     # 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.
+    # The type-embedding lookup itself is also detached inside _build_sec_x1,
+    # so the shared PDG table is shaped only by its Stage-1 conditioning role
+    # and predict_n_sec, not by chasing the Stage-2 decoder's predictions.
     from giant.constants import K_MAX
 
     pdg_emb_weight = stage1_model.pdg_embedding_weight()