Flow-matching and diffusion models have become priors for linear inverse problems, yet most methods remain bespoke to a single class of trajectory prior. Plug-and-play alternatives exist but implicitly decorrelate the guided trajectory, producing broad priors that underperform when observations are sparse. This work introduces β-Flow, a general plug-and-play method indexed by a correlation parameter β that generalizes over flow, diffusion, and decorrelated priors, together with a confidence parameter κ that controls enforcement of the path distribution. Each step requires only one forward evaluation of a denoising network, followed by a Gaussian conditioning step that is an exact Kalman update independent of network gradients. The authors show that many existing guided samplers are instances or approximations of β-Flow, each implicitly using a different trajectory prior and operator in place of the Kalman gain. From this viewpoint they provide a qualitative comparison of methods in terms of their implicit biases and assumptions. Experiments across 7 inverse problems on 4 datasets demonstrate consistently state-of-the-art performance in PSNR and SSIM. Moreover, the authors show 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.
Why Dataset-Driven Priors Matter for Inverse Problems
Linear inverse problems appear across imaging, signal processing, and scientific computing. The task is to recover a clean signal from incomplete or corrupted observations. In recent years, diffusion models and flow-matching models have been adopted as learned priors that encode statistical structure from training data. These learned priors can dramatically improve recovery quality compared to traditional hand-crafted priors such as total variation or Gaussian Markov random fields. However, most existing approaches fix a single trajectory prior—either a diffusion process or a flow—and tune the model specifically for that prior. When the fixed prior does not match the true data distribution or the observation geometry, recovery quality degrades.
Plug-and-play methods have been proposed as a way to use pre-trained denoising networks as priors without retraining. These methods typically alternate denoising steps with data consistency projections. The problem is that many plug-and-play schemes implicitly decorrelate the guided trajectory, resulting in broad priors that cannot leverage strong statistical dependencies in the data. In sparse observation regimes, where the likelihood provides little constraint, such broad priors struggle to produce sharp, accurate reconstructions.
β-Flow: A Generalized Plug-and-Play Framework
The central contribution of this paper is β-Flow, a framework that introduces two knobs: a correlation parameter β and a confidence parameter κ. The correlation parameter controls how strongly the guided trajectory follows a learned prior, spanning from independent (β = 0) to fully correlated (β = 1) priors. The confidence parameter κ modulates how strongly the path distribution is enforced at each step. Together, β and κ allow the practitioner to interpolate between flow-based priors, diffusion-based priors, and decorrelated priors, depending on the observation regime.
Each step of β-Flow consists of two operations. First, a single forward evaluation of a denoising network produces a denoised estimate. Second, a Gaussian conditioning step applies an exact Kalman update. Crucially, this conditioning step does not depend on network gradients; it is a closed-form linear algebraic operation that conditions the current state on the observed measurements. The Kalman update form means that β-Flow can be interpreted as performing exact Bayesian inference under a linear-Gaussian observation model, with the denoising network supplying the prior mean and covariance structure.
The authors demonstrate that many existing guided samplers—including various diffusion sampling schemes and flow-matching variants—are either special cases or approximations of β-Flow, each implicitly selecting a different point in the (β, κ) parameter space. This reinterpretation provides a unified lens through which to compare methods: the differences are not in the overall algorithmic structure but in the choice of trajectory prior and the effective operator substituted for the Kalman gain.
Exact Gaussian Conditioning as a Kalman Update
The Gaussian conditioning step in β-Flow is an exact Kalman update. Given a current state estimate x and an observation y with observation matrix H and noise covariance R, the Kalman update computes
K = P Hᵀ (H P Hᵀ + R)⁻¹
x₊ = x + K (y - H x)
P₊ = (I - K H) P
where P is the current error covariance and x₊, P₊ are the updated state and covariance after incorporating the observation. In the β-Flow context, the denoising network provides an estimate of the prior mean and an implicit covariance, and the Kalman update fuses this prior with the likelihood encoded by the measurements. Because the update is exact and gradient-free, it can be implemented with standard linear algebra routines without backpropagation through the denoising network.
Experimental Evaluation Across Inverse Problems
The authors evaluate β-Flow across 7 inverse problems on 4 datasets. The inverse problems include image reconstruction tasks such as compressed sensing, inpainting, and deblurring, as well as tomographic reconstruction and inverse problems in scientific computing. Each problem presents a different balance between the informativeness of the observations and the complexity of the prior required.
Performance is measured in PSNR and SSIM, two standard metrics for image quality. β-Flow achieves state-of-the-art level performance across all 7 problems, often matching or surpassing specialized methods that are hand-tuned for each specific inverse problem. The results suggest that the generalized prior control afforded by β-Flow is competitive with, and in some cases superior to, problem-specific designs.
Prior Control and the Balance Between Prior and Likelihood
A key analytical finding is the relationship between the correlation parameter β and the informativeness of the observations. When observations are sparse or noisy—i.e., the likelihood provides little information—confident high correlations (β close to 1, effectively strict priors) work best. The prior dominates the posterior, and the denoising network's learned structure is leveraged fully. When observations are relatively informative—i.e., the likelihood already constrains the solution well—low correlation (β close to 0, effectively decorrelated priors) is preferred. In this regime, over-relying on the prior can hurt recovery; the method should trust the data more and let the likelihood dominate.
This observation reflects the complementary roles of the prior and likelihood in Bayesian inference. The β parameter effectively allows the user to schedule the relative trust between the two, adapting the algorithm's behavior to the specific observation regime without changing the underlying denoising network.
Limitations and Practical Considerations
The paper notes that β-Flow relies on a denoising network pretrained for the task at hand. The quality of the prior is bounded by the network's capacity and the diversity of its training data. If the network was trained on a narrow data distribution, the prior may not generalize to the inverse problem at hand. The authors also acknowledge that the Kalman update assumes a linear-Gaussian observation model; for strongly nonlinear problems, the exact conditioning step may need to be approximated or augmented with iterative refinement.
Additionally, the method introduces two hyperparameters (β and κ) that must be set. While the paper provides guidelines—high β for sparse observations, low β for informative observations—the optimal setting may depend on the specific problem and may require cross-validation. The authors suggest that automated schemes for selecting β and κ based on the observed signal-to-noise ratio could be a fruitful direction for future work.
What It Means for Working Developers
For developers working on inverse problems, β-Flow provides a drop-in generalization of existing plug-and-play and diffusion-based reconstruction schemes. Because each step requires only one forward pass of a denoising network plus a gradient-free Kalman update, the computational overhead relative to existing methods is modest. The β and κ parameters can be adjusted to match the observation regime: set β high and κ confident when observations are sparse, and set β low when observations are already strong.
The unified reinterpretation of existing methods also means that developers can evaluate alternative sampling strategies within a common framework. Instead of choosing between a diffusion sampler and a flow sampler, they can experiment with different points in the (β, κ) space and observe the resulting trade-offs in recovery quality and computational cost.
The exact Kalman conditioning step further implies that β-Flow can be combined with any pre-trained denoiser, even those whose architectures were not designed for probabilistic modeling. The gradient-free nature of the conditioning step makes it straightforward to integrate into existing pipelines that use off-the-shelf denoising networks.
In practical terms, the ability to interpolate between strict and decorrelated priors means that a single β-Flow implementation can potentially replace multiple specialized reconstruction pipelines, reducing maintenance burden and code complexity for teams that handle a variety of inverse problems.
A natural ending: β-Flow demonstrates that a two-parameter family of plug-and-play methods can generalize over flow, diffusion, and decorrelated priors, with exact Gaussian conditioning as the unifying computational kernel. The framework provides a quantitative handle on the prior–likelihood trade-off, and its consistent state-of-the-art performance across diverse inverse problems suggests that generalized prior control is a promising direction for learned Bayesian inference.
Read the paper on arXiv