Linear attention vision transformers for end to end mass regression and classification

This project addresses a key challenge in High Energy Physics: efficient and accurate end-to-end reconstruction of particle properties from detector data. In the CMS experiment at the Large Hadron Collider, particle collision events can be represented as high-resolution, multi-channel images. The goal is to simultaneously classify particle types and regress their invariant mass directly from these images. A major limitation of existing Vision Transformer (ViT) approaches is the quadratic computational complexity of softmax attention, which makes them inefficient for large-scale detector data. While linear attention mechanisms reduce this complexity, they often struggle to capture rich global interactions and multi-channel dependencies required for physics tasks. To address this, this project proposes the development of linear-attention-based Vision Transformers tailored for CMS E2E tasks, with a strong focus on extending the DepthViT architecture. The approach consists of three key components: Efficient Representation Learning A linear attention ViT with Masked Autoencoder (MAE) pretraining will be used to learn strong representations from large-scale unlabeled CMS data, improving downstream performance and convergence. Architecture Adaptation (DepthViT Extension) The DepthViT architecture will be adapted for supervised classification and regression by: Introducing dual-task heads for classification and mass regression Preserving depth-channel and rank structure in later layers to avoid information collapse Integrating Hierarchical Attention Permutable (HAP) blocks to enable efficient cross-window and cross-channel information flow while maintaining linear complexity Systematic Benchmarking The proposed models will be rigorously compared against standard and efficient transformer baselines (ViT, Swin, L²ViT, XCiT) using: Physics metrics (accuracy, MAE, R²) Efficiency metrics (FLOPs, memory, throughput)