Jacobian-Aware Posterior Sampling for Inverse Problems

AI summary

Overview

Research area: Computer Vision and generative machine learning, specifically zero-shot image restoration using diffusion models as priors for ill-posed inverse problems.

Technical level: Advanced. The paper leans heavily on stochastic differential equations, Bayes-rule posterior score decomposition, denoiser Jacobians, and Gaussian conjugate analysis, although its central insight can be stated simply.

One-sentence scope: This paper identifies the diffusion denoiser's Jacobian as the key mathematical difference between two families of diffusion-based inverse-problem solvers, then reconciles them into a single sampler (JAPS) that keeps the Jacobian's useful information while suppressing its training-induced imperfections.

What This Paper Is About

Diffusion models can act as powerful generative priors for recovering a clean image from a corrupted measurement, but different solvers approximate the posterior likelihood in different ways. Under a common isotropic Gaussian assumption and (locally) linear measurement operators, the authors show that "direct" methods and "proximal" methods produce likelihood surrogates that differ only by the presence of the denoiser's Jacobian. The goal is to build one sampler that exploits the genuine prior knowledge encoded in that Jacobian while avoiding the artifacts caused by imperfect Jacobians of realistically trained networks.

Key Contributions

1. Jacobian-Aware Posterior Sampling (JAPS): The authors identify the role of the denoiser's Jacobian in posterior sampling and introduce a sampler that bridges direct and proximal solutions, improving reconstruction quality without additional computational overhead.

2. Nonlinear extension: They extend the framework from linear to nonlinear inverse problems by locally linearizing the measurement operator around the unconditional MMSE estimate, and they demonstrate improved performance on representative nonlinear tasks.

3. DDIM reformulation for conditional guidance: They provide a formulation to incorporate the log-likelihood gradient into the DDIM sampler through a conditional noise estimator, preserving DDIM's original schedule and scaling with time re-spacing and varying levels of stochasticity.

4. Empirical diagnosis of Jacobian non-idealities: They measure violations of symmetric and positive semi-definite (PSD) structure in trained denoiser Jacobians on both an ImageNet-256 denoiser pretrained by Dhariwal and Nichol (2021) and a synthetic Gaussian Mixture Model setup with an analytically known ground-truth score.

Main Findings

• Direct and proximal methods differ by the Jacobian: The likelihood score admits the representation ∇_{x_t} log p(y|x_t) = √(ᾱ_t)/(1−ᾱ_t) · (E[x_0|x_t,y] − E[x_0|x_t]). Under a matched parameterization (r_t² = (1−α_t)/√(ᾱ_t)), the direct surrogate u_t and proximal surrogate v_t satisfy u_t = J_t v_t, so the only structural difference is the denoiser Jacobian J_t.

• Ideal Jacobians equal a scaled posterior covariance: For an ideal MMSE denoiser, J_t = √(ᾱ_t)/(1−ᾱ_t) · Cov(x_0|x_t), which is why the direct methods' use of the Jacobian amounts to preconditioning proximal gradients by uncertainty curvature. This motivates the expectation that a good Jacobian should be symmetric and PSD.

• Trained Jacobians violate ideality: On the ImageNet-256 denoiser and the toy GMM diffusion model, the minimum eigenvalue of the trained denoiser's Jacobian is negative (not PSD), and the Jacobian significantly deviates from symmetry. Figure 2 reports mean curves with standard deviations over 1024/50 samples per time-step for the GMM/ImageNet curves.

• Projection guidance is best: Among candidate guidance surrogates g(y,x_t) ∈ {u_t, v_t, P_{v_t}(u_t), P_{u_t}(v_t), (u_t+v_t)/2}, the orthogonal projection of v_t onto u_t (that is, P_{u_t}(v_t)) attained the best performance across all tested tasks after tuning λ_t ∈ {0.5,1,2,3,4,5,6,7,8}, evaluated on three representative tasks over a subset of 30 ImageNet-256 images with a fixed seed.

