PaperlayerView exact revision on arXiv

Source-linked AI summary

Kinetic Langevin Meets Split Gibbs: Accelerated Posterior Sampling for Imaging Inverse Problems with Diffusion Priors

Dai Hai Nguyen, Duc Dung Nguyen

arXiv:2610.10187v1stat.MLcs.LG

TL;DR

Posterior sampling for imaging with denoiser-based priors is difficult because the prior-side conditional in split Gibbs sampling is costly to update. The paper introduces kinetic Langevin updates within split Gibbs sampling and reports faster convergence and high-quality reconstruction across several imaging tasks.

  • Problem

    Split Gibbs sampling makes the data update tractable, but sampling its prior-side conditional is difficult when the prior is defined through a denoiser or generative model.

    Source: §1 · p. 1

  • Method

    RED-KLwSGS keeps the exact Gaussian data-variable update and uses kinetic Langevin diffusions to update the auxiliary variable.

    Source: §1 · p. 2

  • Results

    Across three inverse problems on FFHQ and ImageNet, the proposed methods show faster convergence and high-quality image reconstruction.

    Source: §1 · p. 2, §5 · p. 8

  • Takeaways & Limitations

    The experiments support kinetic Langevin split Gibbs sampling for image reconstruction across the tested tasks and datasets.

    Source: §1 · p. 2

    The theoretical guarantees for the joint variant are left to future work.

    Source: §3.3 · p. 6

Abstract

from arXiv · show

Split Gibbs sampling (SGS) is a popular framework for posterior sampling in Bayesian imaging inverse problems. It decouples a Gaussian data-fidelity term from a complex prior through an auxiliary variable, so the data variable is updated exactly and only the prior-side conditional is hard to sample. Existing samplers treat this conditional in one of two ways. Plug-and-play SGS runs a multi-step diffusion denoiser at every iteration, which is expensive and lacks non-asymptotic guarantees. Langevin-within-SGS takes cheap overdamped Langevin steps but needs many iterations. We propose RED-KLwSGS, which keeps the exact Gaussian update for the data variable and updates the auxiliary variable with underdamped (kinetic) Langevin diffusions driven by a one-shot denoising score, at the same per-iteration cost as Langevin-within-SGS. We prove non-asymptotic Wasserstein-2 convergence in continuous and discrete time for strongly log-concave priors. We also introduce Joint-RED-KLwSGS, which applies kinetic Langevin diffusions to both variables. Experiments with Denoising diffusion probabilistic models as diffusion priors on FFHQ and ImageNet datasets show faster convergence and high-quality image reconstruction.

1 INTRODUCTION

The introduction motivates split Gibbs sampling as a way to separate data fidelity from an implicit prior, while identifying slow or costly updates of the prior-side conditional as a bottleneck. It presents kinetic Langevin updates as the basis for RED-KLwSGS and Joint-RED-KLwSGS, with theoretical guarantees and experiments across imaging tasks.

1 INTRODUCTION

  • MAP estimation returns a single point and cannot quantify uncertainty, while MCMC can access the posterior but is costly for high-dimensional imaging with implicit denoiser-based priors.Split Gibbs sampling addresses this challenge by separating data fidelity from the prior through an auxiliary variable.Source: §1 · pp. 1–2
  • PnP-SGS uses many diffusion-model evaluations per Gibbs iteration, whereas RED-LwSGS uses one but needs many more iterations to converge.The introduction frames these approaches as a trade-off between per-iteration cost and iteration count.Source: §1 · p. 2
  • RED-KLwSGS keeps the exact Gaussian update for the data variable and uses kinetic Langevin diffusions for the auxiliary variable; Joint-RED-KLwSGS applies them to both variables.The proposed methods build on split Gibbs sampling.Source: §1 · p. 2
  • The paper provides theoretical guarantees for RED-KLwSGS in continuous time and its discretization.The introduction identifies kinetic Langevin diffusions as a faster-converging alternative to overdamped Langevin diffusions (Cheng et al., 2018).Source: §1 · p. 2
  • Experiments on FFHQ and ImageNet across Gaussian deblurring, motion deblurring, and super-resolution report faster convergence and high-quality reconstruction.The passage describes these as applications of the proposed methods.Source: §1 · p. 2

2 RELATED WORK

The related-work section presents RED and split Gibbs sampling as tools for posterior sampling with denoiser-based priors. It compares PnP-SGS and RED-LwSGS, then reports a motion-deblurring timing comparison that motivates the proposed sampler.

2.1 Regularization-by-denoising (RED)

  • RED obtains the prior gradient from the denoising residual using a single denoiser evaluation, without differentiating the denoiser.This identity allows denoisers to be embedded in the prior.Source: §2.1 · pp. 2–3

