Dreamguider: Improved Training free Diffusion-based Conditional Generation

## Abstract

Abstract Diffusion models have emerged as a formidable tool for training-free conditional generation.
However, a key hurdle in inference-time guidance techniques is the need for compute-heavy backpropagation through the diffusion network for estimating the guidance direction. Moreover, these techniques often require handcrafted parameter tuning on a case-by-case basis.
Although some recent works have introduced minimal compute methods for linear inverse problems, a generic lightweight guidance solution to both linear and non-linear guidance problems is still missing. To this end, we propose Dreamguider, a method that enables inference-time guidance without compute-heavy backpropagation through the diffusion network. The key idea is to regulate the gradient flow through a time-varying factor. Moreover, we propose an empirical guidance scale that works for a wide variety of tasks, hence removing the need for handcrafted parameter tuning. We further introduce an effective lightweight augmentation strategy that significantly boosts the performance during inference-time guidance. We present experiments using Dreamguider on multiple tasks across multiple datasets and models to show the effectiveness of the proposed modules. To facilitate further research, we will make the code public after the review process.

## 1 Introduction

Generative modeling utilizing Denoising Diffusion Probabilistic Models (DDPMs) [, , , ] has massively improved
over the past few years. Multiple works have extended the use of diffusion models for text-to-image synthesis [, , ], 3D synthesis [, ], video generation [, , ], as well as for conditioning to solve inverse problems. Moreover, like conditional generative adversarial networks (GANs)[, ], DDPMs can be adapted to tasks based on a labels [, ] or visual prior-based conditioning []. However, like conditional GANs [, ], DDPMs also need to be trained with annotated pairs of labels and instructions to obtain satisfactory results. This poses a limitation in many cases where there is a lack of paired data to train large diffusion models. Due to this reason, there has been recent interest in models that can perform conditional generation without the need for explicit training [, , , ].

Progressing towards this direction is prior research in plug-and-play models. First introduced in [], the initial research on plug-and-play models [, ] enabled conditional sampling from GANs trained with unlabeled data. For this, a pre-trained classifier [, ] or a captioning model was used to estimate the deviation between the GAN-generated image and a given label, and based on this deviation, the GAN input noise was modulated until the generated sample satisfied the given text or class label. A similar approach that has been attempted for diffusion models to facilitate conditional sampling from unconditional diffusion models is classifier guidance [, ], where a noise-robust classifier is trained along with the diffusion model to guide the sampling towards a particular direction. However, classifier guidance brings in the computational costs of training a classifier, which is often undesirable. Some recent works have performed conditional generation without explicit training for the condition by utilizing the implicit guidance capabilities of the diffusion model [, , , , ]. Diffusion posterior sampling (DPS) [] proposed a technique of using an $L_{2}$ norm-based loss function to solve linear inverse problems using unconditional diffusion models. However, DPS often requires a large number of sampling steps for photorealistic results. Freedom [], yet another work, proposed the use of general loss functions during sampling to achieve training-free conditional sampling. Some variants of DPS have also been proposed in the literature []. All the aforementioned loss-guided posterior sampling techniques involve a guidance function at each timestep that requires backpropagation through the diffusion UNet.
Recently, [] proposed Manifold Preserving Guided Diffusion Models (MGD) that remove the need for backpropagating through the diffusion U-Net by performing a gradient descent with respect to the Minimum Mean Square Error (MMSE). Although MGD [] works remarkably well for linear tasks that require more guidance towards the start of the guidance process, it may fail in some tasks where guidance happens earlier, for example, face semantics-to-image and sketch-to-image, where stronger guidance is required from a much earlier stage. Moreover, like [, ], MGD also requires a case-by-case handcrafted parameter. Hence, a generic lightweight method that works well for both linear and non-linear guidance functions is still missing. Moreover, the need to find a handcrafted guidance parameter on a case-by-case basis still remains an open challenge.

In this paper, we introduce a new framework that can adaptively perform zero-shot generation using diffusion models without the need for any manual intervention by the user. We found a rather simple fix to the problem during the initial timesteps of diffusion, i.e., by utilizing the gradient with respect to the diffusion output noise in initial steps of inference. Combined with the guidance with respect to the MMSE estimate, we found that the combination generalizes well to tasks that require guidance at very early stages of guidance. Figure 2 presents the visualization of our approach over existing works present in the literature.
Utilizing the correction term along with the correction with respect to the MMSE estimate significantly boosts the performance in non-linear tasks. We present the corresponding results in Section 5.
Moreover, we treat the energy-based inference-time guidance [, ] as a stochastic gradient optimization of the MMSE estimate and the noise present in the image. This formulation enabled us to leverage recent research in parameter-free learning [, ] to develop a dynamic step size schedule. This step size adjusts itself adaptively based on the initial noise seed input of the diffusion model and guidance functions, hence removing the need for manual parameter tuning for inference-time guidance. Moreover, motivated by the effectiveness of differentiable augmentations while training GANs [], we found that utilizing multiple levels of matching differentiable augmentations to the MMSE estimate and guidance reference significantly improves the sampling quality, enabling very high-quality sampling with a low number of guidance steps. We present an overview of the different applications of our method in Figure 1 and an illustration of the difference of dreamguider with existing methods in Table 1.
Namely, we present results using Stable Diffusion [], unconditional diffusion models released by [] for $256\times 256$ guidance, and class-conditional diffusion models for high-resolution $512\times 512$ conditional synthesis. The different functionalities of Dreamguider are tabulated in Figure 2.

Figure: Figure 2: An illustration of the difference between the existing method and our method. Existing works backpropagate through the diffusion network to perform guidance at each timestep, whereas we find the gradients with respect to the MMSE estimate and the predicted noise based on the timesteps, thereby bypassing the expensive backpropagation operation.
Refer to caption: extracted/5640856/figs/comparev2.png

