ReDi: Efficient Learning-Free Diffusion Inference via Trajectory Retrieval

Assuming data follow an unknown true distribution $p({\bm{x}})$, diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2021; Kingma et al., 2021) define a generation process. For any $p({\bm{x}})$, diffusion models learn a denoising process by iteratively adding noise to original data (denoted as ${\bm{x}}_{0}$) from time step $0$ until time step $T$. The forward process adds noise to ${\bm{x}}_{0}$ such that at time step $t$, the distribution of ${\bm{x}}_{t}$ conditioned on ${\bm{x}}_{0}$ is

where $\alpha(t),\sigma(t)$ are real-valued positive functions with bounded derivatives. The signal-to-noise ratio (SNR) is defined as $\gamma=\alpha^{2}(t)/\sigma^{2}(t)$. By making the SNR decreasing, the marginal distribution of ${\bm{x}}_{T}$ approximates a zero-mean Gaussian, i.e.

The noise-adding forward process transforms a data sample from the original distribution to the zero-mean Gaussian $q({\bm{x}}_{T})$, while the backward process randomly samples from $q({\bm{x}}_{T})$ and denoises the sample with a neural network parameterized conditional distribution $q({\bm{x}}_{s}|{\bm{x}}_{t})$ where $s<t$ until it reaches time step $0$.

Kingma et al. (2021) proves that the forward process is equivalent (in distribution) to the following stochastic differential equation (SDE) in terms of the conditional distribution $q({\bm{x}}_{t}|{\bm{x}}_{0})$,

where $\bm{w}$ is the standard Wiener process.

Song et al. (2021) shows that the forward SDE has an equivalent reverse SDE starting from time $T$ with the marginal distribution $q({\bm{x}}_{T})$,

where $\bar{\bm{w}}_{t}$ is the reverse Wiener process and $t$ goes from $0$ to $T$.

In practice, $\nabla_{\bm{x}}\log q_{t}({\bm{x}}_{t})\approx\bm{\epsilon}/\sigma(t)$ is estimated using a noise-predictor function $\bm{\epsilon}_{\theta}({\bm{x}}_{t},t)$, where $\bm{\epsilon}$ is the Gaussian noise added to ${\bm{x}}_{0}$ to obtain ${\bm{x}}_{t}$, i.e.,

To learn $\bm{\epsilon}_{\theta}$, the following objective function is minimized (Ho et al., 2020; Song et al., 2021),

where $\omega(t)$ is a weighting function, and the integral is estimated using random samples.

Song et al. (2021) proves that the following ordinary differential equation (ODE) is equivalent to Equation 6,

When $\nabla_{\bm{x}}\log q_{t}({\bm{x}}_{t})$ is replaced by its estimation, we obtain the neural network parameterized ODE,

The inference process of diffusion models can be formulated as solving Equation 10 given a random sample from the Gaussian distribution $q({\bm{x}}_{T})$. For each iteration in solving the equation, the noise-predictor function $\bm{\epsilon}_{\theta}({\bm{x}}_{t},t)$ is estimated with the trained neural network $\theta$. Therefore, the inference time consumption can be approximated by the number of function estimations (NFEs), i.e., how many times $\bm{\epsilon}_{\theta}({\bm{x}}_{t},t)$ is estimated.

The sample trajectory describes the intermediate steps to generate the final data sample (e.g. an image), so the inference time linearly correlates with its length, i.e., the number of estimations used. While previous numerical solvers work towards enlarging the step size $\delta$, ReDi aims at skipping some intermediate steps to reduce the length of the trajectory. ReDi is able to do so because the first few steps determine only the layout of the image which can be shared by many, and the following steps determine the details (Meng et al., 2021; Wu et al., 2022).

In the following section, we describe the ReDi algorithm with this definition.

Given a trained diffusion model, ReDi does not require any change to the parameters. It only needs access to the sample trajectory of generated images. We show the inference procedure of ReDi in Figure 2. Instead of generating all $T/\delta$ intermediate samples in the trajectory, ReDi skips some of them to reduce number of function estimations (NFEs) and improve efficiency. This is done by generating the first few samples ${\bm{x}}_{T..k}$, using them as a query to retrieve a similar trajectory ${\bm{x}}^{\prime}_{T..k}$, and then starting from ${\bm{x}}^{\prime}_{v}$ of the trajectory retrieved. This way, the NFEs spent to go from time step $k$ to time step $v$ would be unnecessary.

