Tokenised Flow Matching for Hierarchical Simulation Based Inference

1. 문제 정의 (Problem Statement)

**시뮬레이션 기반 추론 (SBI, Simulation Based Inference)**은 우도 함수 (likelihood) 가 해석적으로 구해지지 않지만 시뮬레이터가 존재하는 상황에서 베이지안 파라미터 추정을 가능하게 합니다. 그러나 복잡한 과학적 시뮬레이션 (예: 전염병 모델, 유체 역학) 에서는 시뮬레이션 실행 비용이 매우 높아 추론의 병목 현상이 됩니다.

특히 계층적 (Hierarchical) 설정에서는 전역 파라미터 (Global Parameters, $\theta_g$) 와 여러 사이트 (Site) 별 로컬 파라미터 (Local Parameters, $\eta_s$) 가 존재하며, 관측 데이터 ($y$) 가 여러 사이트에서 생성됩니다. 기존 계층적 SBI 방법들은 사후분포 (Posterior) 를 분해하여 학습하지만, 여전히 각 학습 샘플마다 여러 사이트의 시뮬레이션을 수행해야 하므로 계산 비용이 사이트 수 ($n_s$) 에 비례하여 급증하는 문제가 있습니다.

2. 제안 방법론 (Methodology)

저자들은 **우도 분해 (Likelihood Factorisation, LF)**와 토큰화 (Tokenisation), Flow Matching을 결합한 새로운 방법론인 **TFMPE (Tokenised Flow Matching for Posterior Estimation)**를 제안합니다.

2.1. 핵심 전략: 우도 분해 (Likelihood Factorisation, LF)

기존의 사후분포 분해 (Posterior Factorisation) 대신 우도 분해를 채택합니다.

• 관측 조건부 독립성: 주어진 로컬 파라미터 하에서 관측 데이터는 사이트 간 독립적입니다 ($p(y|\theta_g, \eta) = \prod p(y_s|\theta_g, \eta_s)$).
• 2 단계 학습 전략:
  1. 1 단계 (Surrogate Training): 각 사이트별로 단일 시뮬레이션 ($\theta_g, \eta_s, y_s$) 만 사용하여 시뮬레이터의 신경망 대리 모델 (Neural Surrogate, $q_{\phi_l}$) 을 학습합니다. 이는 우도 함수를 근사하는 것이 아니라, 다중 사이트 관측 데이터를 합성 (Synthesize) 하기 위한 샘플러로 사용됩니다.
  2. 2 단계 (Posterior Estimation): 학습된 대리 모델을 사용하여 합성된 다중 사이트 데이터 ($\hat{y}$) 를 생성하고, 이를 바탕으로 전체 계층적 사후분포 ($q_{\phi_p}(\theta_g, \eta | y)$) 를 추정하는 모델을 학습합니다.
• 효과: 학습 샘플당 시뮬레이션 호출 횟수를 $n_s$번에서 1 번으로 줄여 계산 효율성을 극대화합니다.

2.2. 토큰화 (Tokenisation)

변수 값이 불규칙하거나 함수 형태 (Function-valued) 일 수 있는 계층적 데이터를 처리하기 위해 토큰화 방식을 도입합니다.

• 모든 파라미터와 관측 데이터를 시퀀스 형태의 토큰으로 변환합니다.
• 각 토큰은 변수 식별자 (Variable ID), 그룹 식별자 (Site ID), 위치 인코딩, 함수 입력 (시간/공간 좌표) 등을 포함합니다.
• **그룹 식별자 (Group ID)**의 명시적 입력은 Transformer 어텐션 메커니즘이 사이트 간 관계를 추론하는 부담을 줄여주며, 노이즈가 많은 매핑에서도 성능을 유지하게 합니다.

2.3. Flow Matching for Posterior Estimation

• Flow Matching: 조건부 유동 매칭 (Conditional Flow Matching) 목적 함수를 사용하여 베이스 분포에서 목표 사후분포로 샘플을 이동시키는 벡터 필드를 학습합니다.
• 아키텍처: 인코더 전용 (Encoder-only) Transformer 를 사용하여 토큰 시퀀스를 처리하고, 파라미터 토큰들로부터 벡터 필드를 출력합니다. 이는 Flow Matching 의 안정적인 학습 동역학과 토큰화의 유연성을 결합합니다.

