DAWN-FM: Data-Aware and Noise-Informed Flow Matching for Solving Inverse Problems (2412.04766v2)

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.

• The paper introduces DAWN-FM, a novel stochastic interpolant framework that embeds observed data and noise for improved solving of ill-posed inverse problems.
• It leverages a UNet architecture with dedicated embedding and tailored loss functions to achieve superior image deblurring and tomography performance.
• The method efficiently quantifies uncertainty via ensemble sampling, delivering reliable reconstructions under high noise conditions.

Data-Aware and Noise-Informed Flow Matching for Ill-Posed Inverse Problems

Problem Statement and Motivation

DAWN-FM (2412.04766) addresses the fundamental challenge of ill-posed inverse problems, such as image deblurring and tomography, where observations are incomplete and/or corrupted by significant noise. Inverse problems are ubiquitous in scientific and engineering domains—ranging from medical imaging to remote sensing—and often require regularization to yield stable, physically plausible solutions. Standard deep-learning and generative approaches, including pre-trained diffusion models, are frequently insufficient in highly ill-posed regimes, particularly when noise characteristics differ between training and inference or when forward operators induce large null spaces in the solution.

The core insight of DAWN-FM is to use Stochastic Interpolants (SI), a unified generative framework integrating deterministic flows and stochastic diffusions, trained directly on the forward operator and noise model for each inverse problem instance. Crucially, DAWN-FM advances this paradigm by explicitly embedding both the observed data and associated noise statistics in the learned SI velocity field, yielding marked robustness to varying and high noise levels. This is in contrast to Bayesian approaches that decouple likelihood and prior, potentially introducing errors in the directions of the null space of the forward map.

Methodology: Stochastic Interpolation with Data and Noise Embedding

DAWN-FM operationalizes the SI framework to model the posterior distribution over solutions conditioned on observed data and noise, by learning a velocity field $v_\theta(x_t, f(b), t, \sigma)$, where $x_t$ is an interpolant trajectory between Gaussian prior $x_0$ and target $x_1$ (with $x_1$ sampled from the true data), $f(b)$ is a learned embedding of the observed data (and optionally, the adjoint of the forward operator), $t$ is the continuous trajectory parameter, and $\sigma$ encodes the noise level.

The network architecture is based on a UNet backbone, with dedicated embedding functions for each of the input conditioning components. The learning objective comprises a velocity matching loss and a data misfit regularization, enforcing consistency of the generated solution with both the observed data and the noise model. The integration for inference is performed by deterministic ODE solvers starting from an ensemble of Gaussian samples, thus naturally enabling posterior quantification of both the mean and uncertainty.

Both variants, DAW-SI (data-aware only) and DAWN-SI (data- and noise-aware), are presented to delineate the impact of explicit noise embedding.

Figure 1: Schematic illustration of SI model training for inverse problems, showing embeddings of observed data and noise into the velocity estimation.

Empirical Validation: Image Deblurring and Tomography

DAWN-FM is validated on canonical ill-posed inverse problems using standard benchmarks (MNIST, STL10, CIFAR10, OrganAMNIST, OrganCMNIST). Experiments systematically vary noise levels from $0\%$ to $20\%$ to test robustness beyond the regime where most diffusion-based priors are typically effective.

Image Deblurring:

DAWN-SI demonstrates substantially superior performance—measured by MSE, SSIM, and PSNR—compared to both DAW-SI, diffusion posteriors, and InverseUNetODE baselines, particularly as data noise increases.

Figure 2: MNIST (left) and STL10 (right) deblurring across noise levels; DAWN-SI (bottom) enables far more accurate restoration under high noise.

Figure 3: Across noise levels, DAWN-SI consistently outperforms DAW-SI in reconstruction metrics, with robustness persisting until $p \lesssim 5\%$.

Figure 4: Comparative visualization of deblurring methods on MNIST.

Figure 5: Comparative visualization of deblurring methods on STL10.

Figure 6: Comparative visualization of deblurring methods on CIFAR10.

Tomography:

DAWN-SI also excels on medical image tomography, achieving better structural recovery and misfit across OrganAMNIST/CMNIST tasks. The direct learning of the conditional posterior, coupled with noise adaptation, yields reconstructions with markedly lower uncertainty and sharper anatomical delineation under high noise.

Figure 7: Tomography on OrganAMNIST; DAWN-SI achieves superior reconstructions under increasing sinogram noise.

Figure 8: Tomography on OrganCMNIST, showing DAWN-SI robustness to noise.

Uncertainty Quantification

A central outcome of DAWN-FM is the ability to efficiently estimate epistemic uncertainty via posterior ensemble sampling. By integrating the velocity ODE from multiple initializations, DAWN-FM produces calibrated uncertainty maps which concentrate at ambiguous boundaries or under severe ill-posedness.

Figure 9: Pixelwise uncertainty (std. dev.) obtained from 32 SI samples per data point; uncertainty localizes on ambiguous object boundaries, especially on MNIST edges.

Figure 10: OrganAMNIST tomography samples from DAWN-SI posterior; variability in anatomical features (lobes) is well-quantified.

Theoretical and Practical Implications

DAWN-FM's design asserts that direct posterior learning via problem-specific SI, leveraging explicit data and noise conditioning, is superior to both prior-likelihood factorizations and unconditional generative models, especially for high-noise or rank-deficient forward maps. The explicit incorporation of noise level, and the use of the adjoint of the forward operator in the embedding, represent bold design decisions that challenge several prevailing heuristics in the field.

This approach also advances the feasibility of uncertainty-aware decision making in downstream applications, as multiple solution samples can be generated at negligible marginal cost. The use of a UNet backbone with multi-scale data embedding further ensures architectural compatibility with a broad range of high-dimensional inverse problems.

Future Prospects

DAWN-FM opens avenues for further research:

• Integration with advanced noise/adversarial models, especially for domain shift or out-of-distribution noise.
• Scalability to 3D inverse problems and integration with physical simulation constraints.
• Extension to nonlinear and hybrid inverse operators beyond convolutional blurs and Radon transforms.
• Real-time uncertainty quantification for critical applications in medical imaging (e.g., adaptive diagnostics).
• Synergy with active and optimal experimental design, leveraging the network's uncertainty outputs.

Conclusion

DAWN-FM establishes a robust, generalizable, and uncertainty-aware framework for solving highly ill-posed inverse problems by direct posterior modeling with data and noise embedding in a stochastic interpolation framework. It achieves demonstrably superior performance under high noise and severe ill-posedness, outperforms standard diffusion-based and encoder-decoder baselines, and provides principled uncertainty quantification. The approach has significant ramifications for both the theory and practice of scientific machine learning in settings where reliable solution existence and uncertainty estimates are critical (2412.04766).