As shown in Figure 2, the retrieval of the similar trajectory ${\bm{x}}^{\prime}$ depends on a precomputed knowledge base $\mathcal{B}$. We formally describe the construction of $\mathcal{B}$ in Algorithm 1. ReDi first computes the sample trajectories for the data samples in a dataset $\mathcal{D}$. For every sample trajectory $\{{\bm{x}}_{n\delta}\}_{n=T/\delta,\dots,1,0}$, a sample early in generation ${\bm{x}}_{k}$ is chosen as the key, while a later sample ${\bm{x}}_{v}$ is chosen as the value. The time consumption for computing one key-value pair is similar to that of one generation of the model. The total time consumption is proportional to $|\mathcal{D}|$. With a moderate-sized $\mathcal{D}$, not only can we achieve comparable performance with much less time, we can also perform zero-shot domain adaptation.

The inference process of ReDi is formally described in Algorithm 2. Given a guidance signal $y$ (in our case, a text prompt), we want to generate a data sample (in our case, an image) ${\bm{x}}$ from it. We first generate the first few steps ${\bm{x}}_{T..k}$ to query the knowledge base $\mathcal{B}$ with ${\bm{x}}_{k}$. We find the top-$H$ keys that are closest to $q$ in terms of the distance measure $d$. Then we find out the set of weights $w^{*}$ that make the linear combination of the top-$H$ keys approximate $q$ the best. With $w^{*}$, we linearly combine the value and use that combination as the initialization point for the remaining steps of the sampling process.

One limitation of the ReDi framework described in subsection 4.2 is its generalizability. When the guidance signal $y$ contains out-of-domain information that does not exist in $\mathcal{B}$, it is difficult to retrieve a similar trajectory from $\mathcal{B}$ and run ReDi. Therefore, We propose an extension to ReDi in order to generalize it to out-of-domain guidance. For an out-of-domain guidance signal $y$, we break it into 2 parts - the domain-agnostic $y^{\text{in}}$, and the domain-specific $y^{\text{out}}$. We use $y^{\text{in}}$ to generate the partial trajectory as the retrieval key. After retrieval, we start from the retrieved value with both $y^{\text{in}}$ and $y^{\text{out}}$ as guidance signal.

For example, under the image synthesis setting, when $\mathcal{B}$ contains style-free images that are generated without any style specifier in the prompt, it is difficult for ReDi to synthesize images from a stylistic prompt $y$ because finding a stylistic trajectory from a style-free knowledge base is hard. However, with the proposed extension, ReDi is able to synthesize stylistic images.

When a stylistic guidance signal $y$ is given, the part of $y$ describing the content of the image is the domain-agnostic $y^{\text{in}}$, and the part of $y$ specifying the style of the image is the domain-specific $y^{\text{out}}$. Although it is difficult to find a trajectory similar to one generated by $y=(y^{\text{in}},y^{\text{out}})$, it is feasible to retrieve a trajectory similar to one generated by $y^{\text{in}}$. Therefore, we first use the content description to generate the retrieval key and then use the whole prompt for the following sampling steps to make the image stylistic.

To show ReDi’s ability to sample from trained diffusion models efficiently, we use ReDi to sample from Stable Diffusion, with two different numerical solvers, PNDM and DPM-Solver. We build the knowledge base $\mathcal{B}$ upon MS-COCO (Lin et al., 2014) training set (with 82k image-caption pairs) and evaluate the generation quality on 4k random samples from the MS-COCO validation set. We use the Fréchet inception distance (FID) (Heusel et al., 2017) of the generated images and the real images. In practice, we choose $L^{2}$-norm as our distance measure $d$ and calculate the optimal combination of neighbors using least square estimation.

Note that although the training of Stable Diffusion costs as much as 17 GPU years on NVIDIA A100 GPUs, constructing $\mathcal{B}$ with PNDM only requires 1 GPU day, which is much more efficient compared to learning-based methods like progressive distillation (Meng et al., 2022).

For PNDM, we generate 50-sample trajectories to build the knowledge base. We choose the key step $k$ to be $40$, making the first 10 samples in the trajectory the key, and alternate the value step $v$ from $30$ to $10$. Different combinations of key and value steps determine how many NFEs are reduced. For DPM-solver, we choose the length of the trajectory to be $20$ and conduct experiments on $k=5,v=8/10/13/15$. To compare the performance of ReDi with the existing numerical solvers. We use them to generate images with the same NFE budget and compare their FIDs with ReDi.

