# Uncertainty-Aware Distribution-to-Distribution Flow Matching for Scientific Imaging

arXiv: https://arxiv.org/abs/2603.21717 · HTML: https://mlsanity.com/paper/2603.21717/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Flow matching, time series and science, one of 14 directions in Diffusion models.
- Builds on 44 of its 55 references that are in the map.
- First posted on arXiv 23 Mar 2026.

## Placed in

- Diffusion models / [Flow matching, time series and science](https://mlsanity.com/diffusion/flow-matching-time-series-and-science/)

## Builds on (most cited first, 30 of 44)

- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [SDEdit: Guided Image Synthesis and Editing with Stochastic Differential Equations](https://mlsanity.com/paper/2108.01073/) (2021)
- [Stochastic Interpolants: A Unifying Framework for Flows and Diffusions](https://mlsanity.com/paper/2303.08797/) (2023)
- [Contrastive Learning for Unpaired Image-to-Image Translation](https://mlsanity.com/paper/2007.15651/) (2020)
- [Diffusion Schr"odinger Bridge with Applications to Score-Based Generative Modeling](https://mlsanity.com/paper/2106.01357/) (2021)
- [A Simple Baseline for Bayesian Uncertainty in Deep Learning](https://mlsanity.com/paper/1902.02476/) (2019)
- [Flow Matching Guide and Code](https://mlsanity.com/paper/2412.06264/) (2024)
- [Why Normalizing Flows Fail to Detect Out-of-Distribution Data](https://mlsanity.com/paper/2006.08545/) (2020)
- [Denoising Diffusion Bridge Models](https://mlsanity.com/paper/2309.16948/) (2023)
- [Score-Based Diffusion Models as Principled Priors for Inverse Imaging](https://mlsanity.com/paper/2304.11751/) (2023)
- [Denoising diffusion models for out-of-distribution detection](https://mlsanity.com/paper/2211.07740/) (2022)
- [Guided Flows for Generative Modeling and Decision Making](https://mlsanity.com/paper/2311.13443/) (2023)
- [Principled Probabilistic Imaging using Diffusion Models as Plug-and-Play Priors](https://mlsanity.com/paper/2405.18782/) (2024)
- [Unpaired Image-to-Image Translation via Neural Schr"odinger Bridge](https://mlsanity.com/paper/2305.15086/) (2023)
- [Understanding Failures in Out-of-Distribution Detection with Deep Generative Models](https://mlsanity.com/paper/2107.06908/) (2021)
- [FBNETGEN: Task-aware GNN-based fMRI Analysis via Functional Brain Network Generation](https://mlsanity.com/paper/2205.12465/) (2022)
- [Measurement-conditioned Denoising Diffusion Probabilistic Model for Under-sampled Medical Image Reconstruction](https://mlsanity.com/paper/2203.03623/) (2022)
- [Unsupervised Out-of-Distribution Detection with Diffusion Inpainting](https://mlsanity.com/paper/2302.10326/) (2023)
- [Estimating Epistemic and Aleatoric Uncertainty with a Single Model](https://mlsanity.com/paper/2402.03478/) (2024)
- [How to Trust Your Diffusion Model: A Convex Optimization Approach to Conformal Risk Control](https://mlsanity.com/paper/2302.03791/) (2023)
- [DiffusionAD: Norm-guided One-step Denoising Diffusion for Anomaly Detection](https://mlsanity.com/paper/2303.08730/) (2023)
- [On the Trustworthiness of Generative Foundation Models: Guideline, Assessment, and Perspective](https://mlsanity.com/paper/2502.14296/) (2025)
- [Zero-shot Medical Image Translation via Frequency-Guided Diffusion Models](https://mlsanity.com/paper/2304.02742/) (2023)
- [BayesDiff: Estimating Pixel-wise Uncertainty in Diffusion via Bayesian Inference](https://mlsanity.com/paper/2310.11142/) (2023)
- [DiffGuard: Semantic Mismatch-Guided Out-of-Distribution Detection using Pre-trained Diffusion Models](https://mlsanity.com/paper/2308.07687/) (2023)
- [Diffusion Models with Implicit Guidance for Medical Anomaly Detection](https://mlsanity.com/paper/2403.08464/) (2024)

## Built on by (most cited first, 1 of 1)

- [Optimal Transport Dropout for Structured Predictive Uncertainty](https://mlsanity.com/paper/2609.33377/) (2026)

## Abstract (arXiv)

Distribution-to-distribution generative models support scientific imaging tasks ranging from modeling cellular perturbation responses to translating medical images across conditions. Trustworthy generation requires reliability, or generalization across labs, devices, and experimental conditions, and accountability, or detecting out-of-distribution cases where predictions may be unreliable. We leverage Stochastic Flow Matching (SFM), a marginal-preserving stochastic extension of flow matching for improved generalization under distribution shift. SFM augments deterministic flows with a diffusion term together with a learned score-based drift correction, retaining the learned transport marginals while modeling conditional variability. Building on this SFM framework, we introduce Bayesian Stochastic Flow Matching (BSFM) as a companion uncertainty quantification mechanism and develop AVUQ (Antithetic Variance-reduction Uncertainty Quantification) to approximately estimate epistemic and aleatoric uncertainty via sample-efficient antithetic sampling with approximate posterior inference. We further use AVUQ to yield anomaly scores for unreliable generation detection. Experiments on cellular imaging (BBBC021, JUMP) and brain fMRI (Theory of Mind) across diverse unseen scenarios show that SFM improves generalization while AVUQ provides effective uncertainty-based anomaly scores under practical sampling budgets.