Fast constrained sampling in pre-trained diffusion models

## Abstract

Abstract Diffusion models have dominated the field of large, generative image models, with the prime examples of Stable Diffusion and DALL-E 3 being widely adopted. These models have been trained to perform text-conditioned generation on vast numbers of image-caption pairs and as a byproduct, have acquired general knowledge about natural image statistics. However, when confronted with the task of constrained sampling, e.g. generating the right half of an image conditioned on the known left half, applying these models is a delicate and slow process, with previously proposed algorithms relying on expensive iterative operations that are usually orders of magnitude slower than text-based inference. This is counter-intuitive, as image-conditioned generation should rely less on the difficult-to-learn semantic knowledge that links captions and imagery, and should instead be achievable by lower-level correlations among image pixels. In practice, inverse models are trained or tuned separately for each inverse problem, e.g. by providing parts of images during training as an additional condition, to allow their application in realistic settings. However, we argue that this is not necessary and propose an algorithm for fast-constrained sampling in large pre-trained diffusion models (Stable Diffusion) that requires no expensive backpropagation operations through the model and produces results comparable even to the state-of-the-art tuned models. Our method is based on a novel optimization perspective to sampling under constraints and employs a numerical approximation to the expensive gradients, previously computed using backpropagation, incurring significant speed-ups.

## 1 Introduction

The recent state-of-the-art in image generation has been dominated by diffusion-based, text-to-image models (Rombach et al., ), which excel in translating input text prompts into images. However, by training on millions of text-image pairs, these models have acquired some general knowledge about the natural image space. Thus, besides from text-to-image generation, these models should also be useful in other image-based inference tasks, such as inpainting, super-resolution etc., given that their knowledge about image statistics can be induced.

Such efforts on utilizing text-to-image models on image-based inference tasks have focused on devising ways to adapt pre-trained, text-guided models to the target task. The simplest approach is to fine-tune the text-based model on the new image-based task (Xie et al., ; Wang et al., ), with the downside of fine-tuning being unnecessarily expensive in smaller-scale scenarios. Alternatively, a different set of algorithms has been developed that utilize only the pre-trained text-to-image diffusion model, and modify the sampling process to infer the missing information (Chung et al., ; Rout et al., ; Chung et al., ). These sampling-based approaches however come with a significant increase in the computations needed.

As a practical example, in Figure 1 we select an image from the web and pose the task of inferring the right half given the pixels in the left half. Based on the Stable Diffusion 1.5 model, we first utilize a fine-tuned variant, trained to be conditioned on images and inpainting masks. Although the inference is fast (4s) and relatively accurate with textures being correctly replicated on the right side and no visible seams, this model required 440k steps of additional inpainting training, on top of the initial SD 1.5 weights.

Secondly, we demonstrate the inpainting results of a recent sampling-based method, PSLD (Rout et al., ). This algorithm requires backpropagating through the denoiser network multiple times during inference, which increases the inference time to 5 minutes. Apart from speed, the method fails to generate a realistic right half, with the generated image achieving an overall correct shape but failing to use the right textures on the fur.

With the shortcomings of existing approaches in mind, we propose a new method for sampling from pre-trained diffusion models under constraints which aims to provide both fast inference speeds and high-quality results. We approach inference under a diffusion prior and a constraint from an optimization perspective and find that (a) there is an alternative gradient update to the diffusion latents during sampling that does not coincide with the gradient updates of previous approaches (b) there is a fast numerical approximation that speeds up this optimization significantly. Our work aims to make inference under constraints practical both by improving the synthesized images and reducing the inference times to a reasonable range.

Figure: Figure 1: Half-image inpainting. Our method allows for fast and accurate inference of the missing part of the image using the original pre-trained Stable Diffusion 1.5 model, which was not tuned for this purpose. In contrast, previous inverse problem solvers are slow and fail to understand the long-range correlations between the image pixels, or rely on models trained/tuned for each given problem. The inpainting model used is the Stable Diffusion 1.5-inpainting fine-tuned model.
Refer to caption: extracted/5951884/figures/gt_cat.png

## 2 Background

### 2.1 Denoising Diffusion

