Solving Linear Inverse Problems Provably via Posterior Sampling with Latent Diffusion Models

Litu Rout; Negin Raoof; Giannis Daras; Constantine Caramanis; Alexandros G. Dimakis; Sanjay Shakkottai

Source-linked AI summary

Solving Linear Inverse Problems Provably via Posterior Sampling with Latent Diffusion Models

Litu Rout, Negin Raoof, Giannis Daras, Constantine Caramanis, Alexandros G. Dimakis, Sanjay Shakkottai

arXiv:2307.00619v1Jul 2, 2023cs.LGcs.AIstat.ML

TL;DR

The paper tackles the lack of inverse-problem methods for pre-trained latent diffusion models, which otherwise requires task-specific finetuning. It introduces Posterior Sampling with Latent Diffusion, combining goodness and gluing updates with a theoretical linear-model analysis, and reports superior performance across diverse restoration tasks. The method is analyzed under assumptions enabling provable recovery, while practical limitations remain around non-unique latent or image mappings.

Problem
Existing inverse-problem algorithms do not work with latent diffusion models, requiring task-specific finetuning to use foundation models such as Stable Diffusion.
Method
The paper extends DPS with goodness and gluing updates that guide latent diffusion toward measurement-consistent latents compatible with the decoder-encoder composition.
Results
PSLD experimentally outperforms state-of-the-art baselines across random and block inpainting, denoising, deblurring, destriping, and super-resolution, while achieving provable recovery in a linear two-step setting.
Takeaways & Limitations
The framework enables use of pre-trained latent foundation-model priors for a broad range of linear inverse problems without finetuning.
Takeaways & Limitations
The approach is limited by non-unique pixel-space measurements and many-to-one decoder mappings, which the goodness objective cannot resolve in general.

Abstract

from arXiv · show

We present the first framework to solve linear inverse problems leveraging pre-trained latent diffusion models. Previously proposed algorithms (such as DPS and DDRM) only apply to pixel-space diffusion models. We theoretically analyze our algorithm showing provable sample recovery in a linear model setting. The algorithmic insight obtained from our analysis extends to more general settings often considered in practice. Experimentally, we outperform previously proposed posterior sampling algorithms in a wide variety of problems including random inpainting, block inpainting, denoising, deblurring, destriping, and super-resolution.

1 Introduction

The paper addresses linear inverse problems with pre-trained latent diffusion models, extending posterior-sampling ideas beyond pixel-space diffusion without task-specific finetuning. It introduces a framework with theoretical recovery guarantees and reports improvements across diverse restoration tasks.

Motivation: Existing inverse-problem methods use supervised restoration models or unsupervised generative priors, with unsupervised approaches avoiding task-specific training and leveraging pre-trained models.Generative posterior sampling can also mitigate bias and regression-to-the-mean issues associated with likelihood maximization.
Motivation: Latent diffusion models encode data into a latent space, diffuse there, and decode back to pixels, powering foundation models such as Stable Diffusion.Their broad use contrasts with the lack of existing inverse-problem algorithms that work directly with latent diffusion models.
Contributions: The paper presents the first framework for solving general inverse problems with pre-trained latent diffusion models.The framework is designed to use foundation-model priors without finetuning for each task.
Contributions: The method extends DPS with an additional gradient update that guides diffusion toward measurement-consistent latents whose decoding-encoding map is not lossy.This design targets latent samples that explain the observations while remaining compatible with the latent representation.
Contributions: The authors theoretically analyze the algorithm and prove sample recovery in a linear model setting with two-step diffusion processes.The analysis is presented as a principled treatment of posterior sampling in the latent setting.
Results: The approach outperforms previous approaches for inpainting, block inpainting, denoising, deblurring, destriping, and super-resolution.The reported comparison spans a wide range of inverse problems.

2 Background and Method

The paper develops posterior sampling with latent diffusion models because existing diffusion-based inverse-problem methods do not apply directly to LDMs. PSLD augments the latent measurement update with a gluing objective that aligns latent samples with the decoder–encoder composition.

Latent diffusion pipeline: The latent diffusion procedure starts from a noisy latent, runs reverse diffusion, and decodes the resulting clean latent into an image.
Motivation: Existing inverse-problem algorithms based on diffusion models do not apply directly to latent diffusion models.
Vanilla latent extension: The vanilla DPS extension estimates the clean image by decoding the conditional expectation of the clean latent given the noisy latent.
Vanilla latent extension: Because the encoder is many-to-one, different latents can explain the measurements, causing the vanilla extension’s gradient and diffusion score to point toward different latents.
GML-DPS: GML-DPS adds a goodness term that guides diffusion toward latents explaining the measurements and acting as fixed points of decoder–encoder composition.
PSLD: PSLD adds a gluing update so denoising, measurement matching, and latent consistency point toward the same optimum.

3 Theoretical Results

The theory proves exact posterior sample recovery for PSLD in a two-step linear diffusion model under recoverability assumptions. It further explains why the latent formulation avoids pixel-space ambient-dimension issues and permits arbitrary finite step sizes.

