[論文レビュー] Re-Simulation-based Self-Supervised Learning for Pre-Training Foundation Models

Philip Harris, M. Kagan+3 more|arXiv (Cornell University)|Mar 11, 2024

Reservoir Engineering and Simulation Methods被引用数 4

ひとこと要約

本稿では、確率的シミュレータにおける干渉を用いてドメイン完備のデータ拡張を生成することで、対照的学習を強化する再シミュレーションベースの自己教師あり学習手法RS3Lを提案する。シミュレーションの中間段階で潜在状態を固定し、下流部を再実行することで、物理的に整合性のある多様なデータ拡張を生成するRS3Lは、強力な基礎モデルの事前学習を可能にし、高エネルギー物理学におけるジェットタギングおよび不確実性低減の分野で最先端の性能を達成する。

ABSTRACT

Self-Supervised Learning (SSL) is at the core of training modern large machine learning models, providing a scheme for learning powerful representations that can be used in a variety of downstream tasks. However, SSL strategies must be adapted to the type of training data and downstream tasks required. We propose RS3L ("Re-simulation-based self-supervised representation learning"), a novel simulation-based SSL strategy that employs a method of re-simulation to drive data augmentation for contrastive learning in the physical sciences, particularly, in fields that rely on stochastic simulators. By intervening in the middle of the simulation process and re-running simulation components downstream of the intervention, we generate multiple realizations of an event, thus producing a set of augmentations covering all physics-driven variations available in the simulator. Using experiments from high-energy physics, we explore how this strategy may enable the development of a foundation model; we show how RS3L pre-training enables powerful performance in downstream tasks such as discrimination of a variety of objects and uncertainty mitigation. In addition to our results, we make the RS3L dataset publicly available for further studies on how to improve SSL strategies.

研究の動機と目的

• 確率的シミュレータを用いてドメイン完備のデータ拡張を生成する自己教師あり学習戦略を開発し、表現学習を向上させること。
• 再シミュレーションによる物理駆動型データ拡張を活用し、高エネルギー物理学における基礎モデルの事前学習を可能にすること。
• 粒子シャワーイベントの再シミュレーションによる変種を活用することで、ジェットタギングおよび不確実性低減の下流タスク性能を向上させること。
• 再シミュレーションベースの拡張が、物理学アプリケーションにおける対照的学習において、経験的拡張手法を上回ることを示すこと。
• 今後の研究を支援するため、RS3Lデータセットを公開すること。

提案手法

• 確率的シミュレーションの中間でイベントの潜在状態を固定し、下流部を再シミュレートすることで複数の拡張を生成する。
• 干渉点以降の確率的変数を再サンプリングすることで、同じ上流状態を前提とした多様な実現を生成する。
• オリジナルと拡張されたジェットイベントの表現を一致させつつ、類似しないものを分離するように、対照的学習（SimCLRフレームワーク）を適用する。
• 正例ペア（オリジナルと拡張されたジェット）間のコサイン類似度を最大化し、負例ペア間を最小化することで、バックボーンニューラルネットワークが不変表現を学習する。
• 対照的損失における温度ハイパーパrameter τ を最適化し、正例と負例ペアの学習をバランスさせる。Higgs 対 QCD 分類において τ=0.1 が最適であることが判明した。
• 固定重みおよび可変重み設定を用いて、事前学習済みRS3Lモデルを微調整し、転移性能を評価する。

Figure 1: Illustration of the RS3L setup, including downstream re-simulation, sampling, graph computation, and the construction of positive and negative pairs. These are then used in a contrastive loss function aiming to align positive pairs and push negative pairs apart.

実験結果

リサーチクエスチョン

• RQ1 再シミュレーションベースのデータ拡張は、高エネルギー物理学における対照的自己教師あり学習を向上させるドメイン完備の表現を生成できるか？
• RQ2 RS3Lの事前学習は、JetCLRなどの経験的拡張手法と比較して、ジェットタギング性能で優れているか？
• RQ3 RS3Lは、シミュレーションと実データの間のドメインシフトに起因する不確実性をどの程度低減できるか？
• RQ4 Higgs 対 QCD 分類において、RS3Lフレームワークにおける対照的学習の最適な温度パrameter τ は何か？
• RQ5 RS3Lの転移学習設定において、可変重みを用いた微調整は固定重みよりも優れた下流性能を達成できるか？

主な発見

• τ=0.1 でRS3Lの事前学習を実施した結果、Higgs 対 QCD 二値分類で最適な性能が達成され、特に可変重みによる微調整で最も優れた結果が得られた。
• 正例ペア（オリジナルと拡張されたジェット）間の平均コサイン類似度は、0.74 から 0.89 の間で収束し、拡張された表現の強い一致を示した。
• 学習時間の経過に伴い、Higgs ジェットの表現はQCDの表現から明確に分離され、両クラスの平均ベクトルが埋め込み空間の反対側に移動した。
• RS3L空間は、すべてのQCD過程で優れた性能を示し、微調整時に可変重みを使用した場合に特に高いROC AUCが観察された。
• 繰り返しの学習実行において性能のばらつきが小さく、RS3Lの事前学習および微調整パイプラインの統計的安定性と再現可能性が裏付けられた。
• RS3Lデータセットは、科学的機械学習における自己教師あり学習戦略に関する今後の研究を支援するため、公開された。

Figure 2: (Left) One of eight features derived in the RS3L pre-training. (Right) Jet substructure variable $N_{2}$ .