Nested Sampling with Slice-within-Gibbs: Efficient Evidence Calculation for Hierarchical Bayesian Models

1. 문제 정의 (Problem Statement)

• 배경: 베이지안 추론에서 모델 비교를 위해 필수적인 '증거 (Marginal Likelihood, $Z$)'를 계산하는 것은 고차원 공간에서 매우 어렵습니다. Nested Sampling (NS) 은 이를 직접 추정할 수 있는 강력한 방법론이지만, 고차원 문제에서 제약된 사전 분포 (constrained prior) 에서 샘플링하는 과정이 병목 현상을 일으킵니다.
• 기존 방법의 한계:
  • 기존 NS: 고차원 계층적 모델 (J 개의 그룹, 각 그룹별 국소 파라미터 $\theta_j$, 공유 하이퍼파라미터 $\psi$) 에서 새로운 샘플을 생성할 때, 매번 전역 우도 제약 조건 ($\sum \ell_j > \ell^*$) 을 확인해야 합니다.
  • 복잡도 문제: $J$개의 그룹이 있는 경우, 한 번의 블록 업데이트마다 전체 우도를 재계산하면 $O(J)$의 비용이 들고, 이를 $J$번 반복하는 Gibbs 스윕에서는 $O(J^2)$의 계산 비용이 발생합니다. 전체 알고리즘 복잡도는 $O(J^3)$에 달해 대규모 데이터 (수천 개 그룹) 에 적용하기 어렵습니다.
  • 기울기 기반 방법 (Gradient-based) 의 한계: HMC 나 NUTS 와 같은 방법은 고차원에서 잘 작동하지만, 복잡한 기하학적 구조 (예: Funnel 분포) 를 가진 계층적 모델에서는 재파라미터화 없이는 수렴이 어렵거나, 증거 (Evidence) 를 직접 추정하지 못한다는 단점이 있습니다.

2. 방법론 (Methodology: NS-SwiG)

저자는 Slice-within-Gibbs 커널을 Nested Sampling 의 내부 루프에 통합하여 위 문제를 해결했습니다.

A. 우도 예산 분해 (Likelihood-Budget Decomposition)

