Spanning Diffusion and Flow-based Posterior Sampling via Exact Gaussian Conditioning

Spanning Diffusion and Flow-based Posterior Sampling via Exact Gaussian Conditioning

arXiv:2610.13812v1 [cs.LG] 9 Oct 2026

Teo Deveney University of Bath

Lea Bogensperger University of Zurich

Abstract

Flow-matching and diffusion models are increasingly used as priors for linear inverse problems, however, most methods are bespoke to one class of trajectory prior. Plugand-play methods have been proposed, however these implicitly decorrelate the guided trajectory, resulting in broad priors that can struggle when data is sparse. This work proposes β-Flow: a general plug-and-play method indexed by a correlation parameter β that generalises over flow, diffusion, and decorrelated priors, and a confidence parameter κ that controls how strongly the path distribution is enforced. Each step requires only one forward evaluation of a denoising network, followed by a Gaussian conditioning step. This step is an exact Kalman update that is not dependent on network gradients. We show that many existing guided samplers can be interpreted as instances or approximations of our method, each implicitly using a different trajectory prior and operator in place of the Kalman gain. From this viewpoint, we provide a qualitative comparison of methods in terms of their implicit biases and assumptions. We demonstrate consistently state-of-the-art level performance across 7 inverse problems over 4 datasets in terms of PSNR and SSIM. Moreover, we demonstrate the utility of this prior control, by showing that confident high correlations (strict priors) work best when observations are less informative, while low correlation is preferred when observations are more informative, reflecting the complementary roles of the prior and likelihood in Bayesian inference.

1

Hok Shing Wong LMU Munich

INTRODUCTION

Flow-matching [1–3] and diffusion models [4, 5] generate data by transporting a simple noise distribution to the data distribution along a learned trajectory, deterministically or stochastically. This makes them natural priors for linear inverse problems. Given an observation y of a clean sample x0 y = H(x0 ) + σy η,

η ∼ N (0, I),

(1)

the aim is to draw samples, or compute a point estimate from the posterior p(x0 | y). A large body of work approaches this problem by steering the model’s reverse trajectory towards y one step at a time. Most of these methods appear tailored to one model class, with likelihood score approximations common in diffusion models (DPS [6], ΠGDM [7]), trajectory unrolling tied to flow-matching models (DFlow [8]), and independent re-noising used in plugand-play methods (PnP-Flow [9], Flower [10]). Consequently, a scheme developed for one does not naturally carry over to the other without changing its prior behaviour. In this paper, we propose a general conditioning method for guided sampling that generalises over flow, diffusion and independent re-noising priors. 1.1

Preliminaries

Let x0 ∼ pdata denote a data sample and x1 ∼ N (0, I) independent noise. We consider the linear interpolant xt = αt x0 + σt x1 ,

αt = 1 − t,

σt = t,

(2)

where t ∈ [0, 1] so that t = 0 corresponds to the data, t = 1 to pure noise. Conditional on x0 , each xt is Gaussian with mean αt x0 and covariance σt2 I. A network trained on this interpolant, whether by flow-matching or as a diffusion denoiser, provides at every xt a onestep estimate x̂0 (xt ) ≈ E[x0 | xt ] of x0 . Interpolant (2) defines a single deterministic path for each noise realisation x1 , but a guided sampler needs flexibility to move away from this path. We therefore

Teo Deveney, Lea Bogensperger, Hok Shing Wong

Figure 1: Three guided trajectories on a toy denoising example for varying βt . Each line is an independent trajectory, and the black marker is the observation. At βt = −1 the trajectory is anti-correlated which results in the closest convergence to the observation, at βt = 0 every step re-noises independently, and as βt increases consecutive states become correlated, until at βt = 1 they follow a deterministic flow. model the current sampler state as a separate variable x̃ti+1 , whose correlation with the reference path is controlled by a time dependent parameter βt ∈ [−1, 1], q
x̃ti+1 := αti+1 x0 + σti+1 βti+1 x1 + 1 − βt2i+1 ε , (3) where ε ∼ N (0, I) is independent of x1 . This process has the same marginal as xti+1 for every value of βti+1 . For βti+1 = 1 it fully coincides with xti+1 , for βti+1 = 0 it is conditionally independent of it given x0 , and for βti+1 < 0 they are negatively correlated. Together, the observation model (1) and interpolant (2) result in the graphical model of Figure 2. From here, it is natural to condition on everything that is known, i.e., x̃ti+1 and y. Each step the sampler must draw from p(xti | x̃ti+1 , y), which decomposes as p(xti | x̃ti+1 , y) = Z p(x0 | x̃ti+1 , y) {z } |

clean-signal posterior

βti+1 τ 2

Coincides with

σt2i σt2i (1 − βt⋆2 ) i+1

0 βt⋆i+1 1

0

Independently noised x0 Reverse diffusion SDE Flow-matching ODE

Table 1: Common cases of the kernel (5). The diffusion schedule (βt⋆i+1 := σti αti+1 /(σti+1 αti )) gives a Markovian forward process. Every schedule with βt ∈ [−1, 1] defines a valid kernel with the correct marginals. the kernel coincides with a known construction. We note that for βt ∈ [0, 1] this derivation serves as an alternative construction of the DDIM family [12] (Appendix A, A.1). Negative values of βt lie outside the DDIM family but define an equally valid kernel. true interpolant path

p(xti | x̃ti+1 , x0 ) dx0 . {z } |

x0

(4)

xti

xti+1

transition kernel

βti+1

H, σy

This decomposition admits a natural interpretation. The first factor is the posterior over the ground truth x0 consistent with the data y and current state x̃ti+1 , and the second factor is the transition kernel, a stochastic interpolant [11] between the endpoints x0 and x̃ti+1 that transports this posterior along the prior trajectory. Every guided sampler approximates this integral, in practice by using a network estimate x̂0 ≈ E[x0 | xt ] as a proxy for the unknown x0 . Standard Gaussian conditioning of (3) (Appendix A, A.1) gives the transition kernel in closed form,
xti | x̃ti+1 , x0 ∼ N c x0 + γ x̃ti+1 , τ 2 I , (5) with coefficients γ = βti+1 σti /σti+1 , c = αti − γαti+1 and τ 2 = σt2i (1−βt2i+1 ). This gives the correct marginal of xti given x0 for every βti+1 . Deterministic flowmatching updates and the standard diffusion posterior are both members of this family, differing only in their choice of βt . Table 1 lists three choices of βt for which

y

x̃ti+1

observation

current state

Figure 2: Graphical model for one reverse step. The state x̃ti+1 held by the sampler depends on x0 and is correlated, with coefficient βti+1 , with the reference point xti+1 on the interpolant path (grey, dashed). It is conditionally independent of xti+1 given x0 when βti+1 = 0 and coincides with it when βti+1 = 1. Guided sampling proceeds from x1 to x0 by iteratively modelling the distribution of xti given x̃ti+1 and y. 1.2

Existing approaches

Three families of guidance cover most of the literature. Each can be seen as approximating (4), and Section 3 writes each in the update form derived in this paper.