Denoising diffusion models were proposed for image generation in Ho et al. (), where their ability to generate diverse and high-quality samples was first showcased in the image space. The original formulation views the training and inference process as a hierarchical latent variable model ${\bm{x}}_{T}\rightarrow{\bm{x}}_{T-1}\rightarrow\dots\rightarrow{\bm{x}}_{1}
\rightarrow{\bm{x}}_{0}$, where the final latent is distributed normally ${\bm{x}}_{T}~{}\sim N(\mathbf{0},{\bm{I}})$ and $p({\bm{x}}_{0})$ represents the data distribution. After choosing a noise schedule $a_{t}$ that defines the forward transitions ${\bm{x}}_{t}\rightarrow{\bm{x}}_{t+1}$, usually Gaussian centered at $\sqrt{\frac{a_{t+1}}{a_{t}}}{\bm{x}}_{t}$ and with variance $(1-\frac{a_{t+1}}{a_{t}})$, the model is trained to reverse each individual step in the diffusion process.

Further iterations of denoising diffusion introduced class conditioning (Nichol & Dhariwal, ) and classifier-free guidance (Ho & Salimans, ), which culminated in the development of Latent Diffusion Models (Rombach et al., ). Latent diffusion proposed the large-scale training of text-to-image diffusion models in the latent space of an image autoencoder, ${\bm{x}}_{0}=\mathcal{D}(\mathcal{E}({\bm{x}}_{0}))$. We note that although we use latent diffusion models for all our experiments, we do not require additional steps during inference to ensure consistency between the pixel and latent spaces, as is performed in previous works (Rout et al., ; Chung et al., ).

### 2.2 Gradient descent steps through backpropagation

A number of recent approaches to solving inverse problems in denoising diffusion models have been investigated in hope that pre-trained large models, which required significant investment, can be used for inference directly, without additional tuning for each different inverse problem (Chung et al., ; Rout et al., ; Chung et al., ). The most typical problem formulation is the denoising generation of the sequence ${\bm{x}}_{T},{\bm{x}}_{T-1},...,{\bm{x}}_{1},{\bm{x}}_{0}$ under the constraint on the final signal in the form ${\bm{A}}{\bm{x}}_{0}={\bm{y}}$, or using a relaxed version that requires minimization of $||{\bm{A}}{\bm{x}}_{0}-{\bm{y}}||_{2}^{2}$, possibly as part of the likelihood function $p({\bm{y}}|{\bm{x}}_{0})={\cal N}({\bm{y}};{\bm{A}}{\bm{x}}_{0},\sigma^{2}{\bm
{I}})$. By Tweedie’s formula (Efron, ), denoising diffusion models approximating $\nabla_{{\bm{x}}_{t}}\log p_{t}({\bm{x}}_{t})$ can be used to express the expected value of ${\bm{x}}_{0}$, denoted as $\hat{{\bm{x}}}_{0}$, but the addition of the constraint as if an additional observed variable ${\bm{y}}$ was generated requires the addition of the term $\nabla_{{\bm{x}}_{t}}\log p({\bm{y}}|{\bm{x}}_{0})$.

At any rate, the regular denoising diffusion steps are altered so that at each $t$, the generated latent ${\bm{x}}_{t}$ is moved in the direction reducing the cost

$$ $C({\bm{x}}_{t})=({\bm{A}}\hat{{\bm{x}}}_{0}({\bm{x}}_{t})-{\bm{y}})^{T}({\bm{A }}\hat{{\bm{x}}}_{0}({\bm{x}}_{t})-{\bm{y}}).$ (1) $$

For example, in the case of inpainting applications, matrix ${\bm{A}}$ extracts a subsection of the pixels in image ${\bm{x}}_{0}$ to be compared with a given target ${\bm{y}}$. The estimated expected value of ${\bm{x}}_{0}$ at the end of the chain is provided by the diffusion model as a nonlinear function $\hat{{\bm{x}}}_{0}({\bm{x}}_{t})$. Typically, these moves are gradient descent moves, i.e. moves of ${\bm{x}}_{t}$ in the direction

$$ ${\bm{h}}=-\nabla_{{\bm{x}}_{t}}C({\bm{x}}_{t})=-{\bm{J}}^{T}{\bm{A}}^{T}({\bm{ A}}\hat{{\bm{x}}}_{0}-{\bm{y}})=-{\bm{J}}^{T}{\bm{e}},\quad{\bm{J}}=\nabla_{{ \bm{x}}_{t}}\hat{{\bm{x}}}_{0}({\bm{x}}_{t}),\ {\bm{e}}={\bm{A}}^{T}({\bm{A}} \hat{{\bm{x}}}_{0}-{\bm{y}})$ (2) $$

