Saving Foundation Flow-Matching Priors for Inverse Problems

AI summary

Overview

• Research area: Machine learning for inverse problems (IPs) in imaging — specifically, using pretrained deep generative priors, and in particular flow-matching (FM) models, as plug-in priors for reconstructing images from measurements.
• Technical level: Advanced. The method itself is compact, but the motivation rests on concentration-of-measure results for Gaussian vectors and on the specific structure of flow-matching ODEs.
• Scope: The paper diagnoses why foundation flow-matching priors underperform domain-specific and even untrained priors on inverse problems, and proposes a plug-in framework, FMPlug, that strengthens those priors for both simple image-restoration tasks and few-shot scientific inverse problems.

Paper details: "Saving Foundation Flow-Matching Priors for Inverse Problems" by Yuxiang Wan, Ryan Devera, Wenjie Zhang, and Ju Sun, all of the Department of Computer Science and Engineering, University of Minnesota. arXiv:2511.16520v3 [cs.LG]. Project page: https://sun-umn.github.io/xm-plug/

What This Paper Is About

Inverse problems ask you to recover an unknown image x from measurements y ≈ A(x), where the forward model A is known; because the measurements are insufficient or approximate, many images could explain the same data, so prior knowledge of what a plausible x looks like is essential. Large "foundation" flow-matching models such as Stable Diffusion and FLUX promise to serve as universal, reusable priors for this purpose, avoiding the need to train a model for each image domain. The paper shows that, as used today, these foundation priors actually perform worse than domain-specific models and even worse than untrained priors, and it introduces FMPlug, a plug-in framework that recovers much of the lost performance.

Key Contributions

1. Identifying and documenting the performance gap. The authors show systematically, in a direct comparison table and in a sweep over deblurring difficulty, that foundation FM priors used in existing IP solvers lag behind domain-specific FM priors and even behind the untrained deep image prior (DIP) — a gap the paper says previous work did not explicitly acknowledge.

2. Diagnosing why existing strengthening ideas fail. The paper analyzes two existing prior-strengthening techniques — the D-Flow initialization plus Gaussianity penalty and the FlowDPS text-conditioning idea — and argues they are nearly useless in practice, explaining the failure of the initialization through concentration of measure and the failure of the regularizer through its flatness.

3. Proposing FMPlug. A plug-in framework combining an instance-guided, time-dependent warm-start strategy (starting the flow at a learned time t rather than at t = 0) with a sharp Gaussianity constraint that forces the latent variable onto a thin spherical shell around radius √d. The few-shot variant additionally optimizes a convex combination (via softmax over a simplex) of a few available similar instances.

4. Confirming effectiveness experimentally. FMPlug is reported to achieve the best results among pretrained foundation FM-based solvers on simple-distortion image restoration tasks across three datasets, and it is also demonstrated on three scientific inverse problems where a few similar samples are available.

Main Findings

• Foundation FM priors underperform badly in a controlled comparison. On Gaussian deblurring of AFHQ-Cat at 256×256 resolution with Stable Diffusion V3 as the foundation model, D-Flow with a domain-specific prior reaches PSNR 28.1389 and LPIPS 0.2783, whereas D-Flow with a foundation prior drops to PSNR 25.0120 and LPIPS 0.5335. FlowDPS moves from PSNR 22.1191 (domain-specific) to 22.1404 (foundation), with LPIPS worsening from 0.3850 to 0.5412.

• An untrained prior beats the foundation priors. The untrained DIP plugged into the standard regularized formulation is the second-best method by three of the four reported metrics, sitting just below D-Flow with a domain-specific prior.

• Existing strengthening tricks barely move the needle. Adding the D-Flow style initialization and Gaussianity regularization (labeled FD-S) changes D-Flow from PSNR 25.0120 to 25.1453 and FlowDPS from 22.1404 to 22.0538 — the paper describes the effect as almost useless.

• The failure mode is visible. For Gaussian deblurring on DIV2K at 512×512 resolution, the paper reports that reconstructions from FlowDPS and FlowChef are no better than the blurry measurement itself, if not worse.

• The generator can approximately represent real images. An image regression test on 1000 random images from DIV2K gives D-Flow PSNR 36.187 / LPIPS 0.181 versus FMPlug PSNR 37.924 / LPIPS 0.093, evidence for the approximate surjectivity of the generator that the plug-in formulation relies on.

• FMPlug improves simple-distortion restoration across tasks and datasets. On 4× super-resolution, FMPlug reaches PSNR 30.31 on AFHQ where DIP gets 29.85 and D-Flow 26.37. On Gaussian deblurring of DIV2K, FMPlug reaches PSNR 26.26 versus DIP 25.23 and D-Flow 23.64. On motion deblurring of DIV2K, FMPlug reaches PSNR 27.38 versus DIP 24.75 and D-Flow 25.21. Gains hold across AFHQ, DIV2K, and RealSR for super-resolution, 70% random inpainting, Gaussian deblurring, and motion deblurring.

• Warm-starting alone accounts for most of the gain. FMPlug-W, which uses only the time-dependent warm-start with NFE = 3, is close to the full FMPlug on nearly every reported entry, with the sharp Gaussianity constraint adding smaller additional improvements.

• The method is demonstrated on scientific IPs. The paper focuses on three: linear inverse scattering (LIS), compressed sensing MRI (MRI), and black hole imaging (BKH), reporting that FMPlug, D-Flow, and DIP are compared qualitatively in a figure. Quantitative numbers for these scientific tasks are not reported in the available content.

Methodology in Plain English

