Linear Attention Vision Transformers for CMS End-to-End Jet Classification and Mass Regression
This project develops linear-scale attention vision transformers for CMS End-to-End jet classification and proxy mass regression on 8-channel 99.1% sparse detector images. Five architectures were trained and benchmarked under identical conditions during the test task: XCiT with domain MAE pretraining (AUC 0.9237, 2.42x S/sqrt(B)), XCiT with ImageNet pretraining, XCiT from scratch, a custom L2ViT combining local window and linear global attention built from scratch, and Swin-T as the softmax baseline. All models use a dual head for joint classification and proxy mass regression. 11 analysis studies were completed beyond standard benchmarking: channel importance (Ch0 and Ch3 critical, AUC collapses to 0.47 without either), channel pair interactions, Grad-CAM saliency, t-SNE feature representations, layer-by-layer CKA similarity, data efficiency (5x labelled data savings), jet multiplicity analysis, energy threshold study (detector implications for zero-suppression), S/sqrt(B) significance curves, adversarial robustness, and calibration with temperature scaling. Key findings: domain MAE pretraining outperforms scratch by +0.125 AUC and leads to fundamentally different representations (CKA 0.008-0.015 at every layer), linear attention beats softmax by +0.063 AUC under identical conditions, and block 1 of the pretrained model alone exceeds the fully trained scratch model in feature separation. The summer work extends this to multi-class classification (quark/gluon/top/W/Z) with per-class channel importance ablations, token pruning exploiting 99.1% sparsity, sparse convolution baselines, hybrid CNN+ViT architectures, and comparison against ParticleNet and JEDI-net on the same held-out test set.
Project details
Technologies
Not listed in the archive