Computation of the Jacobian ${\bm{J}}$ would be expensive both in memory and computation, so the gradient is computed using backpropagation through $C({\bm{x}}_{t})$, and we show its mathematical form here for comparison with a different update we use in this paper. In Chung et al. () and Chung et al. () for example, (several) gradient steps of this form are applied to ${\bm{x}}_{t}$ at each step $t$ of generation before moving on to the next stage. As optimization of ${\bm{x}}_{t}$ might reduce the total noise in the image below what the denoising at $t-1$ was trained for, the gradient steps moving ${\bm{x}}_{t}$ towards optimizing $C({\bm{x}}_{t})$ are combined with adding additional noise, which could also be seen as a form of stochastic averaging.

## 3 Our approach: Newton steps based on the inverse function

Suppose that at the current point in denoising, (${\bm{x}}_{t},\hat{{\bm{x}}}_{0}$), the function $\hat{{\bm{x}}}_{0}({\bm{x}}_{t})$ is locally invertible, i.e. there is a unique, although unknown, inverse function ${\bm{x}}_{t}(\hat{{\bm{x}}}_{0}+{\bm{g}})$, in the neighborhood of $\hat{{\bm{x}}}_{0}$ (for small $||{\bm{g}}||$), in which case $\nabla_{\hat{{\bm{x}}_{0}}}{\bm{x}}_{t}=(\nabla_{{\bm{x}}_{t}}{\bm{x}}_{0})^{-
1}={\bm{J}}^{-1}$. This assumption may seem obviously wrong considering that denoising diffusion models are trained assuming that ${\bm{x}}_{t}$ is distributed as a Gaussian variable centered at ${\bm{x}}_{0}$, i.e. for a given ${\bm{x}}_{0}$ from the data distribution, many different ${\bm{x}}_{t}$ samples are feasible. However, $\hat{{\bm{x}}}_{0}$ is a deterministic function represented by a neural network, through which backpropagation can compute the update in (2). I.e., at least for the directions we are interested in, unless the cost is already optimized, the perturbations in $\hat{{\bm{x}}}_{0}$ will yield nonzero perturbations in ${\bm{x}}_{t}$.

Using the first-order Taylor series we can approximate

$$ ${\bm{x}}_{t}(\hat{{\bm{x}}}_{0}+{\bm{g}})\approx{\bm{x}}_{t}(\hat{{\bm{x}}}_{0 })+(\nabla_{\hat{{\bm{x}}}_{0}}{\bm{x}}_{t}){\bm{g}}={\bm{x}}_{t}+{\bm{J}}^{-1 }{\bm{g}}$ (3) $$

for some small perturbation vector ${\bm{g}}$.

Now suppose that, as in the Gauss-Newton method, ${\bm{x}}_{t}$ should move as close as possible to some target ${\bm{x}}_{t}^{\prime}$, i.e. we need to minimize

$$ $C(\hat{{\bm{x}}_{0}}+{\bm{g}})=({\bm{x}}_{t}(\hat{{\bm{x}}}_{0}+{\bm{g}})-{\bm {x}}_{t}^{\prime})^{T}({\bm{x}}_{t}(\hat{{\bm{x}}}_{0}+{\bm{g}})-{\bm{x}}_{t}^ {\prime})\approx({\bm{x}}_{t}+{\bm{J}}^{-1}{\bm{g}}-{\bm{x}}_{t}^{\prime})^{T} ({\bm{x}}_{t}+{\bm{J}}^{-1}{\bm{g}}-{\bm{x}}_{t}^{\prime})$ (4) $$

w.r.t. to the movement ${\bm{g}}$ of $\hat{{\bm{x}}}_{0}$. Setting $\nabla_{\bm{g}}C(\hat{{\bm{x}}}_{0}+{\bm{g}})=\mathbf{0}$ yields the system

$$ $\displaystyle({\bm{J}}^{-1})^{T}({\bm{x}}_{t}^{\prime}-{\bm{x}}_{t})=({\bm{J}} ^{-1})^{T}({\bm{J}}^{-1}){\bm{g}}$ (5) $\displaystyle{\bm{x}}_{t}^{\prime}-{\bm{x}}_{t}={\bm{h}}={\bm{J}}{\bm{g}}$ (6) $$

Thus, locally, within the validity of the first order approximation, ${\bm{g}}$ and ${\bm{h}}$ represent a pair of corresponding movements in $\hat{{\bm{x}}}_{0}$ and ${\bm{x}}_{t}$ respectively, but here they are linked differently than in (2). The semantics of the relationship here is that ${\bm{g}}$ is the best direction in $\hat{{\bm{x}}}_{0}$ to approximate the target move ${\bm{h}}$ in ${\bm{x}}_{t}$ in the least squares sense. If we assume that ${\bm{g}}=-\epsilon{\bm{e}}$, where ${\bm{e}}$ is the error vector on $\hat{{\bm{x}}}_{0}$ as before in (2) and $\epsilon$ is a small constant used to make the first order approximation valid, we get

