lars
e6e0eb22bf
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>
173 lines
7.9 KiB
Markdown
173 lines
7.9 KiB
Markdown
# Phase 2: Secondary Particle Prediction
|
||
|
||
## Context
|
||
|
||
Phase 1 takes `n_sec` (secondary count) and `e_sec` (total secondary energy) as **conditioning inputs**. Phase 2 must instead **predict** them, making the surrogate self-contained for shower rollout. Per Jan's 2026-06-29 decision: hard discrete `n_sec` integer head; escalation to Gumbel-Softmax only if empirically needed.
|
||
|
||
Two-stage factorization:
|
||
- **Stage 1**: existing 9D flow model (reduced conditioning: drop `n_sec` + `log(e_sec)`) + a new discrete `n_sec` classification head
|
||
- **Stage 2**: non-AR flow matching over `K_MAX` secondary slots simultaneously, each slot predicting `(stick_break_logit, dir_local_3D, type_emb)` — conditioned on pre-step state + Stage 1 output; padded slots masked from loss
|
||
|
||
Training: joint, combined loss `L = L_flow_s1 + λ_nsec * L_nsec + λ_s2 * L_flow_s2`.
|
||
|
||
---
|
||
|
||
## Prerequisite: Determine K_MAX
|
||
|
||
Before implementing, run a quick analysis over existing parquet files to find `max(n_sec)` and the 99th percentile. Expected to be 5–20 for EM shower steps. Set `K_MAX` as a constant in `giant/constants.py` (suggest 15 as a starting point, revise from data).
|
||
|
||
---
|
||
|
||
## New Branch
|
||
|
||
```bash
|
||
git checkout -b phase2-secondary-prediction master
|
||
```
|
||
|
||
---
|
||
|
||
## Part A — Data Pipeline
|
||
|
||
### A1. `scripts/steps_to_parquet.py`
|
||
|
||
Extend `_add_secondary_energy` to also collect per-secondary attributes from the spawning tree join:
|
||
- For each `child_track_id`, look up the child's first step → get `pdg`, `pre_E`, `pre_dx/dy/dz`
|
||
- Emit list columns in the parquet: `sec_pdg_list`, `sec_E_list`, `sec_dx_list`, `sec_dy_list`, `sec_dz_list`
|
||
- Lists are sorted **descending by energy** at write time
|
||
- Truncate to `K_MAX` entries if needed (flag if any row truncated)
|
||
|
||
Re-run ROOT→parquet conversion after this change.
|
||
|
||
### A2. `giant/data/loader.py`
|
||
|
||
In `_df_to_dict`: read the five new list columns. Pad each to length `K_MAX` with zeros (energy) / sentinel values (pdg → 0, dir → (0,0,1)). Return as fixed-shape arrays `(N, K_MAX)` / `(N, K_MAX, 3)`.
|
||
|
||
Also return a boolean validity mask `sec_valid` of shape `(N, K_MAX)`: `True` for slots `i < n_sec`.
|
||
|
||
### A3. `giant/data/transforms.py`
|
||
|
||
Add `encode_secondaries(sec_pdg_list, sec_E_list, sec_dir_list, sec_valid, e_sec, pdg_emb_weight, pre_dir, K_MAX)`:
|
||
1. **Direction**: call existing `local_frame_rotation` per slot
|
||
2. **Energy (stick-breaking)**:
|
||
- Slot 0: `f_0 = E_0 / e_sec` → logit `log(f_0/(1-f_0))` (clamped)
|
||
- Slot i: `f_i = E_i / (e_sec - sum(E_0..E_{i-1}))` → logit
|
||
- Last valid slot: logit = large positive constant (takes all remaining budget)
|
||
- Padding slots (beyond `n_sec`): set logit = 0, masked out of loss anyway
|
||
3. **Type embedding**: index into `pdg_emb_weight` (the PDG embedding table weights) to get the target embedding vector for each secondary's `pdg`. Shape `(K_MAX, emb_dim)`.
|
||
|
||
Returns `sec_targets: (K_MAX, 1 + 3 + emb_dim)` and `sec_valid: (K_MAX,)`.
|
||
|
||
Inverse (`decode_secondaries`): sigmoid stick-breaking fractions → energies, inv local frame rotation → world dirs, nearest-neighbor lookup in PDG embedding table → pdg code.
|
||
|
||
### A4. `giant/data/dataset.py`
|
||
|
||
Update `build_features` and `StreamingStepsDataset.__iter__` to also yield `sec_targets` and `sec_valid` alongside the existing `(cond_cont, cond_cat, x1)` batch items.
|
||
|
||
---
|
||
|
||
## Part B — Constants (`giant/constants.py`)
|
||
|
||
- `COND_DIM`: 10 → **8** (remove `n_sec` and `log(e_sec)`)
|
||
- Add `K_MAX: int` (set after data analysis, e.g. 15)
|
||
- Add `SEC_SLOT_DIM: int` (= 4 + `emb_dim` = 20 for default emb_dim=16; 1 stick + 3 dir + 16 type)
|
||
- Add `SEC_DIM: int = K_MAX * SEC_SLOT_DIM` (flattened Stage 2 target dimension)
|
||
- Update `LOCAL_TARGET_NAMES` (Stage 1 only, still 9D)
|
||
|
||
---
|
||
|
||
## Part C — Model (`giant/model/network.py`)
|
||
|
||
### C1. `DenoisingMLP` — Stage 1 (minimal changes)
|
||
|
||
- `ConditionEncoder.cont_dim` drops from 10 to 8 (COND_DIM change propagates automatically)
|
||
- Add `n_sec_head = nn.Sequential(Linear(cond_out_dim, hidden_dim//2), SiLU(), Linear(hidden_dim//2, K_MAX + 1))` applied to `c_emb` (the condition encoding, not the diffused latent)
|
||
- Add method `predict_n_sec(cond_cont, cond_cat) -> Tensor[B, K_MAX+1]` — no diffusion, just encode conditioning and run the head
|
||
|
||
### C2. `SecondaryDecoder` — Stage 2 (new class)
|
||
|
||
Architecture mirrors `DenoisingMLP` but:
|
||
- **Input**: `x_t` of shape `(B, SEC_DIM)` (flattened K_MAX secondary slots)
|
||
- **Conditioning**: pre-step state (8D cont + 2 cat → same ConditionEncoder as Stage 1) concatenated with Stage 1 output (9D normalized target, detached from Stage 1 loss for stability initially). Total cond dim to the ResBlocks: `time_dim + cond_s1_out_dim + 9`
|
||
- **Output**: vector field of shape `(B, SEC_DIM)`
|
||
- Uses same `ResBlock` / `SinusoidalEmbedding` / `ConditionEncoder` building blocks
|
||
|
||
A `SecondaryConditionEncoder` wraps the base `ConditionEncoder` and concatenates the Stage 1 output:
|
||
```python
|
||
class SecondaryConditionEncoder(nn.Module):
|
||
# base: ConditionEncoder(pdg_vocab, mat_vocab, 8, emb_dim, cond_out_dim)
|
||
# stage1_proj: Linear(X_DIM, stage1_cond_dim)
|
||
# mlp: fuses both
|
||
```
|
||
|
||
---
|
||
|
||
## Part D — Loss / Training
|
||
|
||
### `giant/model/schedule.py`
|
||
|
||
Add `flow_matching_loss_masked(model, x1, cond_cont, cond_cat, mask)`:
|
||
- Same as `flow_matching_loss` but divides by `mask.sum()` instead of `B * SEC_DIM`, zeroing out padded slots before averaging. `mask` shape: `(B, K_MAX)`, broadcast over slot dims.
|
||
|
||
### `giant/train.py`
|
||
|
||
Batch now unpacks as `(cond_cont, cond_cat, x1_s1, n_sec_target, x1_s2, sec_mask)`.
|
||
|
||
Combined loss per batch:
|
||
```
|
||
L_s1 = flow_matching_loss(stage1_model, x1_s1, cond_cont, cond_cat)
|
||
L_nsec = cross_entropy(stage1_model.predict_n_sec(cond_cont, cond_cat), n_sec_target)
|
||
L_s2 = flow_matching_loss_masked(sec_decoder, x1_s2, cond_cont, cond_cat, stage1_detached, sec_mask)
|
||
L = L_s1 + lambda_nsec * L_nsec + lambda_s2 * L_s2
|
||
```
|
||
|
||
Config adds `lambda_nsec` (suggest 0.1) and `lambda_s2` (suggest 1.0) under `[train]`.
|
||
|
||
Both `stage1_model` and `sec_decoder` share a single `optimizer` (AdamW over all parameters).
|
||
|
||
Checkpoint saves both `stage1_model.state_dict()` and `sec_decoder.state_dict()`, plus `K_MAX` and `SEC_SLOT_DIM` in `model_config`.
|
||
|
||
### `giant/pipeline.py`
|
||
|
||
- Compute `K_MAX` from data (max `n_sec` over training events) before constructing models
|
||
- Build both `DenoisingMLP` and `SecondaryDecoder`, pass both to `run_training`
|
||
|
||
---
|
||
|
||
## Part E — Sampling (`giant/sample.py`)
|
||
|
||
```python
|
||
def sample_stage1(model, cond_cont, cond_cat, steps=10):
|
||
# Euler ODE → primary sample (9D), + argmax n_sec head
|
||
...
|
||
|
||
def sample_secondaries(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec, steps=10):
|
||
# Euler ODE on SEC_DIM → decode stick-breaking → energies
|
||
# inv_local_frame_rotation → world-frame dirs
|
||
# nearest-neighbor in pdg_emb_weight → pdg codes
|
||
# mask slots >= n_sec
|
||
...
|
||
```
|
||
|
||
---
|
||
|
||
## Part F — Wiring
|
||
|
||
- **`giant/validate.py`**: add secondary-specific marginals (n_sec distribution, species distribution, energy fraction per slot)
|
||
- **`giant/cli.py`**: `predict` command loads both checkpoints, calls both samplers, appends secondary columns to output parquet
|
||
|
||
---
|
||
|
||
## Type embedding design note
|
||
|
||
The type embedding target at training is `pdg_emb.weight[sec_pdg_idx]` (the Stage 1 PDG embedding table rows). Gradients flow into the embedding table from both the conditioning path (input PDG) and the secondary type loss — this is intentional; the shared embedding space is the bridge. At inference, snap: `argmin_k ||pred_emb - pdg_emb.weight[k]||`.
|
||
|
||
---
|
||
|
||
## Verification
|
||
|
||
1. `uv run pytest` — existing tests pass (Stage 1 shape/interface unchanged beyond COND_DIM)
|
||
2. Unit tests for `encode_secondaries` / `decode_secondaries` (round-trip: energies sum to `e_sec`, directions are unit vectors)
|
||
3. Unit test for `flow_matching_loss_masked`: verify padded slots contribute zero gradient
|
||
4. Short training run (1–2 epochs): confirm all three loss components decrease
|
||
5. Sampling smoke test: verify `sum(sec_E) ≈ e_sec` per sample, all directions unit-normed
|