Teo Deveney, Lea Bogensperger, Hok Shing Wong

The first estimates the ground truth x̂0 (xti+1 ) directly, applies a data consistency-based update, and samples xti independently of xti+1 . PnP-Flow [9] does this via a gradient step as the data consistency update while FLOWER [10] takes a proximal step. In both cases the corrected point is returned to the trajectory by re-noising with independent noise. The second family, commonly used in diffusion settings, approximates the conditional score function by adding the score of an approximated likelihood to the unconditional score and integrating. This type of approximation is derived via the gradient of the likelihood evaluated at x̂0 (xt ). DPS [6] and ΠGDM [7] are examples of this type of approach. The update is therefore based on the network’s own local sensitivity to the state, and its step size is set by hand. The third family, most common in flow-matching settings, replaces the single step estimator with an estimate derived through a roll-out of the ODE solver. This is exemplified by D-Flow [8], which optimises the noise x1 that seeds a deterministic sampler, backpropagating a data-fidelity loss on the final output x0 through the entire discretised trajectory. Here the observation enters once, at the boundary, and the update is applied to x1 rather than to a state on the trajectory. Closely related variants also apply this update at intermediate points of the trajectory [13]. 1.3

Contributions

This paper treats a guided step as a Gaussian conditioning based on the graphical model of Figure 2. Our contributions are as follows. • One sampler for flow, diffusion and decorrelated priors. A single correlation parameter βt spans the flow-matching ODE, the reverse diffusion SDE and independent re-noising, so one method serves all three model classes and the trajectory prior becomes a choice rather than implicitly inherited. • Closed-form, gradient-free guidance. Each guided step is an exact Gaussian conditioning update, computed as a Kalman update from a single forward evaluation of the network, with no backpropagation through the network or the trajectory. • A comparison of existing methods within a unified view. DPS, ΠGDM, PnP-Flow, FLOWER and DFlow can all be written as a choice of transport and gain, which makes their implicit assumptions explicit and comparable. We compare methods qualitatively based on the transport and gain they apply. For PnP-Flow and DPS we prove that no fixed step size in their gradient step can reproduce

the calibrated Kalman gain, as both over-weight the observation increasingly as the trajectory approaches the data. • Empirical results. We show state-of-the-art performance against existing solvers, and demonstrate that the optimal choice of trajectory prior differs based on the problem, showing that control over the trajectory prior is practically beneficial.

2

METHODOLOGY

Recall equation (4), which expresses the conditional distribution of the next state xti given the current state x̃ti+1 and the observation y. We note that this can be further decomposed, being proportional to Z p(xti | x̃ti+1 , x0 ) p(x0 | x̃ti+1 ) p(y | x0 ) dx0 . (6) {z } | {z } | {z } | likelihood

prior

transition kernel

|

{z

clean-signal posterior

}

Our approach is to carry out the Bayesian inference at the level of the clean signal and then transport the result along the trajectory. The observation model gives a likelihood, a prior over the clean signal combines with it to give the clean-signal posterior, and the transition kernel of Section 1.1 transports that posterior to the next state, as in (4). Section 2.1 derives the resulting update for a Gaussian prior with a given mean and covariance. It is available in closed form for every βt , costs one forward evaluation of the network per step, and involves no network Jacobian. Section 2.2 then specifies our choice of mean and covariance using Tweedie’s formulas. 2.1

Guidance as closed form Gaussian conditioning

We assume for now that the prior over the clean signal, given the state the sampler holds, is Gaussian with a given mean mpri and covariance Spri , x0 | x̃ti+1 ∼ N (mpri , Spri ).

(7)

These may depend on the current state x̃ti+1 and timestep t + 1. We specify our choice of these in Section 2.2. The observation model (1) gives the likelihood p(y | x0 ) = N (Hx0 , σy2 I). Posterior over the clean signal. A Gaussian prior combined with a linear-Gaussian likelihood gives a Gaussian posterior, p(x0 | x̃ti+1 , y) = N (mpost , Spost ), with mpost = mpri + Spri H ⊤ (HSpri H ⊤ + σy2 I)−1 r, ⊤

Spost = Spri − Spri H (HSpri H

⊤

(8)

+ σy2 I)−1 HSpri .

Teo Deveney, Lea Bogensperger, Hok Shing Wong

The posterior mean is the prior mean plus a linear correction applied to the residual r = y −Hmpri which incorporates the data. The linear operator applied to the residual is the Kalman gain, which trades off the prior covariance against the observation noise σy2 . A confident prior (small Spri ) leaves mpri almost unchanged, and a diffuse prior lets the observation dominate. The posterior covariance is the prior covariance reduced by what the observation resolves. Transport along the trajectory. By (4) the next state is obtained by transporting this posterior through the kernel (5). The kernel is linear in x0 and its noise is independent of x0 and of y, so, given x̃ti+1 and y,
xti | x̃ti+1 , y ∼ N c mpost + γ x̃ti+1 , c2 Spost + τ 2 , (9) where (c, γ, τ 2 ) are the kernel coefficients (5). That is, the posterior mean over x0 is transported along the same path as the underlying kernel (2), and the posterior variance expands the variance of (2) additively. These formulas can also be derived in a more opaque manner by direct Gaussian conditioning over the jointly Gaussian system (x0 , xt , x̃ti+1 , y). The mean in (9) is therefore E[xti | x̃ti+1 , y], the minimummean-squared-error estimate of xti under the prior. In order to make a transition we compute the posterior transition mean µpost = c mpost + γ x̃ti+1 and add i.i.d. noise with standard deviation τ as prescribed by the kernel, giving an update rule (5) xti = µpost + τ ξ,

ξ ∼ N (0, I).

(10)

We have omitted the contribution of the c2 Spost covariance term. The reason for this is simple; the posterior covariance Spost over x0 is a consequence of our choice of prior over x0 (7) and the observation operator H only. It does not necessarily represent the natural geometry of the data distribution and can introduce excess noise and unnatural structure if propagated. To reduce variance we draw M independent copies of ξ, all of which are propagated. This mirrors the approach taken in [10]. Averaging the resulting outputs approximates marginalisation over the residual uncertainty in xti at each step, instead of committing to a single realisation at each step in the trajectory. Algorithm 1 summarises a step of our algorithm. Figure 5 in Appendix C illustrates this step on a twodimensional problem, showing the evolution of prior and posterior over xti at six points of a trajectory for three for different choices of βt . 2.2

The prior over x0

The only intractable term in (6) is the prior p(x0 | x̃ti+1 ). Our method deals with this by using a Gaussian approximation, as is implicit in many diffusion