We present experiments on publicly released models on generic images, face images, and stable diffusion to show the relevance of our method. We focus on the
tasks of (1) Inpainting, (2) Super-resolution, (3) Colorization, (4) Gaussian Deblurring, (5) Semantic label-to-image generation, (6) Face sketch-to-image, (7) ID guidance and identity generation, and beat existing benchmarks that utilize diffusion models for these tasks, obtaining a significant boost in performance over existing methods leveraging loss-guided models. To summarize, our contributions are:

- •
We propose a zeroth-order loss-guided diffusion guidance that is applicable to both linear inverse problems and non-linear inverse problems.
- •
We remove the need for a manually tuned guidance scale for classifier guidance by proposing a scaling function that works for a wide variety of tasks.
- •
We propose a time-varying guidance scale for improving sampling quality.
- •
We propose a differentiable augmentation strategy to improve sampling quality.

**Table 1: Table illustrating the difference over existing methods performing inference-time guidance.**
| Method | Zeroth order | Linear Tasks | Non-Linear Tasks | Automatic scaling |
| --- | --- | --- | --- | --- |
| DPS [] | ✗ | ✓ | ✗ | ✗ |
| $\pi$GDM [] | ✗ | ✓ | ✗ | ✗ |
| Freedom [] | ✗ | ✗ | ✓ | ✗ |
| MGD [] | ✓ | ✓ | ✗ | ✗ |
| OURS | ✓ | ✓ | ✓ | ✓ |

## 2 Background

### 2.1 Training-free Conditional Sampling using Diffusion Models

Recently, there has been a rise in multiple works that propose utilizing unconditional diffusion models for conditional sampling [, , , ]. The earlier works proposed solving linear inverse problems using diffusion models with the help of priors dependent on the inverse transform of degradation. Recently, diffusion posterior sampling [] considered the degradation to be conditioned on a Gaussian distribution given any intermediate timestep and derived an $L_{2}$ norm-based regularization at each intermediate timestep to solve for linear inverse problems. Recent works such as Freedom [] explored an energy-based perspective and extended guidance to non-linear functions using general loss functions. Universal diffusion guidance [] extended this guidance process to stable diffusion and improved the performance by using forward-backward guidance. More recent works, such as manifold-guided diffusion [], further proposed to constrain the manifold space by projecting for the latent space alone.

### 2.2 Perturbed Markovian Kernel for Diffusion Transition

Let us assume that $r(x_{t},y)$ gives a measure of the distance between an intermediate $x_{t}$ and the condition $y$ and is a positive bounded function. Hence, in the reverse process, the diffusion trajectory should proceed through distributions with a higher probability of being closer to the desired cases. We model these trajectory intermediate distributions with

$$ $\hat{p}(x_{t})=p(x_{t})r(x_{t},y).$ (1) $$

Dickenson et al. [] first proposed the use of Markovian kernels to estimate the distribution of diffusion intermediates. Specifically, given the state $x_{t}$ at the equilibrium of the training process for a diffusion model, the intermediate of a diffusion model at a time instant, the distribution at a timestep $t-1$ can be estimated as

$$ $p(x_{t-1})=\int p(x_{t})p_{\theta}(x_{t-1}|x_{t})dx_{t}.$ (2) $$

As we know, the kernel $p(x_{t-1}|x_{t})$ is a Gaussian distribution whose mean can be estimated using the diffusion UNet and $x_{t}$.
To estimate a perturbed kernel $\hat{p}(x_{t-1}|x_{t})$, the perturbed distribution is

$$ $p(x_{t-1})r(x_{t-1},y)=\int r(x_{t},y)p(x_{t})\hat{p}_{\theta}(x_{t-1}|x_{t}) dx_{t}.$ (3) $$

By merging the constant terms in the transition into the normalization factor, the transition step is

$$ $\hat{p}_{\theta}(x_{t-1}|x_{t})=p_{\theta}(x_{t-1}|x_{t})r(x_{t-1},y).$ (4) $$

The proof is given in the supplementary material. Hence, we can see that rather than considering a Gaussian posterior, as in DPS [], any distance or loss function can be used. Similarly, one other valid transition step of the perturbed process is

$$ $\hat{p}_{\theta}(x_{t-1}|x_{t})=p_{\theta}(x_{t-1}|x_{t})\frac{r(x_{t-1},y)}{r (x_{t},y)},$ (5) $$

which adopts the notion of reciprocal distance from the previous timestep.

### 2.3 Inference-time Guidance of Diffusion Models

For conditional generation tasks using an unconditional diffusion model, ideally, the model would predict intermediates closer to the condition. The formulation can be seen in terms of transition probabilities. Consider a pretrained unconditional diffusion model on a specific domain. The problem at hand needs to guide the diffusion model during inference time conditioned with a condition $y$. Dhariwal et al. [] proposed a general strategy to perform this by conditioning on the condition $y$ and finding the resultant marginal distribution

$$ $p(x_{t}|x_{t+1},y)=p(x_{t}|x_{t+1})p(y|x_{t}).$ (6) $$

By assuming the distribution $p(y|x_{t})$ has much lower curvature compared to $p(x_{t}|x_{t+1})$, considering the marginal distribution close to $x_{t}$,

$$ $\displaystyle\text{log }p(y|x_{t})$ $\displaystyle=(x_{t}-\mu)\nabla_{x_{t}}\text{log }p(y|x_{t}),$ (7) $\displaystyle g$ $\displaystyle=\nabla_{x_{t}}\text{log }p(y|x_{t}).$ $$

Plugging back to $\log(p(x_{t}|x_{t+1},y))$,

$$ $\displaystyle\log(p(x_{t}|x_{t+1},y))$ $\displaystyle=(x-\mu-\Sigma g)^{T}\Sigma^{-1}(x-\mu-\Sigma g)+C,$ (8) $\displaystyle p(x_{t}|x_{t+1},y)$ $\displaystyle\sim N(\mu+\Sigma g,\Sigma).$ $$

Hence, the reverse sampling equation becomes,