2.2 Split Gibbs Sampling (SGS)

  • Split Gibbs sampling alternates between conditionals that decouple the likelihood and prior potentials, generating samples asymptotically distributed according to the augmented target.Its x-conditional is easier to sample than the original posterior because it replaces the prior with an isotropic Gaussian centered at z.Source: §2.2 · p. 3, §2.3 · p. 3

2.3 PnP-SGS and RED-LwSGS

  • PnP-SGS uses a multi-step reverse diffusion process for each auxiliary-variable draw, while RED-LwSGS uses a one-evaluation gradient step that is cheaper per iteration but needs more iterations.The proposed RED-KLwSGS is motivated by this per-iteration-cost versus iteration-count trade-off.Source: §2.3 · p. 4
  • RED-KLwSGS reaches 26.5 dB in 31.5 s, versus 46.64 s for RED-LwSGS and 52.2 s for PnP-SGS on one FFHQ motion-deblurring image.The table compares time per iteration, iteration count, and total time to the target PSNR.Source: §2.3 · p. 4, Table 1 · p. 4

3 PROPOSED METHODS

RED-KLwSGS retains the exact Gaussian update for the data variable and uses kinetic Langevin diffusion for the auxiliary variable, with non-asymptotic convergence analysis. Joint-RED-KLwSGS extends kinetic updates to both variables.

3.1 RED-KLwSGS: Kinetic Lavengin-within-Split Gibbs under RED

  • RED-KLwSGS applies kinetic Langevin diffusion to the prior-side conditional while retaining the exact update for x.Each iteration uses one denoiser evaluation and an exact Gaussian draw for X, matching RED-LwSGS per-iteration cost; the added velocity requires O(d) memory and arithmetic.Source: §3.1 · p. 4

3.2 Theoretical analysis

  • The guarantees for RED-KLwSGS are analyzed for smooth potentials, with strong convexity of gRED ensured when the denoiser is contractive.The cited conditions give gRED strong convexity with parameter 1 − ϵ for a denoiser with Lipschitz constant ϵ < 1.Source: §3.2 · pp. 4–5

3.2.1 Convergence of the continuous-time process

  • The continuous-time RED-KLwSGS process converges exponentially to its stationary distribution under the stated conditions.The analysis assumes a smooth, strongly convex prior potential; a contractive denoiser is sufficient for strong convexity of gRED.Source: §3.2 · pp. 4–5, §3.2.1 · p. 5

3.2.2 Convergence of the discrete-time process

  • The discrete-time process has a unique stationary law, and the analysis establishes convergence to that law.The discrete-time contraction result is stated for step sizes h in (0, 1/2].Source: §3.2.2 · p. 5

3.2.3 Discretization bias and algorithmic convergence

  • The discretization analysis bounds the bias between the continuous- and discrete-time stationary laws and uses that result to establish algorithmic convergence.Lemma 3.7 states a bias bound with a constant independent of h for h in (0, 1/2].Source: §3.2.3 · p. 5

3.3 Joint-RED-KLwSGS: applying kinetic Langevin diffusions to both variables

  • Joint-RED-KLwSGS applies kinetic Langevin diffusion to both variables, addressing settings where the exact Gaussian draw is unavailable or costly; its theoretical guarantees are left to future work.The exact draw requires a quadratic data-fidelity term and can be impractical for an unstructured operator because of inversion and storage costs.Source: §3.3 · pp. 5–6

4 NUMERICAL EXPERIMENTS

The experiments compare restoration quality and convergence for three inverse tasks on FFHQ and ImageNet. The kinetic variants converge faster than RED-LwSGS and improve benchmark PSNR and SSIM at identical per-iteration cost.

4.1 Experimental setup

  • The experiments use 256 × 256 RGB images from FFHQ and ImageNet for Gaussian deblurring, motion deblurring, and 4× super-resolution.The diffusion models are pretrained and used without fine-tuning; deblurring uses diagonal noise covariance, while super-resolution uses 40 dB SNR.Source: §4.1 · p. 6

4.2 Acceleration of RED-KLwSGS and Joint-RED-KLwSGS

  • By iteration 400 on Gaussian deblurring, RED-KLwSGS and Joint-RED-KLwSGS reach PSNR ≈27dB, while RED-LwSGS reaches PSNR ≈17dB.The accelerated variants reach their highest quality by iteration 500; RED-LwSGS needs more than 600 iterations to reach comparable quality, with a similar acceleration effect on the other tasks and ImageNet.Source: §4.2 · p. 7, Figure 2 · p. 7

