Differentiable Stochastic Halo Occupation Distribution with Galaxy Intrinsic Alignments

1. 연구 배경 및 문제 제기 (Problem)

• 배경: 은하의 공간 분포와 고유 형태 (Intrinsic Alignments, IA) 는 우주 구조의 형성과 진화에 대한 핵심 정보를 담고 있습니다. 특히 차세대 약중력렌즈 (Weak Gravitational Lensing) 관측 (Rubin Observatory, Euclid, Roman 등) 에서는 은하 군집화 (Clustering) 와 IA 를 정밀하게 모델링하여 체계적 오차를 제거하고 우주론적 매개변수를 추정해야 합니다.
• 기존 방법의 한계:
  • HOD (Halo Occupation Distribution) 모델: 은하와 암흑물질 헤일로 간의 연결을 통계적으로 설명하는 강력한 도구이나, 기존 halotools-IA 와 같은 구현체는 확률적 샘플링 (Stochastic Sampling) 을 사용합니다.
  • 미분 불가능성 (Non-differentiability): 이산적 (Discrete) 인 샘플링 과정 (예: 베르누이, 포아송 분포) 은 기울기 (Gradient) 추적을 차단합니다. 이로 인해 경사 기반 최적화 (Gradient-based Optimization) 나 해밀토니안 몬테카를로 (HMC) 와 같은 효율적인 베이지안 추론 기법을 적용할 수 없습니다.
  • 계산 비용: 기존 MCMC 기반 추론은 고차원 매개변수 공간에서 계산 비용이 매우 큽니다 (예: Stage IV 분석 시 수 년 소요).
  • 에뮬레이터의 한계: 신경망 기반 에뮬레이터 (IAEmu 등) 는 속도는 빠르지만, 통계량 (Summary Statistics) 만을 모델링하므로 은하 카탈로그 수준에서의 미분 가능성이 없으며, 새로운 물리 모델 적용 시 재학습이 필요합니다.

2. 방법론 (Methodology)

저자들은 diffHOD-IA라는 완전히 미분 가능한 (Fully Differentiable) HOD 구현체를 개발하여 IA 모델링을 포함시켰습니다.

• 핵심 프레임워크:

  • HOD 모델: Zheng et al. (2007b) 의 HOD 공식을 기반으로 합니다.
  • IA 모델: Van Alfen et al. (2024) 의 IA 모델을 통합하여, 은하의 헤일로 정렬 (Alignment) 을 모델링합니다.
  • 미분 가능한 샘플링 (Differentiable Sampling):
    • 이산 분포 처리: 베르누이 (중앙 은하) 및 포아송 (위성 은하) 분포와 같은 이산적 샘플링을 Gumbel-Softmax 트릭을 사용하여 연속적인 완화 (Continuous Relaxation) 로 변환합니다. 이를 통해 파라미터에 대한 기울기를 역전파할 수 있게 합니다.
    • Dimroth-Watson 분포: 은하의 헤일로 오정렬 각도 (Misalignment Angle) 를 모델링하는 Dimroth-Watson 분포에서, 역 누적 분포 함수 (Inverse CDF) 샘플링 기법을 사용하여 미분 가능한 방식으로 샘플을 생성합니다.
  • 미분 가능한 상관 함수 (Differentiable Correlation Functions):
    • 은하 위치 - 위치 ($\xi(r)$), 위치 - 방향 ($\omega(r)$), 방향 - 방향 ($\eta(r)$) 상관 함수를 미분 가능한 추정기로 구현했습니다.
    • 은하 쌍 (Pair) 구조는 고정하고, 은하의 방향 벡터에 대한 기울기만 흐르도록 설계하여 효율성을 높였습니다.

• 구현 기술:

  • JAX 라이브러리를 사용하여 자동 미분 (Automatic Differentiation) 을 지원하며, GPU 가속 및 벡터화를 통해 계산 효율성을 극대화했습니다.

3. 주요 기여 (Key Contributions)