$$ $x_{t-1}=\frac{1}{\sqrt{\alpha_{t}}}\left(x_{t}-\frac{1-\alpha_{t}}{\sqrt{1- \bar{\alpha_{t}}}}\epsilon_{\theta}(x_{t})\right)+\sigma_{t}\epsilon+\Sigma \frac{dr(x_{t-1},y)}{dx_{t-1}},\epsilon\sim\mathcal{N}(0,I).$ (9) $$

### 2.4 Shortcomings of the Existing Methods

Although the energy-based guidance theory supports guidance as a function of the current latent estimate, almost all loss-based guidance techniques derive the distance function as a function of $x_{t}$ rather than $x_{t-1}$ and derive the gradient based on the previous sample. Although this approach works for many tasks, it requires backpropagating through the neural network and modeling the score function for the guidance correction term. This limits the use of classifier guidance since existing diffusion architectures that produce photorealistic results are often very bulky. One can see why the existing framework utilizes the derivative with respect to the previous sample works by taking a better look at Equation 5. As we can see, a reciprocal distance over the previous timestep diffusion latent $x_{t}$ is a perfectly valid distance guidance function. In the next section, we elaborate on Dreamguider.

## 3 Proposed Method

Suppose $x_{t-1}$ denotes the current step and $x_{t}$ denotes the previous step in the inference process of the diffusion module. As mentioned in the previous section, existing works utilize the derivative with respect to the previous step for guidance; one reason for this is to use an off-the-shelf auxiliary distance function on the MMSE estimate at each step $\hat{x_{t}}$, which enables the use of general functions defined on image space for guidance. Here, the MMSE estimate is defined as

$$ $\displaystyle\hat{x_{t}}$ $\displaystyle=\frac{x_{t}-\sqrt{1-\bar{\alpha}_{t}}\epsilon_{\theta}(x_{t})}{ \sqrt{\bar{\alpha_{t}}}},$ (10) $$

where $\bar{\alpha}$ denotes the variance schedule of the diffusion process and
$\epsilon_{\theta}(x_{t})$ is the noise estimated by the network. One other observation to note is that finding the derivative with respect to the current step requires finding $\hat{x}_{t-1}$, which again requires an additional propagation through the diffusion network. Hence, the dilemma of backpropagating through the UNet for guidance still remains unresolved.

### 3.1 Time Variant Classifier Guidance

We found a simple yet effective solution for this dilemma; if we take a look at the ODE estimate at each step proposed by Song et al. []. Hence, rather than perturbing the Gaussian kernel at each timestep, we perturb the components $\hat{x_{t}}$ and $\epsilon_{\theta}(x_{t})$ by a small amount. Specifically, we perform the following operations:

$$ $\displaystyle\hat{x_{t}}$ $\displaystyle=\hat{x_{t}}-c\Sigma\frac{dr(\hat{x}_{t},y)}{d\hat{x}_{t}},t>t_{0}$ $\displaystyle\epsilon_{\theta}(x_{t})$ $\displaystyle=\epsilon_{\theta}(x_{t})-d\Sigma\frac{dr(\hat{x}_{t},y)}{d \epsilon_{\theta}(x_{t})},t<t_{0}$ $\displaystyle x_{t-1}$ $\displaystyle=\frac{1}{\sqrt{\alpha_{t}}}\left(x_{t}-\frac{1-\alpha_{t}}{\sqrt {1-\bar{\alpha_{t}}}}\epsilon_{\theta}(x_{t})\right)+\sigma_{t}\epsilon-c_{t} \Sigma\frac{dr(\hat{x}_{t},y)}{d\hat{x_{t}}}-d_{t}\Sigma\frac{dr(\hat{x}_{t},y )}{d\epsilon_{\theta}(x_{t})}$ (11) $$

where $r(\hat{x}_{t},y)$ is a non negative distance function that measures the distance between the MMSE estimate and condition, $\Sigma$ is the variance of the latent estimate at each timestep as in Equation 8. Please note that we perform a double descent here. The intuition behind the double descent is that performing descent on one of the components, say $\hat{x_{t}}$, guides effectively at the end of the diffusion process where $\alpha_{t-1}$ is one and vice versa.
Hence, during the guidance with the gradient w.r.t. $\hat{x_{t}}$, the maximum component of shift that happens to the sample is when we consider the flow of this correction through $\hat{x_{t}}$. Hence, we define the value as the maximum component of $x_{t-1}$ present in $\hat{x_{t}}$.

$$ $c_{t}=c\sqrt{\alpha_{t-1}}.$ (12) $$

Similarly, we define $d_{t}$ as the maximal component of $\epsilon_{\theta}(x_{t})$ in $x_{t-1}$. Hence,

$$ $d_{t}=-d.\frac{1-\alpha_{t}}{\sqrt{\alpha_{t}}\sqrt{1-\bar{\alpha_{t}}}}.$ (13) $$

Hence, this term gives efficient guidance at all timesteps, bypassing the guidance at the later timesteps alone as in MGD []. In the following section, we proceed to propose an effective empirical estimate for $c$ and $d$ that works for a wide range of tasks.

### 3.2 A Gradient-Dependent Scaling Factor Estimate

Recently, Distance over Gradients (DOG) [] was proposed as an effective parameter-free dynamic step size schedule for SGD problems. Given any Stochastic Gradient Descent (SGD) optimization problem, the Distance over Gradient works as an effective learning rate. Recent works [] have found the diffusion process as a stochastic optimization problem and have derived an SGD-based interpretation of the diffusion sampling process. Hence, inspired by both of these works, we attempted an empirical guidance estimate of the form:

$$ $\displaystyle\gamma_{t}=\begin{cases}\frac{1e^{-5}}{\sqrt{g_{T}^{2}}},&\text{ if }t=T\\ \frac{\max_{i>t}|f_{i}-f_{T}|}{\sqrt{\Sigma_{i=i}^{T}g_{t}^{2}}},&\text{ otherwise}\end{cases}$ (14) $$

