[논문 리뷰] Re-Simulation-based Self-Supervised Learning for Pre-Training Foundation Models

Philip Harris, M. Kagan외 3명|arXiv (Cornell University)|2024. 03. 11.

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 τ를 최적화하여 양의 쌍과 음의 쌍 간의 학습 균형을 조절하며, 힉스 대 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 히iggs 대 QCD 분류를 위한 RS3L 프레임워크에서 대비 학습의 최적 온도 파rameter τ는 무엇인가?
• RQ5 RS3L 전이 학습 설정에서 유동 가중치로의 미세조정이 고정 가중치보다 더 뛰어난 최종 성능를 제공하는가?

주요 결과

• τ=0.1로 RS3L 사전 훈련을 수행한 결과, 최적의 힉스 대 QCD 이진 분류 성능를 달성하였으며, 특히 유동 가중치로의 미세조정 시 최고의 성능를 기록함.
• 양의 쌍(원본 및 증강된 제트) 간의 평균 여유코사인 유사도는 0.74에서 0.89 사이로 수렴하여 증강 표현 간의 강력한 정렬을 나타냄.
• 훈련 시간이 지남에 따라 힉스 제트 표현은 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}$ .