[논문 리뷰] CONNECT: A neural network based framework for emulating cosmological observables and cosmological parameter inference

Andreas Nygaard, Emil Brinch Holm외 2명|arXiv (Cornell University)|2022. 05. 30.

Galaxies: Formation, Evolution, Phenomena인용 수 5

한 줄 요약

CONNECT는 CLASS 봄볼츠만 해소기에서 유도된 천체물리 관측량을 모의하는 PyTorch 기반 신경망 프레임워크로, 빠르고 정확한 천체물리 모수 추론을 가능하게 한다. 반복 최적화된 라틴 초체 샘플링을 사용하여 효율적인 학습 데이터를 생성함으로써, CLASS 평가 횟수를 두 배수 감소시키고 MCMC 추론 속도를 수배수로 향상시키며, 전체 CLASS 실행과 비교해 후행 분포의 차이가 표준편차의 0.1 이내로 유지된다.

ABSTRACT

Bayesian parameter inference is an essential tool in modern cosmology, and typically requires the calculation of $10^5$--$10^6$ theoretical models for each inference of model parameters for a given dataset combination. Computing these models by solving the linearised Einstein-Boltzmann system usually takes tens of CPU core-seconds per model, making the entire process very computationally expensive. In this paper we present extsc{connect}, a neural network framework emulating extsc{class} computations as an easy-to-use plug-in for the popular sampler extsc{MontePython}. extsc{connect} uses an iteratively trained neural network which emulates the observables usually computed by extsc{class}. The training data is generated using extsc{class}, but using a novel algorithm for generating favourable points in parameter space for training data, the required number of extsc{class}-evaluations can be reduced by two orders of magnitude compared to a traditional inference run. Once extsc{connect} has been trained for a given model, no additional training is required for different dataset combinations, making extsc{connect} many orders of magnitude faster than extsc{class} (and making the inference process entirely dominated by the speed of the likelihood calculation). For the models investigated in this paper we find that cosmological parameter inference run with extsc{connect} produces posteriors which differ from the posteriors derived using extsc{class} by typically less than $0.01$--$0.1$ standard deviations for all parameters. We also stress that the training data can be produced in parallel, making efficient use of all available compute resources. The extsc{connect} code is publicly available for download at \url{https://github.com/AarhusCosmology}.

연구 동기 및 목표

• 천체물리학에서 베이지안 모수 추론의 높은 계산 비용 문제를 해결하기 위해, 일반적으로 아인슈타인-봄볼츠만 해소기의 평가 횟수가 10^5~10^6에 이르는 데 기인한다.
• 신경망 모의기의 학습에 필요한 고비용의 CLASS 평가 횟수를 줄이기 위해, 새로운 반복적 샘플링 전략을 개발한다.
• 새로운 데이터셋에 대해 재학습 없이도 빠르고 정확한 모수 추론이 가능한 즉시 사용 가능한 프레임워크를 구축한다.
• 기존의 MCMC 샘플러인 MontePython과 Cobaya와의 호환성을 확보하여 기존의 천체물리학 워크플로우에 널리 활용 가능하도록 한다.
• 재현 가능한 오픈소스 솔루션을 제공하며, 사용자 정의 천체물리 출력과 GPU 가속 기능을 지원한다.

제안 방법

• CLASS가 계산한 천체물리 관측량(예: CMB 스펙트럼, 물질 스펙트럼)을 모의하기 위해 완전히 연결된 피드포워드 신경망을 사용한다.
• 고품질의 학습 데이터를 생성하기 위해 반복적 라틴 초체 샘플링 전략을 적용하여, 필요한 CLASS 평가 횟수를 두 배수 감소시킨다.
• 입력 설정 파일을 통해 사용자 정의 가능한 출력 변수를 정의하고, 천체물리 모수와 그에 상응하는 CLASS 출력 결과의 데이터셋을 기반으로 신경망을 학습시킨다.
• CLASS 호출을 신속한 신경망 추론으로 대체하는 플러그인을 통해 MontePython과 Cobaya에 원활하게 통합한다.
• 모듈식 코드 아키텍처를 통해 사용자 정의 손실 함수와 네트워크 아키텍처를 지원하며, 특정 천체물리 모델에 맞게 맞춤화할 수 있다.
• 병렬 방식으로 학습 데이터를 생성하여 고성능 컴퓨팅 자원을 효율적으로 활용한다.

실험 결과

연구 질문

• RQ1 감소된 CLASS 평가 횟수로 학습된 신경망 모의기로 천체물리 관측량을 고정밀도로 모의할 수 있는가?
• RQ2 반복적 샘플링 전략은 얼마나 많은 CLASS 평가 횟수를 줄일 수 있으며, 동시에 모의기 정확도를 유지할 수 있는가?
• RQ3 전체 CLASS 계산과 비교했을 때, 모의기를 사용해 유도한 천체물리 모수의 후행 분포는 얼마나 정확한가?
• RQ4 재학습 없이도 다양한 데이터 조합에 대해 모의기를 재사용할 수 있는가? 이 경우 추론 속도와 정확도에 어떤 영향을 미치는가?
• RQ5 MontePython과 같은 기존의 MCMC 샘플링 파이프라인에 얼마나 쉽게 통합할 수 있으며, 코드 수정 최소화가 가능한가?

주요 결과

• CONNECT는 표준 추론 실행 대비 학습에 필요한 CLASS 평가 횟수를 두 배수 감소시켜 데이터 생성 비용을 크게 낮춘다.
• CONNECT를 사용한 천체물리 모수 추론 결과는 전체 CLASS 실행 결과와 비교해 모든 모수에서 표준편차의 0.01~0.1 이내로 후행 분포의 차이를 보였다.
• MCMC 추론이 수시간 내로 완료되어, 수일 또는 수주일이 걸리던 기존 방식에 비해 빠르게 수행되며, 추론 시간은 주로 가능도 평가에 의해 결정된다.
• 학습된 모의기는 재학습 없이도 다양한 데이터 조합에 대해 완전히 재사용 가능하여, 새로운 데이터가 추가될 경우 신속한 재분석이 가능하다.
• MontePython과 Cobaya와 호환되며, 플러그인 인터페이스는 원본 코드베이스에 대한 수정 없이도 작동한다.
• 학습 데이터는 병렬로 생성 가능하여 가용한 HPC 자원을 최적화하고 확장 가능한 데이터 생성을 가능하게 한다.