where $g_{t}$ is the gradient of the loss function as defined in the equation, $f_{t}$ can be any of ${\hat{x_{t}},x_{t},\epsilon_{\theta}(t)}$ at timestep $t$ and $f_{0}$ is the initial estimate of $f_{t}$. We noticed that this empirical estimate works well for the first-order sampling involving DPS [] as well. We illustrate more results on the effect of this plug-in value for different cases in the appendix. Hence, utilizing Equation 14, we estimate $c$ and $d$ accordingly by substituting $f_{i}$ as $\hat{x}_{t}$ and $\epsilon_{\theta}(x_{t})$

### 3.3 Differential Augmentation Classifier Guidance

A common practice while performing classifier guidance to augment diffusion models with specific regularization for guidance is to use the noisy estimate at timestep $t$ and utilize it to compute the loss function to regularize the current prediction. However, in many cases,
such guidance can give results with artifacts and color shifts, as portrayed in Figure 3 and Figure 5, due to excessive guidance or insufficient guidance at intermediate timesteps that shift the results off manifold or cause color shifts. One effective solution for this is to imitate different versions of artifacts or color shifts on both the source image and the target image and utilize these augmented versions for a boost in performance. Hence, to perform guidance with a much more robust guidance loss, we introduce DiffuseAugment, an augmentation strategy for diffusion guidance during inference time. Specifically, given an intermediate sample $x_{t}$ and condition $y$, we augment $\hat{x}_{t}$ and $y$ with differentiable augmentations denoted by

$$ $\displaystyle\hat{x}^{aug}_{t},y^{aug}=T(\hat{x}^{aug}_{t},y^{aug}).$ (15) $$

We choose three different types of augmentations for $T$ comprising random cutouts, random translations, and color saturations. Please note that the augmentation of $y$ is dependent on the input signal. For label-based conditioning such as identity or text, we do not perform augmentation for $y$. For image space augmentations, we augment $y$ with the same random augmentation as that of $x$. While computing the effective loss, we find the average across all augmentations. We find that DiffuseAugment significantly boosts the sampling fidelity and quality of the reconstructed image. We present these results in Section 5.

## 4 Experiments

Since our method comprises both linear and non-linear inverse tasks, for linear inverse tasks, we follow DPS and evaluate our method utilizing two different benchmarks: (1) ImageNet [] and (2) CelebA []. For non-linear tasks, we follow Freedom and evaluate using the CelebA dataset. For linear tasks, we evaluate our method quantitatively for Super-resolution ($\times 4$), Colorization, Inpainting (Box), and Gaussian deblurring tasks.
For non-linear tasks, we evaluate for Face Sketch guidance, Face Parse maps guidance, and Face ID guidance. Since our method falls into the category of loss-guided diffusion models, we perform all quantitative evaluations using existing methods that follow this kind of sampling. Please note that although we acknowledge the parallel field of research in tackling inverse problems without backpropagation [, ], we excluded these methods for comparison as they tackle solely Linear inverse problems. In contrast, loss-guided models are generic and applicable to a wider range of problems.

Figure: Figure 3: Qualitative comparisons for Linear Tasks on ImageNet for 100 inference steps
Refer to caption: x2.jpeg

Figure: Figure 4: Qualitative comparisons for Linear Tasks on CelebA dataset for 100 inference steps
Refer to caption: extracted/5640856/results/face/deblur/degraded/00252.jpg

### 4.1 Implementation Details

We perform all experiments on NVIDIA A6000 GPUs. For ImageNet [] based tasks, we utilize the unconditional model released by Guided Diffusion. For Linear Tasks involving faces, we use the model trained on the FFHQ dataset [] and perform experiments on the CelebA dataset [] similar to DPS. For non-linear tasks, we follow Freedom and utilize the model trained unconditionally on the CelebA dataset. We evaluate using conditions derived from existing networks. For the high-resolution results presented in Figure 2, we utilized the class-conditional model of resolution $512\times 512$ released by Guided Diffusion. For all experiments, we used 100 sampling steps. For style transfer, we utilized Stable Diffusion [] v1.5. Please note that our sampling method is generic, and any sampler can be used. We fix the number of augmentations in DiffuseAugment for all the experiments to 8. For linear inverse problems we set the value of $t_{0}$ to 5 in Equation 11 to 30 and for linear inverse problems we set $t_{0}$ to 5

**Table 2: Quantitative evaluation of image restoration tasks on CelebA 256$\times$256-1k with $\sigma_{y}=0.05$, We utilize 100 inference steps for all methods**
|  | Inpaint (Box) | Colorization | SR ($\times$ 4) | Gaussian Deblur |  |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Method | PSNR $\uparrow$ | SSIM $\uparrow$ | LPIPS $\downarrow$ | FID $\downarrow$ | Cons $\uparrow$ | SSIM $\uparrow$ | LPIPS $\downarrow$ | FID $\downarrow$ | PSNR $\uparrow$ | SSIM $\uparrow$ | LPIPS $\downarrow$ | FID $\downarrow$ | PSNR $\uparrow$ | SSIM $\uparrow$ | LPIPS $\downarrow$ | FID $\downarrow$ |
| Score-SDE [] | 9.57 | 0.329 | 0.634 | 94.33 | 0.1627 | 0.3996 | 0.6609 | 118.86 | 20.75 | 0.5844 | 0.3851 | 53.22 | 23.39 | 0.632 | 0.361 | 66.81 |
| ILVR [] | - | - | - | - | - | - | - | - | 26.14 | 0.7403 | 0.2776 | 52.82 | - | - | - | - |
| DPS [] | 19.39 | 0.610 | 0.3766 | 58.89 | 0.0069 | 0.5404 | 0.5594 | 55.61 | 17.36 | 0.4969 | 0.4613 | 56.08 | 20.52 | 0.5824 | 0.3756 | 52.64 |
| MGD [] | 27.21 | 0.7460 | 0.2197 | 11.83 | 0.0018 | 0.6865 | 0.4549 | 38.22 | 27.51 | 0.7852 | 0.2464 | 60.21 | 27.23 | 0.7695 | 0.2327 | 51.59 |
| Ours | 28.84 | 0.8491 | 0.1432 | 5.96 | 0.0014 | 0.7775 | 0.3036 | 20.89 | 29.47 | 0.8429 | 0.1757 | 46.95 | 27.30 | 0.7672 | 0.2202 | 42.70 |