$$ ${\bm{h}}=-\epsilon{\bm{J}}{\bm{e}},$ (7) $$

and so moving ${\bm{x}}_{t}$ in this direction would mean moving it to the target ${\bm{x}}_{t}^{\prime}$ that would *require* movement of $\hat{{\bm{x}}}_{0}$ in the direction opposite of the error vector ${\bm{e}}$. For example, in the case of image inpainting, the error vector (per pixel) ${\bm{e}}={\bm{A}}^{T}({\bm{A}}\hat{{\bm{x}}}_{0}-{\bm{y}})$ has zeros for pixels outside of the given (constraint) area. So the move ${\bm{h}}$ in ${\bm{x}}_{t}$ has to be such so that by searching for a compensatory move in ${\bm{g}}$ in $\hat{{\bm{x}}}_{0}$ the resulting move would not alter $\hat{{\bm{x}}}_{0}$ in the areas outside of the constraint.

The direction of optimization we propose in (5) has at least two advantages over the usual gradient descent update. First, the direction ${\bm{h}}$ can be computed numerically to save both on computation (2-fold) and memory (2.5-fold) compared to using backpropagation on cost (1). To derive that update, consider the function $f(s)=\hat{{\bm{x}}}_{0}({\bm{x}}_{t}-s{\bm{e}})$ where the variable $s$ is scalar. Its derivative at $s=0$ is

$$ $\frac{df}{ds}=-{\bm{J}}{\bm{e}},$ (8) $$

and so the direction ${\bm{h}}$ can be approximated numerically using the approximate derivative of $f(s)$ at $s=0$

$$ $\displaystyle f^{\prime}=\frac{1}{\delta}[\hat{{\bm{x}}}_{0}({\bm{x}}_{t}+ \delta{\bm{e}})-\hat{{\bm{x}}}_{0}({\bm{x}}_{t})],$ (9) $\displaystyle{\bm{h}}=-\epsilon{\bm{J}}{\bm{e}}\approx\frac{\epsilon}{\delta}[ \hat{{\bm{x}}}_{0}({\bm{x}}_{t}+\delta{\bm{e}})-\hat{{\bm{x}}}_{0}({\bm{x}}_{t })],$ (10) $$

requiring no backpropagation but instead, two forward passes through the network: one to compute $\hat{{\bm{x}}}_{0}({\bm{x}}_{t})$ and another to compute $\hat{{\bm{x}}}_{0}$ for the perturbed ${\bm{x}}_{t}$ (in the negative direction of the error vector).

The second advantage comes from the asymmetry of the Jacobian ${\bm{J}}$, as alluded to above. The updates (2) and (5) optimize locally for different things unless ${\bm{J}}={\bm{J}}^{T}$. As described in 3.2 on inpainting, (5) is in the direction that would represent the textures in the constrained region everywhere in the rest of the image as much as possible, whereas (5) only leads to copying some of the texture into some parts of the rest of the image.

There is no theoretical reason for the Jacobian ${\bm{J}}$ to be asymmetric in pre-trained denoising diffusion models in general. In fact, it very well could be symmetric as these models can be seen in terms of score matching so that $\hat{{\bm{x}}}_{0}$ approximates the true expectation $E[{\bm{x}}_{0}|{\bm{x}}_{t}]=\frac{1}{\sqrt{\alpha_{t}}}[{\bm{x}}_{t}+{\bm{v}}
_{t}\nabla_{{\bm{x}}_{t}}\log p_{t}({\bm{x}}_{t})]$, and the gradient of this is indeed symmetric:

$$ $\nabla_{{\bm{x}}_{t}}E[{\bm{x}}_{0}|{\bm{x}}_{t}]=\frac{1}{\sqrt{\alpha_{t}}}[ {\bm{I}}+{\bm{v}}_{t}\nabla^{2}\log p_{t}({\bm{x}}_{t})].$ (11) $$

One could argue that a well-trained model $\hat{{\bm{x}}}_{0}({\bm{x}}_{t})$ should have its Jacobian equal to this. But whether the denoising diffusion models are trained with score matching in mind or using the variational method of Ho et al. (), they do not directly optimize to match the real score $\nabla_{{\bm{x}}_{t}}\log p_{t}({\bm{x}}_{t})$ everywhere nor are they constrained to produce symmetric Jacobians, and as shown in 3.1 the Jacobian for Stable Diffusion 1.5 is not symmetric, which we assume is the reasons why our optimization of ${\bm{x}}_{t}$ may be more suitable for some applications, as demonstrated in Section 4.