3.1 Problem Setup: The analysis uses a two-step Ornstein–Uhlenbeck diffusion and proves perfect inverse-problem recovery in a linear model setting.
3.1 Problem Setup: The data lie in a low-dimensional subspace, and positive definiteness of (AS)^T(AS) ensures enough measurement energy for unique reconstruction.
3.2 Pixel-space posterior sampling: Under the linear assumptions, pixel-space DPS samples from the true posterior and exactly recovers the groundtruth when coordinate-wise step sizes satisfy ζ^j_i = 1/2σ_j.
3.3 Latent diffusion posterior sampling: GML-DPS with an LDM also samples from the true posterior and exactly recovers the groundtruth, but requires a specific coordinate-wise step size.
3.3 Latent diffusion posterior sampling: PSLD samples from the true posterior and exactly recovers the groundtruth for any positive finite step sizes under the same linear setting.
Theoretical insights: The analysis identifies that latent-space posterior sampling avoids the pixel-space DPS curse of ambient dimension and supports a single scalar step size across coordinates.

4 Experimental Evaluation

PSLD is evaluated on in-distribution and out-of-distribution inverse problems using Stable Diffusion and LDM-VQ-4, across inpainting, denoising, super-resolution, and deblurring settings. It outperforms or matches DPS and other baselines across diverse tasks and metrics, while generalizing without finetuning.

Evaluation setup: The evaluation covers random and box inpainting, denoising, Gaussian and motion deblurring, arbitrary masking, and super-resolution.Performance is measured with LPIPS, PSNR, SSIM, and FID.
Evaluation caveat: Commercial inpainting comparisons used platforms deployed in May 2023, whose results may change as providers improve their models.The platforms use Stable Diffusion and additional proprietary models to the authors’ knowledge.
Quantitative results: PSLD outperforms DPS in LPIPS, PSNR, and SSIM as the percentage of dropped pixels varies in random inpainting.The comparison is reported for FFHQ 256 samples and is summarized in Figure 5.
Quantitative results: PSLD with Stable Diffusion outperforms previous approaches on random inpainting and is better or comparable on harder tasks.The comparison includes random inpainting, box inpainting, super-resolution, and Gaussian deblurring.
Quantitative results: With LDM-VQ-4 trained on FFHQ 256, PSLD outperforms DPS in box inpainting and has comparable performance in random inpainting.Both methods use diffusion models trained on the same dataset in this comparison.
Generalization: PSLD generalizes to out-of-distribution samples without finetuning, including random inpainting and denoising with Stable Diffusion v-1.5.The experiments use FFHQ 256 validation data and compare against DPS trained on FFHQ 256.

5 Conclusion

The paper presents PSLD as a latent-diffusion framework for general linear inverse problems, combining theoretical analysis with a modified goodness-and-gluing objective. Experiments show strong performance across inverse-problem tasks, while the method remains subject to dataset bias and has not been extended to nonlinear problems.

Contribution: PSLD uses latent diffusion priors to solve general linear inverse problems that earlier approaches addressed only with pixel-space diffusion models.The method is designed to leverage latent-based foundation generative models.
Method: The algorithm is motivated by a principled linear two-step diffusion analysis and uses a modified objective called goodness and gluing.These design choices lead to the Posterior Sampling with Latent Diffusion algorithm.
Results: PSLD experimentally outperforms state-of-the-art baselines across random inpainting, block inpainting, denoising, destriping, and super-resolution.The conclusion summarizes performance across this broad task set.
Limitations: The evaluation relies on Stable Diffusion trained on LAION, so biases in that dataset and foundation model affect the method.The authors expect improved training datasets to mitigate these issues.
Limitations: The paper does not explore using latent-based foundation models for nonlinear inverse problems.The stated method scope is linear inverse problems.

A.1 Proof of Theorem 3.4

The theorem shows that, under the stated linear-model assumptions, Algorithm 1 can sample the true posterior and exactly recover the ground-truth sample using a carefully chosen matrix-valued step size.

Proof construction: The proof represents (AS)^T(AS) with a singular value decomposition using orthogonal matrices U and V and singular values σ.This decomposition organizes the measurement-dependent update.
Theorem statement: Under the theorem’s assumptions, Algorithm 1 samples from the true posterior p(x0|y) and exactly recovers the ground-truth sample.The conclusion is x̂0 = x0.
Step-size choice: The update uses a positive-definite matrix step size ζ whose eigenvalues are selected from the nonzero singular values of (AS)^T(AS).The proof sets ζ_i = 1/(2σ_i) for each measured singular direction.
Recovery conclusion: With this choice, the algebraic update cancels the reverse-process term and reduces to the ground-truth sample x0.The final identities use the assumptions on the subspace and measurement structure.

A.2 Proof of Proposition 3.5

The proposition constructs a linear encoder-decoder pair that minimizes the VAE objective under the subspace assumption by achieving exact reconstruction and zero distribution loss.