The paper builds on the "plug-in" way of using a generative model: instead of interleaving denoising steps with steps that push the image toward fitting the measurements, you treat the whole generative model as a fixed function that maps a latent input to a generated image, and you optimize the latent input so the generated image explains the measurements. The trouble is where you start that optimization.

The first insight concerns the starting latent. A flow-matching model is trained on latents drawn from a standard Gaussian. A classical result about high-dimensional Gaussians says that essentially all such samples lie on an ultra-thin shell at radius roughly √d, not spread through space. So the generator is, in effect, only trained on inputs from that shell. The existing D-Flow initialization mixes the inverted measurement with noise, which pushes the starting point off that shell — into territory the generator never saw. FMPlug instead keeps the latent on the shell and, because the measurement is assumed to be close to the truth, injects the measurement directly into the flow at a learned time t rather than at t = 0. Since the coefficient multiplying the image shrinks toward zero as t → 0, choosing t lets the method control how much the (possibly imperfect) measurement dominates. This both improves results and shortens the flow.

The second insight concerns the regularizer. D-Flow penalizes the negative log-likelihood of the latent's squared norm under a chi-squared distribution. The authors plot this function and observe it is extremely flat: in a range of [62000, 70000] its value changes by ≤ 0.031% relative to its minimum, so it provides almost no pressure to stay in the narrow region where a Gaussian actually concentrates. FMPlug instead enforces a hard constraint — a "shell" constraint — that the latent's norm lies in a narrow band [(1−ε)√d, (1+ε)√d], implemented by a closed-form projection applied after each optimization step. The default setting is ε = 0.025.

For latent flow-matching models, the formulation is extended by wrapping the generator with the pretrained decoder and feeding the encoder's version of the measurement into the warm-start term. When the measurement has a different dimension than the image, the measurement is upsampled (for super-resolution) or mapped by the Moore-Penrose pseudoinverse of the forward operator (for compressed sensing and other linear inverse problems).

For the few-shot scientific setting, where a handful of structurally similar images are available but the measurement may be far from the target, the measurement in the warm-start term is replaced by a learned convex combination of the available instances. A simplex constraint fixes the scale ambiguity and encourages the method to concentrate weight on the most relevant instances; in practice the constraint is enforced by parameterizing the weights with a softmax.

Why This Matters

• Impact on research. The paper reframes a working assumption in the inverse-problems literature: that bigger, more general generative models automatically make better priors. It argues the opposite can hold, shows why the standard fix fails, and demonstrates that a small, principled change to initialization and regularization recovers much of the gap. It also gives a concrete example of concentration-of-measure reasoning being used to diagnose a practical failure in a deep learning pipeline.

• Real-world applications named in the paper:

  • Medical imaging, including compressed sensing MRI.
  • Remote sensing, which the paper notes often involves high dynamic range imagery.
  • Computer vision image restoration: super-resolution, inpainting, and deblurring.
  • Scientific imaging more broadly, including materials science and astronomy, microscopy modalities, and black hole imaging.

• Industry relevance. Domain-specific generative models are expensive to collect data for and to train, and in scientific settings the clean ground truth may be impossible to obtain. FMPlug's claim is that existing foundation models — Stable Diffusion V3 and FLUX.1 are the named examples — can be reused off the shelf for specialized imaging tasks instead. If that holds, it lowers the barrier for domain scientists who cannot afford to build their own generative priors.

Future Directions

• Closing the remaining gap. FMPlug improves on other foundation FM solvers but, based on the reported AFHQ 256×256 deblurring numbers, domain-specific priors and the untrained DIP remain competitive; how much of the gap can ultimately be closed is an open question.
• Theoretically grounding the plug-in formulation. The paper explicitly notes that theoretical results on whether the generator is surjective — that every reasonable image can be produced by some latent — appear to be lacking, and that analyzing the modeling of high-dimensional distributions for this purpose is tricky. The evidence offered is empirical.
• Understanding the warm-start improvement. The paper observes that most of FMPlug's gains come from the time-dependent warm-start alone (FMPlug-W), which raises the question of how much the shell constraint contributes in settings beyond the ones tested, and how the learned t behaves.
• Broadening beyond images and beyond the tested tasks. The formulation is developed for image objects. Whether the same warm-start and shell arguments transfer to other modalities mentioned in the paper, such as video world models, is not addressed.

Target Audience

Researchers and graduate students working on inverse problems, computational imaging, and deep generative priors, particularly those already familiar with diffusion and flow-matching samplers. The paper is also relevant to applied scientists in medical imaging, remote sensing, materials science, and astronomy who have a small number of representative samples and need to solve inverse problems without training a bespoke generative model. Practitioners building image restoration systems on top of foundation models such as Stable Diffusion or FLUX will find the practical prescription — start the flow at a learned time with the measurement injected, and project the latent onto a thin shell — directly applicable.

Authors’ abstract

Foundation flow-matching (FM) models promise universal priors for solving inverse problems (IPs); yet today, they trail behind domain-specific and even untrained priors. \emph{How can we unlock their potential?} We introduce FMPlug, a plug-in framework that redefines how foundation FMs are used in IPs. FMPlug combines an instance-guided, time-dependent warm-start strategy with sharp Gaussianity regularization, adding problem-specific guidance while preserving the Gaussian structures. For evaluation, we consider both simple image restoration tasks and scientific IPs with a few similar samples -- where the prohibitive cost of data collection and model training hinders the development of domain-specific generative models. Our superior experimental results confirm the effectiveness of FMPlug. Overall, FMPlug paves the way for making foundation FM models practical, reusable priors for IPs, especially scientific ones with few similar samples. More details are available at https://sun-umn.github.io/xm-plug/ .