Technical Summary: Saving Foundation Flow-Matching Priors for Inverse Problems

1. Problem Statement

Inverse problems (IPs) aim to recover an unknown object $x$ from measurements $y \approx A(x)$, where the forward model $A(\cdot)$ is often ill-posed due to insufficient data or noise. While deep generative priors (DGPs) have revolutionized IP solving, a significant performance gap exists between domain-specific priors (trained on specific datasets like FFHQ) and foundation priors (general models like Stable Diffusion or Flux trained on web-scale data).

The paper identifies that while foundation Flow-Matching (FM) models are powerful generative tools, they perform poorly as priors for IPs compared to domain-specific models or even untrained priors (e.g., Deep Image Prior). Existing attempts to adapt foundation FM priors to IPs—such as interleaving generation steps with gradient updates or using simple warm-starts—fail to constrain the solution space effectively, often resulting in reconstructions worse than the raw measurements. The core challenge is that foundation models lack the specific semantic and structural constraints required to resolve the ambiguity in ill-posed IPs without overfitting or hallucinating.

2. Methodology: FMPlug

The authors propose FMPlug, a plug-in framework designed to bridge the performance gap between foundation FM priors and domain-specific solutions. The method operates within the plug-in formulation where the optimization variable is the latent code $z$, and the generator $G_\theta$ is fixed:
$$z^* \in \arg \min_z \ell(y, A \circ G_\theta(z)) + \Omega(z)$$

FMPlug introduces two key components to strengthen the foundation prior:

A. Instance-Guided, Time-Dependent Warm-Start

Standard initialization strategies for FM priors often fail due to a distribution mismatch. Training samples for FM models concentrate on an "ultra-thin shell" in the latent space (due to the concentration of measure phenomenon for Gaussian vectors), whereas naive initialization strategies place the starting point outside this shell, leading to unpredictable generator behavior.

To address this, FMPlug introduces a time-dependent warm-start strategy:

• Simple-Distortion Setting: When the target $x$ and measurement $y$ are close (e.g., image restoration), the method approximates the flow trajectory by injecting $y$ directly into the flow at a learnable time step $t \in [0, 1]$. Instead of starting from $t=0$ (pure noise), the optimization starts from $\alpha_t y + \beta_t z$, where $\alpha_t$ and $\beta_t$ are flow coefficients.
• Few-Shot Setting: For scientific IPs where $x$ and $y$ may not be close, but a small set of similar instances $\{x_k\}$ is available, the method uses a linear combination of these instances as the warm-start seed: $\alpha_t (\sum w_k x_k) + \beta_t z$. The weights $w$ are optimized to prioritize instances structurally similar to the target.

B. Sharp Gaussianity Regularization

The authors argue that previous regularization methods (e.g., penalizing the negative log-likelihood of the Gaussian seed) are too "soft" and fail to enforce the strict concentration of measure required for FM priors.

• They introduce a shell constraint $z \in S^{d-1}_\epsilon(0, \sqrt{d})$, forcing the latent vector $z$ to lie within a narrow shell around the sphere of radius $\sqrt{d}$.
• This is implemented via a projection operator at each optimization step, ensuring the latent code remains within the high-probability region where the foundation model was trained, thereby preventing the generator from operating in uncharted regions of the latent space.

3. Key Contributions

1. Identification of the Performance Gap: The paper empirically demonstrates that foundation FM priors significantly underperform compared to domain-specific priors and even untrained priors (DIP) on standard IP tasks like Gaussian deblurring.
2. FMPlug Framework: The proposal of a unified framework combining instance-guided warm-starts and sharp Gaussian regularization to effectively adapt foundation models to IPs.
3. Few-Shot Adaptation for Scientific IPs: An extension of the framework to handle scientific imaging scenarios where domain-specific training is infeasible due to data scarcity, utilizing a few-shot instance-guided strategy.
4. Systematic Evaluation: Comprehensive benchmarks showing that FMPlug achieves state-of-the-art performance among foundation-based solvers, often surpassing untrained priors and approaching the performance of domain-specific models.

4. Experimental Results

The authors evaluated FMPlug on two main categories of tasks:

• Simple-Distortion IPs (Image Restoration):

  • Tasks: Super-resolution, random-mask inpainting, Gaussian deblurring, and motion deblurring.
  • Datasets: DIV2K, RealSR, AFHQ.
  • Results: FMPlug outperformed all other foundation-based methods (D-Flow, FlowDPS, FlowChef) and the untrained Deep Image Prior (DIP) across most metrics (PSNR, SSIM, LPIPS). Notably, interleaving methods (FlowDPS/FlowChef) with foundation priors performed worse than the blurry input itself in some cases, while FMPlug produced high-fidelity reconstructions.

• Few-Shot Scientific IPs:

  • Tasks: Linear Inverse Scattering (LIS), Compressed Sensing MRI, and Black Hole Imaging (BKH).
  • Results: FMPlug significantly outperformed D-Flow and DIP. In LIS and BKH tasks, FMPlug's performance approached that of Red-Diff, a method using domain-specific priors. For MRI, a slight performance gap remained, which the authors attribute to the test set having higher spatial variability than the few-shot guidance instances, violating the assumption of instance closeness.

• Ablation Studies:

  • Removing either the warm-start or the sharp Gaussian regularization led to significant performance drops, confirming both components are necessary.
  • The learned warm-start time $t^*$ correlated inversely with the severity of degradation (larger blur $\to$ smaller $t^*$), validating the theoretical intuition.
  • The method is robust to low numbers of Function Evaluations (NFE), requiring fewer steps than interleaving methods to achieve high fidelity.

5. Significance and Claims

The paper claims that FMPlug provides a practical, reusable path for utilizing foundation Flow-Matching models in inverse problems. By addressing the specific weaknesses of foundation priors (lack of domain specificity and distribution mismatch), the authors demonstrate that these general models can be effectively "saved" and repurposed for scientific and restoration tasks without the prohibitive cost of collecting domain-specific data or training new models.

The work highlights that while foundation models are generative powerhouses, they require specific structural guidance (via warm-starts and strict latent constraints) to function as effective priors. This approach is particularly significant for scientific imaging, where data scarcity often precludes the development of custom generative models, yet high-quality reconstruction is critical for discovery. The authors position FMPlug not as a replacement for domain-specific models where data exists, but as a superior alternative to untrained priors and a viable solution for data-sparse scientific domains.