技術要約：高エネルギー物理学の基盤モデルに向けた集合体上のマスク粒子モデリング

問題提起

高エネルギー物理学（HEP）のデータ解析は、従来、特定のタスクや限定的なラベル付きデータセットに対して訓練された教師あり学習モデルに依存してきました。自己教師あり学習（SSL）に基づく基盤モデル（Foundation Models: FMs）は、自然言語処理（NLP）やコンピュータビジョン（CV）に革命をもたらしましたが、HEPへの適用はまだ初期段階にあります。HEPコミュニティは、これらのパラダイムを適応させる上で特有の課題に直面しています：

1. データ構造： HEPのデータ、特に粒子ジェットは、テキストの逐次的な性質や画像のグリッド構造とは異なり、**無秩序な集合（unordered sets）**の粒子で構成されています。
2. 特徴量の連続性： 粒子の特徴量（運動量、方向など）は連続的であり、言語モデルで使用される離散的な語彙とは対照的です。
3. ラベルの希少性とドメインシフト： 教師あり学習は膨大なシミュレーションデータを必要とすることが多く、それらが実際の実験データと完全に一致しない場合、ドメインシフトによる不確実性が生じます。ラベル付きデータが少ない状況でも、大量のラベルなしデータから汎用的な表現を学習し、ダウンストリームタスクの性能を向上させるモデルが必要です。

手法：マスク粒子モデリング（MPM）

著者らは、連続的な入力の無秩序な集合に適応させた自己教師あり学習スキームである**マスク粒子モデリング（MPM）**を提案しています。この手法は、NLPにおけるマスク言語モデリング（MLM、例：BERT）や、マスク画像モデリング（例：BEiT）に着想を得ています。

コアフレームワーク

1. 入力表現： ジェットは、$N$個の粒子からなる集合 $X = \{x_i\}$ として扱われ、各 $x_i$ は連続的な特徴量を持つベクトルです。
2. マスキング戦略： 各ジェット内の粒子のサブセットがマスクされ、学習可能なベクトルに置き換えられます。
3. トークン化（離散化）： 連続的な特徴量を扱い、分類ベースの損失関数を可能にするために、著者らは**ベクトル量子化変分オートエンコーダ（VQ-VAE）**を採用しています。
   • 学習済みのVQ-VAEエンコーダが、連続的な粒子特徴量を有限のコードブック内の最も近いベクトルへと写像します。
   • コードブックベクトルのインデックスが、離散的な「トークン」として機能します。
   • これにより、モデルは連続値を直接回帰するのではなく、マスクされた粒子の「アイデンティティ（トークンインデックス）」を予測することが可能になります。
   • 代替案： より単純でコンテキストに依存しないトークン化手法として、K-meansクラスタリングも検討されています。
4. モデルアーキテクチャ：
   • バックボーン： Transformerエンコーダ（Normformerに基づく）が粒子の集合を処理します。
   • 置換不変性： ジェットは無秩序な集合であるため、バックボーンは置換等変（permutation equivariant）である必要があります。しかし、すべてのマスクされた粒子に同一のマスクベクトルを使用すると、モデルがマスクされた位置を区別できなくなるという冗長性が生じます。
   • 順序付け戦略： ダウンストリームタスクに対する置換不変性を損なうことなく、この縮退を打破するために、著者らはマスク予測ヘッドへの入力時のみに順序付けを導入しています。粒子は横運動量（$p_T$）によって順序付けられ、予測ステップにおいてのみ学習された位置エンコーディングが割り当てられます。バックボーン自体は置換不変なまま維持されます。
5. 学習目的： モデルは、マスクされていないコンテキストに基づいて、マスクされた粒子の離散的なトークン・アイデンティティを予測するようにクロスエントロピー損失を最小化します。

ファインチューニング戦略

論文では、以下の3つのファインチューニング手法を評価しています：

1. 固定バックボーン（Fixed Backbone）： 学習済みのバックボーンは凍結し、線形分類ヘッドのみを訓練します。
2. ファインチューニング（Fine-Tuned）： バックボーンと分類ヘッドの両方を更新します。
3. スクラッチからの学習（From Scratch）： 同一のアーキテクチャを持つモデルを、ダウンストリームタスクに対してランダムな初期化から訓練します（教師あり学習のベースライン）。

