PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Teo Deveney, Lea Bogensperger, Hok Shing Wong

arXiv:2610.13812v1cs.LGcs.CV

TL;DR

Guidance methods for flow-matching and diffusion models often assume different trajectory priors, limiting their direct transfer across model classes. This paper uses a correlation-controlled Gaussian-conditioning framework to span these priors and reports state-of-the-art PSNR and SSIM in most tasks, with performance within 1% of the best in the others.

  • Problem

    Many guided samplers are tailored to a model class or trajectory prior, so a scheme developed for one does not naturally transfer to another without changing its prior behavior.

    Source: §1 · p. 1

  • Method

    A single correlation parameter βt lets one sampler use flow, diffusion, or independent re-noising trajectory priors.

    Source: §1.3 · p. 3

  • Results

    The method attains state-of-the-art PSNR and SSIM in most tasks and is within <1% of the best performance in the remaining tasks.

    Source: §5 · p. 8

  • Takeaways & Limitations

    In super-resolution, harder tasks perform substantially better with straighter trajectories and stronger prior strength.

    Source: §5 · p. 8

    The Hessian term is omitted because computing it outright is impractical, while Jacobian-vector-product implementations are significantly slower.

    Source: §2.2 · p. 4

Abstract

from arXiv · show

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. Plug-and-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 $\beta$-Flow: a general plug-and-play method indexed by a correlation parameter $\beta$ that generalises over flow, diffusion, and decorrelated priors, and a confidence parameter $\kappa$ 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 INTRODUCTION

Flow-matching and diffusion priors offer a route to posterior inference for linear inverse problems, but existing guidance methods are often tied to particular model classes or trajectory priors. The paper presents a Gaussian-conditioning framework that makes trajectory correlation an explicit choice and compares prior guidance methods within that common view.

1.1 Preliminaries

  • The sampler models the current state as correlated with the reference interpolant by βt, then conditions the next state on the current state and observation.βt = 1 makes the states coincide, βt = 0 makes them conditionally independent given x0, and negative values produce negative correlation.Source: §1.1 · p. 2

1.2 Existing approaches

  • Existing guidance approaches include independent re-noising after data consistency, likelihood-score guidance, and trajectory optimization through an ODE rollout.These families differ in how they update or transport the sample.Source: §1.2 · p. 3

1.3 Contributions

  • A single βt parameter spans flow-matching, reverse-diffusion, and independent re-noising priors, making the trajectory prior an explicit choice across model classes.The trajectory prior is no longer implicitly inherited from the model class.Source: §1.3 · p. 3
  • Each guided step is an exact Gaussian-conditioning Kalman update requiring one network forward evaluation and no backpropagation through the network or trajectory.Source: §1.3 · p. 3
  • DPS, ΠGDM, PnP-Flow, FLOWER, and D-Flow can be compared as choices of transport and gain, exposing their implicit assumptions.For PnP-Flow and DPS, no fixed gradient-step size reproduces the calibrated Kalman gain because they increasingly over-weight observations near the data.Source: §1.3 · p. 3
  • The method reaches state-of-the-art performance against existing solvers and shows that the best trajectory prior depends on the inverse problem.The paper identifies control over the trajectory prior as practically beneficial.Source: §1.3 · p. 3

2 METHODOLOGY

The method performs Bayesian inference on the clean signal, then transports the posterior along the trajectory. It uses a Gaussian prior and linear-Gaussian conditioning to produce an update requiring one network forward pass and no network Jacobian.

2 METHODOLOGY

  • The update uses a prior, incorporates the observation through a clean-signal posterior, then transports that posterior along the trajectory.The resulting step is available in closed form for every β_t and requires one network forward evaluation without a network Jacobian.Source: §2 · p. 3

2.1 Guidance as closed form Gaussian conditioning

  • The posterior mean corrects the prior mean using the residual through a Kalman gain, while observation-resolved information reduces the posterior covariance.A confident prior leaves the prior mean nearly unchanged; a diffuse prior lets the observation dominate.Source: §2.1 · pp. 3–4
  • The posterior is transported through the trajectory kernel, and the update combines its transported mean with kernel noise.The method omits propagation of the posterior covariance term because it may introduce excess noise and unnatural structure.Source: §2.1 · p. 4
  • Averaging M independently propagated noise copies approximates marginalisation over residual uncertainty rather than committing to one trajectory realisation.The authors state that this is used to reduce variance.Source: §2.1 · p. 4

