[논문 리뷰] Dataset for flavour tagging R&D

I. Ochoa, S. B. Klein외 2명|arXiv (Cornell University)|2024. 08. 20.

Electron and X-Ray Spectroscopy Techniques인용 수 5

한 줄 요약

본 논문은 토큰화 제거, 디코더 업그레이드, 다양한 재구성 작업 및 집합-대 집합(set-to-set) 생성 탐구를 통해 제트 물리 학습에 기반 모델(pre-train foundational-model–style backbones)을 제트 데이터에 적용하는 masked particle modeling (MPM)을 개선한다. MPMv2 및 set-to-set flow-matching을 도입하여 OOD 작업을 포함한 강력한 하류 성능을 시연한다.

ABSTRACT

In this work, we significantly enhance masked particle modeling (MPM), a self-supervised learning scheme for constructing highly expressive representations of unordered sets relevant to developing foundation models for high-energy physics. In MPM, a model is trained to recover the missing elements of a set, a learning objective that requires no labels and can be applied directly to experimental data. We achieve significant performance improvements over previous work on MPM by addressing inefficiencies in the implementation and incorporating a more powerful decoder. We compare several pre-training tasks and introduce new reconstruction methods that utilize conditional generative models without data tokenization or discretization. We show that these new methods outperform the tokenized learning objective from the original MPM on a new test bed for foundation models for jets, which includes using a wide variety of downstream tasks relevant to jet physics, such as classification, secondary vertex finding, and track identification.

연구 동기 및 목표

• 레이블이 없는 제트 데이터를 사용하여 고에너지 물리학에 대한 foundation-model 스타일의 사전 학습을 동기 부여한다.
• VQVAE 토큰화 없이 개선된 masked particle modeling (MPMv2)을 개발한다.
• 조건부 생성 방식 등을 포함한 다수의 재구성 작업을 평가한다.
• 집합-대-집합 흐름 매칭(set-to-set flow-matching)을 제트에 대한 경쟁력 있는 사전 학습 패러다임으로 제안한다.

제안 방법

• 반복된 마스킹 토큰을 제거하고 전체 트랜스포머 디코더를 사용하도록 MPM을 수정한다.
• 태스크를 단순화하지 않도록 마스킹된 요소들 사이에서만 위치 인코딩을 제공한다.
• 다섯 가지 연속 특성 재구성 작업과 하나의 범주형 작업(입자 ID)을 조사한다.
• 대안 타깃들을 탐색한다: K-Means 토큰화, CNF (conditional normalizing flow), flow-matching (CFM), 그리고 set-to-set flow-matching (SSFM).
• Delphes로 시뮬레이션된 JetClass 및 BTag 데이터셋을 사용하여 백본 표현을 사전 학습하고 평가한다.
• 디코더 유형, 추가 특성 및 학습 설정을 비교하는 어블레이션 연구(Table 1).

Figure 1 : A comparison of the original MPM encoder-decoder setup (left) and the new model configuration (right). The new model includes multiple reconstruction tasks, swaps the MLP decoder for a transformer, and only encodes the reduced set.

실험 결과

연구 질문

• RQ1 토큰화 제거와 더 강한 디코더 사용이 원래의 MPMv1보다 MPM 성능을 향상시키는가?
• RQ2 대안 재구성 타깃(CNF, K-Means, flow 기반 방식)이 MPM 사전 학습에서 VQVAE 토큰화와 경쟁력이 있는가?
• RQ3 개선된 백본은 분포 내(in-distribution), 약지도(weakly supervised), 분포 외(out-of-distribution) 제트 작업에서 어떤 성능을 보이는가?
• RQ4 집합-대-집합 흐름 매칭이 unordered 제트 구성요소에 대한 실행 가능한 사전 학습 패러다임을 제공하는가?
• RQ5 연장된 학습, 마스크 비율 조정, 추가 특성이 하류 작업에 어떤 영향을 미치는가?

주요 결과

• MPMv2는 트랜스포머 디코더와 입력 원소 수를 줄임으로 MPMv1보다 분류 정확도를 크게 향상시킨다.
• MAE 스타일 디코딩으로 전환하고 위치 인코딩을 제한하면 성능이 개선되고 GPU 메모리 사용이 감소한다.
• 충돌 매개변수 특징 및 입자 ID를 추가하면 정확도가 더 향상된다(예: 회귀 62.2에서 80.4로; k-means 70.2에서 83.0으로 추가된 경우).
• 완전한 트랜스포머 기반 디코더(MAE)는 어블레이션에서 회귀 79.2와 k-means 81.4를 달성한다.
• 더 긴 학습, 더 깊은 디코더, 40% 마스크 비율은 최상의 결과를 낳는다: 회귀 83.3 및 k-means 84.0.
• 사전 학습된 백본은 분포 내, 약지도, 분포 외 작업에서 무작위 초기화보다 우수한 일반화를 보인다.

Figure 2 : A schematic overview of the SSFM model.