Full event interpretation with machine-learning-based particle-flow reconstruction in the CMS detector

本論文は、CMSフレームワーク内でGPU上に実装された、イベント再構成を単一のモデルへと統合しつつ、標準的な手法と同等の物理性能を達成し、ジェットエネルギー分解能を大幅に向上させ、推論時間を5分の1に削減した、最先端の機械学習ベースの粒子フロー（MLPF）アルゴリズムを提示するものである。

要約

CERNの大型ハドロン衝突型加速器（LHC）を、巨大で高速な粒子衝突マシンとして想像してみてください。2つの陽子が衝突すると、単に壊れるだけでなく、数千もの微細で目に見えない破片へと爆発的に散逸します。CMS検出器は、この爆発の様子を撮影しようとする、高度な技術を備えた巨大なカメラです。その役割は、あらゆる破片が正確に何であるか（それは光子か？電子か？陽子の破片か？）そして、どのくらいの速さで動いているのかを解明することです。

長年、CMSは「粒子フロー（Particle-Flow: PF）」アルゴリズムと呼ばれる「レシピ本」を使用してきました。従来のPFアルゴリズムは、パズルを解こうとする人間の探偵チームのようなものです。彼らは検出器のさまざまな部分（トラッカーやカロリメータ）からの手がかりを調べ、「もしトラックがこのような形をしていて、エネルギーの塊があのような形であれば、それらは同じ粒子であるはずだ」といった、長く厳格に書き込まれたルールに従って点と点を結びつけます。これはうまく機能しますが、動作が遅く、硬直的であり、多くの手動調整を必要とします。

この論文は、よりスマートな探偵である**MLPF（機械学習による粒子フロー）**を紹介しています。

新しい探偵：ニューラルネットワーク

厳格なルールブックに従う代わりに、MLPFは、何百万冊もの物理学の教科書を読み、何百万回ものシミュレーションによる爆発を観察してきた学生のようです。これは、高度な言語モデルの背後にある技術である**トランスフォーマー（Transformer）**と呼ばれる人工知能の一種を使用しています。

• 学習方法: チームはこのAIに、何百万回もの「シミュレートされた」衝突を学習させました。AIに生のデータ（トラックやエネルギーの塊）を与え、「これはシミュレーションで実際に起きたことである」と教えました。これにより、AIは人間が作ったルールでは見逃してしまうかもしれないパターンや相関関係を認識することを学んだのです。
• 思考方法: ヒントを一つずつチェックするのではなく、AIは爆発の「全体」を一度に捉えます。すべてのパズルのピースが、他のすべてのピースとどのように関連しているかを同時に理解するのです。

大きな成果

1. はるかに速い（スピードスター）
従来の探偵（標準的なPF）は標準的なコンピュータプロセッサ（CPU）上で動作し、1回の衝突を分析するのに約110ミリ秒かかります。これは、トランプの束を仕分けするのに長い時間がかかるようなものです。
新しいAI探偵（MLPF）は、この種の重い処理を得意とする専用のグラフィックスカード（GPU）上で動作します。MLPFは同じ作業をわずか20ミリ秒で完了します。これは5倍の高速化です。これは、トランプを手作業で仕分けることから、高速マシンを使うことに切り替えるようなものです。LHCがより忙しくなり、より多くの衝突をより短い時間で処理する必要があるため、このスピードは極めて重要です。

2. より正確（シャープシューター）
AIは膨大な事例から学習したため、従来のルールブックよりも詳細を正確に把握できます。

• ジェット・エネルギー分解能: 物理学において「ジェット」とは、粒子がひとつのパッケージのように振る舞う噴流のことです。論文によると、中規模のジェットにおいて、新しいAIは従来のメソッドよりもエネルギーを10〜20%精密に測定します。リンゴの袋の重さを量る場面を想像してください。従来の方法では数オンスの誤差が出るかもしれませんが、新しい方法ではグラム単位まで正確です。
• 中性粒子: AIは「中性ハドロン」（電荷を持たないため追跡が難しい粒子）を見つけることにも特に優れており、間違いを増やすことなく、より多くの粒子を特定します。

3. 柔軟性がある（カメレオン）
従来のルールは特定の検出器の状態に合わせて作られていました。検出器が変化したり、衝突のエネルギーが変わったりすると、ルールを書き直す必要がありました。しかし、AIは物理学の「原理」を学習しました。論文では、トレーニング中に見ていない、わずかに異なる年やエネルギーレベルのデータでテストした場合でも、AIは良好に機能したことが示されています。つまり、完全な刷新を必要とせずに適応できる「汎用性」を備えているのです。