### 3.1 Symmetry of the Jacobian

Figure: Figure 2: Sample pairs of pixels $(i,j),(k,l)$ of the denoiser Jacobian $\nabla_{{\bm{x}}_{t}}\hat{{\bm{x}}_{0}}=\nabla_{{\bm{x}}_{t}}E[{\bm{x}}_{0}|{ \bm{x}}_{t}]$ for different timesteps $t$. We use Stable Diffusion 1.5.
Refer to caption: x1.png

In Figure 2 we verify the claim that the trained diffusion model’s Jacobian is not symmetric. We again employ SD 1.5 and given a random input image from ImageNet, we scale it and add noise to get an intermediate latent of the diffusion process at different timesteps. We then pass the image through the denoiser and compute the gradients $\partial\hat{{\bm{x}}}_{0}^{k,l}/\partial{\bm{x}}_{t}^{i,j}$ and $\partial\hat{{\bm{x}}}_{0}^{i,j}/\partial{\bm{x}}_{t}^{k,l}$ for randomly chosen pixels $(i,j),(k,l)$ using backpropagation. When plotting the gradients we see that the values deviate from $y=x$, which would indicate a symmetric Jacobian.

### 3.2 Difference between gradient directions

Figure: Figure 3: Comparison between 5 gradient updates at $t=800$ using a learning rate $\lambda=1$.
Refer to caption: x2.png

To highlight the difference between the two gradient updates of (2) and (5) we perform the experiment showcased in Figure 3. We create a synthetic black-and-white grid image (Figure 3 (a)), which we will pass to the diffusion model. We compute the diffusion latent at $t=800$ by first blurring and scaling the original image and then adding the appropriate magnitude noise. Then we aim to find how the model intends to change the entire image when asked to add a red square in the middle. We pose this as an inpainting task and compute the two different gradients when given a measurement corresponding to the center patch of the image with an added red square (Figure 3 (b)). Starting from the initial diffusion latent, we perform 5 gradient updates using the two different gradients and visualize the final $\hat{{\bm{x}}}_{0}$ produced by the model.

The resulting images show that the proposed direction (Figure 3 (d)) produces a more coherent image that copies the newly introduced texture to the correct locations (intersections). The direction given by backpropagation (Figure 3 (c))

Figure: Algorithm 1 The proposed algorithm for linear inverse problem solving.

## 4 Experiments

### 4.1 Inverse Problems

We validate our approach on ImageNet (Deng et al., ) by performing inpainting and super-resolution. Following Chung et al. (), we randomly choose 1000 images from the 10k test images of the ctest10k split. For the diffusion model, we use SD v1.5, which is pre-trained on the LAION (Schuhmann et al., ) text-image dataset. All experiments were done on a single NVIDIA RTX A5000 24GB GPU.

We use Stable Diffusion 1.5, a latent diffusion model, which spatially compresses the images by a factor of 8. To demonstrate the versatility of our approach, we approach linear inverse problems from two different angles. For inpainting, instead of decoding the latent into an image and comparing it to the measurement in pixel space, we opt to dilate the pixel-level masks and directly perform inpainting in the latent space. That means we ’discard’ some of the information in the given image by only keeping $8\times 8$ patches that do not overlap with the masked pixels. Consequently, we apply the masking and un-masking operators, ${\bm{A}}$ and ${\bm{A}}^{T}$ respectively, in the latent space and inpaint latent values instead of pixels. Even with fewer pixels, our method performs better than existing methods in inferring the missing information. We utilize the $10-20\%$ free-form masking from Saharia et al. () as the method of masking pixels.

