[논문 리뷰] Uncertainty Quantification for Distribution-to-Distribution Flow Matching in Scientific Imaging

Dongxia Wu, Yuhui Zhang외 3명|arXiv (Cornell University)|2026. 03. 23.

Generative Adversarial Networks and Image Synthesis인용 수 0

한 줄 요약

분포-대-분포(distribution-to-distribution) 이미지 생성에서 알레이터릭 및 에피스테믹 불확실성을 정량화하기 위해 Bayesian Stochastic Flow Matching(BSFM)을 제안하고, 과학 이미징 작업의 신뢰성(일반화)과 책임성(OOD 탐지)을 향상시킨다.

ABSTRACT

Distribution-to-distribution generative models support scientific imaging tasks ranging from modeling cellular perturbation responses to translating medical images across conditions. Trustworthy generation requires both reliability (generalization across labs, devices, and experimental conditions) and accountability (detecting out-of-distribution cases where predictions may be unreliable). Uncertainty quantification (UQ) based approaches serve as promising candidates for these tasks, yet UQ for distribution-to-distribution generative models remains underexplored. We present a unified UQ framework, Bayesian Stochastic Flow Matching (BSFM), that disentangles aleatoric and epistemic uncertainty. The Stochastic Flow Matching (SFM) component augments deterministic flows with a diffusion term to improve model generalization to unseen scenarios. For UQ, we develop a scalable Bayesian approach -- MCD-Antithetic -- that combines Monte Carlo Dropout with sample-efficient antithetic sampling to produce effective anomaly scores for out-of-distribution detection. Experiments on cellular imaging (BBBC021, JUMP) and brain fMRI (Theory of Mind) across diverse scenarios show that SFM improves reliability while MCD-Antithetic enhances accountability.

연구 동기 및 목표

• 과학 이미징을 위한 신뢰할 수 있는 distribution-to-distribution 생성 모델을 신뢰성(일반화)과 책임성으로 구분하여 촉진한다.
• 알레이터릭 및 에피스테믹 구성요소로 불확실성을 분해하는 통합 UQ 프레임워크(BSFM)를 개발한다.
• 분포 간 이동에 대한 일반화를 개선하기 위해 흐름 매칭에 확률적 요소를 도입한다(SFM).
• 효율적인 OOD 탐지 및 불확실성 추정에 대한 확장 가능한 베이지안 접근법(MCD-Antithetic)을 제안한다.
• 세포 이미징 및 fMRI 데이터 세트에서 일반화 및 OOD 탐지의 개선를 입증한다.

제안 방법

• marginals를 보존하는 점수 보정 SDE를 통한 학습 가능한 확산 항을 도입하여 결정론적 흐름 매칭을 확률적 흐름 매칭으로 확장한다.
• 변환과 데이터 변이성(알레이터릭 불확실성) 모두를 포착하기 위해 공동 속도장과 스코어 모델을 학습한다.
• (속도) 및 스코어 파라미터의 포스터리를 근사하기 위해 MC-Dropout으로 모델 파라미터의 베이지안 처리를 적용한다.
• 중첩 샘플링을 통해 총 예측 불확실성을 추정하고 이를 알레이터릭 및 에피스테믹 구성요소로 분해한다.
• SDE 기반 샘플링의 분산을 줄이고 에피스테믹 및 알레이터릭 신호를 분리하여 OOD 탐지의 신뢰성을 높인다.
• 농축된 포스트eriors에서 알레이터릭 불확실성에 대한 MAP 기반 근사를 제공하여 계산량을 줄인다.

실험 결과

연구 질문

• RQ1 분포 간 이미지 생성에서 불확실성을 알레이터릭 및 에피스테믹 구성요소로 어떻게 분해할 수 있는가?
• RQ2 과학 이미징 작업에서 분포 시프트에 따른 일반화를 흐름 매칭으로 개선할 수 있는가?
• RQ3 안정적인 antithetic 샘플링을 갖춘 확률적 접근법이 효과적이고 효율적인 OOD 탐지 신호를 제공하는가?
• RQ4 세포 이미징 및 신경영상 데이터에서 신뢰성(일반화)과 책임성(OOD 탐지) 간의 절충은 어떠한가?
• RQ5 생성 품질과 OOD 감지 모두에서 BSFM이 기존 기준선보다 실질적인 이점을 주는가?

주요 결과

• 확률적 흐름 매칭(SFM)은 세포학 및 fMRI 데이터 세트 전반에서 분포 시프트에 대한 일반화를 개선한다.
• MCD-Antithetic은 OOD 탐지를 위한 효과적이고 샘플 효율적인 에피스테믹 및 알레이터릭 불확실성 추정을 제공한다.
• Antithetic 샘플링은 분산을 감소시키고 에피스테믹 및 알레이터릭 신호를 분리하여 OOD 탐지의 신뢰성을 향상시킨다.
• 보지 않은 섭동, 강도 변화, 보지 못한 세포주, 낮은 ToM 피험군에서 BSFM 기반 방법이 생성 및 OOD 지표에서 대략의 기준선보다 우수하다.
• 심한 시프트 하에서 에피스테믹 불확실성이 우세한 경향이 있으며, 더 온건한 시프트 하에서 알레이터릭 불확실성이 주도하여 보완적인 OOD 신호를 생성한다.