Algorithm 1 One reverse step Require: known state x̃ti+1 , observation y, operator H, noise level σy , network fθ , correlation βti+1 , scale Sti+1 , particles M 1: x̂0 ← fθ (x̃ti+1 , ti+1 ) 2: compute γ, c, τ 2 (5) 3: mpost ← posterior mean of x0 (8) 4: xti ← c mpost + γ x̃ti+1 + τ ξ with ξ ∼ N (0, I) guidance methods [14], which we motivate here using Tweedie’s moment estimates. Given x̃ti+1 we can produce a one-step estimate x̂0 = x̂0 (x̃ti+1 ) = E[x0 |x̃ti+1 ] (e.g. using Tweedie’s formula [15]). We use this as the mean of the Gaussian prior mpri = x̂0 (7). We base the covariance Spri = Sti+1 on Tweedie’s second order formula for the covariance of x0 given x̃ti+1 , σt2i+1  αt2i+1

 I + σt2i+1 ∇2 log p(x̃ti+1 ) .

(11)

In particular we take the data independent component to give a schedule for our prior covariance, and combine this with a tunable confidence parameter κ, giving   σti+1 2 Sti+1 = κ I. (12) αti+1 The removal of the Hessian term is primarily a pragmatic one, since this is impractical to compute outright, and Jacobian-vector product-based implementations are significantly slower. We remark however, that this curvature is primarily negative in high likelihood regions of interest, so the omission of this term amounts to enlarging the variance in directions of high negative curvature rather than restricting it to no longer cover the required ground truth. 2.3

Computational cost and stability

Every quantity in (9) is an explicit function of x̂0 and x̃ti+1 , and x̂0 is obtained from a single forward evaluation of the network that is held fixed for the remainder of the step. Computing xti therefore requires no further evaluation of fθ and no differentiation through it. The network enters only through x̂0 , and its uncertainty enters only through the scalar Sti+1 in the gain. This is the same cost profile as proximal [9] and dataconsistency samplers [10], whereas score-based guidance [6, 7] obtains its update from the network’s gradients, and flow-based rollout methods require gradients over the full trajectory rollout [8]. A summary of network evaluations and gradients required per update between methods is summarised in Table 2. In terms of stability, our approach is constructed at the level of the joint Gaussian system, and therefore

Teo Deveney, Lea Bogensperger, Hok Shing Wong

is not subject to discretisation error inherent to methods framed around integrating an ODE/SDE. Furthermore, since our method is not guided by network gradients, it remains robust in regions of low density where the lower quantity of training data can result in unpredictable gradients.

3

METHOD COMPARISON

The update (9) has three distinct parts. A prior over the clean signal centred at a point estimate x̂0 , a data informed update of that estimate with the residual r := y − H x̂0 , and transport of the updated estimate back along the trajectory. Every method considered here has this structure, and can be written
xnew = T x̂0 + G r; xcur ,

(13)

so that each method is specified by its gain G and its transport T . For our update, FLOWER, PnPFlow, DPS and ΠGDM, xcur = x̃ti+1 , xnew = xti , and T is the kernel (5) for some βt , T (z; x̃ti+1 ) = c z + γ x̃ti+1 + τ ξ. PnP-Flow and FLOWER have this structure by construction, and Appendix A, A.2, shows that the score-based guidance schemes DPS and ΠGDM follow it implicitly. ODE rollout methods that optimise the latent noise x1 have a slightly different (k) (k+1) structure, setting xcur = x1 , xnew = x1 (where k indexes optimisation iterations). Here the point estimate is defined as the endpoint of the integration (k) trajectory x̂0 = F (x1 ), and transport of the updated estimate follows by the reversal of this trajec(k) tory T (z; x1 ) = F −1 (z). For all methods we derive the implicit gains G, which specify exactly how each method updates its estimate of x0 . Table 2 summarises the results, here we briefly summarise. Our update is the case G = G(Sti+1 ), where G(S) := SH ⊤ (HSH ⊤ + σy2 I)−1 is the Kalman gain of (8). It automatically balances the prior’s uncertainty against the observation noise, resulting in a properly normalised gain which stays bounded as the prior uncertainty S grows. The other methods differ in how, and whether, they supply this normalisation. 3.1

FLOWER and PnP-Flow

These methods re-noise with independent noise in each iteration, and therefore implicitly fix βt = 0. FLOWER computes the posterior mean of x0 exactly so G = G(Sti+1 ) with Sti+1 = σt2i+1 /(αt2i+1 + σt2i+1 ). FLOWER can therefore be considered as an instance of our method that uses a different prior approximation. PnP-Flow instead takes a gradient step on the data fidelity, G = ρH ⊤ , with ρ tuned per problem.

3.2

DPS and ΠGDM

DPS adds to the unconditional diffusion step (βt = βt⋆ ) the gradient of the misfit at the network estimate, differentiated through the network. With J := ∂ x̂0 /∂xti+1 and a hand-set step size ζ, its gain is GDPS = λi JH ⊤ with λi := 2ζ/c. For an exact denoiser, Tweedie’s formula gives J = (αti+1 /σt2i+1 ) Cov(x0 | xti+1 ), so DPS approximates the Kalman gain with the prior covariance replaced by a multiple of the network’s own posterior covariance, and the inverse residual covariance (HSH ⊤ + σy2 I)−1 replaced by the scalar λi . ΠGDM replaces the point mass by an isotropic Gaussian with covariance rt2 I, giving a gain proportional to J ⊤ H ⊤ (rt2 HH ⊤ +σy2 I)−1 . It contains the normalisation that DPS lacks, with rt2 I in place of S, but still uses the network Jacobian where the Kalman gain uses S. 3.3

D-Flow

D-Flow optimises the noise x1 that seeds a deterministic flow-matching model, x0 = F (x1 ), by gradient descent on ∥y − HF (x1 )∥2 . With J := ∂F/∂x1 at the current iterate and step size η, one step on x1 changes the point estimate to x̂0 + 2η J¯J ⊤ H ⊤ r, where J¯ is the average of ∂F/∂x1 along the straight line from the current iterate to the next (Appendix A, A.2). To first order J¯ = J , and the gain has the form proportional to J J ⊤ H ⊤ , and so the update can be considered a gradient step preconditioned by J J ⊤ , where J is the Jacobian of the rollout F . 3.4

Comparison summary

Written in this form, the methods differ along three axes. The first is the transport; PnP-Flow and FLOWER use the kernel at βt = 0 and DPS at βt = βt⋆ , while D-Flow transports along the learned flow. The flow and diffusion kernels are based on continuous ODE/SDE dynamics, whilst the PnP-Flow and FLOWER are discontinuous and considerably broader. The second is the gain; our update and FLOWER apply a Kalman gain built from an explicit prior on x0 , whereas PnP-Flow, DPS and D-Flow apply a gradient or Jacobian step whose scalar step size stands in for the residual covariance and so has to be tuned per problem. For PnP-Flow and DPS, Appendix B shows that the Kalman gain cannot be reproduced by choice of step size: as t → 0 the Kalman gain vanishes while theirs does not, so both increasingly over-weight the observation (Proposition 1). The third is the cost of obtaining the gain, with DPS and ΠGDM requiring a gradient through the network at every step and DFlow requiring the derivative of the whole trajectory, whereas a gain built from the prior needs none.

Teo Deveney, Lea Bogensperger, Hok Shing Wong Method This paper FLOWER PnP-Flow DPS ΠGDM D-Flow

Transport T

Point estimate