2.2 The prior over x0

  • The prior mean is the network’s one-step estimate of the clean signal, and its covariance uses a data-independent schedule with tunable confidence κ.The Hessian term is omitted as impractical to compute, although the authors note its curvature is primarily negative in high-likelihood regions.Source: §2.2 · p. 4

2.3 Computational cost and stability

  • One forward network evaluation supplies the estimate for the step; no further evaluation or differentiation through the network is needed.The method is constructed from a joint Gaussian system and is not subject to ODE/SDE integration discretisation error; it also avoids network-gradient guidance in low-density regions.Source: §2.3 · pp. 4–5

3 METHOD COMPARISON

The comparison frames guided samplers by their trajectory transport, update gain, and cost of obtaining that gain. The methods differ in their implicit priors and operators, while the proposed prior-based gain is calibrated without network derivatives.

3.4 Comparison summary

  • The proposed update uses a normalized Kalman gain that balances prior uncertainty against observation noise and remains bounded as prior uncertainty grows.Table 2 summarizes each method’s update in terms of its residual and gain.Source: §3 · p. 5, Table 2 · p. 6

3.1 FLOWER and PnP-Flow

  • FLOWER uses a prior approximation to compute the posterior mean, whereas PnP-Flow applies a data-fidelity gradient step with a problem-tuned scale.Both methods re-noise independently at each iteration, corresponding to β_t = 0.Source: §3.1 · p. 5

3.2 DPS and ΠGDM

  • DPS approximates the Kalman gain using the network’s posterior covariance and a scalar in place of residual-covariance normalisation; ΠGDM adds normalisation but still uses the network Jacobian.The cited comparison says the Kalman gain’s prior covariance cannot be reproduced by choosing a step size for these schemes.Source: §3.2 · p. 5

3.3 D-Flow

  • D-Flow updates the latent noise that seeds a deterministic sampler, with its first-order gain proportional to J J^⊤H^⊤.This is equivalent to a gradient step preconditioned by the rollout Jacobian covariance J J^⊤.Source: §3.3 · p. 5

3.4 Comparison summary

  • Methods differ in transport, gain, and gain-computation cost; prior-based gains need no network derivatives, unlike gradient- or Jacobian-based updates.PnP-Flow and FLOWER use β_t = 0, DPS uses β*_t, and D-Flow transports along the learned flow; PnP-Flow and DPS increasingly over-weight observations as t approaches zero.Source: §3.4 · p. 5

4 NUMERICAL RESULTS

Across natural-image and medical reconstruction tasks, β-Flow performs strongly against comparison methods, while super-resolution ablations show that the preferred trajectory prior changes with downsampling severity.

4 NUMERICAL RESULTS

  • Evaluation covers seven inverse problems, with N=100 held-out test images per problem and averaged PSNR, SSIM, and LPIPS.Hyperparameters were tuned per case, and final reported numbers use test data independent of training and validation.Source: §4 · p. 6

4.1 AFHQ-Cats and CelebA

  • On AFHQ-Cats, β-Flow has the best PSNR in 3 of 5 experiments and is within 0.03 dB of the best in the other two; CelebA shows best PSNR in 3 of 5 cases.For SSIM, it is highest in 3 of 5 AFHQ-Cats experiments and second in the rest by less than 1%; on CelebA it is best in 4 of 5 cases. LPIPS favors other approaches, though β-Flow remains competitive.Source: §4.1 · p. 6, Table 3 · p. 7, Table 4 · p. 7

4.2 Medical MRI/CT reconstruction

  • On MRI, Flower and β-Flow lead PSNR closely, while β-Flow is clearly superior on sparse-view CT; its lower LPIPS is consistent with the perception-distortion tradeoff [22].The medical comparison uses the same pretrained flow-matching prior for both tasks, and β-Flow achieves state-of-the-art performance on each.Source: §4.2 · pp. 6–7, Table 5 · p. 7

