[논문 리뷰] JAX-COSMO: An End-to-End Differentiable and GPU Accelerated Cosmology Library

J.E. Campagne, François Lanusse외 8명|arXiv (Cornell University)|2023. 02. 10.

Galaxies: Formation, Evolution, Phenomena인용 수 4

한 줄 요약

jax-cosmo는 JAX 기반으로 구축된 미분 가능하고 GPU 가속 기능을 갖춘 천체물리학 라이브러리로, 천체물리학적 우도 함수의 엔드 투 엔드 자동 미분을 가능하게 한다. 이는 히어레틱 편차 행렬 계산과 기울기 기반 방법(예: 해밀토니안 몬테 카를로 및 NUTS)을 통한 사후 분포 추론을 가속화하며, DES-Y1 3x2pt 분석에서 전통적인 MCMC 대비 GPU 시간당 약 5배 높은 유효 표본 크기를 달성한다.

ABSTRACT

We present jax-cosmo, a library for automatically differentiable cosmological theory calculations. It uses the JAX library, which has created a new coding ecosystem, especially in probabilistic programming. As well as batch acceleration, just-in-time compilation, and automatic optimization of code for different hardware modalities (CPU, GPU, TPU), JAX exposes an automatic differentiation (autodiff) mechanism. Thanks to autodiff, jax-cosmo gives access to the derivatives of cosmological likelihoods with respect to any of their parameters, and thus enables a range of powerful Bayesian inference algorithms, otherwise impractical in cosmology, such as Hamiltonian Monte Carlo and Variational Inference. In its initial release, jax-cosmo implements background evolution, linear and non-linear power spectra (using halofit or the Eisenstein and Hu transfer function), as well as angular power spectra with the Limber approximation for galaxy and weak lensing probes, all differentiable with respect to the cosmological parameters and their other inputs. We illustrate how autodiff can be a game-changer for common tasks involving Fisher matrix computations, or full posterior inference with gradient-based techniques. In particular, we show how Fisher matrices are now fast, exact, no longer require any fine tuning, and are themselves differentiable. Finally, using a Dark Energy Survey Year 1 3x2pt analysis as a benchmark, we demonstrate how jax-cosmo can be combined with Probabilistic Programming Languages to perform posterior inference with state-of-the-art algorithms including a No U-Turn Sampler, Automatic Differentiation Variational Inference,and Neural Transport HMC. We further demonstrate that Normalizing Flows using Neural Transport are a promising methodology for model validation in the early stages of analysis.

연구 동기 및 목표

• 천체물리학적 파라미터 및 시스템적 요인에 대해 천체물리학적 우도 함수의 자동 미분을 가능하게 하는, 유연한 천체물리학 라이브러리 개발
• JAX의 JIT 컴파일, GPU 가속, 배치 처리를 활용한 고차원 파라미터 공간에서의 천체물리학적 추론 가속화
• 정교한 튜닝 없이도 안정적인 히어레틱 편차 행렬 계산과 기울기 기반 사후 분포 추론(예: 해밀토니안 몬테 카를로 및 변분 추론) 지원
• 다양한 히어레틱 편차 행렬을 통한 설문 조사 최적화 지원으로, Figure of Merit에 대한 기울기 하강 기반 최적화 가능
• NumPyro와 같은 확률적 프로그래밍 언어와의 통합을 통해 천체물리학 분석에서 확장성 있고 고성능의 사후 샘플링을 위한 지원 제공

제안 방법

• JAX의 자동 미분(autodiff)을 활용해 천체물리학적 관측치의 정확한 기울기 계산을 천체물리학적 파라미터 및 시스템적 요인에 대해 수행
• 배경 진화, 선형 및 비선형 스펙트럼(_halofit_ 또는 _Eisenstein & Hu_ 사용), 그리고 리머 근사법을 통한 각도 스펙트럼(Cℓ)과 같은 핵심 천체물리학 컴포넌트의 유연한 버전 구현
• JAX의 JIT 컴파일 및 GPU/TPU 가속을 활용해 우도 평가 및 샘플링 알고리즘의 성능 향상
• NUTS, ADVI, Neural Transport HMC와 같은 기울기 기반 추론 기법을 지원하기 위해 _NumPyro_ 와 같은 확률적 프로그래밍 프레임워크와 통합
• 다중 GPU 또는 CPU를 통해 천체물리학 시뮬레이션 및 파라미터 스윕을 위한 배치 처리를 위해 _vmap_ 활용
• 3x2pt 분석을 위한 유연한 우도 함수 지원으로 이론에서 데이터 비교에 이르기까지 엔드 투 엔드의 유연성 확보

Figure 1: Comparison of the radial comoving distance between CCL and jax-cosmo . More plots are available in the companion notebook \faFileCodeO .

실험 결과

연구 질문

• RQ1 JAX에서의 자동 미분이 천체물리학에서 정교한 튜닝 없이도 안정적인 히어레틱 편차 행렬 계산을 가능하게 할 수 있는가?
• RQ2 실제 천체물리학 분석에서 해밀토니안 몬테 카를로 및 NUTS와 같은 기울기 기반 추론 기법이 기존 MCMC 대비 효율성이 높은가?
• RQ3 유연한 우도 함수를 통해 Figure of Merit에 대한 기울기 하강 기반 설문 조사 설계 최적화가 가능할 수 있는가?
• RQ4 효율성 샘플 크기 단위 시간 기준으로 JAX 기반 샘플러가 최적화된 전통적 MCMC 대비 성능 향상은 어느 정도인가?
• RQ5 노멀라이징 플로우를 활용한 신경 운반 HMC는 초기 단계의 천체물리학 분석에서 모델 검증에 도움이 될 수 있는가?

주요 결과

• DES-Y1 3x2pt 분석에서 JAX 기반 NUTS 샘플러는 잘 최적화된 Cobaya 메트로폴리스-해스팅스 샘플러 대비 GPU 시간당 약 5배 높은 유효 표본 크기를 확보했다.
• jax-cosmo로 계산된 히어레틱 편차 행렬은 정확하고 안정적이며 더 이상 수동 튜닝이 필요 없어져 계산 및 구현 오버헤드가 크게 감소했다.
• 히어레틱 편차 행렬 자체가 유연하게 구현되어 있어, 적색편이 구간 설정이나 관측 전략과 같은 설문 조사 설계 파라미터에 대한 기울기 하강 기반 최적화가 가능하다.
• NUTS 및 ADVI와 같은 기울기 기반 추론 기법은 고차원 파라미터 공간으로의 확장에서 기존 MCMC와 일치하는 사후 샘플을 생성하며 더 효율적인 스케일링을 보였다.
• 노멀라이징 플로우를 활용한 신경 운반 HMC는 초기 단계의 천체물리학 분석에서 모델 검증에 유망한 성능을 보이며, 신속하고 정확한 사후 근사치 제공 가능.
• 이 라이브러리는 천체물리학적 우도 함수의 엔드 투 엔드 유연성 제공으로, 최소한의 코드 수정과 함께 고성능 GPU 하드웨어에서 복잡한 추론 파이프라인을 지원한다.

Figure 2: Comparison of the non-linear matter power spectrum ( halofit function) between CCL and jax-cosmo . Also shown is the linear power spectrum.