From 980b6ae7da337855d10150cdbadf2075d7ff8752 Mon Sep 17 00:00:00 2001
From: Lars Bogner <upyxo@student.kit.edu>
Date: Wed, 8 Jul 2026 11:41:24 +0200
Subject: [PATCH] Weight Stage-2 secondary loss equally between direction and
 type-embedding dims

The masked flow-matching loss for the secondary decoder averaged uniformly
over all 20 per-slot dims, letting the 16 type-embedding dims outvote the
4 physically-interesting ones (stick-break logit + direction). Split the
two blocks and average each over its own width before summing, so they
contribute with equal weight regardless of EMB_DIM.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 giant/constants.py      |  3 +++
 giant/model/schedule.py | 24 ++++++++++++++++--------
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/giant/constants.py b/giant/constants.py
index 8b42522..230d1e3 100644
--- a/giant/constants.py
+++ b/giant/constants.py
@@ -15,6 +15,9 @@ K_MAX = 15
 SEC_SLOT_DIM = 20  # 1 + 3 + 16
 EMB_DIM = 16       # must match model emb_dim default
 
+# Per-slot continuous (non-embedding) width: stick-breaking logit + local dir.
+CONT_SLOT_DIM = SEC_SLOT_DIM - EMB_DIM  # 4
+
 # Flattened Stage-2 target dimension
 SEC_DIM = K_MAX * SEC_SLOT_DIM  # 15 * 20 = 300
 
diff --git a/giant/model/schedule.py b/giant/model/schedule.py
index ed2877c..a5a384c 100644
--- a/giant/model/schedule.py
+++ b/giant/model/schedule.py
@@ -86,8 +86,14 @@ def flow_matching_loss_secondary(
 
     Only valid-slot dimensions contribute to the loss; padded slots are zeroed
     before averaging, so the loss is not diluted by empty slots.
+
+    Each slot packs CONT_SLOT_DIM continuous dims (stick_logit, dir) followed
+    by EMB_DIM type-embedding dims. A flat per-dimension mean would let the
+    16 embedding dims outvote the 4 physically-interesting ones, so the two
+    blocks are each averaged over their own width first and then combined
+    with equal weight — this stays correct if EMB_DIM/CONT_SLOT_DIM change.
     """
-    from giant.constants import SEC_SLOT_DIM
+    from giant.constants import CONT_SLOT_DIM, EMB_DIM, K_MAX, SEC_SLOT_DIM
 
     B = x1.size(0)
     t = torch.rand(B, device=x1.device)
@@ -96,10 +102,12 @@ def flow_matching_loss_secondary(
     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
+    err = ((v_t - u_t) ** 2).view(B, K_MAX, SEC_SLOT_DIM)
+    cont_err = err[..., :CONT_SLOT_DIM].mean(dim=-1)  # (B, K_MAX)
+    emb_err = err[..., CONT_SLOT_DIM : CONT_SLOT_DIM + EMB_DIM].mean(dim=-1)
+
+    mask = sec_mask.float()
+    denom = mask.sum().clamp(min=1)
+    cont_loss = (cont_err * mask).sum() / denom
+    emb_loss = (emb_err * mask).sum() / denom
+    return cont_loss + emb_loss