• Strong perceptual quality on linear tasks: Across CelebA-HQ and ImageNet-256 for super-resolution ×4, Gaussian deblurring, motion deblurring, random inpainting, and box inpainting, JAPS achieved the best or second-best LPIPS in all settings while remaining competitive on PSNR. Notable examples: box inpainting on CelebA-HQ with 27.58 PSNR / 0.874 SSIM / 0.046 LPIPS / 32.26 FID, and super-resolution ×4 on ImageNet-256 with 24.55 PSNR / 0.687 SSIM / 0.195 LPIPS / 40.24 FID.

• Nonlinear tasks also improve: On nonlinear blur (ImageNet-256) JAPS recorded 23.64 PSNR / 0.677 SSIM / 0.213 LPIPS versus 21.87 / 0.598 / 0.345 for ΠGDM, 19.98 / 0.364 / 0.586 for DAPS, and 19.12 / 0.517 / 0.363 for DPS. On high dynamic range (ImageNet-256) JAPS recorded 24.70 / 0.841 / 0.099 versus 24.55 / 0.822 / 0.110 for DAPS and 11.11 / 0.434 / 0.522 for ΠGDM. The nonlinear experiments used 100 CelebA-HQ and ImageNet-256 images with measurement noise σ_y = 0.05, an HDR factor of 2, and the neural blur model of Tran et al. (2021).

• GMM ablation clarifies when the Jacobian helps: With an analytically derived perfect denoiser, the direct surrogate u_t generally performed best, re-validating the information encoded in the Jacobian. With a trained model, the non-ideal Jacobian degraded performance, and JAPS effectively mitigated this effect.

• Robustness checks: JAPS attenuated degradation at high measurement-noise levels and was substantially less sensitive to the stochasticity setting of the diffusion process than alternative approaches.

Methodology in Plain English

The paper starts from the Bayes-rule decomposition of the posterior score into a prior score (which the pretrained diffusion model supplies) and a likelihood score (which must be approximated because only the clean-image likelihood p(y|x_0) is directly available). The authors write out two versions of the likelihood surrogate. One comes from solving a proximal problem in clean-image space using the unconditional Tweedie estimate as a fixed point; that version, called v_t, avoids differentiating through the network. The other comes from the direct approach, where the chain rule places the denoiser's Jacobian J_t in front, producing u_t. By choosing a specific r_t², they arrange these two to satisfy u_t = J_t v_t, isolating the Jacobian as the sole difference.

Because a real trained network only approximates the true MMSE denoiser, and because networks are usually trained for denoising score matching rather than to match higher-order Jacobian behavior, the authors empirically check whether trained Jacobians at least preserve the symmetry and PSD properties expected of a covariance-like operator. They find they do not. Their remedy is simple: they compute both u_t and v_t from the same computational graph (so v_t costs nothing extra) and use the orthogonal projection of the proximal surrogate onto the direct surrogate as the guidance term. This keeps the direction information encoded by the Jacobian while discarding the scale that inherits its non-idealities.

They then plug this guidance into a reformulated DDIM step where the posterior information enters through an explicit time-dependent coefficient γ_t, derived from a "posterior noise estimator" that replaces the unconditional noise prediction. The final update is x_{t−1} = DDIM(x_t) + γ_t λ_t √(1−ᾱ_t) g(y, x_t). For nonlinear operators, they replace the measurement matrix with a local Jacobian J_A evaluated at x̂_0, solve the resulting linear system with Conjugate Gradients, and use vector-Jacobian and Jacobian-vector products obtained by standard auto-differentiation.

Experiments use a DDIM sampler with η = 1 (DDPM-equivalent) and 100 diffusion steps, i.e. 100 NFEs; baselines mostly use 100 NFEs, with DPS and DAPS requiring 1,000 NFEs and DSG also using 1,000 steps for ImageNet-256. Noise is zero-mean i.i.d. Gaussian with σ_y = 0.05 in [0,1] units, multiplied by 2 to match the [-1,1] image normalization. Metrics are PSNR, SSIM, LPIPS (AlexNet variant), and FID, averaged or computed over 1K samples per dataset.

Why This Matters

