By flow-based models, we mean diffusion models (DMs) and flow-matching (FM) models.

There are two families of methods, as shown above: (Left) Interleaving Methods alternate between flow generative steps and measurement-inconsistency reduction steps; (Right) Plug-in Methods directly minimize the measurement inconsistency with the generative priors plugged into the objective:

where \(\mathcal G_{\boldsymbol{\theta}}\) is the whole pretrained generation process, and the optimization is over the latent variable \(\boldsymbol{z}\). These two families of methods are contrasted in the table below (\(\widehat{\boldsymbol{x}}\) is the final estimate)

Generative models for domains such as images and videos are trained on internet-scale data, and hence are foundational enough to be considered as universal generators. Do these powerful generative models lead to powerful generative priors for inverse problems?

In FMPlug, we assume that additional guide samples close to the target \(\boldsymbol{x}\) are available, and propose augmenting the weak foundation flow-based prior on \(\boldsymbol{x}\) by guidance from these samples.

Instance-guided warm-start

Assume \(\boldsymbol{x}_0\) is a sample close to the target \(\boldsymbol{x}\). The ideal flow path in flow-based models is \[ \boldsymbol{x}_t = \alpha_t \boldsymbol{x} + \beta_t \boldsymbol{z} \quad t\in[0,1], \] which can be approximated by \[ \boldsymbol{x}_t \approx \alpha_t \boldsymbol{x}_0 + \beta_t \boldsymbol{z},\quad t\in[0,1]. \] with an approximation error \(\alpha_t \|\boldsymbol{x}_0-\boldsymbol{x}\|\). Since \(\alpha_t\) is a predefined function of \(t\), we can leave \(t\) learnable to be adaptive to \(\|\boldsymbol{x}_0-\boldsymbol{x}\|\) so that \(\alpha_t \|\boldsymbol{x}_0-\boldsymbol{x}\|\) is small.

Plugging the instance-guided warm-start into the DMPlug objective yields our FMPlug framework to save foundation flow-based priors (\(\Omega\) omitted):

Here, the constraint \(\boldsymbol{z} \in S^{d-1}_{\epsilon}(0,\sqrt{d})\) ensures that the seed \(\boldsymbol{z}\) lies in a thin Gaussian shell as dictated by the concentration of measure behavior of Gaussian random vectors (we assume the standard Gaussian source distribution). For simple image restoration problems, we can often take \(\boldsymbol{x}_0 = \boldsymbol{y}\). For linear IPs, we can often take \(\boldsymbol{x}_0 = \mathcal{A}^\dagger \boldsymbol{y}\), where \(\mathcal{A}^\dagger\) is the Moore-Penrose pseudoinverse of \(\mathcal{A}\).

Many scientific domains are data-scarce and hence cannot afford to train domain-specific flow-based generative models. However, their adjacent domains with similar object types might already have foundation generative models, which could be repurposed for the scientific domain under consideration, e.g., generative models of natural images can be repurposed for microscopy imaging. Moreover, since scientific IPs are typically highly nonlinear, precluding guidance from \(\mathcal{A}^\dagger \boldsymbol{y}\) or \(\boldsymbol{y}\) itself as in the simple IPs above.

We consider a set of \(K\) guide samples \(\{\boldsymbol{x}_k\}_{k=1}^{K}\), some of which are close to the target \(\boldsymbol{x}\), and consider the combined guide \(\sum_{k=1}^{K}w_k \boldsymbol{x}_k\) with learnable weights \(\boldsymbol{w}\), leading to

where \(\Delta_K \doteq \{\boldsymbol{w}\in\mathbb{R}^K: w_k\geq 0, \sum_{k=1}^{K}w_k=1\}\) is the probability simplex. The simplex weights \(\boldsymbol{w}\) select and combine the most relevant few-shot instances from \(\{\boldsymbol{x}_k\}_{k=1}^{K}\).

Visual results of few-shot FMPlug on scientific IPs