• 핵심 아이디어: 전역 우도 제약 조건을 각 그룹별 '예산 (Budget)'으로 분해합니다.
  • 전체 제약: $\sum_{j=1}^J \ell_j(\theta_j, \psi) > \ell^*$
  • 이를 $k$번째 그룹의 파라미터 $\theta_k$를 업데이트할 때의 조건으로 재배열하면:
    $\ell_k(\theta'_k, \psi) > B_k$, 여기서 $B_k = \ell^* - \sum_{j \neq k} \ell_j$
• 캐싱 전략: 현재 상태의 전체 로그 우도 합 ($S$) 과 각 그룹별 로그 우도 ($\ell_j$) 를 캐시합니다.
  • 새로운 제안이 받아들여지면, 전체 합 $S$를 $O(1)$ 시간에 업데이트하고, 해당 그룹의 예산 $B_k$를 즉시 계산할 수 있습니다.
  • 이로 인해 제약 조건 확인 비용이 $O(J)$에서 $O(1)$로 감소합니다.

B. 제약된 Slice-within-Gibbs 커널

• 하이퍼파라미터 업데이트: $\psi$를 조건부 분포에서 슬라이스 샘플링으로 업데이트합니다 (이때는 모든 그룹의 우도를 재계산하여 $O(J)$ 비용 발생).
• 국소 파라미터 업데이트 (Inner Sweep): 각 그룹 $j$에 대해, 해당 그룹의 예산 $B_j$를 만족하는 범위 내에서 $\theta_j$를 슬라이스 샘플링합니다.
  • 각 그룹 업데이트 비용: $O(1)$ (예산 확인만 하면 됨).
  • 전체 $J$개 그룹에 대한 스윕 비용: $O(J)$.
• 결과: 한 번의 Gibbs 스윕 비용이 기존 $O(J^2)$에서 $O(J)$로 감소했습니다.

C. 구현

• JAX 와 BlackJAX 프레임워크를 사용하여 벡터화된 형태로 구현하여 여러 입자 (particles) 를 병렬 처리합니다.
• Hit-and-Run 슬라이스 샘플링을 사용하여 국소 블록 내에서 방향을 무작위로 선택합니다.

3. 주요 기여 (Key Contributions)

1. 우도 예산 분해: 전역 제약 조건을 캐싱된 그룹별 기여도를 기반으로 $O(1)$로 업데이트 가능한 국소 예산으로 변환하는 방법을 제시했습니다. 이는 마르코프 구조를 가진 잠재 변수 (예: 확률적 변동성 모델) 로도 자연스럽게 확장 가능합니다.
2. 제약된 Slice-within-Gibbs 커널 설계: 하이퍼파라미터와 국소 파라미터를 번갈아 업데이트하는 블록 Gibbs 커널을 Nested Sampling 에 적용하여, $J$개의 그룹에 대해 $O(J)$의 선형 스케일링을 달성했습니다.
3. 실증적 검증: 수천 차원 (최대 $d \approx 2500$) 에 이르는 벤치마크에서 기존 방법 (NSS, NUTS, SMC-HMC) 과 비교하여 정확도와 확장성을 입증했습니다.

4. 실험 결과 (Results)

논문의 4 가지 벤치마크를 통해 다음과 같은 결과를 도출했습니다.

• 계층적 가우시안 모델 (Hierarchical Gaussian):
  • 그룹 수 $J$가 증가함에 따라 NS-SwiG 는 $O(J)$의 평가 횟수 (full-likelihood equivalents) 를 보인 반면, 기존 Joint Nested Sampling (NSS) 은 $O(J^2)$의 비용이 발생했습니다.
  • $J=1000$ (차원 1001) 에서도 정확한 증거 추정 ($\log Z$) 을 제공하며, 계산 시간이 기존 방법 대비 수백 배 단축되었습니다.
• Funnel 분포 (Neal's Funnel):
  • 하이퍼파라미터와 국소 파라미터 간의 강한 결합으로 인해 기하학적으로 병목이 발생하는 'Funnel' 문제에서, NUTS 는 재파라미터화 없이는 수렴에 실패하거나 발산 (divergent transitions) 이 발생했습니다.
  • NS-SwiG 는 재파라미터화 없이도 전체 기하 구조를 정확하게 포착하고 정확한 증거를 추정했습니다.
• 맥락 효과 모델 (Radon Data):
  • 85 개 카운티, 946 개 관측치를 가진 계층적 회귀 모델에서 NS-SwiG 는 NUTS 와 유사한 사후 분포 정확도를 유지하면서 증거를 추정했습니다.
  • SMC-HMC 는 중심화 (centered) 파라미터화에서 편향된 증거 추정을 보였습니다.
• 확률적 변동성 모델 (Stochastic Volatility):
  • 마르코프 구조를 가진 시계열 데이터 ($J=2516$) 에서도 NS-SwiG 가 확장 가능함을 보였습니다. 국소 예산을 이웃과 공유하도록 확장하여 적용했습니다.

5. 의의 및 결론 (Significance & Conclusion)

• 계산 효율성: NS-SwiG 는 고차원 계층적 모델에서 Nested Sampling 의 계산 복잡도를 $O(J^3)$에서 $O(J^2)$로 줄였습니다 (입자 수 $m$을 고정할 때). 이는 천문학, 우주론 등에서 수천 개의 천체나 관측 데이터를 다루는 대규모 계층적 추론을 가능하게 합니다.
• 강건성: 기울기 (Gradient) 가 필요 없으며, 재파라미터화가 어려운 복잡한 기하학적 구조 (Funnel, Multimodal 등) 에서도 견고하게 작동합니다.
• 모델 비교: MCMC 기반 방법들이 주로 사후 분포 샘플링에 집중하는 반면, NS-SwiG 는 모델 비교에 필수적인 증거 (Evidence) 를 직접적이고 정확하게 추정할 수 있다는 장점이 있습니다.
• 한계 및 향후 과제: 우도 함수가 분해 가능한 구조를 가져야 한다는 전제가 필요합니다. 또한, 고차원에서는 정확한 추정을 위해 필요한 입자 수 ($m$) 가 증가하여 메모리 및 계산 비용이 여전히 문제가 될 수 있습니다.

요약하자면, NS-SwiG 는 계층적 모델의 구조적 특성을 활용하여 Nested Sampling 의 확장성을 획기적으로 개선한 알고리즘으로, 대규모 데이터에 대한 베이지안 모델 비교 및 추론을 위한 실용적인 도구로 자리 잡을 것으로 기대됩니다.