PaperlayerView exact revision on arXiv

Source-linked AI summary

Improving Diffusion Models for Inverse Problems using Manifold Constraints

Hyungjin Chung, Byeongsu Sim, Dohoon Ryu, Jong Chul Ye

arXiv:2206.00941v3cs.LGcs.AIcs.CVstat.ML

TL;DR

Existing diffusion solvers for inverse problems can move sampling paths away from the data manifold, causing accumulated error while relying on costly projection-based corrections. The paper introduces MCG, an additional manifold-constrained correction based on Tweedie denoising that works with existing measurement-consistency steps. Across inpainting, colorization, and sparse-view CT, the method is reported to outperform prior methods, while remaining stochastic and slow and showing failures in some low-measurement or extreme-inpainting settings.

  • Problem

    Existing inverse-problem diffusion solvers can produce suboptimal results because their sampling paths leave the data manifold and errors accumulate.

  • Method

    MCG adds a Tweedie-denoising-inspired correction that uses measurement information to keep the gradient on the data manifold alongside reverse diffusion and data consistency.

  • Results

    The method improves diffusion-based inverse-problem reconstruction and outperforms prior methods in inpainting, colorization, and sparse-view CT.

  • Takeaways & Limitations

    MCG complements the score’s normal direction with a tangent-direction correction, helping prevent generation from falling off the manifold and reducing accumulated error.

  • Takeaways & Limitations

    The stochastic method can fail to produce high-quality reconstructions at low measurement dimensions, may artifact on extreme inpainting, and remains slow to sample.

Abstract

from arXiv · show

Recently, diffusion models have been used to solve various inverse problems in an unsupervised manner with appropriate modifications to the sampling process. However, the current solvers, which recursively apply a reverse diffusion step followed by a projection-based measurement consistency step, often produce suboptimal results. By studying the generative sampling path, here we show that current solvers throw the sample path off the data manifold, and hence the error accumulates. To address this, we propose an additional correction term inspired by the manifold constraint, which can be used synergistically with the previous solvers to make the iterations close to the manifold. The proposed manifold constraint is straightforward to implement within a few lines of code, yet boosts the performance by a surprisingly large margin. With extensive experiments, we show that our method is superior to the previous methods both theoretically and empirically, producing promising results in many applications such as image inpainting, colorization, and sparse-view computed tomography. Code available https://github.com/HJ-harry/MCG_diffusion

1 Introduction

Diffusion models can solve inverse problems without problem-specific training, but naive projection-based sampling may produce artifacts and require costly repeated noising-denoising iterations. The paper uses Tweedie-based denoising to motivate MCG, which improves reconstruction quality by keeping sampling near the data manifold.

  • Motivation: Diffusion models use pretrained unconditional score functions to solve inverse problems without problem-specific training.Examples include super-resolution, inpainting, compressed-sensing MRI, and sparse-view CT.
  • Limitations of existing solvers: Naive inverse-problem solvers can produce unsatisfactory reconstructions, including boundary artifacts in inpainting.High-quality reconstruction may require more than 10 back-and-forth noising and denoising iterations per iteration, increasing computational cost.
  • Motivation: Tweedie’s formula shows that Bayes-optimal denoising can be performed in one step when the gradient of the log density is known.The paper leverages this single-step denoising result to construct its correction.
  • Contribution: The proposed MCG correction keeps the sampling path closer to the data manifold and significantly improves diffusion-based reconstruction.The method combines denoising-based correction with the existing sampling procedure; Figure 1 reports high-fidelity inpainting relative to score-SDE.

2 Related Works

The paper situates MCG among unconditional diffusion, projection-based conditional sampling, spectral-domain solvers, plug-and-play methods, GAN approaches, and Tweedie-based denoising. These methods seek inverse-problem reconstructions while avoiding retraining conditional score models for every measurement condition.

  • Diffusion models: Continuous diffusion transforms data into an approximately spherical Gaussian distribution, while the reverse SDE generates samples using the score function.In practice, the score is replaced by a trained denoising score-matching model.
  • Inverse problems: Inverse problems seek unknown x from measurements y = Hx + ϵ, with conditional samples drawn from p(x|y).The measurement model uses H ∈ R^m×n and noise ϵ ∈ R^m.
  • Conditional diffusion: Conditional diffusion methods retain an unconditional score and impose measurement conditions through a projection or other contraction mapping.This avoids retraining a conditional score whenever the measurement conditions change.
  • Related approaches: Prior approaches include spectral-domain diffusion solvers, plug-and-play denoisers with data consistency, and GAN-based inverse-problem solvers.A spectral-domain method reported feasible solutions with as few as 20 diffusion steps.
  • Tweedie’s formula: Tweedie’s formula connects posterior-mean denoising with Gaussian noise and extends to broader noise distributions, motivating its use in the proposed method.The paper rewrites the formula for the discrete diffusion setting.

3 Conditional Diffusion using Manifold Constraints

