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>
This commit is contained in:
2026-06-29 11:34:31 +02:00
parent c627142135
commit e6e0eb22bf
18 changed files with 1174 additions and 234 deletions
+4 -2
View File
@@ -39,7 +39,9 @@ def test_denoising_mlp_gradients_flow():
t = torch.rand(B)
cond_cont = torch.randn(B, COND_DIM)
cond_cat = torch.zeros(B, 2, dtype=torch.long)
loss = model(x_t, t, cond_cont, cond_cat).sum()
loss.backward()
# Both paths must be exercised to get gradients through all parameters.
flow_loss = model(x_t, t, cond_cont, cond_cat).sum()
nsec_loss = model.predict_n_sec(cond_cont, cond_cat).sum()
(flow_loss + nsec_loss).backward()
for name, p in model.named_parameters():
assert p.grad is not None, f"no grad for {name}"