GPU 가속을 위한 벡터화된 중첩 샘플링: Nested Slice Sampling

초록

본 논문은 하드 제약을 갖는 중첩 샘플링을 GPU 친화적으로 구현한 Nested Slice Sampling(NSS)을 제안한다. 핵심 아이디어는 Hit‑and‑Run 슬라이스 샘플링을 내부 업데이트로 사용하고, 슬라이스 폭을 차원에 따라 최적화하는 간단한 규칙을 도출함으로써 연산 비용을 예측 가능하게 만든다. 실험 결과, NSS는 복잡하고 다중모드인 베이지안 모델에서 정확한 증거(evidence) 추정과 고품질 사후 샘플을 제공하며, 기존의 템퍼드 SMC 대비 안정성을 보인다.

상세 분석

논문은 먼저 기존 Nested Sampling(NS)의 구조를 외부 루프(증거를 1차원 적분으로 변환)와 내부 제약 샘플링(제한된 영역에서의 표본 생성)으로 분리한다. 외부 루프는 m개의 라이브 파티클을 유지하면서 매 반복마다 가장 높은 에너지 k개의 파티클을 제거하고, 남은 파티클을 복제·변이시켜 새로운 파티클을 채워 넣는다. 이때 k/m 비율이 압축 속도를 결정하고, 병렬화 수준을 조절한다. 기존 CPU‑중심 구현은 이러한 단계들을 순차적으로 수행하거나 MPI 수준에서만 병렬화했지만, 저자는 JAX와 XLA를 활용해 전체 연산을 배치 형태로 GPU에 전달한다.

핵심 기술은 내부 제약 샘플링으로 Hit‑and‑Run Slice Sampling(HRSS)을 채택한 점이다. HRSS는 무작위 방향을 선택한 뒤, 해당 방향에 대한 1차원 슬라이스를 정확히 샘플링한다. 슬라이스 폭 w는 “슬라이스 폭 최적화” 이론을 통해 w∗≈c·d⁻¹/² 형태로 도출되며, 여기서 d는 차원, c는 상수이다. 이 규칙을 적용하면 스텝아웃·수축 연산 횟수가 차원에 대해 거의 일정한 분산을 보이며, 따라서 배치 실행 시 워프 발산이 최소화된다. 논문은 타원형 레벨셋에 대해 평균 코드 길이 ℓ_d∼Θ(d⁻¹/²)임을 증명하고, 최적 폭에서 연산 비용이 정규분포에 가까워짐을 실험적으로 확인한다.

또한 저자는 HRSS와 기존의 반사 기반 혹은 무작위 워크 기반 제약 MCMC를 비교한다. 고차원에서 반사 기반 방법은 경계 근처에서 재현성이 떨어지고, 차원 증가에 따라 수렴 속도가 급격히 저하되는 반면, HRSS는 차원에 거의 독립적인 혼합 속도를 유지한다. 이러한 특성은 NS가 매 단계마다 점점 작아지는 제약 영역을 다루어야 하는 상황에서 매우 중요하다.

GPU 구현 측면에서는 에너지 평가가 가장 비용이 많이 드는 연산임을 강조한다. JAX의 JIT와 vmap을 이용해 에너지 함수를 배치 처리함으로써 메모리 대역폭과 연산량을 동시에 최적화한다. 또한 라이브 파티클의 복제·재배치를 벡터화하고, HRSS의 스텝아웃·수축 루프를 고정 반복 횟수(예: 최대 10회)로 제한해 SIMD 흐름을 보장한다. 결과적으로 per‑step 연산 시간이 예측 가능해져, 대규모 파티클 수와 높은 차원에서도 효율적인 스케일링이 가능하다.

실험에서는 (1) 다중모드 합성 베르누이·가우시안 혼합, (2) 100차원 이상의 베이지안 회귀·신경망 하이퍼파라미터 마진, (3) 고차원 가우시안 프로세스 커널 추정 등 세 가지 벤치마크를 제시한다. 모든 경우에서 NSS는 증거 추정 오차가 1–2% 수준으로 유지되며, 사후 샘플의 ESS(effective sample size)도 템퍼드 SMC 대비 30% 이상 우수했다. 특히 다중모드 문제에서 SMC는 파티클 소실 현상이 발생해 증거가 크게 편향되는 반면, NSS는 제약 영역을 직접 탐색하므로 이러한 현상이 거의 없었다.

마지막으로 저자는 오픈소스 구현(NSS 라이브러리)을 제공하고, JAX 기반 확장성을 강조한다. 이는 다른 확률 프로그래밍 프레임워크와 쉽게 통합될 수 있어, 물리학·천문학·머신러닝 등 다양한 분야에서 대규모 베이지안 모델링에 바로 적용 가능하다.