기술 요약: HIcosmo – 우주론적 추론을 위한 미분 가능한 JAX 기반 프레임워크

1. 문제 정의

다가오는 Stage IV 우주론적 관측(예: Euclid, LSST, DESI, SKA)은 전례 없는 규모의 관측 데이터를 생성할 것입니다. 배경 수준의 우주론적 추론(거리, 거리 비율, 거리 지수 등을 포함하는 과정)은 일반적으로 저차원 파라미터 공간($d \sim 5\text{--}7$)에서 작동하지만, 마르코프 연쇄 몬테카를로(MCMC) 방법에서 요구되는 방대한 양의 가능도(likelihood) 계산으로 인해 계산 비용이 병목 현상이 되고 있습니다.

전통적인 파이프라인(예: CosmoMC, Cobaya, MontePython)은 무작위 보행 메트로폴리스-헤이스팅스(Metropolis–Hastings, MH) 샘플러에 의존합니다. 이러한 방식은 차원($d$)에 따라 선형적으로 증가하는 느린 혼합 시간(mixing times) 문제를 겪으며, 유한 차분 근사(finite-difference approximations)를 위해 추가적인 계산($O(d)$번의 추가 평가)이 필요하고 단계 크기(step-size) 선택에 민감하다는 단점이 있습니다. 또한, 기존의 성숙한 코드들은 미분 가능한 프로그램으로 구현되어 있지 않아, No-U-Turn 샘플러(NUTS)와 같은 경사도 기반 샘플러 및 피셔 예측(Fisher forecasting)을 위한 자동 미분(AD)을 사용하는 것을 방지합니다. JAX 기반 도구들이 섭동 수준(perturbation-level) 분석(예: jax-cosmo, CosmoPower)을 위해 등장하기는 했으나, 가능도, NUTS 샘PLING, 그리고 AD 기반 피셔 예측을 단일 아키텍처 내에 통합하는 체계적으로 검증된 엔드 투 엔드(end-to-end) 미분 가능 프레임워크는 존재하지 않았습니다.

2. 방법론 및 아키텍처

본 논문은 JAX로 완전히 구축된 오픈 소스 프레임워크인 HIcosmo(High-performance Inference for Cosmology)를 제시합니다. 이 프레임워크는 엔드 투 엔드 미분 가능성 원칙을 중심으로 설계되었으며, 이를 통해 모든 연산이 우주론적 파라미터로부터 스칼라 로그 가능도에 이르기까지 JAX의 자동 미분 시스템에 의해 추적될 수 있도록 보장합니다.

2.1 계층적 아키텍처

HIcosmo는 세 가지 주요 계층으로 구성됩니다:

1. 모델 계층 (Model Layer): 이론적 예측을 계산합니다. $\Lambda$CDM, $w$CDM, $w_0w_a$CDM (CPL 매개변수화), 그리고 상호작용하는 암흑 에너지 모델($I\Lambda$CDM)의 네 가지 배경 모델을 구현합니다. 결정적으로, 비미분적인 적응형 적분기(adaptive integrators)를 대신하여 조밀한 적색편이 격자(dense redshift grids) 상의 **고정 격자 사다리꼴 구적법(fixed-grid trapezoidal quadrature)**을 사용합니다. 이를 통해 유한 차분 없이 연쇄 법칙(chain rule)을 통한 정확한 경사도 계산이 가능해집니다.
2. 가능도 계층 (Likelihood Layer): 모델 예측을 관측 데이터와 비교합니다. 다음을 지원합니다:
   • Type Ia 초신성: Pantheon+, DES-SN5YR, Union3.
   • 중입자 음향 진동 (BAO): DESI DR1/DR2, SDSS, BOSS, 6dFGS.
   • CMB: Planck 2018 거리 사전 확률(단축된 $R$, $l_A$, $\omega_b$ 형태).
   • 국소 $H_0$: SH0ES 거리 사다리.
   • 강한 중력 렌즈: H0LiCOW 시간 지연.
     모든 가능도는 순수 JAX로 구현되어 JIT 컴파일 및 벡터화를 지원합니다.
3. 추론 계층 (Inference Layer): 로그 가능도와 그 도함수를 소비합니다. 다음을 제공합니다:
   • NUTS 샘플링: JAX로부터 직접 얻은 경사도를 활용하여 혼합 시간을 개선하는 NumPyro를 통해 구현되었습니다(혼합 성능이 $d^{1/4}$로 스케일링됨).
   • 피셔 예측 (Fisher Forecasting): 중첩된 자동 미분을 통해 헤시안(Hessian) 행렬을 계산하며, 이를 통해 단계 크기 선택의 필요성을 제거하고 유한 차분 오차를 없앱니다.

2.2 구현 세부 사항

• 미분 가능성: 고정 단계 적분과 파라미터에 독립적인 양(예: 역공분산 행렬)의 사전 계산을 통해, 프레임워크는 역방향 모드 AD(reverse-mode AD)를 통해 경사도와 헤시안을 직접 얻을 수 있음을 보장합니다.
• 하드웨어 가속: 프레임워크는 JAX의 통합 CPU/GPU/TPU 실행을 활용합니다. JIT 컴파일은 가능도를 머신 코드로 최적화하여 MCMC 실행 중 발생하는 파이썬 인터프리터 오버헤드를 우회합니다.
• 교정 (Calibration): 물리적 정확성을 보장하기 위해, 프레임워크는 유도된 양(예: 음향 지평선 $r_d$)을 CAMB 볼츠만 솔버와 대조하여 잔차를 $\sim 0.01\%$ 수준으로 줄여 교정합니다.

