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

TL;DR

DAWN-FM introduces a data-aware, noise-informed flow matching framework for inverse problems, improving robustness and uncertainty quantification in noisy, ill-posed scenarios across applications like imaging and tomography.

Contribution

It develops a novel flow matching method that explicitly incorporates data and noise information, tailored for each inverse problem, enhancing stability and uncertainty estimation.

Findings

Outperforms existing methods in noisy inverse problems

Provides accurate solutions with uncertainty quantification

Demonstrates robustness in image deblurring and tomography

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 pretrained 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. The code is available at: https://github.com/ahxmeds/DAWN-FM.git.