kernel, βt ∈ [−1, 1] kernel, βt = 0 kernel, βt = 0 kernel, βt⋆ kernel, βt⋆ flow, F −1

one step one step one step one step one step (k) integrated, F (x1 )

Gain G on the residual
G κ(σti+1 /αti+1 )2
G σt2i+1 /(αt2i+1 + σt2i+1 ) ρ H⊤ λi JH ⊤ ∝ J ⊤ H ⊤ (rt2 HH ⊤ + σy2 I)−1 2η J¯J ⊤ H ⊤

Gradients through fθ none none none every step every step whole trajectory

Table 2: Every method’s update written in the form of (13), with r = y−H x̂0 and G(S) = SH ⊤ (HSH ⊤ +σy2 I)−1 . For D-Flow, J¯ is the average of J along the step. Appendix A, A.2, derives the DPS and D-Flow entries.

4

NUMERICAL RESULTS

We evaluate the method across various inverse problems, including 4× super-resolution, denoising, random inpainting, Gaussian deblurring, box inpainting, MRI, and CT reconstruction using a trained flowmatching models as the denoiser, and tested on a full held-out set of N =100 test images per problem. For AFHQ-Cats and CelebA we use the experimental setup and datasets, models, and comparison methods as [10], which is the current state-of-the-art on these problems. For the undersampled MRI and sparse-view CT reconstruction experiments we train a new flowmatching model on the FastMRI knee dataset [16, 17] and re-use a trained flow-matching model from [18] trained on the AAPM dataset for CT [19] and compare all methods discussed in Section 3. Results reported for our method are all based on a 100 step discretisation, and hyperparameters were tuned on a per-case basis for the highest PSNR using the leftout validation set. Comparison methods also all used 100 steps with the exception of two pre-existing baselines which used 200∗ and 500† steps on the AFHQCats dataset consistent with [10] (indicated by the corresponding superscripts in the table). The final reported numbers report the performance on the test dataset, which is independent of the train and validation sets. We report averaged PSNR, SSIM [20] and LPIPS [21] over these test samples, highlighting the best method in bold and the second best underlined. 4.1

AFHQ-Cats and CelebA

Following the set-up of [9, 10], we evaluate on five inverse problems using the AFHQ-Cats and CelebA datasets. Results are shown in Tables 3 and 4. For AFHQ-Cats our method has the best PSNR in 3 of 5 experiments, and is second (≤ 0.03dB from the best) in the remaining two. A similar story holds for SSIM, where our method is highest in 3 of 5 experiments and second in the remaining ones by < 1%. For LPIPS our method is less performant, though very much competitive with the other approaches. This is not surpris-

ing, since our construction marginalises uncertainty to produce MMSE estimates, which (correctly) smooth over the unidentifiable fine structure that LPIPS is sensitive to. This is a classic demonstration of the perception-distortion tradeoff [22]; our method excels in accuracy, and therefore must sacrifice some perceptual capability. The results for CelebA are similar, with our method being best in 3 of 5 cases in PSNR and 4 of 5 in SSIM and second by a < 1% margin in the remaining cases. Again LPIPS favours different approaches for state-of-the-art performance, but our method still remains competitive. 4.2

Medical MRI/CT reconstruction

We evaluate on medical image reconstruction problems, sparse-view CT and accelerated MRI, using for each problem a single pretrained general purpose flowmatching prior with no problem-specific retraining. For MRI, we use the FastMRI single-coil knee dataset (128 × 128, complex-valued, stored as 2 real channels). The forward operator is an orthonormal 2D FFT followed by a variable-density random Cartesian mask retaining 20% of phase-encode lines corresponding to an effective acceleration factor of R = 5. We evaluate on a held-out test split of 100 images. For CT, we use the AAPM CT dataset (256 × 256, single channel) [19]. The forward operator is a sparse-view Radon transform with 60 projection angles evenly spaced over [0◦ , 180◦ ), normalised to unit operator norm, with additive Gaussian noise of standard deviation σ = 0.01 on the normalised sinogram. We evaluate on 50 test slices, evenly spaced across the AAPM test volume. We compare our method, β-Flow, against five baselines using the same pretrained flow-matching prior. All method-specific hyperparameters were tuned on a held-out validation split, disjoint from the test images reported here. The quantitative results are shown in Table 5. β-Flow achieves state-of-the-art performance on both medical imaging tasks, demonstrating that our approach generalises well beyond natural images. On MRI reconstruction, Flower per-

Teo Deveney, Lea Bogensperger, Hok Shing Wong

Table 3: AFHQ-Cat performance. Most results were using 100 steps, with the exception of some pre-existing PnP-Flow/Flower baselines which used 500∗ and 200† steps in certain tasks. Denoising

Deblurring

Super-resolution

Random inpainting

Box inpainting

Method

PSNR

SSIM

LPIPS

PSNR

SSIM

LPIPS

PSNR

SSIM

LPIPS

PSNR

SSIM

LPIPS

PSNR

SSIM

LPIPS

Degraded

20.00

0.314

0.509

23.94

0.517

0.444

11.70

0.208

0.873

13.36

0.223

1.081

21.80

0.740

0.198

PnP-GS DiffPIR OT-ODE D-Flow Flow-Priors PnP-Flow Flower β-Flow (ours)

32.58 30.58 30.03 26.13 29.41 31.43 32.35 32.55

0.894 0.835 0.815 0.574 0.763 0.864 0.891 0.894

0.072 0.189 0.076 0.175 0.153 0.168 0.116 0.089

27.91 27.56 27.06 27.82 26.47 28.19∗ 28.97 29.18

0.753 0.728 0.713 0.721 0.700 0.766∗ 0.784 0.790

0.349 0.342 0.123 0.164 0.181 0.332∗ 0.283 0.290

24.15 23.65 25.91 24.64 23.51 27.37∗ 26.57∗ 28.38

0.632 0.624 0.716 0.601 0.570 0.774∗ 0.750∗ 0.797

0.362 0.402 0.108 0.190 0.272 0.183∗ 0.282∗ 0.186

29.42 31.70 29.40 32.20 32.37 33.75† 33.70† 33.88

0.836 0.881 0.839 0.894 0.906 0.929† 0.927† 0.928

0.126 0.062 0.090 0.040 0.047 0.048† 0.045† 0.057

– – 24.62 26.26 26.20 26.68 26.88 26.86

– – 0.875 0.842 0.818 0.901 0.922 0.915

– – 0.085 0.077 0.118 0.120 0.066 0.083

Table 4: Results on 100 test images of the dataset CelebA. Denoising

Deblurring

Super-resolution

Random inpainting

Box inpainting

Method

PSNR

SSIM

LPIPS

PSNR

SSIM

LPIPS

PSNR

SSIM

LPIPS

PSNR

SSIM

LPIPS

PSNR

SSIM

LPIPS

Degraded

20.00

0.348

0.372

27.83

0.740

0.126

10.26

0.183

0.827

11.95

0.196

1.041

22.27

0.742

0.214

PnP-GS DiffPIR OT-ODE D-Flow Flow-Priors PnP-Flow Flower β-Flow (ours)