MCG adds a manifold-based correction to unconditional reverse diffusion while retaining measurement consistency. The correction uses Tweedie-denoised estimates and measurement information to keep the gradient on the data manifold, complementing the reverse-diffusion dynamics.

  • Method: The method adds an additional constraint to the unconditional score-based reverse diffusion process.This supplements the existing measurement-consistency step rather than replacing it.
  • Manifold Constrained Gradient: MCG is defined using Tweedie’s Bayes-optimal denoising result so the measurement-term gradient remains on the data manifold.The paper presents this preferred condition both empirically and theoretically.
  • Combined update: The combined update applies the manifold correction and then data consistency to address potential measurement deviation.The resulting reverse-diffusion scheme integrates both constraints.
  • Mechanism: Although Tweedie denoising itself does not use the measurement, the MCG gradient incorporates y to constrain the measurement term on the manifold.This distinction enables the denoising estimate and measurement correction to play complementary roles.
  • Scope: MCG generalizes a related gradient method beyond temporal imputation and super-resolution to arbitrary linear inverse problems.The connection arises when the MCG term is combined with the data-consistency update.

4 Geometry of Diffusion Models and Manifold Constrained Gradient

The paper models diffusion as transitions across noisy-data manifolds and shows that measurement corrections can leave these manifolds. Its manifold constrained gradient instead projects measurement information onto the data manifold’s tangent space, while score-estimation errors are least consequential near high noise.

  • Geometric assumptions: The analysis assumes a locally linear, low-dimensional data manifold M, with the data distribution uniform on M.The manifold dimension l is much smaller than the ambient dimension n.
  • Noisy-data geometry: Noisy data concentrates on intermediate manifolds M_i, so diffusion steps can be viewed as transitions between the clean manifold and a pure-noise hypersphere.These interpolating manifolds describe transport from the data manifold to the hypersphere through the diffusion process.
  • Noisy-data geometry: Score functions are trained on points concentrated near noisy-data manifolds, making their application away from those manifolds a source of inaccurate inference.This geometric view explains why deviations caused by other sampling steps can matter.
  • Manifold constrained gradient: The score function captures normal directions, while measurement fidelity supplies discrimination along tangent directions needed to reconstruct the original signal.The tangent-space projection is represented locally by Q_i, which maps noisy samples to denoised estimates.
  • Manifold constrained gradient: Theorem 1 states that the manifold constrained gradient correction does not leave the data manifold and projects the data-fidelity gradient onto T_x̂0M.By contrast, naive measurement fidelity can push the inference path away from the manifolds and lead to inaccurate reconstruction.
  • Caveat: The method’s practical caveat is that score-estimation errors are concentrated near t ∼1, where Tweedie inference is uninformative, but MCG magnitudes are small; as t →0, estimation becomes exact.Thus, the paper argues that score-matching suboptimality has minimal impact on MCG implementation in the problematic regime.

5 Experiments

Experiments evaluate MCG across inpainting, colorization, and sparse-view CT reconstruction, with ablations examining measurement consistency, inference cost, NFE, and the correction strength α. Across these settings, MCG generally outperforms diffusion baselines and is competitive with supervised methods, while combining MCG with projection preserves data consistency.

  • Inpainting: MCG outperforms diffusion-model baselines by a large margin on inpainting and is competitive with or better than fully supervised methods.The comparison uses FFHQ and ImageNet with shared score functions for the diffusion methods, differing only in inference method.
  • Colorization: MCG outperforms other methods on both PSNR and LPIPS for LSUN-bedroom colorization and achieves strong performance on FFHQ.The evaluation compares score-SDE, DDRM, cINN, and pix2pix.
  • Sparse-view CT reconstruction: MCG outperforms score-CT by a large margin on sparse-view CT reconstruction, while supervised and analytical baselines fall behind.The evaluation uses PSNR and SSIM and reports high-fidelity, detailed reconstructions.
  • Ablation studies: Combining the MCG and projection steps provides perfect data consistency with further improved reconstruction, whereas MCG alone improves LPIPS but introduces measurement error.Gradient steps without Tweedie denoising substantially degrade performance, indicating that using the denoised estimate x̂_0 is important.
  • NFE and runtime: MCG remains in the best LPIPS regime across NFEs except below 100, while score-SDE worsens as NFE increases from several hundred to 1000.MCG takes longer than score-SDE at the same NFE because its MCG step requires backpropagation.
  • Ablation studies: α values from 0.1 to 1.0 produce satisfactory results; smaller values approach projection-only behavior, while larger values cause exploding gradients and saturation.The experiments normalize the squared norm and vary α′ to avoid cluttered notation.

6 Conclusion

