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