Similarity Pairing with Energy Mover's Distance for Self-Supervised Pre-Training at the LHC

技術要約：LHCにおける自己教師あり学習のためのEnergy Mover's Distanceを用いた類似性ペアリング

問題提起
自己教師あり学習（SSL）は、高エネルギー物理学（HEP）、特に大型ハドロン衝突型加速器（LHC）のデータにおける基盤モデルの学習において、不可欠な要素となっています。標準的なSSLワークフローは、データの拡張（augmentation）に依存してイベントの複数の「ビュー」を作成し、特定の歪みや対称性に対するモデルの不変性を学習させます。しかし、HEPにおいては、コンピュータビジョンとは異なり、任意の拡張を適用することは厳格な物理法則や検出器の制限によって制約されます。手作業による拡張は、物理的対称性を損なったり、イベントの忠実性を侵害したりするリスクがある一方で、追加のシミュレーションイベントのバリアントを生成することは計算コストが高くなります。その結果、物理の内容を保持しつつ、学習に十分な変化を提供する堅牢な拡張セットを選択することが、大きな課題となっています。

手法
著者らは、手作業による拡張に代わるデータ駆動型の代替案として、Energy Mover's Distance (EMD) による類似性ペアリングを提案しています。単一のイベントを歪ませて拡張されたビューを作成する代わりに、この手法は、幾何学的にエネルギー分布が類似している2つの異なる実衝突イベントをペアにします。

• EMDペアリング: 最適輸送（optimal transport）指標であるEnergy Mover's Distanceは、あるイベントのエネルギー分布を別のイベントへと変換するために必要な「仕事量」を定量化します。小さなEMD値は、2つの異なるイベントが幾何学的に類似していることを示しており、これは実質的に「ゆらぎを除けば同一のイベントである」ことを意味します。EMD値が低いイベントをペアにすることで、モデルはイベントの物理的内容を変更したり人工的な歪みを導入したりすることなく、真の物理的ゆらぎに対する不変性を学習します。
• 事前学習フレームワーク: 著者らは、このペアリング戦略を、iBOTアーキテクチャに基づく自己蒸留法である jBOT フレームワークに統合しています。この設定では：
  • 2つの異なるジェットが、EMDの類似性に基づいてペアにされます。
  • 教師・生徒（teacher-student）エンコーダ・アーキテクチャが採用されています。生徒はマスクされたバージョンのジェットを処理し、教師は完全なジェットを処理します。
  • 目的関数は、粒子レベルおよびジェットレベルの両方において、教師と生徒の出力分布間のクロスエントロピーを最小化することです。
  • このペアリングは拡張ステップを置き換えるものであり、ビューのペアを必要とする限り、特定のSSLアルゴリズムに依存しない手法となっています。

実験設定
本手法は、13 TeVでの陽子・陽子衝突ジェットを含む JetNetデータセット を用いて評価されました。このデータセットには、5つのクラス（軽クォーク $q$、グルオン $g$、$W$ ボソン、$Z$ ボソン、トップクォーク $t$）が含まれています。

• 事前学習: モデルは、Transformerエンコーダを用い、QCDジェット（$q$ および $g$）のみを用いて事前学習されました。
• ペアリング・スキーム: 以下の3つのペアリングモードを探索しました：
  1. Same-class（同クラス）: $q$ と $q$、$g$ と $g$ をペアにする。
  2. Different-class（異クラス）: $q$ と $g$、またはその逆をペアにする。
  3. Random（ランダム）: クラスの制約なしに、EMDの類似性のみに基づいてペアリングを行う。
• ダウンストリーム・タスク: 事前学習された埋め込み（embedding）は、ファインチューニングを行わずに、k最近傍法（kNN）、コサイン類似度、およびマハラノビス距離を用いて、異常検知（QCD背景事象に対する $W, Z, t$ ジェットの識別）のためにプローブされました。

主な結果

• 埋め込み構造: t-SNE可視化により、EMDペアリング手法が明確なクラスタリング構造を持つ埋め込みを生成することが明らかになりました。特筆すべきは、モデルが事前学習に使用されなかったジェットクラス（$W, Z, t$）を正常に分離できたことであり、これは学習された表現が高次のセマンティクスを捉えていることを示しています。QCDジェット（$q$ と $g$）は、ベースラインとなる拡張ベースのアプローチと比較して、より密度が高く局在化したクラスターを形成しました。
• 異常検知性能:
  • ペアリング・スキームの中で、ランダム・ペアリング・モード（クラス制約のないもの）が最も高い性能（Combined Signal AUC: $0.8188 \pm 0.0076$）を示し、これは同クラス・モード（$0.8142 \pm 0.0090$）と統計的に互換性がありました。これは、クラス制約を取り除くことで、不変性学習の空間が広がることを示唆しています。
  • ベースラインとの比較: 拡張ベースのjBOTベースラインと比較した際、EMDペアリング法は同等または優れた性能を達成しました。具体的には、コサイン類似度スコアを用いた場合、EMDペアリングモデルは、結合信号（combined signal）に対して高いAUC（$0.8309$対$0.8269$）を達成し、$W$信号クラスに対しても高い値を示しましたが、$Z$および$t$ クラスについては、使用する指標に応じて同等またはわずかに低い性能となりました。

意義と主張
本論文は、EMDによる類似性ペアリングが、高エネルギー物理学における自己教師あり事前学習のための、手作業による拡張に代わる実行可能なデータ駆動型の手法であることを主張しています。その主な意義は、以下の能力にあります：

1. 物理の保存: 人工的な歪みを避けることで、物理法則や検出器の制約の完全性を維持する。
2. 計算コストの削減: 追加のシミュレーションイベントのバリアントを生成する必要性を排除する。
3. 表現の向上: 物理的に類似したイベント間のゆらぎに対して鈍感なセマンティックなジェット埋め込みを生成し、従来の拡張ベースの手法に匹敵するかそれを上回るダウンストリームの識別力を実現する。

著者らは、この技術が基盤モデルの学習において有効であると結論付けており、他のデータタイプ（画像ベースの検出器データなど）への将来的な応用を示唆していますが、この範囲を超えた具体的な実装の詳細については述べていません。