Tokenised Flow Matching for Hierarchical Simulation Based Inference

Motivation

• Known: 기존 hierarchical SBI 방법들은 posterior를 global/local parameter로 분해(factorised inference)하여 시뮬레이션 효율을 개선하지만, 학습 샘플당 여전히 여러 사이트(최대 nmax개)를 시뮬레이션해야 한다. SimFormer와 같은 tokenised SBI는 function-valued observation을 유연하게 다룰 수 있고, flow matching은 posterior estimation에서 score matching보다 안정적인 학습을 제공한다.
• Gap: Posterior factorisation 기반 방법들은 여전히 학습 샘플당 다중 사이트 시뮬레이션을 요구하여 시뮬레이터 호출 비용이 큰 병목이며, global posterior 추정 오차가 local 추정으로 전파되는 문제가 있다. 또한 계층적 SBI 방법들을 체계적으로 비교할 표준 벤치마크가 부재하다.
• Why: 과학 분야의 계층적 시뮬레이터(감염병 모델, 유체역학 모델 등)는 사이트 수가 늘어날수록 파라미터와 시뮬레이션 비용이 함께 증가하여 실질적인 추론 병목이 되므로, 시뮬레이션 효율을 근본적으로 개선하는 것은 계산 비용이 큰 과학적 캘리브레이션 문제 전반에 중요한 영향을 미친다.
• Approach: Likelihood factorisation(LF)을 통해 사이트당 단일 시뮬레이션으로 per-site neural surrogate를 학습한 뒤 합성 다중 사이트 관측치를 생성하고, 이를 이용해 tokenised flow matching(TFMPE)으로 전체 계층적 posterior를 amortise하여 추정한다.

Achievement

1. LF sampling 제안: likelihood를 factorise하여 학습 샘플당 시뮬레이션 요구량을 단일 사이트로 축소하는 훈련 전략을 도입했다.
2. TFMPE 제안: function-valued observation을 지원하고 flow matching의 안정적 학습 특성을 활용하는 tokenised flow matching 기반 hierarchical SBI 기법을 개발했다.
3. Hierarchical SBI 벤치마크 구축: SBI benchmark를 계층적 설정에 맞게 확장한 최초의 벤치마크를 제시하여 TFMPE와 기존 hierarchical/non-hierarchical baseline을 체계적으로 비교했다.
4. 실제 응용 검증: 계절성 SEIR 감염병 모델(함수형 관측치 포함)과 Computational Fluid Dynamics(혈류역학) 모델에서 TFMPE가 잘 보정된(calibrated) posterior를 제공하면서 계산 비용을 절감함을 확인했다.

How

• 계층적 모델을 global parameter θg와 site별 local parameter η_s로 분리하고, likelihood factorisation을 통해 per-site neural surrogate q_φl(y_s | η_s, θg)를 단일 시뮬레이터 호출로 학습한다.
• 학습된 surrogate로부터 다중 사이트에 대한 합성 관측치(ŷ)를 생성하여, 별도의 posterior estimator q_φp(θg, η | y)를 학습시켜 전체 계층적 posterior를 amortise 추론한다.
• 파라미터(θg, η)와 관측치(y)를 값(value), 변수 식별자, 변수 내 위치, 그룹(사이트) 식별자, 함수형 입력 좌표(ξ, Gaussian Fourier feature로 임베딩) 등의 필드를 갖는 토큰 시퀀스로 표현하는 tokenisation 스킴을 설계한다.
• Encoder-only transformer 아키텍처 위에서 Conditional Flow Matching objective(FMPE)로 vector field v_t를 학습시켜 posterior 또는 likelihood를 근사한다.
• posterior consistency는 ℓ-C2ST 지표로 평가하고, SEIR 및 CFD(혈류역학) 실사례에서 posterior predictive check와 pairplot을 통해 보정 여부를 확인한다.

Originality

• 기존 hierarchical SBI가 posterior factorisation에 의존하는 것과 달리, likelihood factorisation을 통해 학습 샘플당 시뮬레이션 요구를 단일 사이트로 줄이는 새로운 훈련 패러다임(LF sampling)을 제시했다.
• Tokenisation과 flow matching을 결합해 function-valued(functional) observation을 계층적 SBI 맥락에서 다루는 최초의 접근인 TFMPE를 제안했다.
• 계층적 SBI 전용 벤치마크를 최초로 구축하여 향후 연구의 표준 비교 기반을 마련했다.

Limitation & Further Study

• LF sampling은 per-site neural surrogate가 실제 likelihood를 정확히 근사한다는 가정에 의존하므로, surrogate 근사 오차가 합성 관측치 생성 및 최종 posterior 추정에 전파될 위험이 있다.
• 논문 발췌에는 전체 실험 세부사항(예: 다양한 사이트 수 ns, 고차원 θg/η에서의 스케일링, baseline 대비 정량적 성능 표 등)이 제한적으로 제시되어 있어 방법의 일반화 성능 평가가 추가로 필요하다.
• 벤치마크가 SBI benchmark를 계층적으로 확장한 초기 버전이므로, 더 다양한 실제 응용(예: 더 많은 사이트 수, 이질적 사이트 구조)에서의 검증이 후속 연구로 요구된다.
• Encoder-only transformer 기반 아키텍처의 계산 비용과 토큰 수 증가에 따른 확장성(scalability)에 대한 논의가 추가로 필요하다.

Evaluation

총평: Likelihood factorisation이라는 참신한 관점 전환을 통해 계층적 SBI의 핵심 병목인 시뮬레이션 비용을 효과적으로 줄이고, tokenised flow matching과 결합하여 function-valued observation까지 지원하는 실용적이고 체계적으로 검증된 기법을 제시한 우수한 연구이다.