Uncertainty Quantification for Distribution-to-Distribution Flow Matching in Scientific Imaging

2026-03-23

AI Summary

This work addresses the challenge in scientific imaging that distribution-to-distribution generative models often lack reliable uncertainty quantification and struggle to simultaneously achieve cross-condition generalization and effective anomaly detection. To this end, we propose the Bayesian Stochastic Flow Matching (BSFM) framework, which, for the first time in flow matching models, disentangles epistemic and aleatoric uncertainties and incorporates a diffusion term to enhance generalization. By integrating Monte Carlo Dropout with antithetic sampling, we introduce the MCD-Antithetic method, significantly improving the efficiency and reliability of uncertainty estimation. Experiments demonstrate that BSFM achieves superior out-of-distribution detection performance and strong generalization to unseen conditions on both cellular imaging datasets (BBBC021, JUMP) and brain fMRI data.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Relational Probabilistic Models
• Computer Vision: Diffusion Models for Vision

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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.

Problem

• Uncertainty Quantification
• Distribution-to-Distribution Generative Models
• Scientific Imaging
• Out-of-Distribution Detection
• Reliability and Accountability

Innovation

• Uncertainty Quantification
• Stochastic Flow Matching
• Bayesian Deep Learning
• Out-of-Distribution Detection
• Generative Modeling