JetBench: Benchmarking Vision Models for Jet Observables' Classification in Heavy-Ion Physics

ICLR 2026 Conference Submission14583 Authors

Keywords: Multi-Task Learning, Multi-Parameter Classification, Benchmark Datasets, Vision Transformers, State Space Models, Loss Weighting in Multi-Task Learning, Model Calibration and Generalization, Training Dynamics and Optimization, Scientific Machine Learning, Relativistic Heavy Ion Collisions, Quark–Gluon Plasma, Physics-Informed Deep Learning

TL;DR: We study multi-task learning on simulated heavy-ion collisions, benchmarking CNNs, Transformers, and state-space models. Models reach ~100% (energy), ~95% ($\alpha_s$), ~78% ($Q_0$). Loss weighting highlights inter-task trade-offs.

Abstract: Relativistic heavy-ion collisions provide a window into quark--gluon plasma formation, but extracting parameters such as the energy loss mechanism, strong coupling $\alpha_s$, and virtuality scale $Q_0$ has traditionally required costly Bayesian inference. We introduce \textbf{JetBench}, a benchmark for multi-parameter classification of heavy-ion events using the ANONYMIZED dataset. Each event is encoded as a $32\times32$ jet image with three targets: energy loss module, $\alpha_s$, and $Q_0$. We systematically evaluate CNNs (EfficientNetV2, ConvNeXt V2), Transformers (ViT-CoMer, Swin V2), and state space models (Mamba) under unified training. Results show saturated performance on energy loss ($\sim$100\%), strong accuracy on $\alpha_s$ ($\sim$95\%), and up to $78\%$ on $Q_0$, with ViT-CoMer achieving the best joint accuracy (74.5\%). Loss-weight ablations reveal trade-offs between tasks, with $Q_0$ emphasis improving recall at modest cost to $\alpha_s$. Probability calibration confirms errors follow physics continuity (e.g., $\alpha_s=0.2/0.3$, $Q_0=2.0/2.5$). These findings establish JetBench as a scalable complement to Bayesian approaches. Code and preprocessing scripts are available at ANONYMIZED url.

Primary Area: applications to physical sciences (physics, chemistry, biology, etc.)

Submission Number: 14583