3. 주요 기여 (Key Contributions)

1. LF 샘플링 전략 도입: 학습당 시뮬레이션 횟수를 사이트 수에 관계없이 1 번으로 줄이는 새로운 훈련 전략을 제안했습니다.
2. TFMPE 모델 개발: 계층적 SBI 를 위한 토큰화된 Flow Matching 접근법을 제안하여, 불규칙한 함수형 관측 데이터 (Function-valued observations) 를 효과적으로 처리하고 안정적인 학습을 가능하게 했습니다.
3. 계층적 SBI 벤치마크 구축: 기존 SBI 벤치마크를 계층적 구조로 확장하여, 샘플 효율성과 확장성을 체계적으로 평가할 수 있는 표준 벤치마크를 개발했습니다.
4. 실제 적용 사례 검증: 계절적 SEIR 전염병 모델과 1 차원 혈류 역학 (Haemodynamics) CFD 모델을 통해 실제 복잡한 문제에서의 성능을 입증했습니다.

4. 실험 결과 (Results)

4.1. 벤치마크 성능

• 샘플 효율성: TFMPE 는 대부분의 계층적 태스크에서 기존 방법 (NPE, SNPE, Posterior Factorisation 등) 보다 훨씬 적은 시뮬레이션 횟수로 더 정확한 사후분포를 추정했습니다. 특히 사이트 수가 많을 때 (Large $n_s$) 성능 우위가 두드러졌습니다.
• 비교: 비계층적 방법들은 시뮬레이션 비용이 증가함에 따라 성능이 정체되는 반면, TFMPE 는 추가 예산에 대해 선형적으로 효율적으로 반응했습니다.
• 예외 사항: "Two Moons"와 같은 비선형성이 강한 문제에서는 대리 모델의 근사 오차가 누적되어 Posterior Factorisation 방법이 더 좋은 성능을 보였으나, 이는 우도 모델이 단순한 경우 (Gaussian Mixture 등) 에 TFMPE 가 압도적으로 우월함을 반증합니다.

4.2. 실제 적용 사례

• 계절적 SEIR 모델 (전염병): 100 개 이상의 사이트 (지역) 에서 불규칙한 시간 간격으로 관측된 데이터를 처리했습니다. MCMC 는 5 개 사이트 이상에서 수렴에 실패하거나 계산 비용이 과도하게 증가했으나, TFMPE 는 단일 사이트 시뮬레이션만으로 잘 보정된 (Well-calibrated) 사후분포를 생성했습니다.
• 혈류 역학 (Haemodynamics) 모델: CFD 시뮬레이션은 한 번 실행에 약 11.88 초가 소요되지만, 학습된 TFMPE 대리 모델은 6.41ms 만에 합성 데이터를 생성하여 약 1,852 배의 속도 향상을 이루었습니다. 전체 파이프라인에서는 2 배 이상의 속도 향상을 달성하면서도 16 명의 환자 데이터에 대해 95% 신뢰구간 커버리지를 유지했습니다.

5. 의의 및 결론 (Significance)

이 논문은 시뮬레이션 비용이 높은 계층적 베이지안 추론 문제에 대한 획기적인 해결책을 제시합니다.

• 계산 비용 절감: 시뮬레이션 호출 횟수를 사이트 수에 비례하지 않게 줄여, 대규모 계층적 모델 (전염병, 기후, 공학 등) 에 대한 실용적인 추론을 가능하게 합니다.
• 유연성: 토큰화 기법을 통해 시간/공간적으로 불규칙한 데이터나 함수형 데이터를 자연스럽게 처리할 수 있습니다.
• 확장성: Flow Matching 의 안정성과 대리 모델 기반의 효율성을 결합하여, 기존 방법론들이 다루기 어려웠던 고차원 및 대규모 계층적 문제에 대한 새로운 표준을 제시합니다.

결론적으로, TFMPE 는 시뮬레이션 비용이 추론의 주요 병목인 현실적인 과학적 문제들에서 정확성과 효율성을 동시에 달성할 수 있는 강력한 프레임워크임을 입증했습니다.