32.64 31.20 30.54 26.04 29.34 32.30 33.14 33.14

0.910 0.885 0.859 0.607 0.768 0.911 0.926 0.926

0.035 0.060 0.032 0.092 0.134 0.056 0.038 0.035

34.03 32.77 33.01 31.25 31.54 34.80 35.67 35.63

0.924 0.912 0.921 0.854 0.858 0.940 0.954 0.954

0.041 0.060 0.029 0.038 0.056 0.047 0.032 0.036

31.31 31.52 31.46 30.47 28.35 31.49 33.09 33.87

0.892 0.895 0.907 0.843 0.713 0.906 0.932 0.943

0.064 0.033 0.025 0.026 0.102 0.056 0.040 0.029

29.22 31.74 28.68 33.67 32.88 33.98 33.95 34.52

0.875 0.917 0.871 0.943 0.871 0.953 0.953 0.959

0.070 0.025 0.051 0.015 0.019 0.022 0.020 0.022

– – 29.40 30.70 30.07 31.09 31.87 31.58

– – 0.920 0.899 0.858 0.940 0.952 0.949

– – 0.038 0.026 0.048 0.043 0.023 0.023

Intuitively, higher values of β encourage trajectories to remain closer to flow paths, thus high β values regUndersampled MRI Sparse-view CT ularise by biasing reconstructions towards determinMethod PSNR SSIM LPIPS PSNR SSIM LPIPS istic flows. This effect can be seen in Table 6 and DPS 36.19 0.926 0.037 28.79 0.740 0.107 Figure 6, where the optimal β shows clear correlation ΠGDM 36.12 0.924 0.039 35.40 0.838 0.154 with the super-resolution downsampling scale. SimiD-Flow 34.19 0.871 0.158 27.42 0.740 0.157 larly, κ regularises behaviour by controlling the relaPnP-Flow 36.30 0.907 0.094 25.70 0.695 0.186 Flower 37.70 0.942 0.091 33.41 0.905 0.076 tive importance of prior and observations at each step β-Flow (ours) 37.76 0.943 0.102 36.44 0.898 0.079 throughout the trajectory. A lower value of κ implies a more confident prior, thus regularises behaviour by biasing it towards unconditional dynamics. This trend is clearly reflected in Table 7 and is shown for images forms very closely to our method, with the two leadin Figure 7, where the optimal κ decreases monotoniing the comparison in PSNR, whereas on sparse-view cally as downsampling scale increases. Both effects are CT, β-Flow yields clearly superior results in PSNR, visualised on the 2d toy example in Figure 4. which we attribute to the sparse-view CT problem being a more ill-conditioned problem. LPIPS perfor×2 ×3 ×4 ×6 mance is lower, consistent with our other experiments, which we attribute to the perception-distortion tradeβt PSNR SSIM PSNR SSIM PSNR SSIM PSNR SSIM off as discussed above [22]. Our method also compares −0.6 33.08 .909 29.87 .838 27.14 .762 22.00 .609 favourably in terms of runtime as shown in Table 9. −0.3 33.07 .908 29.93 .839 27.36 .768 22.40 .620 0.0 33.04 .907 29.99 .840 27.60 .772 22.84 .632 Fig. 3 compares reconstructions for all datasets. Table 5: FastMRI and AAPM CT datasets results.

4.3

βt and κ super-resolution scale ablations

Here we investigate the impact of prior hyperparameters κ and β , which both regularise the reconstruction in different ways. We apply super-resolution with the AFHQ-Cats dataset with increasingly severe downsampling scale and assess the impact of these prior hyperparameters on reconstruction quality.

0.3 32.96 .904 30.01 .839 27.76 .776 23.40 .645 0.6 32.76 .896 29.96 .836 27.85 .777 23.93 .657

Table 6: PSNR (dB) and SSIM on AFHQ-Cat superresolution across the correlation βt , with κ = 1 fixed. The optimal βt increases with the downsampling factor, showing reconstructions favour more strongly correlated trajectories at higher scale factors.

Teo Deveney, Lea Bogensperger, Hok Shing Wong

Figure 3: Qualitative reconstructions across datasets and problems. All methods used 100 steps (100 updates in the case of D-Flow), and hyperparameters tuned for each dataset and task. Columns 2-4 represent reconstructions from each class of model. In Column 2 the likelihood score methods OT-ODE/ΠGDM are applied to the upper/lower two rows respectively. These excel at producing realistic looking fine detail that does not necessarily match the ground truth. In Column 3 is the ODE rollout method D-Flow, which also produces some fine details, but appears sensitive to the degradation process as can be seen in the noise artifacts in row 2 and vertical smearing in row 3. Column 4 shows the independent-renoising method Flower, which achieves excellent accuracy when measurements are informative, but struggles at 100 steps with the sparser 4x super-resolution problem. In Column 5 β-Flow achieves excellent accuracy across all problems by adapting the prior to the task. ×2

×3

×4

×6

κ PSNR SSIM PSNR SSIM PSNR SSIM PSNR SSIM 0.1 32.09 0.5 33.00 1.0 33.08 3.0 32.99

.882 .907 .909 .904

30.15 30.19 30.01 29.49

.845 .846 .839 .823

28.40 28.22 27.85 26.47

.798 .790 .777 .735

25.57 25.12 23.93 20.66

.712 .696 .657 .553

Table 7: PSNR (dB) and SSIM on AFHQ-Cat superresolution for varying confidence scale κ. For each column βt is chosen to be the PSNR-optimal value from Table 6 (see Figure 7). The optimal κ decreases as the downsampling factor increases, favouring a more confident prior when observations are less informative.

5

CONCLUSION

We have introduced a general plug-and-play technique for guided flow sampling for inverse problems that

makes the prior behaviour itself an explicit choice. Our method allows an arbitrary denoising generative model to define trajectory priors ranging from deterministic flows to decorrelated, and even anticorrelated paths with controllable strength, while being simply interpreted as standard conditioning in a Gaussian graphical model. We illustrated the practical relevance of this control in super-resolution by demonstrating that harder tasks perform substantially better with straighter trajectories and stronger prior strength. Our method attains state-of-the-art performance in PSNR and SSIM in most tasks, and is < 1% from the best performance in all tasks where it is not. Moreover, we demonstrated that many existing approaches can be interpreted as approximations to our approach, and presented a qualitative comparison across model classes that reveals the implicit assumptions contained in three common model classes.

Teo Deveney, Lea Bogensperger, Hok Shing Wong