Proposition: Under Assumption 3.1, the encoder E(x0;ϕ) = S^T x0 and decoder D(z0;ω) = Sz0 minimize the VAE training objective.The construction uses the linear subspace representation of the data.
Reconstruction loss: The reconstruction term is analyzed by substituting the linear encoder and decoder and using that x0 lies in a linear subspace.The subspace assumption supplies the key identity in the reconstruction proof.
Distribution loss: The distribution loss is zero because the encoded distribution E♯p equals N(0, I_k), matching the VAE’s target Gaussian.The resulting KL divergence is KL(N(0, I_k), N(0, I_k)) = 0.

A.3 Proof of Theorem 3.6

The proof analyzes latent diffusion training in a VAE latent space and derives a normalized closed-form solution. Under the stated construction, the learned model recovers the true latent subspace and reconstructs samples through the decoder.

The regression solution yields θ∗ = √(1−β)I_k, which normalizes to the desired identity solution.
The resulting solution recovers the true subspace of the data distribution.
Latent diffusion training is performed in the latent space of a pre-trained VAE, where the objective reduces under Assumption 3.1.
Composing the latent generator with the VAE decoder produces the reconstructed sample from the recovered latent.

A.4 Proof of Theorem 3.7

The proof establishes exact posterior recovery for GML-DPS in the linear setting only with a uniquely matched step size. It shows that the measurement update already minimizes the goodness objective, while approximation and step-size errors remain important limitations.

GML-DPS combines latent denoising, pixel-space measurement matching, and goodness-objective regularization.
The exact analysis uses a closed-form solution, whereas solving the goodness minimization in practice can be difficult and may require gradient descent.
Under the theorem’s assumptions, a unique step size makes GML-DPS sample from the true posterior and exactly recover the groundtruth sample.
The measurement-consistency update is shown to produce a solution that is already a minimizer of the goodness objective.
Because the goodness objective cannot correct measurement-update error, GML-DPS requires the exact step size for posterior sampling.

A.5 Proof of Theorem 3.8

The PSLD proof replaces GML-DPS’s goodness objective with a gluing objective and derives exact recovery in the linear model. Unlike GML-DPS, PSLD remains exact for arbitrary positive step sizes under the stated assumptions.

PSLD replaces the goodness objective with a gluing objective and analyzes its corresponding latent-space minimization problem.
In general, the gluing minimization is difficult and is typically approximated with gradient descent, although the linear setting admits a closed-form solution.
The theorem permits any fixed variance β ∈ (0, 1) and any positive step sizes η_j.
Under the theorem’s assumptions, PSLD samples from the true posterior and exactly recovers the groundtruth sample.
The proof derives a closed-form latent solution that decodes to the groundtruth sample.
PSLD exactly recovers the groundtruth irrespective of step-size choice, whereas GML-DPS requires its uniquely matched step size.

B.1 Implementation Details

The implementation glues observed pixels into inpainting reconstructions while using direct decoding for other inverse problems. Experiments use Stable Diffusion and related image sources across the documented figures and examples.

For inpainting, PSLD generates missing regions while preserving consistency with the known image portions through its gluing step.
Unlike DPS, PSLD glues observations onto the corresponding generated image regions rather than generating the entire image without exact observation matching.
For motion deblur, Gaussian deblur, and super-resolution, the final gluing step is unnecessary because these tasks are not box inpainting.
The implementation uses the same measurement operator and transpose as the DPS code, while adjusting kernel sizes for the different image resolutions.
The documented examples include web images, an ImageNet 256 fisherman, and a celebrity face from the web.
The web application example uses 512 × 512 images generated by Stable Diffusion v-2.1.

B.2 Additional Experimental Evaluation

Additional experiments evaluate PSLD across inpainting, super-resolution, deblurring, and destriping on both in-distribution and out-of-distribution images. The reported results show close perceptual recovery on FFHQ super-resolution, stronger general-domain performance than DPS, and lower LPIPS for destriping.

Web demonstration: The web demo applies PSLD to user-defined inpainting masks at 512 × 512 resolution using a generic foundation model.The authors recommend finetuning the foundation model on a target image class for better task-specific performance.
Super-resolution: PSLD closely captures ground-truth perceptual quality in 4× super-resolution on in-distribution FFHQ 256 images, unlike DPS.The passage describes DPS samples as far from the ground truth, while PSLD samples preserve the identity-related perceptual content.
Deblurring: PSLD clearly outperforms DPS on motion and Gaussian deblurring for out-of-distribution ImageNet samples.The comparisons use Stable Diffusion v1.5 on ImageNet validation images.
General-domain reconstruction: The evaluation also covers random inpainting, super-resolution, and destriping on general-domain samples beyond the ImageNet deblurring tests.Figures 14, 15, and 16 report these additional reconstruction tasks on general-domain or web samples.
Destriping: LPIPS is 0.244 for PSLD versus 0.613 for DPS in horizontal destriping, and 0.255 versus 0.597 in vertical destriping.Both comparisons are reported on out-of-distribution web samples at 256 × 256 resolution.