4.3 βt and κ super-resolution scale ablations

  • Across AFHQ-Cat super-resolution, optimal βt rises with downsampling severity, while optimal κ falls, favoring more correlated trajectories and a more confident prior as observations weaken.Figure 3 depicts β-Flow achieving excellent accuracy across problems by adapting the prior to the task.Source: Table 6 · p. 7, §4.3 · p. 7, Table 7 · p. 8

5 CONCLUSION

β-Flow makes trajectory-prior behavior an explicit choice, with controllable correlation and strength. The conclusion reports state-of-the-art PSNR and SSIM in most tasks and near-best performance elsewhere.

5 CONCLUSION

  • The method attains state-of-the-art PSNR and SSIM in most tasks and is < 1% from the best performance in all others.The conclusion also reports that existing approaches can be interpreted as approximations and qualitatively compared by their implicit assumptions.Source: §5 · p. 8
  • An arbitrary denoising model can define trajectory priors from deterministic flows to decorrelated or anti-correlated paths, with controllable strength.The conclusion reports that harder super-resolution tasks perform substantially better with straighter trajectories and stronger prior strength.Source: §5 · p. 8
  • Many existing approaches are interpretable as approximations to this method, enabling qualitative comparison of their implicit assumptions across three common model classes.Source: §5 · p. 8

A DERIVATIONS

The derivations connect the Gaussian kernel family to flow and diffusion updates, including negatively correlated transitions. They also express DPS and D-Flow through the paper’s transport-and-gain formulation.

A.1 The kernel family.

  • Conditioning jointly Gaussian states yields the kernel update, whose β choice recovers the diffusion posterior and whose β = 1 case gives deterministic flow updates.The family also covers the DDIM range for β in [0, 1].Source: §A.1 · p. 11
  • Negative β defines valid transitions with active negative correlation, beyond the DDIM family described here.Source: §A.1 · p. 11

A.2 DPS and D-Flow in the form of (13).

  • DPS can be written with the diffusion-kernel choice β⋆ and a Jacobian-based gain; ΠGDM instead uses an isotropic covariance for residual weighting while retaining J.Source: §A.2 · p. 11
  • D-Flow’s update on the seeded noise induces the paper’s transport-and-gain form, with the gain involving an average Jacobian along the optimization step.To first order, the average Jacobian is the Jacobian at the step start; beyond first order it depends on the endpoint.Source: §A.2 · p. 12

B STEP-SIZE GAINS VERSUS THE KALMAN GAIN

The section contrasts calibrated Kalman corrections with fixed-step gradient corrections by tracking how each changes the observation residual. It argues that PnP-Flow and DPS increasingly overweight observations near the data endpoint without step-size scheduling.

B STEP-SIZE GAINS VERSUS THE KALMAN GAIN

  • The Kalman step removes a residual share that balances prior and observation uncertainty; removing more risks fitting observation noise.The residual combines projected prediction error and observation noise.Source: B STEP-SIZE GAINS VERSUS THE KALMAN GAIN · p. 12
  • Kalman gains, including the proposed gain, tend to zero as t approaches zero, whereas PnP-Flow retains a fixed residual correction share.The DPS comparison assumes its denoiser Jacobian tends to the identity in the no-noise limit.Source: B STEP-SIZE GAINS VERSUS THE KALMAN GAIN · p. 12
  • PnP-Flow and DPS continue reducing the residual as t approaches zero, over-weighting observations where the true conditional mean treats the remaining residual as noise.No fixed step size reproduces the Kalman gain along the trajectory, so these methods require scheduling.Source: B STEP-SIZE GAINS VERSUS THE KALMAN GAIN · p. 12

C TWO-DIMENSIONAL ILLUSTRATION

The two-moons illustration visualizes how correlation and confidence shape trajectories and conditioning at successive times. The implementation section describes inherited or trained models used for the experiments.

