Comparison of image processing models in quark–gluon jet classification

Abstract

Quark–gluon discrimination provides a useful test case for studying how different machine-learning architectures learn the spatial structure of QCD radiation. In this work, we compare convolutional neural network (CNN), vision transformers (ViT), and hierarchical Swin transformers using the same three-channel jet-image representation, consisting of charged-particle momentum, neutral-particle momentum, and charged-particle multiplicity from PYTHIA 8 jets. We study their performance for different training-set sizes and fine-tuning configurations, with particular attention to the role of local and global information in the jet images. CNN and Swin models consistently perform better than ViT in the cases studied. Since both CNN and Swin retain a strong local component in their architectures, this suggests that local jet substructure plays an important role in quark–gluon discrimination. The performance of the hierarchical Swin model also suggests that combining local features over larger spatial scales is useful. Block-wise fine-tuning improves the performance of the transformer models, although the improvement becomes smaller and the training less stable as more blocks are unfrozen. We also find that self-supervised momentum contrast pretraining improves the model initialization, particularly when the amount of labeled training data is limited. Based on these observations, we developed a smaller Swin model adopted to the jet-image representation used in this study. It achieves comparable performance with substantially fewer parameters. The results show that it is important to adapt the model architecture and training procedure to the specific input characteristics of high energy physics (HEP) data when applying vision models in HEP.