Learned harmonic mean estimation of the marginal likelihood for multimodal posteriors with flow matching

기술 요약: 다중 모드 사후 분포를 위한 플로우 매칭 기반 학습된 조화 평균 추정

문제 정의
주변 가능도(marginal likelihood, 또는 베이즈 증거)로 표기되는 $z = p(d|M)$는 베이즈 모델 비교를 위한 근본적인 양으로, 베이즈 인자(Bayes factor)를 통한 정량적 평가를 가능하게 한다. 그러나 이 양을 계산하려면 전체 파라미터 공간에 대한 다차원 적분을 수행해야 하며, 이는 복잡하고 고차원적인 모델의 경우 계산적으로 불가능에 가깝다. 마르코프 연쇄 몬테카를로(MCMC) 및 해밀토니안 몬테카를로(HMC) 방법은 고차원에서 사후 샘플을 생성하는 데 탁월하지만, 본질적으로 주변 가능도를 제공하지는 않는다. 반대로, 네스티드 샘플링(nested sampling)은 증거를 제공하지만 특정 중첩 계층 구조에 샘플링 전략을 제한한다.

"학습된 조화 평균 추정량(learned harmonic mean estimator)"은 사용된 샘플링 방식에 관계없이 사후 샘서들만을 사용하여 주변 가능도를 추정하기 위해 이전에 도입되었다. 이 접근 방식은 사후 분포 내에 포함되면서도 꼬리(tails)가 더 얇은 내부 타겟 밀도 $\phi(\theta)$를 학습하는 것에 의존한다. 그러나 이전에 사용된 이산 정규화 흐름(discrete normalizing flows)이나 특수 밀도 추정기들은 매우 다중 모드(multimodal)인 사후 분포를 처리하는 데 어려움을 겪었다. 이러한 아키텍처들은 휴리스틱한 수정(예: 기저 분포를 위해 가우시안 혼합 모델을 수동으로 구성) 없이 복잡한 위상(topology)을 포착하는 데 실패하거나, 밀도 브릿지(density bridges) 부분에 무시할 수 없는 확률 질량을 할당하여 추정량의 분산을 무한대로 만드는 문제를 겪었다.

방법론
본 연구는 플로우 매칭(Flow Matching)을 학습된 조화 평균 프레임워크에 통합함으로써 이전 밀도 추정기들의 한계를 해결한다. 방법론은 다음과 같다:

1. 학습된 조한 평균 추정량: 추정량은 사후 샘플 $\theta_i \sim p(\theta|d)$와 학습된 타겟 밀도 $\phi(\theta)$를 사용하여 역수 주변 가능도 $\rho = z^{-1}$를 근사한다:
   $$\hat{\rho} = \frac{1}{N} \sum_{i=1}^{N} \frac{\phi(\theta_i)}{L(\theta_i)\pi(\theta_i)}$$
   추정량이 유한한 분산을 갖기 위해서는, $\phi(\theta)$가 사후 분포의 고확률 영역 내에 집중되어야 하며 꼬리 부분에서 더 빠르게 감소해야 한다.

2. 플로우 매칭 아키텍처: 저자들은 이산 정규화 흐름 대신, 플로우 매칭을 통해 훈련된 **연속 정규화 흐름(Continuous Normalizing Flows, CNFs)**을 채택하였다.

   • 정의: CNF는 단순한 기저 분포(표준 가우시안)와 타겟 사후 분포 사이의 변환을 시간 의존적 속도장(velocity field) $v_t(\theta)$를 통해 매개변수화한다.
   • 훈련 목적 함수: 네트워크는 시뮬레이션이 필요 없는 회귀 목적 함수를 통해 훈련된다. 조건부 확률 경로(conditional probability path)는 노이즈 샘플 $\theta_0$와 사후 샘플 $\theta_1$ 사이의 직선으로 정의된다: $\theta_t = (1-t)\theta_0 + t\theta_1$. 네트워크는 예측된 속도장과 실제 조건부 벡터장 $(\theta_1 - \theta_0)$ 사이의 평균 제곱 오차를 최소화한다.
   • 장점: 이 접근 방식은 아키텍처 제약이 없는 고용량 신경망(Swish 활성화 함수를 사용하는 다층 퍼셉트론)을 사용할 수 있게 하며, 초기 CNF 방법에서 요구되었던 비용이 많이 드는 ODE 시뮬레이션을 피할 수 있게 한다.

