[논문 리뷰] Track Seeding and Labelling with Embedded-space Graph Neural Networks

Nicholas Choma, Daniel Murnane외 17명|arXiv (Cornell University)|2020. 06. 30.

Advanced Graph Neural Networks참고 문헌 16인용 수 11

한 줄 요약

이 논문은 그래프 신경망(GNN)과 학습된 임bedding을 사용한 엔드 투 엔드 입자 궤적 재구성 프레임워크를 제시한다. 임bedded 공간 내 비선형 거리 구조를 활용해 효율적인 근접성 질의와 군집화를 수행함으로써, 높은 시딩 효율성(>93%)과 순도(>99%)를 달성하였으며, HL-LHC 데이터에 대해 기존 방법보다 빠르고 확장 가능한 성능을 보였다. 바렐 영역에서 TrackML 점수는 0.932를 기록하였다.

ABSTRACT

To address the unprecedented scale of HL-LHC data, the Exa.TrkX project is investigating a variety of machine learning approaches to particle track reconstruction. The most promising of these solutions, graph neural networks (GNN), process the event as a graph that connects track measurements (detector hits corresponding to nodes) with candidate line segments between the hits (corresponding to edges). Detector information can be associated with nodes and edges, enabling a GNN to propagate the embedded parameters around the graph and predict node-, edge- and graph-level observables. Previously, message-passing GNNs have shown success in predicting doublet likelihood, and we here report updates on the state-of-the-art architectures for this task. In addition, the Exa.TrkX project has investigated innovations in both graph construction, and embedded representations, in an effort to achieve fully learned end-to-end track finding. Hence, we present a suite of extensions to the original model, with encouraging results for hitgraph classification. In addition, we explore increased performance by constructing graphs from learned representations which contain non-linear metric structure, allowing for efficient clustering and neighborhood queries of data points. We demonstrate how this framework fits in with both traditional clustering pipelines, and GNN approaches. The embedded graphs feed into high-accuracy doublet and triplet classifiers, or can be used as an end-to-end track classifier by clustering in an embedded space. A set of post-processing methods improve performance with knowledge of the detector physics. Finally, we present numerical results on the TrackML particle tracking challenge dataset, where our framework shows favorable results in both seeding and track finding.

연구 동기 및 목표

• 고광도 LHC(HL-LHC) 입자 추적의 계산적 확장성 문제를 해결하기 위해, 막대한 데이터 부피와 높은 충돌 밀도로 인한 도전 과제를 다루기 위해.
• 기계 학습, 특히 그래프 신경망(GNN)을 활용해 전통적인 조합적 비용이 큰 알고리즘을 대체함으로써, 궤적 시딩 및 레이블링 성능을 향상시키기 위해.
• 학습된 히트 임bedding, 임bedded 공간 내 그래프 구축, GNN 기반 엣지 분류를 통합하는 엔드 투 엔드 학습 파이프라인을 개발하기 위해.
• 임bedded 공간 내 효율적 군집화와 검출기 물리학 지식을 통합한 후처리를 통해 성능을 향상시키기 위해.
• 실제 데이터 제약 조건을 반영한 바렐 영역에 집중하여 TrackML 입자 추적 도전 과제 데이터셋에서 최신 기술 수준의 성능을 입증하기 위해.

제안 방법

• 히트를 저차원 유클리드 공간으로 투영하여 효율적인 근접성 질의를 가능하게 하는 학습된 임bedding 공간에서 k-최근접 이웃(KNN)을 사용해 그래프를 구축한다.
• 메트릭 학습을 통해 신경망을 훈련시켜 동일한 입자에서 온 히트들이 임bedded 공간 내에서 더 가까이 오도록 임bedding을 학습시킨다. 이는 비선형 거리 구조를 형성한다.
• 검출기 기하학 및 임bedding에서 유도된 노드 및 엣지 특징을 사용해, GNN을 적용하여 엣지가 진짜인지(feat)인지 분류한다 (즉, 두 히트가 동일한 궤적에 속하는지 여부).
• GNN의 엣지 예측 결과를 바탕으로 고순도 궤적 시딩을 생성한 후, 임bedded 공간에서 DBSCAN 군집화를 적용해 히트를 전체 궤적으로 그룹화한다.
• 검출기 물리학 지식을 반영한 후처리 단계를 도입하여 성능을 향상시키며, 바렐 영역(|η| < 2.1)에 국한하고 조각(프래그먼트)을 필터링한다.
• 빠른 처리와 병렬화를 위해 파이프라인 최적화를 수행하여, 약 5초 내로 이벤트당 처리가 가능하게 하여 HL-LHC 실시간 처리에 적합하게 하였다.

실험 결과

연구 질문

• RQ1 비선형 거리 구조를 가진 학습된 임bedding 공간에서 입자 궤적 재구성에 대한 그래프 구축의 정확도와 효율성을 향상시킬 수 있는가?
• RQ2 GNN 기반 엣지 분류기가 이중체 및 삼중체 시딩에서 기존 수작업 기반 기준을 얼마나 뛰어나게 성능을 낼 수 있는가?
• RQ3 기존 군집화 파이프라인과 비교했을 때, 임bedded 공간 내 군집화가 궤적 탐색에서 순도와 효율성 측면에서 어떤가?
• RQ4 바렐 영역에 국한하고 엔드캡 프래그먼트를 제거하는 것이 최종 TrackML 점수 및 재구성 성능에 어떤 영향을 미치는가?
• RQ5 최소한의 수작업 히ュ리스틱에 의존하면서도 전체 파이프라인이 고성능을 달성할 수 있는가? 이는 HL-LHC에 대한 확장 가능한 엔드 투 엔드 솔루션을 가능하게 하는가?

주요 결과

• 바렐 영역(|η| < 2.1)에 국한하고 엔드캡 프래그먼트를 제외한 경우, TrackML 데이터셋에서 0.932의 TrackML 점수를 기록하여 재구성 가능성 제약 조건 하에서 뛰어난 성능을 보였다.
• 시딩 효율성은 93% 이상, 순도는 99% 이상을 기록하여 정확도와 계산적 확장성 측면에서 기존 방법을 크게 능가하였다.
• 이중체 GNN은 0.805의 TrackML 점수를, 삼중체 GNN은 0.846의 점수를 기록하여 GNN 기반 엣지 분류기가 궤적 시딩에 있어 효과적임을 입증하였다.
• 임bedded 공간 그래프에 대해 DBSCAN 군집화를 적용한 결과 강력한 성능을 기록하였으며, 점수 저하의 주요 원인은 효율성 부족에서 비롯되었고 순도 문제에서는 아니었기 때문에 모델의 일반화 능력이 뛰어남을 시사하였다.
• 파이프라인은 약 5초 내로 이벤트당 처리가 가능하여, HL-LHC 데이터 속도에 적합한 빠르고 병렬 처리 가능한 처리 성능을 확보하였다.
• 전통적인 군집화 및 칼만 필터 기반 방법과 비교해도 유사한 성능을 보였으며, 특히 고밀도 이벤트 처리에서 조합 복잡도를 줄이는 데서 유리함을 입증하였다.