1. diffHOD-IA 프레임워크 개발: 은하 카탈로그 생성부터 요약 통계량 (Summary Statistics) 계산까지 End-to-End 미분 가능한 HOD+IA 모델을 처음 구현했습니다.
2. 미분 가능한 IA 샘플링: Dimroth-Watson 분포에 대한 미분 가능한 역 CDF 샘플링 알고리즘을 개발하여, IA 파라미터 ($\mu_{cen}, \mu_{sat}$) 에 대한 정확한 기울기를 제공합니다.
3. 성능 검증:
   • 정확도: 비미분 가능한 기준 코드 (halotools-IA) 와 비교하여 은하 수, 2 점 상관 함수 ($\xi, \omega, \eta$) 에서 샘플 변동성 (Sample Variance) 이내의 높은 일치도를 보였습니다.
   • 기울기 정확도: 유한 차분법 (Finite Difference) 과 자동 미분 (Autodiff) 결과 비교를 통해 기울기 계산의 정확성을 입증했습니다.
4. 과학적 적용 사례:
   • 경사 기반 최적화: 모의 관측 데이터 (Mock Data) 에서 IA 파라미터를 성공적으로 복원했습니다.
   • HMC 추론: 기존 MCMC 대비 압도적인 속도 향상 (GPU 기준 5 분 vs CPU 150 코어 기준 1 일) 을 보여주며, 정확한 사후 분포 (Posterior) 를 추정했습니다.

4. 결과 (Results)

• 모델 일치도:
  • 은하 수 (Number Counts) 는 두 구현체 간 약 0.17% 미만의 차이를 보였습니다.
  • 2 점 통계량 ($\xi(r), \omega(r), \eta(r)$) 에서 평균 편차는 각각 0.28%, 0.16%, 1.77% 로 매우 높은 정확도를 보였습니다.
• 기울기 검증:
  • HOD 파라미터 ($\log M_{min}, \sigma_{\log M}$ 등) 와 IA 파라미터에 대한 기울기가 유한 차분법과 거의 완벽하게 일치함을 확인했습니다.
• 파라미터 복원 (Optimization):
  • 모멘트 매칭 (Moment-matching) 과 상관 함수 매칭을 통한 최적화 실험에서, 무작위 초기값에서 시작하여도 tng300 시뮬레이션의 참값 (Ground Truth) 으로 수렴함을 확인했습니다.
• HMC 성능 비교:
  • diffHOD-IA (HMC): NVIDIA A100 GPU 에서 약 5 분 소요.
  • IAEmu (HMC): 약 1 분 소요 (에뮬레이터 특성상 더 빠름).
  • halotools-IA (MCMC): 150 개 CPU 코어에서 약 1 일 소요.
  • diffHOD-IA 는 에뮬레이터보다 느리지만, 카탈로그 수준의 미분 가능성과 유연성을 제공하며, 기존 MCMC 대비 약 2,800 배 이상 빠른 추론 속도를 달성했습니다.

5. 의의 및 향후 전망 (Significance & Future Work)

• 차세대 관측 대비: 차세대 약중력렌즈 관측 (Stage IV) 은 수백 개의 시스템 파라미터를 처리해야 하므로, 기존 MCMC 방식은 계산적으로 불가능합니다. diffHOD-IA 는 경사 기반 추론 (Gradient-based Inference) 을 가능하게 하여 이러한 고차원 문제를 해결할 수 있는 길을 열었습니다.
• 유연성: 에뮬레이터와 달리, 새로운 물리 모델 (예: 어셈블리 바이어스, 환경 의존적 효과) 이나 HOD 공식 변경 시 재학습 없이 모델에 직접 적용 가능합니다.
• 확장성:
  • 현재는 3 점 상관 함수까지 미분 가능하지만, 향후 2 차원 투영 통계량 (Weak Lensing 관측과 직접 비교 가능) 및 고차 통계량으로 확장 가능합니다.
  • 미분 가능한 입자 메쉬 솔버 (jaxPM) 및 헤일로 찾기 (jFoF) 와 결합하여 우주론 파라미터부터 은하 필드까지 End-to-End 미분 가능한 전체 파이프라인 구축의 기반이 됩니다.

결론적으로, 이 논문은 우주론적 시뮬레이션과 관측 데이터 간의 간극을 해소하고, 차세대 대규모 관측 데이터를 효율적으로 분석하기 위한 미분 가능한 은하 형성 모델링의 새로운 표준을 제시했습니다.