This work provides a clean theoretical account of why two popular classes of diffusion-based inverse-problem solvers behave differently, and it turns that account into a concrete, zero-extra-cost correction. For the research community, it offers a diagnostic framework (Jacobian symmetry and PSD checks) and a DDIM-based conditional-guidance formulation that explicitly carries the time-step coefficient γ_t that earlier direct methods omitted and earlier proximal methods handled heuristically. It also supplies an extension path to nonlinear operators via local linearization and Conjugate Gradient solves.

Real-world applications that build on diffusion-prior inverse problems include:

• Medical imaging reconstruction such as accelerated MRI, where measurement noise and undersampling make strong priors essential.
• Computational photography and mobile imaging, including deblurring of motion-degraded photos, burst super-resolution, and HDR correction.
• Surveillance and restoration of archival imagery, where randomized motion-blur kernels and inpainting masks resemble the tasks in Table 1.
• Scientific and remote sensing imaging, where nonlinear forward models (for example sensor responses and high-dynamic-range effects) are common.

Industry relevance is significant because JAPS improves perceptual quality at the same sampling budget as its baselines (100 NFEs in most settings), while methods like DPS and DAPS require 1,000 NFEs and DSG needs 1,000 steps for ImageNet-256. Lower inference cost at equal or better LPIPS matters for any deployment where diffusion sampling time is the bottleneck.

Future Directions

• Better Jacobian quality: Since the analysis shows trained Jacobians are neither symmetric nor PSD, an open question is whether training objectives or architectural choices can make higher-order derivatives of the denoiser more faithful to the ideal covariance-like structure, which could let direct methods recover their theoretical advantage.

• Stronger nonlinear approximations: The nonlinear extension relies on local linearization with a curvature penalty that is negligible only for weakly nonlinear operators; extending the analysis to strongly nonlinear measurement models remains open.

• Broadening the surrogate family: The ablation compares only five candidates for g(y,x_t); whether richer parameterizations of Jacobian-aware guidance can outperform the single orthogonal projection is not yet established, and the paper's additional experiments on measurement-noise levels and stochasticity suggest further systematic sensitivity studies are natural next steps.

• Scaling and heavier data regimes: The evaluation covers CelebA-HQ and ImageNet-256, 100 images for nonlinear tasks, and 30 images for one ablation; whether the gains hold on larger datasets, higher resolutions, and more diverse degradation models is not yet reported.

Target Audience

This paper is best suited to graduate students and researchers working on diffusion models, score-based generative modeling, or inverse problems in imaging, especially those already familiar with DDIM and posterior sampling. Practitioners who deploy diffusion priors in medical imaging, computational photography, or restoration pipelines will find the JAPS guidance rule directly implementable, and reviewers or engineers evaluating the perception–distortion trade-off will benefit from the head-to-head PSNR/SSIM/LPIPS/FID tables. Beginners can follow the high-level bridging intuition but will likely need background in SDEs and Bayes-rule score decomposition to engage with the derivations.

Authors’ abstract

Diffusion models provide powerful generative priors for solving inverse problems by sampling from a posterior distribution conditioned on corrupted measurements. Existing methods primarily follow two paradigms: direct methods, which approximate the likelihood term, and proximal methods, which incorporate intermediate solutions satisfying measurement constraints into the sampling process. Under standard Gaussian approximations and locally-linear measurements, we demonstrate that these approaches differ fundamentally in their treatment of the diffusion denoiser's Jacobian within the likelihood term. While this Jacobian encodes critical prior knowledge of the data distribution, training-induced non-idealities can degrade performance in zero-shot settings. In this work, we bridge direct and proximal approaches by proposing a principled Jacobian-Aware Posterior Sampler (JAPS). JAPS leverages the Jacobian's prior knowledge while mitigating its detrimental effects through a corresponding proximal solution, requiring no additional computational cost. Additionally, we integrate our guidance into DDIM sampling, with a corrected conditional factor that has been missing in previous works. Our method enhances reconstruction quality across diverse linear and nonlinear noisy imaging tasks, outperforming existing diffusion-based baselines in perceptual quality while maintaining or improving distortion metrics.