For superresolution, we cannot work in the latent space as the operation of downsampling pixels does not correspond to downsampling latents. Therefore, in order to compute the error direction ${\bm{e}}$ we backpropagate the pixel-level constraint cost $({\bm{A}}\mathcal{D}({\bm{x}}_{t})-{\bm{y}})^{T}({\bm{A}}D({\bm{x}}_{t})-{\bm{
y}})$ through the decoder network and utilize the resulting gradient w.r.t. the latent as the error direction. We note that by setting ${\bm{e}}=-\partial C(\mathcal{D}({\bm{x}}_{t})/\partial{\bm{x}}_{t}$ we only require backpropagation through the decoder model, which is significantly less expensive than backpropagating through the denoiser network, as is done in previous works. For both degradations, we also include additive white Gaussian noise with $\sigma_{\bm{y}}=0.05$.

In both experiments, we also find it useful to perform warm restarts of our algorithm. After running Algorithm 1 from $t=1000$ to $t=0$, we reset the inferred $x_{0}$ by adding the appropriate noise to
When superresolving, we add an additional perturbation to the gradient to avoid local minima. We find that our approach quickly converged to blurry images, which satisfy completely the superresolution constraint, but do not contain the desired details and texture. A simple perturbation of the gradient with random noise around the current $\hat{{\bm{x}}}_{0}$ seems to be adequate in solving this issue. For inpainting, this is not a problem, as the textures and overall fidelity of the inferred image are dictated by the given regions.

In Figure 4 we showcase the qualitative performance of our algorithm compared to similar recent works. In the case of inpainting our results show greater coherence with the rest of the image. For superresolution, our method seems to introduce more new textures in the generated image, which may not always align perfectly with the shown content (e.g. artifacts). In Table 1 we present quantitative results by measuring PSNR, LPIPS and FID scores. Inpainting shows a clear advantage for our approach, whereas superresolution struggles to improve significantly. It is important to consider that for our results, the inference time is around 2 minutes, including the warm restarts, whereas competing methods always require more than 4 times longer for a single image.

Figure: Figure 4: Comparison between our method and other works. We directly use the images and results from (Chung et al., ) since there is no code available to replicate their method. We run our algorithm on the same ground truth images for inpainting the given ’gray region’ masks and $\times 8$ super-resolution.
Refer to caption: x3.png

**Table 1: Quantitative evaluation (PSNR, LPIPS, FID) of free-form inpainting and superresolution ImageNet.**
|  | Inpaint (Freeform) |  | SR ($\times 8$) | Time |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Method | PSNR↑ | LPIPS↓ | FID↓ |  | PSNR↑ | LPIPS↓ | FID↓ | (approx.) |
| Ours | 22.20 | 0.275 | 30.45 |  | 22.29 | 0.428 | 73.05 | 2min |
| P2L (Chung et al., ) | 21.99 | 0.229 | 32.82 |  | 23.38 | 0.386 | 51.81 | 30min |
| LDPS | 21.54 | 0.332 | 46.72 |  | 23.21 | 0.475 | 61.09 | 8min |
| PSLD (Rout et al., ) | 20.92 | 0.251 | 40.57 |  | 23.17 | 0.471 | 60.81 | 12min |

### 4.2 Image layer inference

Our approach allows for fast inference of missing regions in a given image. With that in mind, we go beyond inpainting and propose a new inference problem for pre-trained diffusion models, layer inference. Given an input image ${\bm{x}}_{0}$, we aim to generate two new images ${\bm{x}}_{0}^{1}$, ${\bm{x}}_{0}^{2}$ and a pixel-wise mask ${\bm{m}}$ that will satisfy ${\bm{x}}_{0}={\bm{m}}{\bm{x}}_{0}^{1}+(1-{\bm{m}}){\bm{x}}_{0}^{2}$. We want the model to generate a possible ’decomposition’ of the input image into two layers, e.g. foreground and background along with the blending mask.

We propose a simple algorithm to perform this task. The assumption is that pixels originating from the same layer should have stronger correlations than pixels from different layers. To probe this correlation given a soft mask ${\bm{m}}$, we generate multiple possible images by sampling a binary mask from the (0,1) values of ${\bm{m}}$ and performing inpainting. After running multiple iterations of inpainting, we can then compute the likelihood of each pixel in the original image to belong in one of the layers. We choose to model a layer image as Gaussian with per-pixel means and variances dictated by the generated samples in that layer.

Instead of performing the full sampling process, we find that we can substitute each layer image with the predicted $\hat{{\bm{x}}}_{0}({\bm{x}}_{t})$. To induce more variety to the ${\bm{x}}_{0}$ estimates, we perturb ${\bm{x}}_{t}$ and help generate multiple inpainting variations without running the full inference process. The mask is randomly initialized with values in (0,1) and we choose to sample $K=5$ images per layer.

In Figure 5 we provide qualitative results by running the layer inference algorithm on images from the web. To guide each layer we also condition the denoiser $\hat{{\bm{x}}}_{0}({\bm{x}}_{t})$ on layer-specific text prompts that describe the image content we want to select. The results show that we can control the separation between the layers and get two images that describe different characteristics of the source image. The key is the speed of our method in inferring the missing pixels; this inference would have been computationally infeasible to perform with previous sampling-based approaches.

Figure: Figure 5: Given an input image we sample two new images which when blended with an inferred mask, reconstruct the original input. Our approach allows us to perform decomposition of random images from the web, guided by simple text prompts.
Refer to caption: x4.png

## 5 Conclusion

In this work, we presented a new algorithm for fast inference in pre-trained diffusion models under constraints. Our novel approach exploits a different gradient update that has different qualitative results and requires no expensive backpropagation operations through the model. Our method produces results comparable to the state-of-the-art with significantly less inference time. We also introduce a new layer inference task, which is enabled by our fast and high-quality constrained sampling algorithm. We believe that a fast and accurate method to sample from large pre-trained generative image models, under any condition, can have the potential to enable countless downstream applications that rely on utilizing a strong image prior.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: While the method shows promising results on inpainting and super-resolution tasks, how well would it perform for broader inverse problems? 

What are the failure cases of this approach, and how sensitive is it to hyperparameter choices?
soundness: 3
strengths: The paper's primary strength lies in its technical contribution that offers a practical solution to a significant problem in diffusion models. By developing a mathematically rigorous alternative to traditional backpropagation-based methods, the authors achieve impressive speed improvements, while maintaining or improving quality. This makes constrained sampling more practical for applications.

Another strength is the method's ability to work directly with pre-trained models without requiring task-specific fine-tuning. The authors validate their approach thoroughly, providing both quantitative metrics and qualitative results across multiple tasks, particularly showing strong performance in inpainting applications. They also ground their practical results in solid theoretical analysis, examining the Jacobian's asymmetry in diffusion models.

The introduction of the "layer inference" task further demonstrates the method's practical value, showing how the faster sampling enables new applications that weren't previously feasible due to computational constraints.
summary: The paper introduces a faster approach for constrained sampling in pre-trained diffusion models. The key innovation lies in the reformulation of the gradient update step: instead of using expensive backpropagation through the model to optimize constraints, the authors derive an alternative update direction by treating the denoiser's output as locally invertible. This enables approximation using just two forward passes, making it faster than traditional methods. The algorithm iteratively updates latent variables at each diffusion timestep using this efficient numerical approximation, achieving comparable or better results than existing methods while being faster. Using Stable Diffusion 1.5, they demonstrate the method's effectiveness on tasks like inpainting and super-resolution.
weaknesses: The paper's main weakness lies in its limited scope of validation. While the method shows strong performance on inpainting tasks, its effectiveness on super-resolution is less impressive. This suggests potential limitations in the method's general applicability that are not fully explored or explained.

From a theoretical perspective, while the authors provide interesting analysis of the Jacobian's asymmetry, they do not fully explain why this leads to better results in practice, leaving some theoretical questions unanswered. Additionally, the lack of direct comparisons for their novel layer inference task makes it difficult to assess the relative improvement over potential baseline approaches.

---

confidence: 4
contribution: 4
initial_rating: 3
presentation: 3
questions: * The examples show in for the teasers lack creativity and do not demonstrate robustness of the approach. If you really want to show the verstality of the method it would be really enlightening to try to frankenstein two different images together and see if the method is able to create an image on the natural image prior. That is to say if I take an image an outline the left third with a dog head, and the right third with say a cat head, can it connect the images in a convincing way?

* How many objects can the layer inference handle? It would be useful if you could use it to segment multiple objects or part of objects from an image at once.

* FID is evaluated, but it does not clarify how many images it is evaluated on. Note: FID is a biased metric, so it's value will vary depending on how many images you run it on (ie running it on more images will lower the value). If I am reading it correctly, This is only calculated on for $FID_1000$ and $FID_10000$ is usually the smallest value reported in the literature. Otherwise, one should use Kernal Inception Distance (Demystifying MMD GANs) as it's far more reliable on such a small number of samples. I am concerned too because such a small sampling of images from ImageNet will not even cover all the ImageNet classes.

All in all, this paper shows a lot of promise, but I worry the empirical results are entirely qualitative without a sufficient number of examples. The quantitative numbers in Table 1 provide a very mixed bag of results as well.
soundness: 3
strengths: * The paper fairly clearly explains the method. The approach also implements multiple baselines on the problems for infilling. The paper's methodology is rather original and I do not recall seeing it in the literature. The method could help to use diffusion models to solve inverse methods without retraining on task specific tasks or do expensive backprop.
summary: The paper describes a method for fast constrained sampling in pre-trained diffusion models, such as Stable Diffusion, without the need for expensive backpropagation operations. The key idea is to reformulate the inference process using an optimization perspective that speeds up constrained sampling. This approach focuses on efficiently estimating gradient updates through numerical approximations rather than full backpropagation, significantly reducing the amount of computation.
weaknesses: * Figure 2 is a bit hard to read
* Tasks are somewhat simplistic. Showing the results on more challenging inverse problems aside from infilling would be make the paper stronger
* It would be helpful to using a Deep Image Prior for these or classical infilling methods that do not require task specific training. The convolutional prior of the Stable Diffusion UNet could be doing a lot of heavy lifting. Although this method approach is orthogonal, it could serve as helpful baseline.
* Why isn't GPU memory analyzed here? if no backpropogation is used, I would expect significant memory savings that would allow a much larger batch size and therefore make the speed up vs the baselines even more dramatic for a fixed memory budget in table 1.
* Figure 4.2 is lacking any baselines.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 1
questions: * Does Equation (5) and (6) correspond? (5) contains $J^{-1}$ while (6) contains $J$.
* (Line 182) It will be better if additional arguments on the computational costs. (For instance, big-O notation, wall-clock time, FLOPs or memory constraints.)

(Minor typos and grammar check)
* (LaTeX) $\hat{x_0}$ and $\hat{x}_0$ are mixed up in the manuscript. I recommend unifying to $\hat{x}_0$.
* The paragraph in Line 134 is confusing: too many parentheses and division of phrases. For example, \
\
Suppose that at the current point in denoising, $(x_t, \hat{x}_0)$, the function $\hat{x}_0 (x_t)$ is locally invertible, i.e, there is a unique, although unknown, inverse function $x_t(\hat{x}_0 + g)$, in the neighborhood of $\hat{x}_0$ ...
\
can be refined to
\
Suppose that the denoiser function $\hat{x}_0$ is locally invertible over an $\varepsilon$-ball centered at $x_t$ ... \
\
Moreover, this denoising function should have time $t$ as its argument, of sub (or super)-script of the function, for a finer understanding.
* (Line 198) whereas (5) $\to$ whereas (2)

---

In conclusion, this paper is not ready for publication in the current stage of manuscript, mainly because of the difficulty of understanding the writing. Even though the paper seems to contain enough novelty, the manuscript seems to be written in a hurry, yielding a draft rather than the full paper.
soundness: 2
strengths: * By introducing the non-symmetry of Jacobian in the inverse problem setting, solving $h=Jg$ instead of $h=J^\top g$, after proposing some motivating examples on comparing the symmetric Jacobians.
* By excluding the gradient descent in the algorithm, this method achieved faster and memory-efficient solving of diffusion inverse problem, with better performance for inpainting and super-resolution tasks, and enabled the image layer inference.
summary: This paper proposes a method that mitigates the "heavy" backpropagation of constrained sampling (inverse problem) with diffusion models by first-order approximation to the gradients to use Gauss-Newton methods. By doing this, the paper found that using Jacobian (rather than the Jacobian transpose) works for inverse problem, while these two are different empirically in general inverse problem and constrained sampling settings. The algorithm consists of three steps:
1. Find an error vector $e=A^\top (A \top{x}_0 (x_t)  - y)$, which is the "direction" of error of $x_0$.
2. Based on $e$, find the error vector for noisy $x_t$, $h$.
3. Propagate $x_t$ to the direction of $h$.
weaknesses: * The paper is hard to read, by having lots of typos, grammar errors and unreleased arguments.
* More derivation on the Jacobian computation is required: too much part is omitted.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: * Figure 1 indicates that the Inpainting model is faster (4s) compared to your method (17s). Does this suggest that the default Inpainting method remains more efficient?

* Regarding soundness, is Formula 1's assumption of linear multiplication form realistic and expressive enough?

* In applying Newton method, does it require second order gradient? 

* In Algorithm 1, concerning the calculation of e, what does 'depends on the tasks' imply? Can you provide further elaboration on this?
soundness: 2
strengths: This method is ambitious as it aims to replace computationally expensive gradient operations for acceleration.
summary: The paper examines generative modeling for in-painting and introduces a sampling algorithm for rapid in-painting using pre-trained Stable Diffusion models. Their approach utilizes numerical gradient approximations to accelerate sampling, leading to faster in-painting results in empirical evaluations.
weaknesses: * The writing should be more polished and structured for better coherence.

* In experiments, consider conducting additional experiments (Table 1) and ensuring that baseline methods are comprehensive and representative. Since pre-trained Stable Diffusion models were mentioned, it would be beneficial to include evaluations for text-to-image tasks.