Reproduce “Uncertainty-Aware Distribution-to-Distribution Flow Matching for Scientific Imaging” on OpenResearch

About this paper

Dongxia Wu, Yuhui Zhang, Serena Yeung-Levy, Emma Lundberg, Emily B. Fox

Stochastic Flow Matching (SFM) is introduced as a marginal-preserving stochastic extension to flow matching, enhancing generalization in distribution-to-distribution generative models for scientific imaging. Complementary to SFM, Bayesian Stochastic Flow Matching (BSFM) with Antithetic Variance-reduction Uncertainty Quantification (AVUQ) provides an efficient framework for detecting unreliable and out-of-distribution generations, demonstrating improved generalization performance and superior anomaly detection across cellular and fMRI imaging tasks.

Key insights

• Introducing stochasticity via a learned score-based drift correction in flow matching significantly improves generalization by explicitly modeling conditional variability without compromising the integrity of learned data marginals.
• Decomposing and efficiently estimating both aleatoric and epistemic uncertainties provides crucial, complementary signals for robustly detecting out-of-distribution inputs and unreliable model generations in scientific imaging applications.
• Antithetic sampling effectively reduces Monte Carlo variance during SDE-based uncertainty estimation, making Bayesian inference for high-dimensional generative models computationally tractable even with limited sampling budgets.

Summary by alphaXiv.