4.3 Comparison with existing methods

  • RED-KLwSGS achieves the best PSNR and SSIM in five of six task-rows on both datasets, improving over RED-LwSGS by up to 1.8 dB at identical per-iteration cost.Joint-RED-KLwSGS ranks second on PSNR/SSIM in nearly every row but has the best LPIPS in three of six rows, consistent with a perception-distortion trade-off.Source: Table 2 · p. 8, Table 3 · p. 8, §4.3 · p. 8
  • The comparison includes methods that provide point estimates and samplers that provide posterior distributions for uncertainty quantification.The methods are evaluated using PSNR, SSIM, and LPIPS averaged over 1000 test images.Source: §4.3 · p. 8

5 CONCLUSION

The conclusion presents RED-KLwSGS as a kinetic Langevin sampler within RED split Gibbs, pairing an exact Gaussian data-variable update with an auxiliary-variable update. It reports convergence guarantees and faster reconstruction, while identifying extensions to the analysis and setting as future work.

5 CONCLUSION

  • RED-KLwSGS pairs the exact Gaussian update for the data variable with a kinetic Langevin update for the auxiliary variable.Joint-RED-KLwSGS applies momentum to both variables.Source: §5 · p. 8
  • The paper establishes nonasymptotic Wasserstein-2 convergence guarantees for RED-KLwSGS in continuous time and for its discretization.Source: §5 · p. 8
  • On FFHQ and ImageNet across three inverse problems, RED-KLwSGS improves PSNR and SSIM over RED-LwSGS at identical per-iteration cost and converges in fewer iterations.Joint-RED-KLwSGS performs competitively on perceptual metrics.Source: §5 · p. 8
  • Future work includes extending convergence analysis to the joint variant, relaxing RED conditions, and studying forward operators beyond the quadratic-Gaussian setting considered.Source: §5 · p. 8

CHECKLIST

The checklist marks the paper’s explanations of assumptions and theoretical results, as well as its reporting of experimental reproducibility information, as provided.

CHECKLIST

  • The checklist marks clear explanations of model assumptions and complete proofs of theoretical results as provided.Source: CHECKLIST · p. 10
  • It marks code, data, and reproduction instructions, training details, and descriptions of metrics and error bars as provided.Source: CHECKLIST · p. 10
  • The checklist marks computing infrastructure information as provided, while several asset and human-subjects items are marked not applicable.Source: CHECKLIST · p. 10

A DIFFUSION MODELS AND SPLIT GIBBS SAMPLERS

The appendix describes a DDPM as both a learned score and denoiser, then relates its noise levels to the prior-side conditional in split Gibbs. It contrasts using the diffusion model as a multi-step sampler with using it as a one-evaluation score source.

A.1 Diffusion models as stochastic denoisers

  • A DDPM forward chain gradually corrupts clean images with Gaussian noise, while its backward chain converts noise back to data.The model uses a fixed noise schedule.Source: §A.1 · p. 11
  • The trained network predicts added noise, and one network evaluation yields both a score and an MMSE estimate of the clean image.The noise predictor’s minimizer is the conditional mean of the added noise.Source: §A.1 · p. 11

A.2 Connection between diffusion models and Split Gibbs Samplers

  • The split-Gibbs prior-side conditional is a denoising posterior for recovering z from x = z + ρn under prior e−βg.For β = 1, the matching noise level corresponds to the diffusion model’s trained denoising posterior.Source: §A.2 · p. 12
  • PnP-SGS (Coeurdoux et al., 2024) samples the conditional by running a reverse chain, whereas RED-LwSGS (Faye et al., 2024) uses one network evaluation at a fixed noise level for a Langevin step.For β = 1/ν², the RED-LwSGS residual corresponds to the score of the prior smoothed at noise level ν.Source: §A.2 · p. 12

B PROOF OF THEOREM 3.3

The proof couples two continuous-time processes using shared noise and bounds their differences. A Lyapunov argument and Gronwall’s inequality establish contraction under the stated spectral bounds.

B PROOF OF THEOREM 3.3

  • The proof uses synchronous coupling, with shared Gaussian noise and Brownian motion, to control the distance between two processes.The initial states are drawn independently from their respective distributions.Source: B PROOF OF THEOREM 3.3 · p. 12
  • The coupled differences obey an evolution equation involving the position and velocity differences and the Hessian term Ht.The proof cancels shared noise when forming the velocity difference.Source: B PROOF OF THEOREM 3.3 · pp. 12–13
  • With γ = 2 and the stated spectral bound, a Lyapunov-functional argument yields a contraction estimate via Gronwall’s inequality.The argument bounds the minimum eigenvalue of Rt below by 1/(2κ).Source: B PROOF OF THEOREM 3.3 · p. 13

C PROOF OF THEOREM 3.5

The proof establishes discrete-time contraction by bounding the update matrix through trace and determinant conditions. It uses step sizes h ∈ (0, 1/2] and controls the worst case over the normalized curvature range.

