Implement set prediction

[jpata]

Summary

Adds an optional set-prediction MLPF path for CLD/CLIC hit datasets. The model uses learned particle slots, cross-attention, and Hungarian matching to predict a compact unordered particle collection while preserving elementwise MLPF as the default.

It also extends validation with comprehensive, scheme-independent particle and jet metrics, and runs the first length-matched comparison of heuristic PF, learned PF on tracks+clusters, and hit-based elementwise/set MLPF (notebooks/studies/20260911_cld_pf_hits_comparison, archived at hf://buckets/jpata/particleflow-studies).

Done

• Added compact set targets with independent padding, masks, and absolute kinematic transforms.
• Added backward-compatible output_mode and validated set-decoder configuration.
• Implemented fixed and input-conditioned queries, decoder layers, output heads, Hungarian matching, set losses, cardinality and auxiliary decoder-layer losses.
• Routed set mode through training, validation, prediction, and parquet serialization.
• Added scheme-independent particle metrics (Hungarian matching efficiency/purity/F1, duplicates, per-class efficiency/PID, event energy/pT/MET closure) and comprehensive jet metrics (one-to-one Hungarian jet matching with angular and response-qualified precision/recall/F1/fake rate).
• Added distributed partial-batch handling, seeded scenario launchers (Flatiron, Tallinn, LUMI), and the cld_pf_hits_comparison / hit-backbone comparison scenarios.
• Added unit, regression, permutation, numerical-stability, and 10k-hit integration tests.
• Ran and archived the 40k-step, seed-12345, batch-512, 8xH100 comparison campaign with all three trained runs completed.

Results (step 40k, interleaved 512-event validation)

	Heuristic PF	Learned PF (tracks+clusters)	Elementwise (hits)	Set (hits)
Particle F1	0.596	0.745	0.602	0.625
Count bias / event	+3.2	−0.4	−10.4	−4.2
Energy response	0.989	0.994	0.854	0.947
MET abs. error [GeV]	4.3	4.8	12.3	7.7
Matched ΔR mean	0.015	0.014	0.027	0.029
Matched rel. pT error	0.097	0.086	0.163	0.157
ttbar jet IQR / match fraction	0.078 / 0.934	0.066 / 0.964	0.239 / 0.896	0.190 / 0.921
Walltime (40k steps)	—	1.9 h	20.9 h	39.6 h

The heuristic reference is scored with the exact same jet_matching_metrics and compute_validation_particle_metrics code as the trained runs (eval_pf_baseline.py), on 520 events per sample.

Diagnosis

1. Set overtakes elementwise at matched training length and is still improving at 40k (F1 0.579 → 0.617 → 0.625 over the last three checkpoints, ttbar IQR 0.203 → 0.190). Elementwise plateaued from 30k and its checkpoints are noisy (count bias oscillates between −11.9 and −17.6). The earlier 20k-step verdict against set was a schedule artifact.
2. Elementwise's dominant failure is missed particles, not bad kinematics. It predicts 60 particles for 70.7 targets; the deficit appears coherently as 15% low energy response, 13% low scalar-pT response, ttbar jet median response 0.91, and the largest angular→response-qualified jet F1 drop (0.885 → 0.840 on qq). The per-hit presence classifier is under-firing.
3. Set's weak spot is class identity and duplicates, not closure. Highest duplicate fraction (0.192) and poor PID on rare classes: 0.49 (neutral hadron), 0.30 (electron), 0.37 (muon) vs 0.53/0.52/0.63 for elementwise and 0.76/0.42/0.68 for learned PF. Lowest photon (0.52) and neutral-hadron (0.33) efficiency, but decisively better than elementwise on charged hadrons (0.725 vs 0.513) and muons (0.84 vs 0.66): strong on track-seeded objects, weak on calorimeter-only ones.
4. Both hit models share a kinematic-resolution floor. Matched ΔR and relative pT/energy errors are ~2x worse than learned PF, which directly produces jet IQR 0.19-0.24 vs 0.066. This is independent of the output formulation and points at the hit encoder / regression head.
5. Learned PF beats the heuristic on every particle metric except MET (4.8 vs 4.3 GeV) and is still improving at 40k. The heuristic's low F1 (0.596, purity 0.583) under strict ΔR<0.1 and rel-pT<0.5 matching should be sanity-checked against the target definition.

Next steps

Elementwise (hits):

• Sweep the presence threshold on the 40k checkpoint (analogous to calibrate_set_inference.py); the −10 count bias may be calibration rather than capacity.
• Break down missed targets by class and by tracker- vs calorimeter-representative hit.
• Verify the hit-to-particle target alignment: particles whose hits are all owned by another particle's representative hit are unlearnable under this scheme.

Set (hits):

• Train longer or extend the cosine schedule; not converged at 40k.
• Add duplicate suppression (presence threshold + ΔR NMS sweep) and report the F1/duplicate trade-off.
• Rebalance PID: class-weighted PID loss or higher matcher.pid cost (currently 1.0 vs geometry 2.0); inspect the PID confusion matrix on accepted matches.
• Revisit tracker_query_fraction=0.6; 40% calorimeter-seeded queries may be too few for photons and neutral hadrons.

Both hit models:

• Plot matched pT resolution and ΔR vs true pT, η, and class to localize the resolution floor.
• Expose summed calorimeter-hit energy in the attention neighborhood as a feature and regress energy relative to it.
• Report jet response vs η (barrel vs endcap).
• Implement packed memory-efficient cross-attention with a CPU fallback; the hit backbones process ~75x more elements per event than track/cluster.

Study hygiene:

• Run at least two more seeds for the set/elementwise pair before drawing architectural conclusions.
• Evaluate the heuristic on exactly the trained runs' 512 test events and check heuristic purity under looser matching.
• Re-evaluate elementwise 30k/35k/40k on a fixed event list to rule out validation-ordering noise.
• Add explicit set-mode ONNX handling.

40k steps training comparison on CLD:
hit_output_comparison.slides.pdf
https://huggingface.co/buckets/jpata/particleflow-studies/tree/20260911_cld_pf_hits_comparison/output/hit_output_comparison.slides.pdf

References

• Better Queries, Cheaper Attention: Adapting Transformers for Efficient Sparse Reconstruction
• HEPTv2: End-to-End Efficient Point Transformer for Charged Particle Reconstruction

🤖 Generated with Claude Code