3. 온도 스케일링(Temperature Scaling): 타겟 밀도 $\phi(\theta)$가 사후 분포보다 더 얇은 꼬리를 갖도록 보장하기 위해, 플로의 기저 분포 분산을 조절하는 온도 파라미터 $T \in (0, 1)$가 적용된다. 이는 플로 분포를 집중시켜 조화 평균 추정량의 포함 요구 사항을 충족시킨다.

주요 기여

• 새로운 아키텍처: 내부 밀도 추정기로서 플로우 매칭으로 훈련된 연속 정규화 흐름을 도입하였다.
• 다중 모드 처리: 이 아키텍처가 기저 분포에 대한 휴리스틱한 수정이나 미세 조정 없이도 복잡하고 고도로 다중 모드인 사후 구조를 정확하게 포착할 수 있음을 입증하였다.
• 확장성: 샘플링 전략에 구애받지 않으며(MCMC, HMC, NUTS, SBI 등과 호환 가능), 중등도-고차원 공간에서 작동할 수 있는 방법을 제공한다.

실험 결과
저자들은 두 가지 도전적인 벤치마크를 통해 접근 방식을 검증하였다:

1. 라스트리진 함수 (2차원):

   • 많은 국소 모드(local modes)를 가진 표준 벤치마크이다.
   • 플로우 매칭 모델(6개 층, 256개 은닉 유닛)은 온도 $T=0.98$에서 모든 모드를 성공적으로 포착하였다.
   • 100회의 실험 실행 결과, 추정된 주변 가능도 값은 정확했으며 수치 적분을 통해 얻은 참값과 밀접하게 일치하였다. 오차 추정치 또한 신뢰할 수 있었다.

2. 20차원 가우시안 혼합 모델:

   • 서로 겹치지 않는 모드와 상관관계가 도입된 20차원의 5개 가우시안 성분 혼합 모델이다.
   • 모델(10개 층, 256개 은닉 유닛)은 40,000개의 샘플로 훈련되었다.
   • 학습된 타겟 분포는 사후 분포 내에 적절히 집중되어 도전적인 위상을 포착함을 보여주었다.
   • 100회 실행에 걸친 추정된 주변 가능도는 해석적 참값과 일치하였으며, 이는 이 복잡한 중등도 차원 시나리오에서도 방법론의 정확성과 견고한 오차 추정 능력을 확인시켜 준다.

의의 및 주장
저자들은 플로우 매칭으로 훈련된 CNF의 표현력을 활용함으로써, 학습된 조화 평균 추정량이 매우 복잡하고 다중 모드인 사후 분포에 대한 증거를 견고하게 추정할 수 있다고 주장한다. 본 연구의 주요 의의는 다중 모드를 처리하기 위해 이전에 필요했던 휴리스틱한 조정(예: 수동으로 선택된 가우시안 혼합)에 대한 의존성을 제거했다는 점에 있다. 저자들은 이 접근 방식이 임의의 가속화되거나 확장 가능한 샘플링 방법을 사용하는 샘플을 통해 정확한 증거 계산을 가능하게 하며, 이를 통해 천체 물리학과 같은 데이터 집약적인 분야에서 전통적인 네스티드 샘플링의 계산적 병목 현상이나 기존 조화 평균 추정량의 불안정성을 극복할 수 있다고 단언한다.

저자들이 언급한 향후 과제에는 중력파 천체 물리학과 같은 실제 문제에 이 방법론을 적용하는 것과, 영상 처리 및 필드 수준의 우주론적 추론과 같은 매우 높은 차원의 작업을 위한 추가적인 아키텍처 개선을 탐구하는 것이 포함된다.