C PROOF OF THEOREM 3.5

  • For a real symmetric 2 × 2 matrix, positive semidefiniteness is equivalent to nonnegative trace and determinant.The proof uses this criterion to bound the update matrix by cI.Source: C PROOF OF THEOREM 3.5 · p. 14
  • For h ∈ (0, 1/2], the update matrix has trace T(h, q) and determinant D(h) as explicitly given.The proof shows that the trace decreases with q, so the worst case over q ∈ [1/κ, 1] occurs at q = 1/κ.Source: C PROOF OF THEOREM 3.5 · p. 14
  • The proof shows β1 and β2 are negative throughout h ∈ (0, 1/2] and q0 ∈ (0, 1], completing the matrix-bound argument.The β2 maximum is at q0 = 1, while β1 is handled by separate cases.Source: C PROOF OF THEOREM 3.5 · p. 15

D PROOF OF BIAS LEMMA 3.7

The bias proof compares an idealized chain with the discretized update using synchronous coupling. It bounds the resulting bias through moment and error estimates, then closes the estimate with a contraction argument.

D PROOF OF BIAS LEMMA 3.7

  • The idealized kernel preserves Πρ by evolving (Z, V) with X fixed and then drawing a new X from its Gaussian conditional.The frozen-x kinetic SDE has stationary distribution p(z|x) ⊗ N(0, uId).Source: D PROOF OF BIAS LEMMA 3.7 · p. 17

D.1 Bounding E

  • The D.1 estimates bound moments and increments of the interpolated process to control discretization errors.The argument assumes F is strongly convex and smooth, and uses a uniform moment bound.Source: §D.1 · pp. 18–20

D.2 Bounding E

  • The D.2 error bound uses exact invariance to keep the idealized chain distributed as Πρ at discrete times.Under this stationary law, the velocity has distribution N(0, uId).Source: §D.2 · p. 21

D.3 Closing the bias estimate

  • The D.3 argument combines error bounds with Grönwall’s inequality and the discrete contraction to close the bias estimate.The proof starts both processes from Πρ and accounts for the Lipschitz map recovering X from Z.Source: §D.3 · pp. 21–22

E PROOF OF THE CONVERGENCE THEOREM 3.8

The convergence proof combines the discrete contraction with the bias estimate. Choosing the step size and iteration count appropriately gives a Wasserstein-2 error at most δ.

E PROOF OF THE CONVERGENCE THEOREM 3.8

  • Choosing h so that √(C_bias h) ≤ δ/2 and enough iterations to make the contraction term at most δ/2 yields W2(Φ̄_k^h μ, Πρ) ≤ δ.The two terms are controlled separately before combining them into the claimed guarantee.Source: E PROOF OF THE CONVERGENCE THEOREM 3.8 · p. 23

F.1 Image restoration tasks

The experiments use shared image observations across three inverse problems on FFHQ and ImageNet, with task-specific forward operators and a Gaussian data-side conditional. Implementation details specify the diffusion priors, coupling schedule, iteration budget, and momentum resets.

F.1 Image restoration tasks

  • Experiments use 256 × 256 RGB images from FFHQ and ImageNet, and each image's noisy observation is generated once and shared by all methods.All forward operators act independently on each color channel.Source: §F.1 · p. 23
  • Gaussian deblurring uses a symmetric 61 × 61 Gaussian convolution, while motion deblurring uses a non-symmetric 61 × 61 camera-path kernel that blurs anisotropically.The motion kernel is generated from a random camera path; its noise setting matches Gaussian deblurring.Source: §F.1 · pp. 23–24
  • Super-resolution blurs with a 9 × 9 Gaussian kernel, then subsamples each image by a factor of four per direction to 64 × 64 observations.The noise is white, with variance chosen per image for 40 dB SNR.Source: §F.1 · p. 24
  • The data-side conditional is Gaussian and sampled exactly by solving one linear system, except Joint-RED-KLwSGS, which instead updates X using the data-fidelity gradient.The exact draw is up to linear-solver tolerance; the Joint variant requires applying A and A⊤ but does not form or solve with Q.Source: §F.1 · p. 24

F.2 Implementation details of RED-LwSGS, RED-KLwSGS and Joint-RED-KLwSGS

  • All three methods use pretrained diffusion models and a decreasing coupling parameter initialized at 0.5 with decay rate 0.85.The FFHQ and ImageNet models are from Choi et al. (2021) and Dhariwal and Nichol (2021), respectively; their noise schedules follow PnP-SGS.Source: §F.2 · p. 25
  • Each method runs 800 iterations with 20 burn-in iterations, and the kinetic methods reset momentum to zero when running-average PSNR stops improving.The posterior-mean estimator is approximated by the empirical mean after burn-in; the running PSNR uses the five most recent iterates.Source: §F.2 · p. 25
Loading 2610.10187v1…