[논문 리뷰] All-in-one simulation-based inference

Manuel Gloeckler, Michael Deistler외 3명|arXiv (Cornell University)|2024. 04. 15.

Simulation Techniques and Applications인용 수 5

한 줄 요약

Simformer는 joint parameter-data 입력에 대해 트랜스포머를 사용하는 확산-모델 기반의 암모터라이즈된 SBI 방법으로, 임의의 조건부(포스터리어, 가능도, 그 외) 샘플링 및 비정형 데이터, 함수값 파라미터, 관측 간격을 다룰 수 있다.

ABSTRACT

Amortized Bayesian inference trains neural networks to solve stochastic inference problems using model simulations, thereby making it possible to rapidly perform Bayesian inference for any newly observed data. However, current simulation-based amortized inference methods are simulation-hungry and inflexible: They require the specification of a fixed parametric prior, simulator, and inference tasks ahead of time. Here, we present a new amortized inference method -- the Simformer -- which overcomes these limitations. By training a probabilistic diffusion model with transformer architectures, the Simformer outperforms current state-of-the-art amortized inference approaches on benchmark tasks and is substantially more flexible: It can be applied to models with function-valued parameters, it can handle inference scenarios with missing or unstructured data, and it can sample arbitrary conditionals of the joint distribution of parameters and data, including both posterior and likelihood. We showcase the performance and flexibility of the Simformer on simulators from ecology, epidemiology, and neuroscience, and demonstrate that it opens up new possibilities and application domains for amortized Bayesian inference on simulation-based models.

연구 동기 및 목표

• irregular/structured 데이터 및 함수값 파라미터와 함께 작동하는 더 유연하고 효율적인 암모터라이즈드 SBI의 필요성 제시.
• 파라미터와 데이터의 결합 분포의 임의의 조건부를 샘플링할 수 있는 단일 모델 개발.
• 의존성 인코딩 및 조건부/추론된 출력을 가능하게 하는 확산 모델과 주의(attention) 마스킹이 포함된 트랜스포머 활용.
• 생태학, 역학, 신경과학의 다양한 시뮬레이터에서 방법을 시연하여 광범위한 적용성 보이고자 함.

제안 방법

• 트랜스포머의 입력으로 모든 변수(파라미터 및 데이터)에 대한 토큰 기반 표현 구성.
• denoising score matching을 이용해 결합 분포 p(theta, x)에 대해 점수 기반 확산 모델 학습.
• 의존성 구조를 인코딩하고 임의의 조건부 샘플링을 가능하게 하는 주의 마스크 M_E 도입.
• guided diffusion을 통해 관측 간격 및 제약 조건에 따라 샘플링하도록 조건부를 가능하게 함으로써 제약된 분포에서 샘플링.
• conditioning 마스크의 무작위화를 통해 결합 분포의 임의의 조건부를 출력하도록 모델을 학습.
• 단일 학습된 네트워크에서 포스터리어, 가능도 및 기타 조건부를 샘플링 가능.

실험 결과

연구 질문

• RQ1 단일 확산-트랜스포머 모델이 결합 분포 p(theta, x)의 임의의 조건부를 정확하게 샘플링할 수 있는가?
• RQ2 의존성 구조 마스크를 도입하면 SBI 작업의 정확도와 효율성이 향상되는가?
• RQ3 비정형 데이터, 결측값, 함수값 파라미터(무한 차원)를 다룰 수 있는가?
• RQ4 guided diffusion이 관측 간격이나 추가 제약 조건에 대한 조건부를 어떻게 가능하게 하는가?
• RQ5 생태학, 역학, 신경과학의 다양한 시뮬레이터에서 접근 방식이 효과적인가?

주요 결과

• Simformer는 주어진 시뮬레이션 예산에서 정확도 측면에서 벤치마크 SBI 작업에서 신경망 포스터리어 추정보다 우수한 성능을 보임.
• 도메인에 기반한 주의 마스크를 도입하면 특히 의존성 구조가 더 희박한 경우 정확도가 향상됨.
• 평균적으로 Simformer는 작업과 관측에 대해 NPE보다 약 10배 더 적은 시뮬레이션으로 결과를 얻음.
• 모델은 포스터리어 및 가능도 등을 포함한 임의의 조건부를 샘플링할 수 있으며 비정형 데이터 및 함수값 파라미터를 처리할 수 있음.
• Guided diffusion은 관측 간격 및 다른 제약 조건에 대한 조건부를 가능하게 하여 포스트리어 및 포스트리스트-예측 샘플에 영향을 줌.
• 이 방법은 무정형 관측을 가진 Lotka-Volterra, 무한 차원 파라메터 공간을 가진 SIRD, 에너지 제약을 가진 Hodgkin-Huxley에 대해 정확한 추론을 달성함.