We show the results of our experiments in Table 1 and Table 2. Some samples generated by ReDi are listed in Figure 3. For both PNDM and DPM-Solver, we report the FID scores of the sampler without ReDi, ReDi with top-1 retrieval, and ReDi with top-2 retrieval. The results indicate that with the same number of NFEs, ReDi is able to generate images with a better FID. They also indicate that when ReDi is combined with numerical solvers, we can achieve better performance with only 40% to 50% of the NFEs. In terms of clock time, 50-step PNDM takes 2.94 seconds, while the 30-step ReDi takes 1.75 seconds to generate one sample. The precomputation time for building the vector retrieval index and retrieving top-$h$ neighbors, when amortized by 4000 inferences, is 0.0077 seconds, which only takes 0.4% of the inference time. Therefore, ReDi is capable of keeping the generation quality while improving inference efficiency.

We further validate the effectiveness of ReDi with more ablation studies and baselines in Appendix A.6.

One major difference between ReDi and previous retrieval-based samplers is the keys we use. Instead of using text-image representations like CLIP embeddings, we use an early step from the sample trajectory as the key. Theoretically, this is the optimal solution, because we are directly minimalizing the bound from 5.35.3 by minimalizing $|{\bm{x}}_{k}-{\bm{y}}_{k}|$. In this section, we show empirically that sample trajectory is indeed the better retrieval key.

We compare CLIP embeddings and trajectories by using them alternatively as the key. To use CLIP embeddings as key, we build a knowledge base similar to the one we use in subsection 6.1 with them, where the keys are CLIP embeddings and the values are later samples in the sample trajectory. We run the same inference process except with a different knowledge base.

To evaluate the performance of retrieval keys, we use two criteria - the distance between the actual value and the retrieved value $d({\bm{x}}_{v},\hat{\bm{x}}_{v})$, and the FID scores of the generated images. We report the results in Table 3. Under all NFE settings, the trajectories serve a better purpose as the key for ReDi than CLIP embeddings. This corresponds well to Theorem 5.3.

We use the extended ReDi framework from subsection 4.3 to generate domain-specific images. In particular, we conduct experiments on image stylization (Fu et al., 2022; Feng et al., 2022) with the style-free knowledge base $\mathcal{B}$ from 6.1. To generate an image with a specific style, we do not build a knowledge base from stylistic images. Instead, we use the content description $y^{\text{content}}$ to generate the partial trajectory as key. After retrieval, we change the prompt from $y^{\text{content}}$ to the combination of $y^{\text{content}}$ and $y^{\text{style}}$, where $y^{\text{style}}$ is the style description in the prompt.

In our experiments, we transfer the validation set of MS COCO to three different styles. The content description $y^{\text{content}}$ directly comes from the image captions, while the style description $y^{\text{style}}$ is appended as a suffix to the prompt. The style descriptions for the three styles are “a Chinese painting”, “by Claude Monet”, and “by Vincent van Gogh”.

Unlike other experiments, in this experiment, we choose $K=47$ and $V=40$. Because from the preliminary experiments, we find that the style of the image is determined much earlier than its detailed content. For 100 randomly sampled captions from MS COCO validation, we generate the corresponding images in all three styles. To evaluate ReDi’s performance, we compare them with images generated by the original Stable Diffusion with PNDM solver.

We asked human evaluators from Amazon MTurk to evaluate the generated images. They are paid more than the local minimum wage. Every generated image is rated in three aspects, text faithfulness, style consistency, and layout similarity. Every aspect is rated on a scale of 1 to 5 where 5 stands for the highest level. Textual faithfulness represents how faithful the image is depicting the content description. Style consistency represents how consistent the image is to the specified style. Layout similarity represents how similar the layout of the stylistic image is to the layout of the style-free image.

We report the results of the human evaluation of the 3 different styles in Figure 4. In terms of text faithfulness and style, ReDi’s evaluation is comparable to Stable Diffusion. In terms of layout similarity, the images generated by ReDi have significantly more similar layouts. This indicates that by retrieving the early sample in the trajectory, ReDi is able to keep the layout unchanged while transferring the style. This finding is also demonstrated by the qualitative examples in Figure 5.

Furthermore, for stylistic prompts that can not be explicitly decomposed into $y^{\text{content}}$ and $y^{\text{style}}$, we propose to use prompt engineering techniques from the natural language processing community by asking a large language model to conduct the decomposition for us. We include an example of such method in Appendix A.7.