[논문 리뷰] CosmicRIM : Reconstructing Early Universe by Combining Differentiable Simulations with Recurrent Inference Machines

Chirag Modi, François Lanusse외 3명|arXiv (Cornell University)|2021. 04. 26.

Computational Physics and Python Applications참고 문헌 17인용 수 5

한 줄 요약

CosmicRIM는 우주론에서 초기 조건을 재구성하기 위해 미분 가능한 N-body 시뮬레이션(FlowPM)과 학습된 최적화를 결합한, 미분 가능한 순환 추론 기계(RIM) 프레임워크를 제안한다. 이는 ADAM과 LBFGS의 냉각 기법을 사용하는 것보다 40배 빠른 수렴 속도를 달성하면서도, 물리학에 기반한 냉각 스케줄을 암묵적으로 학습함으로써 대규모 스케일에서 재구성 품질을 유사하거나 향상시킨다.

ABSTRACT

Reconstructing the Gaussian initial conditions at the beginning of the Universe from the survey data in a forward modeling framework is a major challenge in cosmology. This requires solving a high dimensional inverse problem with an expensive, non-linear forward model: a cosmological N-body simulation. While intractable until recently, we propose to solve this inference problem using an automatically differentiable N-body solver, combined with a recurrent networks to learn the inference scheme and obtain the maximum-a-posteriori (MAP) estimate of the initial conditions of the Universe. We demonstrate using realistic cosmological observables that learnt inference is 40 times faster than traditional algorithms such as ADAM and LBFGS, which require specialized annealing schemes, and obtains solution of higher quality.

연구 동기 및 목표

• 우주론적 설문 데이터로부터 초기 밀도장 재구성이라는 고차원적이고 비선형적인 역문제를 해결하기 위해.
• ADAM과 LBFGS와 같은 전통적인 최적화 방법이 수백 번의 비용이 많이 드는 N-body 시뮬레이션을 필요로 하여 계산적으로 비현실적인 문제를 해결하기 위해.
• 수동 조정 없이도 물리학 기반 냉각 스케줄을 암묵적으로 모방하는 효율적인 추론 기법을 학습하기 위해.
• 다음 세대의 우주론적 설문에서 반복적인 재구성이 요구되는 빠르고 확장 가능한 추론을 가능하게 하기 위해.

제안 방법

• 이 방법은 初기 조건에 대한 전방 모델의 기울기를 계산하기 위해 미분 가능한 N-body 시뮬레이터(FlowPM)를 사용한다.
• RIM은 은닉 상태, 현재 재구성 결과, 로그우도의 기울기를 조합하는 순환 업데이트 규칙을 갖는다.
• RIM은 로그우도의 음수를 최소화하도록 훈련되어, 물리학적 냉각을 암묵적으로 모방하는 최적화 경로를 학습한다.
• 아키텍처는 메모리 상태와 재구성 및 은닉 상태에 대한 별도의 업데이트 함수를 포함하여 반복적인 정밀화를 가능하게 한다.
• 초기 조건에 대해 정규 분포 사전을 사용하고, 가우시안 노이즈를 가감하는 우도 모델을 정의하여 우도를 설정한다.
• 고정된 기준 스펙트럼을 가진 시뮬레이션 데이터를 사용하여 추론 네트워크의 엔드 투 엔드 최적화를 수행한다.

실험 결과

연구 질문

• RQ1 학습된 최적화 기법이 ADAM과 LBFGS와 같은 전통적인 기울기 기반 방법보다 초기 밀도장 재구성에서 더 우수한 성능을 내는가?
• RQ2 RIM 프레임워크는 명시적인 설계 없이도 물리학에 기반한 냉각 스케줄을 암묵적으로 학습하는가?
• RQ3 수렴 속도와 재구성 정확도 측면에서 RIM의 성능은 특수화된 냉각 기법과 비교해 어떻게 되는가?
• RQ4 RIM 프레임워크는 다양한 初기 조건과 설문 구성에 대해 일반화하여 훈련될 수 있는가?

주요 결과

• CosmicRIM는 대규모 스케일에서 ADAM과 LBFGS의 냉각 기법을 유사하거나 뛰어나게 하는 재구성 품질을 달성한다.
• 이 방법은 단지 10회의 추론 단계만으로도 냉각 기법을 사용하는 ADAM과 LBFGS(400단계 필요)에 비해 40배 빠른 수렴 속도를 확보한다.
• 표준 재구성 결과로 초기화된 후에도 CosmicRIM는 성능을 향상시키며, 초기화에 대한 강건성을 보여준다.
• RIM은 대규모 구조를 먼저 재구성하는 최적화 경로를 암묵적으로 학습하여 물리학 기반 냉각 기법의 행동을 모방한다.
• 최종 RIM 출력에서 시작해 ADAM이나 LBFGS로 추가 최적화를 수행해도 향상되지 않으며, 이는 RIM이 고품질의 MAP 추정치에 도달했음을 시사한다.
• 이 프레임워크는 확장 가능한 추론을 가능하게 하며, 메모리 비용을 줄이기 위해 하위 네트워크의 순차적 훈련을 통해 확장 가능성이 있다.