Acknowledgements This work was supported by the Maths4DL EPSRC Programme on the Mathematics of Deep Learning (EP/V026259/1), the UKRI Strength in Places Fund MyWorld Project (SIPF00006/1), and the Munich Center for Machine Learning (MCML). References [1] Yaron Lipman, Ricky T. Q. Chen, Heli BenHamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. In International Conference on Learning Representations (ICLR), 2023. [2] Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In International Conference on Learning Representations (ICLR), 2023. [3] Michael S. Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. In International Conference on Learning Representations (ICLR), 2023. [4] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems (NeurIPS), volume 33, 2020. [5] Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations (ICLR), 2021. [6] Hyungjin Chung, Jeongsol Kim, Michael Thompson Mccann, Marc Louis Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. In International Conference on Learning Representations (ICLR), 2023. [7] Jiaming Song, Arash Vahdat, Morteza Mardani, and Jan Kautz. Pseudoinverse-guided diffusion models for inverse problems. In International Conference on Learning Representations (ICLR), 2023. [8] Heli Ben-Hamu, Omri Puny, Itai Gat, Brian Karrer, Uriel Singer, and Yaron Lipman. D-Flow: Differentiating through flows for controlled generation. In Proceedings of the 41st International Conference on Machine Learning (ICML), volume 235 of Proceedings of Machine Learning Research, 2024. [9] Ségolène Martin, Anne Gagneux, Paul Hagemann, and Gabriele Steidl. PnP-Flow: Plug-andplay image restoration with flow matching. In

International Conference on Learning Representations (ICLR), 2025. [10] Mehrsa Pourya, Bassam El Rawas, and Michael Unser. Flower: A flow-matching solver for inverse problems. In International Conference on Learning Representations (ICLR), 2026. [11] Michael S. Albergo, Nicholas M. Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797, 2023. [12] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations (ICLR), 2021. [13] Alexander Denker, Moshe Eliasof, Zeljko Kereta, and Carola-Bibiane Schönlieb. Trajectory stitching for solving inverse problems with flow-based models. arXiv preprint arXiv:2602.08538, 2026. [14] Xinyu Peng, Ziyang Zheng, Wenrui Dai, Nuoqian Xiao, Chenglin Li, Junni Zou, and Hongkai Xiong. Improving diffusion models for inverse problems using optimal posterior covariance. In Proceedings of the 41st International Conference on Machine Learning (ICML), volume 235 of Proceedings of Machine Learning Research, pages 40347–40370, 2024. [15] Bradley Efron. Tweedie’s formula and selection bias. Journal of the American Statistical Association, 106(496):1602–1614, 2011. [16] Jure Zbontar, Florian Knoll, Anuroop Sriram, Matthew Muckley, Mary Bruno, Aaron Defazio, Marc Parente, Krzysztof J. Geras, Joe Katsnelson, Hersh Chandarana, Zizhao Zhang, Michal Drozdzal, Adriana Romero, Michael G. Rabbat, Pascal Vincent, James Pinkerton, Duo Wang, Nafissa Yakubova, Erich Owens, C. Lawrence Zitnick, Michael P. Recht, Daniel K. Sodickson, and Yvonne W. Lui. fastmri: An open dataset and benchmarks for accelerated mri. ArXiv, abs/1811.08839, 2018. URL https://api. semanticscholar.org/CorpusID:53759905. [17] Florian Knoll, Jure Zbontar, Anuroop Sriram, Matthew J. Muckley, Mary Bruno, Aaron Defazio, Marc Parente, Krzysztof J. Geras, Joe Katsnelson, Hersh Chandarana, Zizhao Zhang, Michal Drozdzalv, Adriana Romero, Michael Rabbat, Pascal Vincent, James Pinkerton, Duo Wang, Nafissa Yakubova, Erich Owens, C. Lawrence Zitnick, Michael P. Recht, Daniel K. Sodickson, and Yvonne W. Lui. fastmri: A publicly available raw k-space and dicom dataset of knee images for accelerated mr image reconstruction using machine learning. Radiology: Arti-

Teo Deveney, Lea Bogensperger, Hok Shing Wong

ficial Intelligence, 2(1):e190007, 2020. doi: 10. 1148/ryai.2020190007. URL https://doi.org/ 10.1148/ryai.2020190007. PMID: 32076662. [18] Alexander Denker, Johannes Hertrich, and Sebastian Neumayer. A stability benchmark of generative regularizers for inverse problems. arXiv preprint arXiv:2605.10076, 2026. [19] Cynthia H McCollough, Adam C Bartley, Rickey E Carter, Baiyu Chen, Tammy A Drees, Phillip Edwards, David R Holmes III, Alice E Huang, Farhana Khan, Shuai Leng, et al. Lowdose ct for the detection and classification of metastatic liver lesions: results of the 2016 low dose ct grand challenge. Medical physics, 44(10): e339–e352, 2017. [20] Zhou Wang, Alan C Bovik, Hamid R Sheikh, and Eero P Simoncelli. Image quality assessment: from error visibility to structural similarity. IEEE transactions on image processing, 13(4):600–612, 2004. [21] Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness of deep features as a perceptual metric. In 2018 IEEE/CVF conference on computer vision and pattern recognition, pages 586– 595. IEEE, 2018. [22] Yochai Blau and Tomer Michaeli. The perceptiondistortion tradeoff. In 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 6228–6237. IEEE, 2018. [23] Mukund Sundararajan, Ankur Taly, and Qiqi Yan. Axiomatic attribution for deep networks. In Proceedings of the 34th International Conference on Machine Learning (ICML), volume 70 of Proceedings of Machine Learning Research, 2017.

Spanning Diffusion and Flow-based Posterior Sampling via Exact Gaussian Conditioning: Supplementary Materials

A

DERIVATIONS

A.1 The kernel family. Given x0 , the point xti = αti x0 + σti x1 on the interpolant and the state x̃ti+1 of (3) are affine in the same noise x1 , plus independent noise, so they are jointly Gaussian,       σt2i I βti+1 σti σti+1 I αti x0 xti x0 ∼ N , , βti+1 σti σti+1 I σt2i+1 I αti+1 x0 x̃ti+1 where the cross-covariance comes from the shared component x1 . Further conditioning on x̃ti+1 gives 

 σ ti x̃ti+1 − αti+1 x0 , σt2i 1 − βt2i+1 I , xti | x̃ti+1 , x0 ∼ N αti x0 + βti+1 σti+1

(14)

which is (5) with γ = βti+1 σti /σti+1 , c = αti − γαti+1 and τ 2 = σt2i (1 − βt2i+1 ). To recover the diffusion prior from this we note that DDPM is defined using the Markov forward step xti+1 = a xti + b ξ, with ξ ∼ N (0, I) independent of xti , a = αti+1 /αti and b2 = σt2i+1 − a2 σt2i so that the marginals are preserved, has the reverse posterior q(xti | xti+1 , x0 ) = N

a σ2 ti x0 + 2 ti xti+1 , 2 σti+1 σti+1

 b2 α

b2 σt2i  I . σt2i+1

