[论文解读] 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，一种基于重仿真（re-simulation）的自监督学习方法，通过在随机模拟器中干预以生成领域完备的数据增强，从而增强对比学习。通过在模拟过程中固定潜在状态并重新运行下游组件，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框架）对原始与增强喷射流事件的表征进行对齐，同时拉大不相似样本之间的距离。
• 通过最大化正样本对（原始与增强喷射流）之间的余弦相似度，最小化负样本对之间的余弦相似度，训练主干神经网络以学习不变表征。
• 优化对比损失中的温度超参数τ，以平衡正负样本对的学习，发现在Higgs与QCD分类任务中τ=0.1为最优值。
• 在预训练的RS3L模型上微调单个Sigmoid头，采用固定权重与浮动权重两种配置，以评估迁移性能。

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 在RS3L框架中，针对Higgs与QCD分类任务，对比学习的最优温度参数τ是多少？
• 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}$ .