実世界でのテスト

チームは単にコンピュータ・シミュレーション上でテストしただけではありません。実際に2024年にCMS検出器によって収集された実データを用いて実行しました。彼らは、実データの衝突に対してAIの出力を標準的な手法と比較しました。その結果、物理学的な成果についてはほぼ同一であり、AIが実世界で使用できる準備ができていることを証明しました。

なぜこれが重要なのか（論文による記述）

論文によれば、これはLHCの未来に向けた大きな一歩です。衝突炉がさらに混雑した衝突を扱うようにアップグレードされる際（高輝度LHCと呼ばれるフェーズ）、従来のルールベースの手法では、あまりにも遅く、複雑になりすぎて管理できなくなります。

MLPFアルゴリズムは、複雑で手作業で作られた物理ルールを、以下の特性を持つ単一の統合されたAIモデルに置き換えられることを証明しています。

• より速い（現代的なGPU上で効率的に動作する）。
• よりスマート（測定の精度を向上させる）。
• スケーラブル（将来の膨大なデータ量に対応可能）。

要するに、CMS実験は、チェックリストに従う一対の人間による探偵から、宇宙の深淵なる秘密をより深く見ることを可能にする、全体像を一瞬で捉える超知能AIへと、その「目」をアップグレードしているのです。

技術概要

技術要約：CMS検出器における機械学習ベースの粒子フロー再構成を用いた全イベント解釈

問題提起
粒子フロー（PF）アルゴリズムは、トラッカー、カロリメータ、およびミューオンシステムの信号を統合し、衝突に関するグローバルな粒子レベルの記述を提供することを目的とした、CERNのLHCにおけるCMS実験の基盤となる手法である。標準的なPFの実装は、物理学に基づいたヒューリスティック（経験則）と反復的なルールベースの手法（例：近接性に基づく結合）に依存している。これらの手法は効果的ではあるものの、検出器の複雑化や、高輝度LHC（HL-LHC）時代に向けて増加するパイプアップ（同時プロトン・プロトン衝突）に伴い、スケーラビリティと適応性の面で課題に直面している。従来のクラスタリングタスクは、しばしば高い計算複雑性（$O(N^3)$ または $O(N \log N)$）を示し、現代のハードウェア向けに最適化することが困難な、手作業によるヒューリスティックを必要とする。さらに、現在のPFフレームワークは、全イベント再構成の加速に向けたGPUでの効率的な実行には本来適しておらず、GPUによる加速の可能性を制限している。

手法
本論文では、CMSソフトウェアフレームワーク内における、機械学習ベースの粒子フロー（MLPF）再構成アルゴリズムの実装について述べる。このアプローチは、複数のモジュール化された再構成ステップを、シミュレーションデータ上で直接学習された単一の統一モデルに置き換えるものである。

• アーキテクチャ: MLPFアルゴリズムは、自己注意（self-attention）メカニズムを採用したトランスフォーマーベースのアーキテクチャを利用している。CMSでテストされた従来のグラフニューラルネットワーク（GNN）アプローチとは異なり、このモデルは、アテンションメカニズムに伴う二次的なメモリおよび計算のスケーリングを緩和するために、FLASHATTENTIONカーネルによって加速されたトランスフォーマー層を使用している。モデルは、再構成されたトラックとカロリメータクラスターを入力特徴ベクトルとして受け取り、すべての最終状態粒子を単一の推論ステップで同時に予測する。
• 学習ターゲット: 重要な要素は、「粒子レベルの学習ターゲット」の定義である。このターゲットは、検出器に（直接または崩壊生成物を介して）検出可能な信号を残すすべての粒子を含むように構築されており、イベント生成器（PYTHIA）によって生成された「真の（truth）」粒子（アクセプタンス外に落ちる可能性があるもの）と区別される。ターゲットにはパイプアップの起源を示すバイナリラベルが含まれるが、パイプアップの除去はPUPPIアルゴリズムを用いてダウンストリームで行われる。
• 損失関数: モデルは、以下の3つのコンポーネントからなるマルチタスク損失関数を用いてエンドツーエンドで学習される。
  1. 二値分類 ($L_{cls-binary}$): 入力されたトラックまたはクラスターが、ターゲット粒子の主要な要素であるかどうかを判定する。
  2. 多クラスPID ($L_{cls-PID}$): 粒子種別（光子、電子、ミューオン、荷電ハドロン、中性ハドロン）を分類する。クラスの不均衡に対処するためにフォーカルロス（focal loss）を使用する。
  3. 回帰 ($L_{reg}$): 四元運動量成分（$p_T, \eta, \phi, E$）を予測する。数値的安定性を向上させるため、回帰ターゲットは変換（主要な入力に対する対数比）されており、高エネルギー粒子を強調するための$p_T$依存の重み付けが追加されている。
