Inverse Problem Sampling in Latent Space Using Sequential Monte Carlo

Idan Achituve, Hai Victor Habi, Amir Rosenfeld, Arnon Netzer, Idit Diamant, Ethan Fetaya

TL;DR: An approach for solving inverse problems in the latent space of diffusion models by augmenting the model with auxiliary observations. Sampling from the posterior is done using a novel SMC procedure.

Abstract: In image processing, solving inverse problems is the task of finding plausible reconstructions of an image that was corrupted by some (usually known) degradation operator. Commonly, this process is done using a generative image model that can guide the reconstruction towards solutions that appear natural. The success of diffusion models over the last few years has made them a leading candidate for this task. However, the sequential nature of diffusion models makes this conditional sampling process challenging. Furthermore, since diffusion models are often defined in the latent space of an autoencoder, the encoder-decoder transformations introduce additional difficulties. To address these challenges, we suggest a novel sampling method based on sequential Monte Carlo (SMC) in the latent space of diffusion models. We name our method LD-SMC. We define a generative model for the data using additional auxiliary observations and perform posterior inference with SMC sampling based on a backward diffusion process. Empirical evaluations on ImageNet and FFHQ show the benefits of LD-SMC over competing methods in various inverse problem tasks and especially in challenging inpainting tasks.

Lay Summary:

It is often the case that an image we witness has undergone some kind of loss or damage. Figuring out what the original image might have looked like before that happened is termed inverse problem. The goal of our paper is to propose a method for recovering or "guessing" the original, high-quality image from a damaged, incomplete, or low-quality version of it. Common examples are (1) Deblurring, where we are given a blurry photo (maybe from a shaky camera) and the goal is to recover the sharp, original image; and (2) Inpainting, where an image has missing or damaged parts (like holes, scratches, or objects removed) and the goal is to fill in the missing parts in a natural, realistic way.

Recently generative models (e.g., diffusion models) have been proposed to help solve this task by guiding towards more probable image reconstructions. Most inverse problem methods were developed for generative models that operate in the original pixel space. However, state-of-the-art image generative models are designed to work in some latent, low-dimensional, space. In our paper, we bridge that gap. We developed a method for solving inverse problems that leverage diffusion models operating in a latent space.

We propose to generate auxiliary images based on the distorted image which are then used to guide the diffusion process to generate a clean version of it. The injection of this information along the sampling process of diffusion models allows us to remain faithful to the information in the distorted image while making the resulting image look more natural. We ground our method using a well-established technique in the literature called sequential Monte Carlo which enjoys theoretical guarantees. Empirically our method generates plausible and natural image reconstructions in various inverse problem tasks.

Link To Code: https://github.com/ssi-research/LD-SMC

Primary Area: Deep Learning->Generative Models and Autoencoders

Keywords: Inverse problems, sequential Monte Carlo, Latent diffusion

Submission Number: 6854