Matching the coefficient of (14) gives βti+1 = aσti /σti+1 = σti αti+1 /(σti+1 αti ) = βt⋆i+1 , and at this value the coefficient of x0 and the variance agree as well. For the DDIM family we note that for marginals N (αt x0 , σt2 I), the DDIM kernel with noise scale s ∈ [0, σti ] is q  xt − αti+1 x0 2  qs (xti | xti+1 , x0 ) = N αti x0 + σt2i − s2 i+1 , s I . σti+1 q This is (14) with βti+1 = 1 − s2 /σt2i , so βti+1 ∈ [0, 1] covers the DDIM family, and s = 0 (βti+1 = 1) is the deterministic update used in flow-matching. Negative βti+1 is not covered by DDIM, but nonetheless defines a valid kernel which introduces transitions that are actively negatively correlated. A.2 DPS and D-Flow in the form of (13). A DPS step proceeds by taking an unconditional reverse diffusion step and adding a gradient step based on the data misfit to the estimated ground truth. Taking βt = βt⋆ , this can be written as xti = c x̂0 (xti+1 ) + γ xti+1 + τ ξ − ζ ∇xti+1 ∥y − H x̂0 (xti+1 )∥2 | {z } Unconditional step ⊤

⊤

(15)

= c x̂0 + λi J H r + γ xti+1 + τ ξ, where the second line uses ∇xti+1 ∥y − H x̂0 ∥2 = −2J ⊤ H ⊤ r with J := ∂ x̂0 /∂xti+1 , and λi := 2ζ/c. This is (13) with the kernel at βt⋆ and the gain λi J ⊤ H ⊤ . For an exact denoiser J = (αt /σt2 ) Cov(x0 | xt ), though a trained network satisfies this only approximately. ΠGDM’s isotropic covariance rt2 I changes the residual weighting to (rt2 HH ⊤ + σy2 I)−1 but keeps J in place of S.

Teo Deveney, Lea Bogensperger, Hok Shing Wong (k)

(k)

D-Flow. With x̂0 = F (x1 ) and J := ∂F/∂x1 at x1 , a gradient step on ∥y − HF (x1 )∥2 with step size η, and the update it induces on the point estimate, are x1

(k+1)

= x1 + 2η J ⊤ H ⊤ r,

(k+1)

= x̂0 + 2η J¯J ⊤ H ⊤ r,

F x1 where J¯ :=

R1

(k)

(k)

(16) (k+1)

J (x1 +s ∆x1 ) ds is the average of J along the step ∆x1 := x1

(k)

−x1 . The second line is exact: (k+1) the fundamental theorem of calculus along the straight line between the iterates gives F (x1 ) − x̂0 = J¯ ∆x1 , −1 −1 the identity behind Integrated Gradients [23]. Applying F gives (13) with T = F and the gain of Table 2. To first order J¯ = J , and J J ⊤ is the covariance that a unit Gaussian perturbation of x1 induces on F (x1 ). (k+1) Beyond first order J¯ depends on x1 , and so on r, so the D-Flow gain is not a fixed operator.

B

0

STEP-SIZE GAINS VERSUS THE KALMAN GAIN

Every method in Table 2 corrects the point estimate by Gr, where r = y − H x̂0 is the residual. The step changes the residual to r′ = y − H(x̂0 + Gr) = (I − HG) r,

(17)

so HG is the part of the residual the step removes. The residual can also be written as r = σy η+H(x− x̂0 ), thus it is a combination of the prediction error (projected into observation space), which is signal, and observation noise. The Kalman step removes a proportion of the residual that balances the uncertainty in the prior and observations based on Gaussian conditioning. A step that removes more of the residual over-weights the observations, either by treating the data as more certain, or the prior as less certain, leading to overfitting to the noise. Setting. Let H be any linear operator, let Ct := Cov(x0 | xt ) be the posterior covariance at the current time t = ti+1 , and assume ∥Ct ∥ → 0 as t → 0. Write G(C) = CH ⊤ (HCH ⊤ + σy2 I)−1 for the Kalman gain built from a covariance C. Ours is G(Sti+1 ) with Sti+1 of (12), and Gt := G(Ct ) is the Kalman gain with the true posterior covariance. PnP-Flow applies gt = ρH ⊤ , and DPS applies gt = λi Jt⊤ H ⊤ with λi = 2ζ/c, where Jt := ∂ x̂0 /∂xt is the Jacobian of the network. The following proposition shows that without tuned step size schedules, PnP-Flow and DPS cannot reproduce the calibrated step of the Kalman gain. Proposition 1. Let r be a residual with H ⊤ r ̸= 0, and let t → 0. (i) Kalman gains, including ours. HGt → 0 and HG(Sti+1 ) → 0, for every κ > 0. (ii) PnP-Flow. Hgt = ρHH ⊤ does not depend on t, so the step removes the share ρHH ⊤ r of the residual at every step. (iii) DPS. Assuming Jt → I as t → 0 (the denoiser output converges to its input in the limit of no noise) ∥Hgt r∥ ≥ 2ζ∥HH ⊤ r∥ (1 + o(1)), so the residual removed does not vanish. PnP-Flow and DPS therefore overweight the observations by continuing to reduce the residual as t → 0, where the true conditional mean recognises models the remaining residual as noise. In particular, no fixed step size reproduces the Kalman gain along the trajectory, and these methods require scheduling as a result. Proof. (i) With A := σy−2 HCH ⊤ ⪰ 0, HG(C) = A(I + A)−1 , so ∥HG(C)∥ ≤ ∥A∥ ≤ ∥H∥2 ∥C∥/σy2 , which tends to zero for C = Ct and for C = Sti+1 = κ(σt2 /αt2 )I. (ii) HH ⊤ r ̸= 0 since H ⊤ r ̸= 0. (iii) c = αti − γαti+1 ≤ 1 for the diffusion kernel (γ ≥ 0), so λi ≥ 2ζ, and HJt⊤ H ⊤ r → HH ⊤ r when Jt → I.

Teo Deveney, Lea Bogensperger, Hok Shing Wong

C

TWO-DIMENSIONAL ILLUSTRATION

In this appendix we provide some further 2-dimensional visualisations in order to further inspect the mechanisms of the approach, which is difficult to see directly for the high-dimensional image problems of Section 4. We apply our construction to a two-moons denoising problem (H = I), so that the priors and posterior at each step are two-dimensional and can be plotted directly. Figure 4 compares the impact of the correlation and confidence parameters β, κ on the trajectories generated, and Figure 5 visualises the prior and posterior distributions that are modelled during a single transition at various stages of sampling. Both figures share the same trained flow-matching network, the same true signal x0 and observation y.

Figure 4: Trajectories for a two-moons denoising problem under for trajectory priors in β ∈ {1, βt∗ , 0} and endpoint-uncertainty scale κ ∈ {0, 0.1, 1, 5}. For κ = 0 the Kalman gain vanishes, so the particles follow the unconditional prior unaffected by the observations. For κ > 0 the prior weakens and lets the Kalman gain pulls the particles towards observation with the magnitude of the pull determined by κ. Trajectories in the top row are deterministic given the initial noise and observations, as expected of an ODE update, whereas those in the lower two rows take random transitions at each step.

Teo Deveney, Lea Bogensperger, Hok Shing Wong

Figure 5: The conditioning step at six times along a trajectory (columns) for three choices of βt (rows): the decorrelated kernel (βt = 0), the diffusion schedule (βt = βt⋆ ) and the flow-matching kernel (βt = 1). Each panel shows the prior over xti before conditioning on y (dashed orange) and the posterior (9) (solid green), together with the true signal (red diamond) mapped into current time point (i.e. the red diamonds are showing αt y for the time associated to that column, as this is what the time conditional posterior is targeting). Larger βt leaves less fresh noise τ 2 in the kernel, so both prior and posterior contract, and at βt = 1 the conditioning step makes only a small correction to a nearly deterministic update.

