DAWN-FM: Data-Aware and Noise-Informed Flow Matching for Solving Inverse Problems

TL;DR

DAWN-FM tackles ill-posed inverse problems by learning a data-aware, noise-informed flow that maps a Gaussian reference to the posterior conditioned on measurements. By embedding the observed data and noise level into the velocity field, it trains a problem-specific interpolant that directly targets the posterior $\pi(x_1|b)$, enabling sampling of multiple likely solutions. The framework yields uncertainty quantification through posterior sampling and demonstrates improved performance on image deblurring and tomography compared to baselines, particularly under higher noise. This approach offers robust, data-informed reconstructions and practical uncertainty maps for decision-making in imaging.

Abstract

Inverse problems, which involve estimating parameters from incomplete or noisy observations, arise in various fields such as medical imaging, geophysics, and signal processing. These problems are often ill-posed, requiring regularization techniques to stabilize the solution. In this work, we employ Flow Matching (FM), a generative framework that integrates a deterministic processes to map a simple reference distribution, such as a Gaussian, to the target distribution. Our method DAWN-FM: Data-AWare and Noise-informed Flow Matching incorporates data and noise embedding, allowing the model to access representations about the measured data explicitly and also account for noise in the observations, making it particularly robust in scenarios where data is noisy or incomplete. By learning a time-dependent velocity field, FM not only provides accurate solutions but also enables uncertainty quantification by generating multiple plausible outcomes. Unlike pre-trained diffusion models, which may struggle in highly ill-posed settings, our approach is trained specifically for each inverse problem and adapts to varying noise levels. We validate the effectiveness and robustness of our method through extensive numerical experiments on tasks such as image deblurring and tomography.

Figures (11)

• Figure 1: Example images from MNIST (left) and STL10 (right) datasets showing the recovery of deblurred images from the blurred (noisy) data (top panel) at different levels of noise (0%, 1%, 5%, 10%, 20%). Our methods DAW-FM (middle panel) incorporates the embedding for the blurred data within the network, while DAWN-FM (bottom panel) incorporates embedding for both data and noise levels. DAWN-FM is superior to DAW-FM at deblurred image recovery, especially at higher levels of noise.
• Figure 2: Generation of the distribution of two Gaussians by FM without and with data.
• Figure 3: Schematic of the training process for the FM model for solving inverse problem, where the forward model is given by ${\bf A}$. This figure specifically represents the deblurring inverse problem. The network with parameters ${\boldsymbol \theta}$ can either have noise embedding (DAWN-FM) or not (DAW-FM). The two loss terms $\mathcal{L}_1$ and $\mathcal{L}_2$ represent the error in prediction of velocity and misfit, respectively. In this figure, the transformation $f$ for generating the data embedding was chosen as $f = {\bf A}^\top$.
• Figure 4: Computing uncertainty in the solutions obtained by FM for some example images from MNIST (left), STL10 (middle) and CIFAR10 (right) for the deblurring task. The reconstruction (posterior mean) is computed by averaging over runs from 32 randomly initialized ${\bf x}_0$. The uncertainty in the posterior mean is computed by evaluating its standard deviation for the predictions from 32 runs.
• Figure 5: Assessing the sensitivity of the metrics for the recovered deblurred image as a function of the level of noise ($p\%$) added to the data. The DAWN-FM is more robust to noise across all three datasets, MNIST, STL10 and CIFAR10 than DAW-FM, with the crossing point happening at values of $p$$\lesssim$$5\%$ for all metrics.

Theorems & Definitions (1)

• Example 2.1