The paper presents a general framework for improving diffusion-model inverse-problem solvers and reports applications across inpainting, colorization, and sparse-view CT. It argues that MCG keeps generation near the data manifold and complements the score function, while acknowledging stochasticity and reconstruction limitations.

  • The proposed framework enhances diffusion-model solvers for inverse problems across inpainting, colorization, and sparse-view CT reconstruction.
  • MCG prevents generation from falling off the data manifold, reducing errors that might accumulate across sampling steps.
  • The method is stochastic, can fail for very low measurement dimensions or extreme inpainting masks, and remains slow because it inherits diffusion-model sampling costs.

Checklist

The checklist reports that the paper addresses claims, limitations, ethics, assumptions, proofs, reproducibility, and asset attribution, while noting that repeated experiments and error bars were not provided.

  • The authors report that they described their contributions, scope, limitations, and potential negative societal impacts.
  • The paper states that it gives the full assumptions and complete proofs for its theoretical results, with proofs in supplementary material.
  • The authors report including code, data, instructions, and training details for reproducing the main experiments, with code release planned after publication.
  • The experiments did not report error bars because limited resources prevented running multiple experiment sets.
  • The paper cites existing dataset creators, does not state dataset licenses, and uses publicly available datasets without consent or personally identifiable-information discussions.

A Proofs

The proofs characterize noisy data and denoising estimates under a strong linear-manifold assumption, then establish that manifold-constrained corrections remain tangent to the data manifold. The section also describes diffusion discretizations, inverse-problem updates, and experimental algorithm choices.

  • Assumptions: The theory assumes that the data manifold is a low-dimensional linear structure coinciding with its tangent space.
  • Propositions: Noisy data concentrates near a higher-dimensional manifold determined by the scaled data manifold and the noise level.
  • Propositions: Under the linear-manifold assumption, the denoising map Qi acts as a locally orthogonal projection onto the data manifold.
  • Theorems: Theorem 1 states that a manifold-constrained gradient correction does not leave the data manifold and projects the data-fidelity term onto the tangent space.
  • Geometric interpretation: Mixing MCG with POCS pushes updates toward tangent directions and is expected to reduce deviation from the manifold caused by POCS.
  • Algorithms: The section specifies VP-SDE and VE-SDE sampling, ancestral inpainting updates, colorization as spectral-space imputation, and ART for sparse-view CT.
  • Algorithms: For sparse-view CT, repeated projection/FBP was found numerically unstable, especially with limited training data, motivating the ART-based update.

D Generative process of the proposed method

The proposed generative process adds MCG correction steps to the alternating-projection process. In the comparison, MCG corrects boundary errors that otherwise emerge and evolve during score-SDE sampling.

  • Score-SDE uses alternating projections, whereas the proposed method uses MCG as additional correcting steps during generation.
  • Without correction, an unnatural boundary forms between masked and unmasked regions and evolves toward t → 0; the additional gradient step corrects these boundary errors.

E Limitations

The proposed algorithm has a dataset-specific inpainting limitation on ImageNet: half-mask reconstructions often exhibit unrealistic symmetries, unlike FFHQ results.

  • ImageNet half-mask inpainting often produces unrealistic symmetries in reconstructions.This effect is not observed in the reported FFHQ experiments.

F.1 Implementation details

The implementation uses pretrained or separately trained score and comparison models, with task-specific sampling configurations and substantial computation for diffusion-based inference.

  • The inpainting experiments use pretrained FFHQ and ImageNet score functions available online.
  • The CT experiment trains an NCSN++ model with VE-SDE defaults for 200 epochs on one RTX 3090 GPU, taking about one week.
  • Inpainting reconstruction takes about 90 seconds per 256×256 image at 1000 NFE, while colorization and CT take about 600 seconds at 4000 NFE.
  • MCG without its manifold-constraint step is identical to score-SDE and uses the same score function, parameters, and sampler.
  • RePAINT uses the same score function and sampler as MCG, with N = 200 and U = 10 denoising-noising iterations per update.
  • DDRM uses the proposed method’s score function with σ_y = 0, η = 0.85, η_b = 1, and 20 sampling steps for noiseless inverse problems.
  • The comparison baselines include supervised, adversarial, transformer, invertible-network, and optimization-based methods with task-specific training or default configurations.

G Further Experimental Results

Further experiments compare MCG with established methods across inpainting and sparse-view CT, and show multiple stochastic reconstructions for inpainting and colorization.

  • Comparative results: The comparison study covers inpainting on FFHQ and ImageNet and sparse-view CT reconstruction on AAPM.
  • Comparative results: The inpainting figures compare MCG with DSI, LaMa, DDRM, score-SDE, and RePAINT, with IAGAN additionally included for FFHQ.
  • Comparative results: The sparse-view CT figure compares FBP, FISTA-TV, cGAN, SIN-4c-PRN, Score-CT, MCG, and ground truth.
  • Multimodal reconstructions: MCG generates three stochastic samples for 128×128 box-region inpainting and 92% pixel-missing imputation on FFHQ, LSUN-bedroom, and ImageNet.
  • Multimodal reconstructions: Colorization results on FFHQ and LSUN-bedroom show three different reconstructions sampled for each measurement.
Loading 2206.00941v3…