**Table 3: Quantitative evaluation of image restoration tasks on ImageNet 256$\times$256-1k with $\sigma_{y}=0.05$. Bold: best, We utilize 100 inference steps for all methods**
|  | Inpaint (Box) | Colorization | SR ($\times$ 4) | Gaussian Deblur |  |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Method | PSNR $\uparrow$ | SSIM $\uparrow$ | LPIPS $\downarrow$ | FID $\downarrow$ | Cons $\uparrow$ | SSIM $\uparrow$ | LPIPS $\downarrow$ | FID $\downarrow$ | PSNR $\uparrow$ | SSIM $\uparrow$ | LPIPS $\downarrow$ | FID $\downarrow$ | PSNR $\uparrow$ | SSIM $\uparrow$ | LPIPS $\downarrow$ | FID $\downarrow$ |
| Score-SDE [] | 9.66 | 0.2087 | 0.7375 | 133.54 | 0.1723 | 0.3105 | 0.8197 | 194.87 | 14.07 | 0.2468 | 0.6766 | 129.91 | 15.39 | 0.3158 | 0.620 | 134.67 |
| ILVR [] | - | - | - | - | - | - | - | - | 15.51 | 0.4033 | 0.5253 | 64.13 | - | - | - | - |
| DPS [] | 15.23 | 0.4261 | 0.6087 | 97.90 | 0.021 | 0.3774 | 0.8011 | 106.25 | 14.94 | 0.3258 | 0.6594 | 87.26 | 17.19 | 0.3980 | 0.5817 | 84.74 |
| MGD [] | 21.94 | 0.6920 | 0.2410 | 40.30 | 0.0057 | 0.5809 | 0.5427 | 73.75 | 23.12 | 0.6025 | 0.3936 | 70.83 | 23.13 | 0.6092 | 0.3695 | 61.49 |
| Ours | 23.49 | 0.7271 | 0.2001 | 30.72 | 0.0055 | 0.6804 | 0.3362 | 52.76 | 24.23 | 0.6818 | 0.2884 | 43.00 | 23.31 | 0.6157 | 0.3566 | 58.38 |

### 4.2 Qualitative Analysis

We present results on Gaussian Deblurring, super-resolution, and colorization. As we can see, DPS fails since 100 steps of diffusion are used, and the DPS scaling factor is not strong enough to perform proper guidance within 100 steps of diffusion. We set the amount of posterior noise for the measurement as $0.05$ in all experiments. MGD works remarkably well for the deblurring and inpainting tasks; however, it fails for colorization since early guidance is required for the flow of natural colors.

For ImageNet tasks, the performance of DPS falls more because the problem is more ill-posed. This can be seen in the eagle diagram, where the method is unable to reconstruct the eagle properly. In contrast, our method performs relatively better, producing much more realistic images. We highlight the performance improvement on colorization since we argue that these results are obtained because of the early flow of gradients. For non linear invere problems, as we can see, Freedom is able to produce realistic-looking results for even the difficult task of Parse Maps to Faces. We argue that this is because backpropagation through the UNet purifies the gradient flow; hence, the generated images look much more naturalistic.

### 4.3 Quantitative Analysis

We utilize Dreamguider and quantitatively evaluate CelebA and ImageNet datasets. The results for face restoration tasks are shown in Table 2 and Table 3. We evaluate these tasks utilizing four different metrics. SDEdit [] fails for the task of face inpainting and colorization as a single perturbation in the noisy domain throws the image off the manifold. DPS requires more inference steps for proper guidance. ILVR is originally designed for super-resolution. Hence, we quantitatively evaluate ILVR [] only for the task of super-resolution. Since DPS and MGD are applicable to all cases, we evaluate with these methods.
As we can see, our approach obtains better results than the baselines because of the flow of gradients, which allows for better reconstruction quality. For faces, the difference is much more highlighted in the task of colorization, where we get a significant boost of 18 FID score above the baseline. General linear inverse problems in ImageNet are much more complex than in faces; hence, there is an overall drop in metrics for the natural domain images in ImageNet. In our case, DiffAugment purifies the gradient; hence, we look for much better realistic-looking images. However, MGD does not produce realistic results for sketch-to-image and anime-to-face synthesis.

Figure: Figure 5: Qualitative comparisons for Non-linear Tasks on CelebA dataset for 100 inference steps
Refer to caption: extracted/5640856/results/nonlinear/anime/degraded/00392.jpg

**Table 4: Non-linear tasks. Best results out of zeroth-order optimization algorithms are highlighted.**
|  | Semantic Parsing | ID Guidance | Face Sketch |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Method | Distance$\downarrow$ | LPIPS$\downarrow$ | FID$\downarrow$ | Distance$\downarrow$ | LPIPS$\downarrow$ | FID$\downarrow$ | Distance$\downarrow$ | LPIPS$\downarrow$ | FID$\downarrow$ |
| First-order |  |  |  |  |  |  |  |  |  |
| Freedom [] | 1864.51 | 0.6030 | 66.89 | 0.3767 | 0.7058 | 81.40 | 39.05 | 0.6583 | 86.51 |
| Zeroth-order |  |  |  |  |  |  |  |  |  |
| MGD [] | 2698.27 | 0.6995 | 104.32 | 0.4291 | 0.7178 | 92.61 | 39.34 | 0.6576 | 70.42 |
| Ours | 2722.51 | 0.6199 | 79.42 | 0.3780 | 0.5932 | 82.70 | 39.03 | 0.5509 | 69.51 |

## 5 Ablation Studies