D

IMPLEMENTATION DETAILS

We base our implementation on the github repository for Flower https://github.com/mehrsapo/Flower which is a fork of the PnP-Flow repository https://github.com/annegnx/PnP-Flow. We inherit the model architecture, weights, and datasets for the AFHQ-Cats and CelebA from these works. For the FastMRI experiment we trained a model using the same architecture, and for the AAPM model we used the pre-trained model from [18]. A fully reproducible implementation with pre-trained weights and dataset splits is planned for released in the near future. The following subsections describe the problem-specific implementation details. D.1

AFHQ-Cats and CelebA

We inherit the experimental setup, pre-trained model weights, datasets/splits, and baseline hyperparameters directly from [10], which in turn builds on [9]. We evaluate on two datasets. CelebA consists of 128 × 128 human face images; AFHQ-Cat consists of 256 × 256 cat images. All images are normalised to the range [−1, 1]. Models are trained on the full training splits of both datasets. All metrics are reported on 100 test images per dataset using a consistent split with the comparison methods. All flow-matching methods (including β-Flow and the baselines OT-ODE, D-Flow, Flow-Priors, and PnP-Flow) share the same U-Net backbone trained with mini-batch optimal-transport flow matching and a Gaussian latent prior. Pre-trained weights are taken from [9], trained with learning rate 10−4 : on CelebA for 200 epochs (batch size 128) and on AFHQ-Cat for 400 epochs (batch size 64). We evaluate on five restoration tasks: denoising with additive Gaussian noise (σn = 0.2); Gaussian deblurring with a 61 × 61 kernel (σb = 1.0 for CelebA, σb = 3.0 for AFHQ-Cat) and additive noise (σn = 0.05); super-

Teo Deveney, Lea Bogensperger, Hok Shing Wong

resolution at 2× for CelebA and 4× for AFHQ-Cat (σn = 0.05); random inpainting with 70% of pixels removed (σn = 0.01); and box inpainting with a centred rectangular mask (40 × 40 for CelebA, 80 × 80 for AFHQ-Cat, σn = 0.05). We compare against OT-ODE, D-Flow, Flow-Priors, PnP-Flow, PnP-GS, and DiffPIR. Hyperparameters for all baselines are adopted from [10], which performed a grid search on the validation set to maximise PSNR. The number of steps was fixed to 100 for all methods; the Flower and PnP-Flow results that used 200† or 500∗ steps on AFHQ-Cat are quoted directly from [10] and are indicated by the corresponding superscripts in Table 3. For β-Flow, βt was selected from a uniform grid on [−0.9, 0.9] at spacing 0.1, augmented with ±0.95 and ±0.99 near the boundaries, and κ from {0.1, 0.5, 1, 5, 10}, both tuned independently per dataset and task on the held-out validation split. The resulting values are reported in Table 8.

Table 8: β-Flow hyperparameters (βt , κ) per task, selected on the validation split to maximise PSNR. CelebA

AFHQ-Cat

βt

βt

κ

Denoising 0.1 1.0 0.4 Deblurring −0.99 1.0 −0.8 Super-resolution −0.9 1.0 0.5 Random inpainting −0.9 10.0 −0.9 Box inpainting 0.9 0.1 0.8

1.0 5.0 0.1 10.0 0.1

Task

D.2

κ

Medical MRI/CT reconstruction

Every method was tuned on a separate validation set, and the number of steps was fixed to 100 (except for D-Flow). Per-image runtimes for both medical tasks are reported in Table 9. For the fastMRI dataset, images are 128 × 128, complex-valued (2 real channels). The forward operator retains 20% of phase-encode lines (yielding an effective R=5 acceleration), with Gaussian noise σ = 0.05 added directly in k-space. DPS used ζ = 1.0, ΠGDM used a guidance scale of 0.5, with the conjugate-gradient solver run for up to 100 iterations (tolerance 10−5 ), PnP-Flow used 5 sample averaging and ρ = 1.5, kept constant throughout, Flower was run with 5 sample averaging (Flower5), D-Flow used an initialization blend of 0.7, a regularization weight of 10−8 , 15 outer LBFGS iterations with 20 inner iterations each, and a 6-step Euler discretization of the flow ODE, β-Flow used β = −0.9 and κ = 1.0 and also used 5 sample averaging. For the CT dataset, images are 256×256, single-channel. The forward operator is a sparse-view Radon transform with 60 projection angles over [0◦ , 180◦ ), with Gaussian noise σ = 0.01 on the normalized sinogram. DPS used ζ = 5.0, ΠGDM used a guidance scale of 1.5, with the same conjugate-gradient settings as above, PnP-Flow used 5 sample averaging and ρ = 2.0, kept constant throughout, Flower was again run with 5 sample averaging, D-Flow used an initialization blend of 0.8, with the same remaining settings as on MRI, β-Flow used β = −0.7 and κ = 5 and also used 5 sample averaging.

Table 9: Wall-clock time per image, in seconds, for the MRI and CT results of Table 5, under the same fixed step budget of 100 used there (except D-Flow, which uses its own optimisation budget). Method

MRI

CT

DPS 8.92 7.43 ΠGDM 9.02 21.75 D-Flow 312.52 283.87 PnP-Flow 15.64 15.95 Flower 23.48 48.76 β-Flow (ours) 6.67 27.10

Teo Deveney, Lea Bogensperger, Hok Shing Wong

E

βt AND κ ABLATION VISUALISATION

Figure 6: Qualitative effect of βt on AFHQ-Cat super-resolution with fixed κ = 1. Leftmost column is the observation, and Columns 2-6 show reconstructions with different βt = β parameters. As the downsampling factor increases the observation constrains x0 less, and reconstructions benefit more from increased βt provides. At ×4 and ×6, βt = −0.6 have visible grid artefacts and washed out tone that improve when βt increases, demonstrating the trend in Table 6.

Teo Deveney, Lea Bogensperger, Hok Shing Wong

Figure 7: Qualitative effect of κ on AFHQ-Cat super-resolution using the same test image as Figure 6. For each scale factor βt is chosen to be its own PSNR-optimal βt (×2: −0.6, ×3: 0.3, ×4, ×6: 0.6) from Figure 6. A larger κ widens the prior and lets the observation dominate the update; by κ = 3 this produces blurring and a strong grid artefact, most pronounced at ×4 and ×6 where the observation is least informative, matching the trend in Table 7.

Related documents

• Spanning Diffusion and Flow-based Posterior Sampling via Exact Gaussian Conditioning#1298308
• Exact Global MCMC with Denoising Diffusion#819735
• Exact Global MCMC with Denoising Diffusion#809739
• Order-Optimal Sample Complexity of Rectified Flows#1232651
• Score-based diffusion models for severely ill-posed problems in diffuse optical tomography#1014605