Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Saving Foundation Flow-Matching Priors for Inverse Problems

Yuxiang Wan, Ryan Devera, Wenjie Zhang, Ju Sun

TL;DR

FMPlug addresses the gap that foundation FM priors underperform domain-specific and untrained priors in solving inverse problems. It introduces a plug-in strategy that combines an instance-guided, time-dependent warm-start with a sharp Gaussianity constraint to inject problem-specific cues while preserving Gaussian structure. The approach is complemented by a mean-variance calibration and an extension to few-shot settings, enabling reuse with limited data. Across simple-distortion IPs and few-shot scientific IPs, FMPlug achieves consistent, state-of-the-art improvements over baselines, bridging the performance gap toward domain-specific priors and demonstrating practical viability of foundation FM priors.

Abstract

Foundation flow-matching (FM) models promise a universal prior for solving inverse problems (IPs), yet today they trail behind domain-specific or even untrained priors. 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 a sharp Gaussianity regularization, adding problem-specific guidance while preserving the Gaussian structures. This leads to a significant performance boost across image restoration and scientific IPs. Our results point to a path for making foundation FM models practical, reusable priors for IP solving.

Saving Foundation Flow-Matching Priors for Inverse Problems

TL;DR

FMPlug addresses the gap that foundation FM priors underperform domain-specific and untrained priors in solving inverse problems. It introduces a plug-in strategy that combines an instance-guided, time-dependent warm-start with a sharp Gaussianity constraint to inject problem-specific cues while preserving Gaussian structure. The approach is complemented by a mean-variance calibration and an extension to few-shot settings, enabling reuse with limited data. Across simple-distortion IPs and few-shot scientific IPs, FMPlug achieves consistent, state-of-the-art improvements over baselines, bridging the performance gap toward domain-specific priors and demonstrating practical viability of foundation FM priors.

Abstract

Foundation flow-matching (FM) models promise a universal prior for solving inverse problems (IPs), yet today they trail behind domain-specific or even untrained priors. 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 a sharp Gaussianity regularization, adding problem-specific guidance while preserving the Gaussian structures. This leads to a significant performance boost across image restoration and scientific IPs. Our results point to a path for making foundation FM models practical, reusable priors for IP solving.

Paper Structure

This paper contains 32 sections, 1 theorem, 18 equations, 9 figures, 12 tables.

Table of Contents

  1. Introduction
  2. Background and challenges in current FM-based IP solving
  3. Flow matching (FM)
  4. Pretrained FM priors for IPs
  5. Foundation FM priors for IPs
  6. Foundation FM priors $\ll$ domain-specific or even untrained ones
  7. Current ideas to strength foundation FM priors do not quite work
  8. Method
  9. An instance-guided & time-dependent warm-start strategy
  10. Why is the initialization strategy in D-Flow problematic?
  11. Our time-dependent warm-up strategy
  12. Additional mean-variance calibration
  13. A sharp Gaussianity regularization
  14. Why is the Gaussian regularization in D-Flow problematic?
  15. Our sharp Gaussian regularization via an explicit constraint
  16. ...and 17 more sections

Key Result

Theorem 3.1

For a $d$-dimensional Gaussian random vector $\boldsymbol z \sim \mathcal{N}(\boldsymbol 0, \boldsymbol I_d)$, $\mathbb P [|\norm{\boldsymbol z}_2 - \sqrt{d}| \ge t] \le 2e^{-ct^2}$ for a universal constant $c > 0$.

Figures (9)

  • Figure 1: Visual illustration of the difference between the interleaving approach (left) and the plug-in approach (right) to IPs with pretrained FM priors
  • Figure 2: Comparison between foundation FM, domain-specific FM, and untrained priors for Gaussian deblurring the on AFHQ-Cat (resolution: $256 \times 256$). DS: domain-specific FM; FD: foundation FM; FD-S: strengthened foundation FM; DIP: deep image prior. Bold: best, & underline: second best, for each metric/column. The foundation model is Stable Diffusion V3 here.
  • Figure 3: Comparison between foundation FM, domain-specific FM, and untrained priors for Gaussian deblurring with varying kernel size (Gaussian sigma) and hence varying difficulty level. Notations are the same as in \ref{['tab:comp-prior-strength']}.
  • Figure 4: Plot of the function $h(\boldsymbol z_0)$ (after a change of variable $u = \norm{\boldsymbol z_0}_2^2$). An ideal regularization function should blow up sharply away from the narrow concentration region in orange to promote Gaussianity effectively.
  • Figure 5: Visual comparison of results in Gaussian deblurring.
  • ...and 4 more figures

Theorems & Definitions (1)

  • Theorem 3.1: Concentration of measure in Gaussian random vectors vershynin2018high