We perform extensive ablation studies with respect to the effect of DiffuseAugment as well as the effect of each guidance term. For the ablation experiments, rather than utilizing the whole testing dataset of 1000 images, we utilize 100 images and report the average LPIPS value.

### 5.1 Effect of DiffuseAugment

We notice that for linear tasks, even for low values of $T$ such as $T=20$, just by increasing the number of augmentations at the output to 8, the perceptual quality drastically improves, matching that of diffusion inference with $T=50$ with just 2 augmentations. Further, we notice that although the effect of augmentations is very significant for linear tasks, the
performance is not that significant or rather drops in some cases for low $T$
such as $T=20$; this is because with 20 diffusion steps, most intermediate MMSE estimates remain noisy, and hence the guidance network ArcFace [] cannot handle such input and hence returns irregular gradients affecting the quality. However, we can see that as $T$ increases and when there are enough gradient steps, DiffuseAugment plays a significant role in boosting the performance.

Figure: Figure 6: Ablation analysis on linear and non-linear tasks. FaceID guidance & ImageNet superresolution
Refer to caption: extracted/5640856/graphs/compo_linear_updated_ablation100ab.png

### 5.2 Effect of Different Components of Guidance

We present the ablation analysis of the effect of different terms of guidance in Figure 6. Please note that for this experiment, we set the number of augmentations from DiffuseAugment as 1. We also turn off time travel sampling for this experiment. For this experiment, we perform guidance with respect to $\epsilon_{\theta}(t)$ until $t_{0}$ and perform guidance with respect to $\hat{x_{t}}$ for $t>t_{0}$. Here $t=100$ represents pure gaussian noise and $t=0$ represents the image. As we can see, guidance with $\hat{x}_{t}$ alone faces a drop in performance initially for a low number of inference steps for non linear cases. We argue that this is because the guidance flow through the MMSE estimate is weak during the earlier steps of diffusion. Although time travel sampling helps to alleviate this issue, careful parameter tuning is required to obtain satisfactory results. We also notice that guiding utilizing the gradients of the output noise of the network closer to the start of the generation process produces better results.

## 6 Limitations and Future Works

Although we illustrated the working across various tasks for pixel space diffusion models, the direct approach cannot be used for latent diffusion models for the task of linear inverse problems, and one might have to apply multiple steps of time travel sampling to fix this issue, making a large computational overhead of the overall sampling time. We emphasize that this problem arises due to the reconstruction error in the VAE that encodes the image to the latent space. In the future, we will attempt to improve upon this with better optimization techniques. Moreover, although the proposed empirical estimate based on distance over gradients works for most tasks and shows the existence of an optimal parameter estimate, a thorough mathematical evaluation and the most optimal parameters are still missing. We leave this problem up to future works to estimate the optimal guidance parameter.

## 7 Conclusion

In this paper, we proposed an improvement to existing loss-guided techniques for zero-shot conditional generation with an unconditional diffusion model. Specifically, we proposed a sampling technique that removes the need to backpropagate through the diffusion U-Net in order to tackle sampling for general inverse problems. We also present an empirical function for automatic scaling parameters that removes the need for manual scaling parameter tuning, which was previously a huge hurdle in using classifier-free guidance. The newly proposed scaling parameter also removes the need for model-specific tuning of start and end guidance steps. We also introduced a differentiable data augmentation method that significantly improves the sampling fidelity. We illustrated the working of our method across 4 linear and 3 non-linear tasks across faces and real image domains. Our sampling technique produces photorealistic samples with much lower sampling time and higher fidelity than existing methods.

## 8 Algorithm of Dreamguider

We present the over algorithm of dreamguider without time travel sampling and the parameter estimation algorithm in Algorithm 1

Figure: Algorithm 1 Dreamguider

## 9 Proof for perturbed Markovian kernel equation

In the main paper, we emphasized that any positive distance function can be utilized for performing conditional generation using the perturbed Markovian kernel equation. hHre we proceed to derive the perturbed transition step. For the proof we closely follow the work from Dickenson et al []. Given a unconditional transition distribution $p_{\theta}(x_{t-1}|x_{t})$ and a distance function $r(.,y)$, where y is the condition provided
Please note that we assume $r(.,y)$ has relatively small variance compared to $p_{\theta}(x_{t-1}|x_{t})$, We know that at equilibrium state, the distribution at any timestep $t$ ina diffusion model can be written as

$$ $p(x_{t-1})=\int p(x_{t})p_{\theta}(x_{t-1}|x_{t})dx_{t}.$ (16) $$

To estimate a perturbed transition kernel $\hat{p}(x_{t-1}|x_{t})$,we start the perturbed distribution as

$$ $p(x_{t-1})r(x_{t-1},y)=\int r(x_{t},y)p(x_{t})\hat{p}_{\theta}(x_{t-1}|x_{t}) dx_{t}.$ (17) $$

By simple algebraic manipulations, taking $r(x_{t-1},y)$ to the other side, we get

$$ $p(x_{t-1})=\int\frac{r(x_{t},y)}{r(x_{t-1},y)}p(x_{t})\hat{p}_{\theta}(x_{t-1} |x_{t})dx_{t}.$ (18) $$

By comparing Equation 16 and Equation 18 we can see that one solution for the transitional distribution is

$$ $\hat{p}_{\theta}(x_{t-1}|x_{t})=p_{\theta}(x_{t-1}|x_{t})\frac{r(x_{t-1},y)}{r (x_{t},y)}.$ (19) $$

Also since normalization constants doesn’t affect the score function or transition step, Absorbing $x_{t}$ to the normalization factor of $p_{\theta}(x_{t-1}|x_{t})$, another valid perturbed transition kernel is

$$ $\hat{p}_{\theta}(x_{t-1}|x_{t})=p_{\theta}(x_{t-1}|x_{t})\frac{r(x_{t-1},y)}{Z}.$ (20) $$

Please note that the term $Z$ does not affect the transition step in the reverse process when the variance of $r(.,y)$ is small.