主な貢献と実験結果

1. トークン化と順序付けのデザイン

JetClassデータセット（1億ジェット、10クラス）を用いた実験により、以下が明らかになりました：

• 順序付け： 粒子を（バックボーンではなく）予測ヘッドでのみ順序付ける手法が最も高い精度（56.8%）を示し、バックボウンの入力として順序付けする場合や順序付けを行わない場合よりも優れた性能を発揮しました。これにより、バックボーンの置換不変性を維持しつつ、精密な粒子ごとの予測のための対称性を打破できます。
• 離散化 vs 回帰： 連続的な特徴量を直接回帰するよりも、離ックなトークン（VQ-VAE経由）を用い分類損失を使用する方が、大幅に高い性能を示しました。回帰は事後分布の平均を学習する傾向があり、マルチモーダルな分布を捉えることに失敗しますが、分類はトークン上の完全な事後分布を学習します。
• VQ-VAE vs K-means： VQ-VAEのトークンはK-meansのトークンよりも優れた性能を示しました。これは、おそらくVQ-VAEがコンテキスト依存の特徴を捉えられるためです。ただし、K-meansも比較的容易な学習で妥当な性能を提供しました。

2. ダウンストリームタスクの性能

• インコンテキスト学習（同一クラス）： JetClassテストセットにおいて、ラベル付きデータが少ない場合（<10kジェット）、事前学習済みモデルは「スクラッチ」モデルを大幅に上回りました。ラベル付きデータが増加するにつれ、教師ありモデルはファインチューニングされた事前学習モデルの性能に収束しました。
• アウトオブコンテキスト学習（新クラス）： 6つのクラス（HiggsおよびQCD）で事前学習を行い、未知の4つのクラス（Top, W, Z）でファインチューニングを行いました。事前学習済みモデルは、限られたデータにおいて教師ありベースラインを再び上回り、学習された表現が汎用的かつ転移可能であることを示しました。
• アウトオブドメイン学習（新しいデータセット）： JetClassで事前学習し、異なる検出器シミュレーションやジェットクラスタリングパラメータを持つRODEMデータセットでファインチューニングした結果、強いデータ効率が示されました。事前学習済みモデルは、小さなラベル付きセットを用いてRODEMデータセットにおいて高い性能を維持しており、ドメインシフトに対する堅牢性を示唆しています。
• 弱教師あり分類： シグナルとバックグラウンドのイベントを混合した「ノイズの多いラベル」を用いたタスクでテストされました。事前学習済みのバックボーンは、たとえ線形ヘッドのみをファインチューニングした場合でも、スクラッチで訓練されたモデルと比較して、純粋なシグナルとバックグラウンドのデータセットを分離する有意性を大幅に向上させました。

意義と主張

本論文は、**マスク粒子モデリング（MPM）**が、高エネルギー物理学における基盤モデルを構築するための実行可能かつ効果的な戦略であることを主張しています。その主な意義は以下の通りです：

1. 集合体へのSSLの適応： 無秩序な連続データの集合に対して、NLPやCVからのマスキング戦略を正常に一般化し、HEP特有の構造的課題に対処しました。
2. データ効率： 事前学習された「バックボーン」は強力な初期化を提供し、スクラッチでの訓練と比較して、はるかに少ないラベル付きデータでダウンストリームタスクの高い性能を達成することを可能にします。
3. 汎用性： 学習された表現は、事前学習中に見た特定のクラスやデータセットに限定されません。それらは、新しい粒子クラスや新しいデータドメイン（異なるパラメータを持つシミュレーション）に対して効果的に転移します。
4. 実データへの事前学習への道： 著者らは、ラベルなしの実の実験データ（豊富に存在する）で直接事前学習を行い、より小さなシミュレーションデータセットでファインチューニングできる可能性を強調しています。このアプローチは、不完全なシミュレーションから生じるドメインシフトの不確実性を軽減できる可能性があり、これはHEPにおける極めて重要な問題です。

本研究は、MPMがHEPにおける大規模な自己教師あり学習に向けた基礎的な一歩であることを位置づけており、事前学習データとモデルサイズをスケールアップすることで、ドメイン適応の課題をさらに克服できる可能性を示唆しています。