3. 주요 결과

3.1 검증 및 정확도

HIcosmo는 $\Lambda$CDM, $w$CDM, $w_0w_a$CDM 모델에 대해 참조 구현체(Cobaya) 및 공식 가능도 코드와 체계적으로 검증되었습니다.

• 가능도 정밀도: $\chi^2$ 값은 참조 코드와 절대 차이 $10^{-6}$에서 $10^{-2}$ 사이에서 일치합니다.
• 파라미터 제약: HIcosmo와 Cobaya의 주변화된 제약 조건은 테스트된 모든 분석에서 0.2$\sigma$ 미만의 차이를 보입니다. SN+BAO+CMB 결합 분석의 경우, 일치도는 0.04$\sigma$ 이내입니다.
• 일관성: 프레임워크는 결합 분석에 대해 $H_0 = 68.24 \pm 0.38$ km/s/Mpc 및 $\Omega_m = 0.307 \pm 0.005$와 같은 확립된 제약 조건을 재현함으로써, 미분 가능한 근사가 유의미한 편향을 도입하지 않음을 확인했습니다.

3.2 성능 벤치마크

• Scipy 대비 속도 향상: 거리 계산에 대해, HIcosmo는 JIT 컴파일 후 scipy 적응형 구적법보다 140배 빠른 속도를 달고 합니다. 전체 Pantheon+ 가능도 파이프라인의 경우, CPU에서 16배 더 빠릅니다.
• 샘플링 처리량: 8코어 CPU에서 Cobaya(MH 샘플러 사용)와 비교했을 때, HIcosmo(NUTS 사용)는 8.7배 더 높은 유효 샘플 크기(ESS) 처리량(400 ESS/s vs 45.8 ESS/s)을 달성합니다. 이는 더 빠른 가능도 평가와 더 우수한 NUTS의 혼합 성능($\tau \approx 1.2$ vs $\tau \approx 11.4$) 덕분입니다.
• GPU 가속: 작은 데이터셋(Pantheon+)에 대한 단일 평가는 오버헤드로 인해 완만한 GPU 이득을 보이지만, GPU 가속은 데이터셋 크기에 따라 확장됩니다. 합성 데이터셋($N_{SN} = 20,000$)의 경우, GPU는 CPU 대비 20.8배의 속도 향상을 달성합니다.
• 피셔 행렬: AD 기반 피셔 행렬 계산은 안정적이며 단계 크기 튜닝이 필요 없으며, 작은 단계 크기에서 유한 차분 방법을 괴롭히는 부동 소수점 반올림 오차를 피할 수 있습니다.

3.3 응용 분야

논문은 다음과 같은 사례를 통해 프레임워크의 유용성을 입증합니다:

• 다중 프로브 결합 제약: Pantheon+, DESI BAO, Planck 사전 확률을 결합하여 $H_0$ 및 $\Omega_m$에 대한 제약을 강화합니다.
• 암흑 에너지 제약: DESI 2024 결과와 일치하는 $w_0$ 및 $w_a$에 대한 제약을 도출합니다.
• 21 cm 강도 매핑 예측: 6개의 조사(SKA1, MeerKAT, BINGO, Tianlai, CHIME)에 대한 피셔 예측을 수행합니다. 프레임워크는 SKA1 Band 1이 암흑 에너지 상태 방정식에 대해 $\sigma(w) \approx 0.046$의 제약을 가할 수 있음을 예측합니다.

4. 의의 및 주장

저자들은 HIcosmo가 체계적으로 검증된, 엔드 투 엔드로 미분 가능한 JAX 네이티브 배경 우주론 파이프라인을 제공함으로써 우주론적 추론 지형의 중요한 공백을 메운다고 주장합니다.

• 효율성: 경사도 기반 샘플링(NUTS)과 JIT 컴파일된 가능도를 결합하면 전통적인 무작위 보행 MCMC보다 계산 속도와 통계적 효율성(ESS) 측면 모두에서 크게 앞설 수 있음을 보여줍니다.
• 확장성: 프레임워크는 GPU 가속을 활용하여 20배에 달하는 속도 향상을 달성할 수 있으므로, 미래의 조사(Stage IV)에서 발생하는 거대한 데이터를 처리할 수 있는 독보적인 위치에 있습니다.
• 모듈성 및 확장성: 모델, 가능도, 추론 계층을 분리함으로써, 미분 코드를 다시 작성하지 않고도 새로운 모델이나 데이터셋을 쉽게 추가할 수 있습니다.
• 섭동 수준 추론을 위한 토대: 현재는 배경 우주론(볼츠만 솔버 제외)에 국한되어 있지만, 저자들은 배경 수준에서의 성공적인 엔드 투 엔드 검증이 견고한 토대를 마련했다고 주장합니다. 이 아키텍처는 모델 계층을 미분 가능한 볼츠만 에뮬레이터(예: CosmoPower-JAX)로 교체함으로써 고차원 Stage IV 문제를 위한 경사도 기반 추론이 가능하도록 설계되었습니다.

요약하자면, HIcosmo는 기존 도구들에 대한 재현 가능하고 고성능인 대안을 제공하며, 자동 미분과 현대적 하드웨어 가속이 표준적인 우주론적 추론 작업에 효과적으로 적용될 수 있음을 증명합니다.