Figure: (a) Face Colorization $c=1.11$, $d=99.30$
Refer to caption: extracted/5640856/supplementary/supp_section_2/grayscale.jpg

**Table 5: Non-linear tasks ablation analysis on time taken, the value is represented in seconds**
| Method | Freedom | Dreamguider(1) | Dreamguider(2) | Dreamguider(3) |
| --- | --- | --- | --- | --- |
| Sketch to Face | 24.95 | 17.55 | 27.04 | 35.09 |
| FaceID to Face | 24.94 | 20.45 | 31.89 | 41.80 |
| FaceParse to Face | 56.25 | 48.35 | 75.43 | 107.02 |

## 10 Time comparison for Dreamguider with timetravel sampling and Freedom(First order) for non linear tasks

We present the time taken by Freedom, a first order algorithm for one step of time travel sampling [, ] in Table 5

## 11 Estimated parameter value for different tasks

In this section, we present the result and the parameter estimated by our approach for different tasks. For this experiment, we use 100 timesteps of diffusion and present the value at the 100th timestep. Here we define $d$ as the scaling factor of the scaling constant of the the loss derivative relative to $\epsilon_{\theta}(x_{t})$ and c as that of
$\hat{x}_{t}$ as in the main paper . The corresponding results are shown in Figure 7

## 12 Non cherry picked results for different tasks.

Figure: Figure 8: Figure illustrating Non cherry picked results for ImageNet colorization
Refer to caption: x30.jpeg

Figure: Figure 9: Figure illustrating Non cherry picked results for ImageNet superresolution
Refer to caption: x32.jpeg

Figure: Figure 10: Figure illustrating Non cherry picked results for Gaussian deblurring on ImageNet
Refer to caption: x34.jpeg

Figure: Figure 11: Figure illustrating Non cherry picked results for face colorization
Refer to caption: extracted/5640856/supplementary/supp_noncherry/graydegraded_grid.jpg

Figure: Figure 12: Figure illustrating Non cherry picked results for face superresolution
Refer to caption: extracted/5640856/supplementary/supp_noncherry/srdegraded_grid.jpg

Figure: Figure 13: Figure illustrating Non cherry picked results for Gaussian Deblurring
Refer to caption: extracted/5640856/supplementary/supp_noncherry/srdeblurdegraded_grid.jpg

Figure: Figure 14: Figure illustrating Non cherry picked results for face inpainting
Refer to caption: extracted/5640856/supplementary/supp_noncherry/inpaintdegraded_grid.jpg

Figure: Figure 15: Figure illustrating Non cherry picked results for sketch to face synthesis
Refer to caption: extracted/5640856/supplementary/supp_noncherry/animedegraded_grid.jpg

Figure: Figure 16: Figure illustrating Non cherry picked results for Face ID guidance
Refer to caption: extracted/5640856/supplementary/supp_noncherry/iddegraded_grid.jpg

Figure: Figure 17: Figure illustrating Non cherry picked results for Face Parse Guidance
Refer to caption: extracted/5640856/supplementary/supp_noncherry/parsedegraded_grid.jpg

## Peer Reviews

confidence: 5
contribution: 1
initial_rating: 3
presentation: 1
questions: Please refer to the weakness for my questions.
soundness: 2
strengths: 1. The paper tackles a relevant and challenging problem in the realm of conditional (controllable) generation for unconditional diffusion models.
2. It proposed two useful tricks to improve the generation quality: 1) an empirical guidance scale and 2) differentiable augmentation for classifier guidance.
summary: The authors propose a generic lightweight guidance solution, named Dreamguider, which enables inference-time guidance without the need for backpropagation through the entire diffusion network. Dreamguider can address both linear and non-linear guidance problems. The authors also introduce an empirical guidance scale strategy to removing the need for handcrafted parameter tuning. This paper showcases experiments across multiple tasks, datasets, and models to demonstrate the effectiveness of the proposed methods. The key contributions include: 1) Dreamguider, a zeroth-order loss-guided diffusion guidance applicable to both linear and non-linear inverse problems, 2) a time-varying guidance scale, and 3) a differentiable augmentation strategy.
weaknesses: 1. Overall, the contributions of this paper are limited. The proposed Dreamguider is built upon MGD [1], an earlier proposed zeroth-order loss-guided guidance method that also eliminates the need for backpropagation through the diffusion network. The authors claim that Dreamguider can address non-linear inverse problems, unlike MGD. However, they should provide a detailed description and evidence to support this claim. In the current manuscript, I cannot find the intuition or evidence needed to substantiate this.

2. The paper should clarify the intuition behind the methodological design discussed in Section 3.1, which currently lacks depth and clarity. Given that Dreamguider is closely related to MGD, a comprehensive comparison of each designed component should be included to enhance understanding.

3. MGD is capable of handling latent diffusion models, so it is ridiculous that an improved method like Dreamguider loses this important capability. The authors should provide a detailed explanation for this shortcoming, and consider addressing it rather than deferring the issue to future work.

4. The showcased results are not appealing, as the performance differences between Dreamguider and MGD on some tasks are marginal. If possible, incorporating a user study could significantly enhance the experimental validation.

Typos:

In Table 2, some second-rank values are incorrectly bolded.

[1] Yutong He, et al. "Manifold preserving guided diffusion" ICLR 2024.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 2
questions: It would be great if the authors can address the weakness mentioned above.
In addition, I am wondering if the authors have used techniques such as time traveling/repaint to stabilize the sampling process?
soundness: 2
strengths: Almost all prior works in training-free diffusion guidance are very sensitive to hyperparameters, especially the step size, and currently there is no principled way to determine these hyperparameters. In fact, I think this is one of the biggest problems with these methods that prevent them from being applied widely. This paper provides a concrete way to decide these hyperparameters, which is extremely valuable for this field. The authors also propose a very creative way to perform the guidance by modifying the noise/score estimation directly (as opposed to the Tweedie’s clean data estimation only), which I find to be very interesting as well. The authors also provide a very effective way to perform data augmentation in this diffusion guidance task, which is also very useful.

