[논문 리뷰] Particle Transformer for Jet Tagging

H. Qu, C. Li외 1명|arXiv (Cornell University)|2022. 02. 08.

Astrophysics and Cosmic Phenomena인용 수 45

한 줄 요약

이 논문은 JetClass라는 대규모 제트 태깅 데이터셋과 ParT를 소개하며, 쌍별 상호작용에 민감한 어텐션을 갖춘 Transformer 기반 모델로 JetClass에서 이전 방법들을 능가하고 미세조정 후 벤치마크를 개선합니다.

ABSTRACT

Jet tagging is a critical yet challenging classification task in particle physics. While deep learning has transformed jet tagging and significantly improved performance, the lack of a large-scale public dataset impedes further enhancement. In this work, we present JetClass, a new comprehensive dataset for jet tagging. The JetClass dataset consists of 100 M jets, about two orders of magnitude larger than existing public datasets. A total of 10 types of jets are simulated, including several types unexplored for tagging so far. Based on the large dataset, we propose a new Transformer-based architecture for jet tagging, called Particle Transformer (ParT). By incorporating pairwise particle interactions in the attention mechanism, ParT achieves higher tagging performance than a plain Transformer and surpasses the previous state-of-the-art, ParticleNet, by a large margin. The pre-trained ParT models, once fine-tuned, also substantially enhance the performance on two widely adopted jet tagging benchmarks. The dataset, code and models are publicly available at https://github.com/jet-universe/particle_transformer.

연구 동기 및 목표

• 대규모 공용 제트 태깅 데이터셋의 필요성을 제시하고, 콜라이더 물리학의 딥러닝 방법을 발전시키려는 동기를 제시한다.
• 받아들일 Transformer 기반 아키텍처(ParT)를 제안하여 어텐션에 쌍별 입자 상호작용을 통합하고 태깅 성능을 향상시킨다.
• JetClass에서의 포괄적 평가를 제공하고 미세조정(fine-tuning)을 통해 기존 제트 태깅 벤치마크에서의 전이 이득을 입증한다.

제안 방법

• JetClass를 100,000,000제트에 걸쳐 10종의 제트 유형으로 정의하되, 이전에 탐색되지 않은 범주를 포함한다.
• ParT를 도입하는 Transformer 변형으로 두 입력을 사용한다: 각 입자 특징과 각 쌍의 상호작용 특징.
• 사전 소프트맥스 어텐션 가중치에 쌍별 상호작용 바이어스 U를 추가한 particle self-attention(P-MHA)을 보강한다.
• 최종 분류를 위한 글로벌 제트 표현을 생성하기 위해 클래스 어텐션 메커니즘을 구현한다.
• 입자 4-벡터에서 도출된 쌍별 상호작용 특징을 인코딩하고 8차원 상호작용 임베딩으로 변환한다.
• ParT를 엔드투엔드로 학습하고 PFN, P-CNN, ParticleNet과 비교한다; 또한 다른 데이터세트에서의 사전 학습 + 미세조정도 평가한다.

실험 결과

연구 질문

• RQ1 쌍별 입자 상호작용 정보를 가진 Transformer 기반 아키텍처가 대규모 제트 데이터셋에서 기존 제트 태깅 모델보다 우수한가?
• RQ2 어텐션에서 쌍별 상호작용 바이어스(P-MHA)의 기여도는 전체 성능에 어떠한 영향을 주는가?
• RQ3 JetClass에서 ParT를 사전 학습한 후 미세조정을 통해 기존 제트 태깅 벤치마크의 성능이 개선될 수 있는가?

주요 결과

• ParT가 JetClass에서 ParticleNet 및 다른 베이스라인보다 더 높은 정확도와 AUC를 달성한다.
• P-MHA를 통한 쌍별 입자 상호작용의 도입은 일반적인 Transformer 대비 유의한 이득을 가져온다.
• JetClass에서 사전 학습하고 다른 제트 태깅 벤치마크에서 미세조정할 때 ParT가 상당한 개선을 제공한다.
• 여러 신호 유형에 걸쳐 ParT가 배경 거부력을 크게 개선하며, 일부 제트 유형에서는 ParticleNet에 비해 큰 이득을 보인다.
• 절제된 아블레이션에서 P-MHA를 표준 MHA로 교체하면 정확도와 배경 거부력이 감소하는데, 이는 쌍별 상호작용 모델링의 가치를 확인한다.