C TWO-DIMENSIONAL ILLUSTRATION

  • For κ = 0 the Kalman gain vanishes and particles follow the unconditional prior; for κ > 0 the prior weakens and the observation pull grows with κ.The top-row trajectories are deterministic, while the lower two rows make random transitions at each step.Source: C TWO-DIMENSIONAL ILLUSTRATION · p. 13
  • Larger β_t leaves less fresh kernel noise, contracting both prior and posterior; at β_t = 1 conditioning makes only a small correction to a nearly deterministic update.Figure 5 compares six trajectory times across three β_t choices and displays prior and posterior distributions.Source: C TWO-DIMENSIONAL ILLUSTRATION · p. 14, Figure 5 · p. 14

D IMPLEMENTATION DETAILS

  • The experiments inherit AFHQ-Cats and CelebA architecture, weights, and datasets; FastMRI uses a trained model, while AAPM uses a pretrained model.A reproducible implementation with pretrained weights and dataset splits was planned for release.Source: D IMPLEMENTATION DETAILS · p. 14

D.1 AFHQ-Cats and CelebA

The experiments use CelebA faces and AFHQ-Cat images across five restoration tasks, with shared flow-matching models and tuned baselines. β-Flow’s βt and κ are selected independently for each dataset and task on a held-out validation split.

Experimental setup

  • The experimental setup, pretrained weights, dataset splits, and baseline hyperparameters are inherited from [10], which builds on [9].Source: §D.1 · p. 14
  • Five restoration tasks cover denoising, Gaussian deblurring, super-resolution, random inpainting, and box inpainting on CelebA and AFHQ-Cat.The datasets contain 128 × 128 face images and 256 × 256 cat images; evaluation uses 100 test images per dataset.Source: §D.1 · pp. 14–15

Models and baselines

  • All flow-matching methods use 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].Source: §D.1 · p. 14
  • The comparisons include six methods, with baseline hyperparameters adopted from [10] after validation-set grid search for maximum PSNR.Methods use 100 steps, except Flower and PnP-Flow results on AFHQ-Cat quoted from [10] with 200† or 500∗ steps.Source: §D.1 · p. 15

Hyperparameter selection

  • β-Flow’s βt and κ are tuned independently per dataset and task on the held-out validation split to maximise PSNR.The validation search uses a uniform βt grid on [−0.9, 0.9], augmented by ±0.95 and ±0.99, and κ values {0.1, 0.5, 1, 5, 10}.Source: §D.1 · p. 15, Table 8 · p. 15

D.2 Medical MRI/CT reconstruction

The medical reconstruction experiments evaluate MRI and CT using task-specific forward operators, noise levels, and sampler settings. Per-image runtimes are reported under a fixed 100-step budget except for D-Flow.

MRI reconstruction

  • MRI uses 128 × 128 complex images, 20% of phase-encode lines (effective R=5), and k-space noise σ = 0.05.β-Flow uses β = −0.9 and κ = 1.0 with five-sample averaging.Source: §D.2 · p. 15

CT reconstruction

  • CT uses 256×256 single-channel images, 60-angle sparse-view Radon measurements, and normalized-sinogram noise σ = 0.01.β-Flow uses β = −0.7 and κ = 5 with five-sample averaging.Source: §D.2 · p. 15

Runtime

  • Per-image wall-clock times for MRI and CT are reported in seconds under the 100-step budget, except D-Flow, which uses its own optimization budget.Source: Table 9 · p. 15, §D.2 · p. 15

E βt AND κ ABLATION VISUALISATION

The visualizations show how βt and κ affect AFHQ-Cat super-resolution reconstructions. Increasing βt benefits reconstructions more as downsampling increases, while larger κ can produce blur and grid artefacts.

βt ablation

  • At ×4 and ×6, βt = −0.6 produces visible grid artefacts and washed-out tone that improve as βt increases.The reported trend is that reconstructions benefit more from increased βt as the downsampling factor increases and the observation constrains x0 less.Source: Figure 6 · p. 16

κ ablation

  • At κ = 3, the wider prior lets the observation dominate the update, producing blurring and a strong grid artefact, especially at ×4 and ×6.Source: Figure 7 · p. 17
  • Figure 7 uses the same test image as Figure 6 and selects each scale factor’s PSNR-optimal βt: −0.6 at ×2, 0.3 at ×3, and 0.6 at ×4 and ×6.Source: Figure 7 · p. 17

Related papers

Loading 2610.13812v1…