Overall I really like the methods proposed in this paper, and I really want to accept it. However, there are certain aspects in this paper that I think need to be improved before it gets accepted.
summary: In this paper, the authors propose a new method for training-free diffusion posterior guidance. In particular, instead of only modifying the Tweedie’s (MMSE) estimation in DDIM sampling, they propose to also guide the noise/score prediction to obtain a better guided sample. They also provide a way to perform data augmentation and a closed form solution to the step size of these guidance, which is a critical contribution to the field as so far there is no principle way to determine these hyperparameters. Their experiments show superior performance in comparison to baselines across many different tasks.
weaknesses: 1. The writing should really be improved: there are a lot of undefined notations and unjustified algorithm components.
2. The main method (Section 3.1 & 3.2) is not very well explained. I think guiding the $\epsilon$ is a very interesting idea and would love to see more theoretical justifications about why this is effective. Section 3.2 is not very informative at all, I would love to see the deduction process of Eq 14.
3. Eq 1 is a very strong assumption, maybe it is better to assume proportional to rather than equals.
4. Also distance (r) doesn’t necessarily implies relationships to $p(y \mid x_t)$, so it may not be connected to Eq 6 and hence you don’t necessarily get Eq 9. What exactly is the assumption made here? 
5. $\Sigma$ is not clearly defined when it’s first mentioned in Eq 8, since the authors use Tweedie’s/DDIM (MMSE) estimation for $\hat{x}_t$, is it just $\sqrt{1-\alpha_t}I$?
6. I think Eq 11 (Line 256-261) is quite confusing. Do you change $\hat{x_t}$ and $\epsilon_\theta(x_t)$ and then plug them into the $x_{t-1}$ formula? Or the $x_{t-1}$ formula is a summary of the changes you would do in the $\hat{x_t}$ formula and $\epsilon_\theta(x_t)$ formula? Seems like the latter one is consistent with the algorithm provided in the appendix, but it is very confusing to read here.
7. I think Eq 12 should be $c_t = \frac{c}{\sqrt{\alpha_t}}$ and in Eq 13 $d_t$ shouldn’t have the negative sign?
8. The authors claim that Eq 14 also works very well for DPS, do they have any experimental results to show that?
9. Writing suggestion: I don’t think the formulation in Section 2.3 is very different from DPS if you use the Tweedie’s (MMSE) estimation like in Section 3.1, so maybe just introducing the problem following DPS or MPGD is ok.
10. MPGD can also perform nonlinear tasks (error in Table 1)

Minor:
1. Line 256 -261, it should be partial derivatives $\partial$ instead of $d$.
2. Eq 9 should be $-\Sigma\frac{\partial r(x_{t-1},y) }{\partial x_{t-1}}$ instead of $+$ because usually the smaller the distance is the larger the conditional probability would be. This is probably connected to Weakness 7 the problem with Eq 13 and Weakness 4 the assumptions about the distance function and the conditional probability $p(y \mid x_t)$.
3. The method He et. al (2023) proposed is called MPGD not MGD.
4. Citation error in Line 382 & 389.

I am willing to accept this paper if the authors can improve the writing and provide more justifications as suggested above.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: - What is the main difference between the proposed method and MGD? Both algorithms compute the gradient $\nabla_{\hat{x}_t} r(\hat{x}_t, y)$ and add it to $x\_{t-1}$. The proposed algorithm also adds $\nabla\_{\epsilon\_{\theta}(x\_t)} r(\hat{x}_t, y)$ but $\hat{x}_t$ and $\epsilon\_{\theta}(x\_t)$ are related by Eq. 10. It seems like that the gradient $\nabla\_{\epsilon\_{\theta}(x\_t)} r(\hat{x}_t, y)$ is not providing any extra guidance/information for the inference process but is rather used as a proxy to correctly weigh the change of $x\_{t-1}$ to match the constraint at timestep $t-1$. Given the above, could someone reformulate your method as MGD with different weighting?
soundness: 3
strengths: - The proposed algorithm seems to overall improve upon the results attained by previous methods. When comparing with the similar work done in MGD, the results show that there is a clear advantage for the proposed algorithm, both quantitatively and qualitatively.
- The augmentation guidance scheme proposed, could be of wider applicability and improve other training-free conditional inference methods in the future. Since the augmentation is not dependent on the method used it has the potential to be adopted by other works in the future.
summary: The paper proposes a training-free algorithm for conditional inference in diffusion models. The authors improve upon previous works by coming up with a guidance schedule that balances better the conditioning signal during the inference process. They also propose a differentiable augmentation strategy that improves the guidance by averaging it over different augmentations of the input and the condition. They validate their method experimentally on a set of linear and non-linear inverse problems, across different diffusion models.
weaknesses: - The novelty is somewhat limited. The proposed algorithm seems to be an extension of the previous work of MGD (also a zeroth-order guidance scheme) and the contributions are not clear enough. Out of the four contributions mentioned in the introduction the ones that are clearly presented are regarding the tuning of the constraint guidance and the differentiable augmentations.
- In Table 2 and Table 3 the citation for MGD is incorrect (points to DPS).

---

confidence: 2
contribution: 3
initial_rating: 5
presentation: 3
questions: See weaknesses.
soundness: 3
strengths: - The topic looks promising. Solving the inverse problem with less compute is important.
- The performance gap between existing method looks good.
summary: This paper propose Dreamguider which solve both linear and non-linear inverse problem. While all the previous work requires to contain back propagation path including the diffusion model, this paper find the way to compute MMSE estimates without diffusion network.
weaknesses: - The motivation of proposed method is "fast sampling", but speed comparison to existing method is missing for table 2,3,4. I can only see Figure 6 about the speed, but this is self ablation study. I want see the figure that x-axis = speed / y-axis=performance compared to existing method.

- Can you apply this method on "solving inverse problem with unconditional flow-matching model"?