[Paper Review] 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 Methods4 citations

TL;DR

This paper proposes RS3L, a re-simulation-based self-supervised learning method that enhances contrastive learning by generating domain-complete data augmentations through intervention in stochastic simulators. By fixing latent states mid-simulation and re-running downstream components, RS3L produces diverse, physics-consistent augmentations, enabling powerful foundation model pre-training that achieves state-of-the-art performance in jet tagging and uncertainty mitigation in high-energy physics.

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.

Motivation & Objective

• To develop a self-supervised learning strategy that generates domain-complete data augmentations using stochastic simulators for improved representation learning.
• To enable the pre-training of foundation models in high-energy physics using physics-driven data augmentation via re-simulation.
• To improve downstream performance in jet tagging and uncertainty mitigation by leveraging re-simulated variations of particle shower events.
• To demonstrate that re-simulation-based augmentation outperforms empirical augmentation methods in contrastive learning for physics applications.
• To release the RS3L dataset to support further research on self-supervised learning in scientific machine learning.

Proposed method

• Intervene in the middle of a stochastic simulation to fix the latent state of an event, then re-simulate downstream components to generate multiple augmentations.
• Use a simulator to re-sample stochastic variables after the intervention point, producing diverse realizations conditioned on the same upstream state.
• Apply contrastive learning (SimCLR framework) to align representations of original and augmented jet events while pushing apart dissimilar ones.
• Train a backbone neural network to learn invariant representations by maximizing cosine similarity between positive pairs (original and augmented jets) and minimizing it for negative pairs.
• Optimize the temperature hyperparameter τ in the contrastive loss to balance positive and negative pair learning, finding τ=0.1 optimal for Higgs vs. QCD classification.
• Fine-tune the pre-trained RS3L model with a single sigmoid head on top, using both fixed and floating weight configurations to evaluate transfer performance.

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.

Experimental results

Research questions

• RQ1 Can re-simulation-based data augmentation generate domain-complete representations that improve contrastive self-supervised learning in high-energy physics?
• RQ2 How does RS3L pre-training compare to empirical augmentation methods like JetCLR in jet tagging performance?
• RQ3 To what extent can RS3L mitigate uncertainties arising from domain shift between simulations and real data?
• RQ4 What is the optimal temperature parameter τ for contrastive learning in the RS3L framework for Higgs vs. QCD classification?
• RQ5 Does fine-tuning with floating weights yield better downstream performance than fixed weights in the RS3L transfer learning setup?

Key findings

• RS3L pre-training with τ=0.1 achieved optimal Higgs vs. QCD binary classification performance, with the best results observed using floating weight fine-tuning.
• The average cosine similarity between positive pairs (original and augmented jets) converged to values between 0.74 and 0.89, indicating strong alignment of augmented representations.
• Over training time, Higgs jet representations were increasingly separated from QCD representations, with the average vectors of the two classes moving to opposite ends of the embedding space.
• The RS3L space demonstrated robust performance across all QCD subprocesses, with superior ROC AUC observed when using floating weights during fine-tuning.
• Repeated training runs showed low variance in performance, indicating statistical stability and reproducibility of the RS3L pre-training and fine-tuning pipeline.
• The RS3L dataset has been released publicly to support further research into self-supervised learning strategies in scientific machine learning.

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