• 統合: 学習されたモデルは、Open Neural Network Exchange（ONNX）形式にエクスポートされ、ONNXRUNTIMEインターフェースを介してCMSオフラインソフトウェアに統合される。これにより、上流のトラッキングおよび下流のジェットクラスタリングモジュールには手を加えることなく、標準的なPF再構成モジュールを置き換えることが可能となる。

主な貢献

1. 初のデータ検証済みML再構成: 本研究は、ハドロン衝突実験（CMS Run 3, 2023–2024）によって収集されたデータによって検証された、MLベースの全イベント再構成パイプラインの最初の適用である。
2. 統一トランスフォーマーモデル: 反復的なヒューリスティックに依存することなく、検出条件や衝突エネルギーを一般化して学習可能な、全イベント再構成を行うトランスフォーマーベースのアーキテクチャの導入。
3. GPU加速: アルゴリズムのGPUへの展開成功。これにより、複雑なグローバルイベントの再構成をCPUからアクセラレータへオフロードできることを示し、大幅なスピードアップを実現した。
4. 性能検証: MLアプローチが、標準的なPFと同等、あるいは一部の指標において標準的なPFよりも優れた物理性能を達成していることを示す包括的な評価。

結果
MLPFアルゴリズムの性能は、シミュレーションサンプル（トップクォーク・反トップクォークおよびQCD多重ジェットイベント）およびRun 3データの小規模なサブセットを用いて評価された。

• 物理性能:
  • ジェットエネルギー分解能: Run 3条件下でのシミュレーションされたトップクォーク・反トップクォークイベントにおいて、ジェットの横運動量（$p_T$）が30–100 GeVの範囲にあるジェットのエネルギー分解能は、標準的なPFと比較して**10–20%**改善した。
  • 粒子識別: MLPFは、中性ハドロンに対して改善された効率を達成しつつ、標準的なPFと同じ誤識別率を維持した。荷電ハドロンおよび光子については、MLPFは高い効率を示したが、わずかに高い誤識別率を示しており、これは調整が可能である。
  • 欠損横運動量 ($p_T^{miss}$): 再構成された $p_T^{miss}$ 分布は、概ね標準的なPFと一致していたが、MLPFはテイル部分においてやや硬いスペクトルを再構成した。
  • データ検証: Run 3データにおいて、MLPFによって再構成されたダイジェット非対称性とリーディングジェットの $p_T$ 分布は、標準的なPFによるものと一致していることが確認され、現実的な物理出力を生成するモデルの能力が裏付けられた。
• 計算性能:
  • 推論時間: Nvidia L4 GPU上で、MLPFアルゴリズムはイベントあたり中央値20 msの推論時間を達成した。これは、CPU上で動作する標準的なCMS PF再構成の約110 msと比較して大幅な短縮である。
  • スケーラビリティ: FLASHATTENTIONの使用により、複雑なイベントに対しても安定した狭いランタイム分布が可能となり、単一のGPU上で64の並列推論ストリームがサポートされた。アルゴリズムは、反復的な標準PFと比較して、イベントの複雑さが増加した場合のランタイムのスケーリングにおいて優れた特性を示した。

意義
本論文は、MLPFアルゴリズムが、従来のヒューリスティックベースの粒子フロー再構成を、同等または改善された物理性能と優れた計算スケーラビリティを提供する統一モデルによって置き換えられることを実証したと主張している。この研究は、パイプアップの増加と検出器の複雑化が増大するHL-LHCに向けた、効率的なGPU加速ソリューションの基礎を築くものである。本アプローチを実データで検証したことにより、著者らは、MLベースの再構成が単なるシミュレーション上の演習ではなく、プロダクションレベルの物理解析のための実行可能な経路であることを示した。論文は、現在の結果はRun 3の条件に限定されているものの、その手法が高輝度アップグレードおよび将来の衝突器に向けたスケーラブルなパイプラインを提供し、検出器最適化のための微分可能プログラミングを可能にする可能性があることを強調している。