[論文レビュー] Dataset for flavour tagging R&D

I. Ochoa, S. B. Klein+2 more|arXiv (Cornell University)|Aug 20, 2024

Electron and X-Ray Spectroscopy Techniques被引用数 5

ひとこと要約

この論文はトークン化を排除し、デコーダを強化し、多様な再構成タスクと集合対集合の生成を探索することで、ジェット物理学のマスクド粒子モデリング（MPM）を改善し、 foundation-model風のバックボーンをジェットデータで事前学習する。これにより MPMv2 と集合対集合フローマッチングを導入し、OODタスクを含む強力な下流性能を示す。

ABSTRACT

In this work, we significantly enhance masked particle modeling (MPM), a self-supervised learning scheme for constructing highly expressive representations of unordered sets relevant to developing foundation models for high-energy physics. In MPM, a model is trained to recover the missing elements of a set, a learning objective that requires no labels and can be applied directly to experimental data. We achieve significant performance improvements over previous work on MPM by addressing inefficiencies in the implementation and incorporating a more powerful decoder. We compare several pre-training tasks and introduce new reconstruction methods that utilize conditional generative models without data tokenization or discretization. We show that these new methods outperform the tokenized learning objective from the original MPM on a new test bed for foundation models for jets, which includes using a wide variety of downstream tasks relevant to jet physics, such as classification, secondary vertex finding, and track identification.

研究の動機と目的

• Foundation-modelスタイルの事前学習を高エネルギー物理学で unlabeled jet data を用いて推進する。
• VQVAE トークン化なしで改良された masked particle modeling (MPMv2) を開発する。
• 条件付き生成アプローチを含む複数の再構成タスクを評価する。
• セット対セットのフローマッチングをジェットの競争力ある事前学習パラダイムとして提案する。

提案手法

• MPM を改良し繰り返し出現するマスク済みトークンを排除し、完全なトランスフォーマーデコーダを使用する。
• trivial を避けるために masked 要素のみに位置エンコーディングを提供する。
• 5つの連続特徴再構成タスクと1つのカテゴリタスク（粒子ID）を調査する。
• K-Means トークン化、CNF（conditional normalizing flow）、flow-matching（CFM）、set-to-set flow-matching（SSFM）など代替ターゲットを探索する。
• JetClass と Delphes でシミュレートした BTag データセットを用いて backbone 表現を事前学習・評価する。
• デコーダタイプ、追加機能、および訓練設定のアブレーション研究（Table 1）。

Figure 1 : A comparison of the original MPM encoder-decoder setup (left) and the new model configuration (right). The new model includes multiple reconstruction tasks, swaps the MLP decoder for a transformer, and only encodes the reduced set.

実験結果

リサーチクエスチョン

• RQ1 トークン化を排除し、より強力なデコーダを使用することで、元の MPMv1 より MPM の性能が改善されるか？
• RQ2 代替再構成ターゲット（CNF、K-means、フローに基づく手法）は MPM の事前学習において VQVAE トークン化と競合し得るか？
• RQ3 改良されたバックボーンは分布内タスク、弱教師付き、分布外タスクでどう性能を示すか？
• RQ4 集合対集合フローマッチングは unordered な jet constituents の妥当な事前学習パラダイムを提供できるか？
• RQ5 extended training、マスク率の調整、追加機能が下流タスクに与える影響は？

主な発見

• MPMv2 はトランスフォーマーデコーダと入力基数を削減することで MPMv1 より分類精度を大きく向上させる。
• MAEスタイルのデコードへ切替え、位置エンコーディングを制限することで性能が向上し、GPU メモリ使用量が抑制される。
• 影響パラメータ特徴量と粒子IDを追加するとさらに精度が向上する（例：回帰 62.2 から 80.4、k-means 70.2 から 83.0 の追加後）。
• 完全にトランスフォーマーベースのデコーダ（MAE）を用いたアブレーションで回帰 79.2 と k-means 81.4 を達成。
• より長い訓練、より深いデコーダ、および 40% のマスク率は最良の結果をもたらす：回帰 83.3 と k-means 84.0。
• 事前学習済みのバックボーンは distribution 内・弱教師付き・分布外タスクのすべてでランダム初期化より優れており、良好な一般化を示す。

Figure 2 : A schematic overview of the SSFM model.