[論文レビュー] Track Seeding and Labelling with Embedded-space Graph Neural Networks

Nicholas Choma, Daniel Murnane+17 more|arXiv (Cornell University)|Jun 30, 2020

Advanced Graph Neural Networks参考文献 16被引用数 11

ひとこと要約

本稿では、グラフ構築に学習された埋め込みを用いたグラフニューラルネットワーク（GNN）を用いたエンドツーエンドの粒子線跡再構成フレームワークを提示する。埋め込み空間内の非線形なメトリクス構造を活用して効率的な近隣探索とクラスタリングを行うことで、本手法は高いシーディング効率（>93%）と純度（>99%）を達成し、バーレル領域におけるTrackMLスコアは0.932を記録。HL-LHCデータの計算スケーラビリティと速度面で、従来手法を上回る性能を発揮する。

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）を用いて、従来の組み合わせ的に高コストなアルゴリズムに代わる、トラックのシーディングとラベル付けの性能を向上させる。
• 学習されたヒット埋め込み、埋め込み空間におけるグラフ構築、GNNベースのエッジ分類を統合したエンドツーエンドの学習パイプラインを開発する。
• 埋め込み空間における効率的なクラスタリングと、検出器物理学の知識を組み込んだ後処理により、性能を向上させる。
• 実際のデータ制約を反映した条件下で、バーレル領域に焦点を当てたTrackML粒子追跡チャレンジデータセットにおいて、最先端の結果を示す。

提案手法

• ヒットを低次元ユークリッド空間に射影することで、効率的な近隣探索を可能にする学習された埋め込み空間におけるk近傍（KNN）を用いてグラフを構築する。
• メトリクス学習を用いてニューラルネットワークを訓練し、同じ粒子由来のヒットが埋め込み空間内で近接するように埋め込みを学習する。これにより、非線形なメトリクス構造が形成される。
• 検出器の幾何学的特徴と埋め込みから得られるノードおよびエッジ特徴を用いて、GNNを用いてエッジを「本物」または「偽物」（2つのヒットが同じトラックに属するか）に分類する。
• GNNのエッジ予測結果をもとに高純度のトラックシードを生成し、その後、埋め込み空間でDBSCANクラスタリングを適用してヒットを完全なトラックにグループ化する。
• 検出器物理学を組み込んだ後処理ステップを導入し、バーレル領域（|η| < 2.1）に制限し、断片をフィルタリングすることで性能を向上させる。
• パイプラインの高速化と並列処理を最適化し、1イベントあたり約5秒の処理時間を達成。HL-LHCのリアルタイム処理に適した性能を実現。

実験結果

リサーチクエスチョン

• RQ1 非線形メトリクス空間における学習された埋め込みは、粒子線跡再構成のためのグラフ構築の精度と効率を向上させ得るか？
• RQ2 GNNベースのエッジ分類器は、二重ヒットおよび三重ヒットのシーディングにおいて、従来の手作業で設計された基準をどの程度上回るか？
• RQ3 埋め込み空間におけるクラスタリングは、従来のクラスタリングパイプラインと比較して、トラック探索における純度と効率の面で優れているか？
• RQ4 バーレル領域への制限およびエンドキャップ断片の除外が、最終的なTrackMLスコアおよび再構成性能に与える影響は何か？
• RQ5 完全なパイプラインは、手作業で設計されたヒューリスティクスにほとんど依存せず、HL-LHCに適したスケーラブルでエンドツーエンドのソリューションを実現できるか？

主な発見

• バーレル領域（|η| < 2.1）に制限しエンドキャップ断片を除外した場合、TrackMLデータセット上で0.932のTrackMLスコアを達成。再構成可能性の制約下でも優れた性能を示した。
• シーディング効率は93%以上、純度は99%以上を達成。両方の観点で、従来手法に比べて顕著に優れた精度と計算スケーラビリティを実現。
• 二重ヒット用GNNは0.805のTrackMLスコア、三重ヒット用GNNは0.846のスコアを達成。GNNによるエッジ分類の有効性が裏付けられた。
• 埋め込み空間のグラフにDBSCANクラスタリングを適用した結果、頑健な性能を発揮。スコアの低下は主に効率性の欠如に起因し、純度の低下ではないため、モデルの一般化能力が優れていることが示された。
• パイプラインは1イベントあたり約5秒で実行可能であり、高速かつ並列処理可能な処理が可能で、HL-LHCのデータレートに適した性能を発揮した。
• 特に高密度イベントの処理において、従来のクラスタリング手法やカルマンフィルタベースの手法と比較して、組み合わせ的複雑性を低減する点で顕著な優位性を示した。