S²-DMs: Skip-Step Diffusion Models

Generative models, especially deep generative models, play a foundational role in the machine learning domain (Karras et al., 2020); (Oord et al., 2016). Architectures like Variational Autoencoders (VAEs; (Kingma & Welling, 2013)) and Autoregressive models(Van den Oord et al., 2016);(Brown et al., 2020); (Salimans et al., 2017), Generative Adversarial Networks (GANs; (Goodfellow et al., 2014); (Yu et al., 2017); (Hjelm et al., 2017);(Fedus et al., 2018)), and Restricted Boltzmann Machines (RBMs; (Hinton, 2012)) have been at the forefront. VAEs, while providing a structured probabilistic framework, occasionally yield blurry samples. GANs, acclaimed for their prowess in generating high-resolution images, can face training instabilities (Adler & Lunz, 2018); (Gulrajani et al., 2017); (Karras et al., 2019). RBMs, though seminal, find themselves overshadowed by more recent architectures in scalability and performance. Against this backdrop, diffusion models, Denoising diffusion probabilistic models(DDPMs; (Ho et al., 2020); (Sohl-Dickstein et al., 2015); (Song et al., 2020b)), have emerged as a compelling alternative, exhibiting unmatched capabilities in generating superior samples in diverse domains, from image synthesis to molecule design (Bengio et al., 2014).

However, diffusion models do come with challenges. Their inherently slow sampling speed, driven by the multitude of necessary sampling steps, remains a significant concern. Recent research has honed in on this computational bottleneck, with the goal of optimizing the sampling process (Jolicoeur-Martineau et al., 2021), (Nichol & Dhariwal, 2021). A significant breakthrough in this area is the Denoising Diffusion Implicit Models (DDIMs; (Song et al., 2020a)). DDIMs utilize a subset sampling strategy, achieving faster performance by sampling from a smaller subset instead of the entire set of steps. This method, due to its omission of certain steps, is coined ”skip-step sampling.” Yet, this acceleration introduces an inconsistency between training and sampling. During training, the model undergoes every step, but during sampling, some steps are selectively skipped, posing a risk of information loss. Although DDIMs, with sampling algorithm adjustments, have lessened the adverse effects of this approach compared to DDPMs, they haven’t specifically addressed and optimized for the missing intermediate information. Consequently, this challenge persists, leading to the suboptimal performance of diffusion models during expedited sampling and hindering the generation of high-quality samples.

Driven by these observations, our research proposes a method of integrating skip-step sampling as an optimization target into the training process. The model trained in this way can consider and adapt to the information lost during the sampling period due to skip-step sampling. Our model not only maintains efficient generation when using other accelerated sampling algorithms but also improves the quality of the generated samples, thereby achieving a balance of efficiency and performance. Specifically, during the training process, the original loss function is retained, and at the same time, a skip-step loss is introduced, measuring the difference between the model’s current step prediction and the skip-step result. This skip-step loss is combined with the original loss using a weighted mechanism. Figure 1 shows an overview of the diffusion models.

Empirical results demonstrate that when the skip-step loss is incorporated into the loss objective, the S² -DMs achieves superior performance in unconditional generation on the CIFAR10 (Krizhevsky et al., 2009) and CELEBA datasets, under the same sampling algorithms (DDIMs, PNDMs (Liu et al., 2022a), and DEIS (Zhang & Chen, 2022)), compared to models trained with the original loss objective. Importantly, our study shows that our model achieves better results than the original model at different steps, with an improvement ranging from 3.27%-14.06% (CIFAR10) and 8.97%-27.08% (CELEBA), under the same sampling algorithms. We also conducted qualitative experiments and ablation studies, along with related analyses.

The key contributions of this work are:

1. 1.

   Modeling for Training-Sampling Discrepancy: To our knowledge, our study proposes the first method specifically designed to mitigate the inherent mismatch between training and sampling in diffusion models. This strategy consistently demonstrates superior performance.

2. 2.

   Innovative Skip-Step Loss: We introduce a trailblazing skip-step loss, embedding the selective sampling modality directly within the training process. This method empowers models to proactively navigate potential sampling information deficits, enhancing the quality of the samples.

3. 3.

   Simplicity of Implementation: The S²-DMs approach stands out not just for its efficacy but also its simplicity. With minimal code alterations required, it offers a convenient solution for both researchers and practitioners. Crucially, it’s adaptable to a range of sampling algorithms.

This study is based on DDPMs (Ho et al., 2020) and DDIMs (Song et al., 2020a), so a brief review is in order. DDPMs specifies a prior Markov forward diffusion process, which gradually adds noise to the data over $T$ steps. Refer to the background description of (Watson et al., 2021). Following the notation of (Ho et al., 2020),

where $q(x_{0})$ represents the data distribution and $1-\alpha_{t}$ signifies the variance of the Gaussian noise added at step $t$. For each $t$, we have $\alpha_{t}=1-\beta_{t}$ and $\bar{\alpha}_{t}=\prod_{s=1}^{t}\alpha_{s}$. To facilitate the transformation of noise back into data, DDPMs are trained to invert equation 1 with a model $p_{\theta}(x_{t-1}|x_{t})$. This model is trained by optimizing a (possibly reweighted) evidence lower bound (ELBO),

DDPMs explicitly select the model for parameterization as

In this framework, optimizing the ELBO corresponds to minimizing denoising score matching goals as explained by (Vincent, 2011). (Song et al., 2020a) introduced the DDIMs concept, a set of ELBOs complemented by forward diffusion processes and sampling mechanisms. These ELBOs, having similar marginals as DDPMs, offer flexibility in determining posterior variances (Chen et al., 2020). (Song et al., 2020a) emphasized crafting alternative ELBOs using a subset of original timesteps $S\subset\{1,...,T\}$ with consistent marginals. This results in $q_{S}(x_{t}|x_{0})=q(x_{t}|x_{0})$ for every $t$ in $S$, permitting faster sampling processes compatible with pre-trained models by integrating new timesteps. Their work also suggests the feasibility of creating a vast range of non-Markovian processes, denoted as $\{q_{\sigma}:\sigma\in[0,1]^{T-1}\}$, with each $q_{\sigma}$ maintaining marginals aligned with the original progression,

and where the posteriors are defined as

In their research, (Song et al., 2020a) observed that the special case of employing all-zero variances, termed as DDIMs($\eta=0$), persistently enhances the quality of samples in the short-step domain. When amalgamated with an apt choice of timesteps for assessing the modeled score function, known as strides, DDIMs($\eta=0$) sets a new benchmark in the realm of few-step diffusion model sampling, especially with minimal inference step allocations. A pivotal advancement we bring is the enhancement of sample quality by introducing skip information (i.e., the aforementioned subset) during the training phase, ultimately establishing a novel diffusion model training paradigm. For a more comprehensive discussion on the S²-DMs, please refer to Section 3.

Acceleration approaches often employ skip-step sampling as a strategy for acceleration. However, this approach inherently introduces non-smooth denoising, leading to a potential decline in performance. This observation prompted us to re-evaluate the entire training and sampling workflow. Intriguingly, we identified an asymmetry between the training and sampling phases: the former proceeds in single steps, while the latter uses skip-step.

To improve the quality of skip-step sampling, we devised a novel yet straightforward objective function for the training phase. By incorporating the skip-step objective into our original loss function (Section 3.2), we achieved a symmetrical training effect. Ultimately, the performance of the model trained by this method met our expectations. (Section 4).

In this section, we demonstrate that the S²-DMs outperform DDIMs (Song et al., 2020a), PNDMs (Liu et al., 2022a) and DEIS (Zhang & Chen, 2022) in image generation, achieving this with an identical number of steps and the same sampling algorithm. Moreover, the latent variables in the images generated by the S²-DMs retain a high level of image features, allowing for interpolation within the latent space.

In consideration of computational resources, our experiments utilized the CIFAR10 dataset with a resolution of 32×32 and the CelebA dataset with a resolution of 64×64. For the training setup, we adopted the same architecture (He et al., 2016); (Ronneberger et al., 2015); (Kingma & Ba, 2014) as provided in the official DDIMs repository. We also ensured that all parameters were kept consistent, guaranteeing the reproducibility of our experiments. For hardware, we trained both datasets on two NVIDIA A100 GPUs, with 600K steps for CIFAR10 and 400K steps for CelebA. Model performance was evaluated based on the FID (Heusel et al., 2017); (Jolicoeur-Martineau et al., 2020). Specifically, Our evaluation method followed the DDIMs repository, wherein we sampled 50,000 images and computed the FID against real images. To further guarantee reproducibility, we fixed the random seed in our experiments, ensuring that all results are replicable. (It is noteworthy that in the experiments, the models are trained based on the same parameters and methods. The results are then obtained using different sampling algorithms. Therefore, some of the results may not be consistent with those reported in original works. This is reasonable, as the parameters trained in each report are different. Our primary comparison focuses on the diffusion models trained using our method versus those trained with the original method, specifically comparing the effects of leap-step sampling under different sampling algorithms.)

In Table 2 and 2, we evaluate the quality of samples generated by models trained on the CIFAR10 and CelebA datasets, measured using the FID as the evaluation metric. We default to $skip=10$ and compare our results with DDIMs. As expected, by incorporating skip-step information, the model is able to capture a broader scope of knowledge, leading to an improved sample quality. We observed that the S²-DMs consistently produces higher quality samples than original model with the same sampling steps and the same algorithm. This demonstrates that the diffusion model enhanced by our training algorithm contributes to the improvement of sample quality. Thus, other models only need a few lines of training code modifications to benefit from the performance boost this method offers, without any additional changes to the sampling algorithm.

Denoising Diffusion Probabilistic Models (DDPMs; (Ho et al., 2020)) and Noise Conditional Score Networks (NCSNs; (Song & Ermon, 2019)) are notable for their sample quality, comparable to GANs. While DDPMs optimize a variational lower bound, NCSNs target score matching over a Parzen density estimator. Both models employ a denoising autoencoder across noise levels and use Langevin dynamics for sampling. The shared approach requires multiple iterations for optimal sample quality. Recent advancements aim to decrease DDPMs’ inference steps through SDE solvers and programming algorithms. However, challenges like the disparity between log-likelihood reduction and FID (Heusel et al., 2017), (Szegedy et al., 2016) remain in some models.

DDIMs (Song et al., 2020a) emerges as an implicit generative model, wherein samples are uniquely defined by latent variables. This lends DDIMs properties akin to GANs and invertible flows, including the capability to produce semantically meaningful interpolations. Conceived from a purely variational standpoint, DDIMs sidesteps the constraints of Langevin dynamics, potentially explaining its superior sample quality compared to DDPMs in fewer iterations. The sampling paradigm of DDIMs also echoes the concepts found in neural networks with continuous depth. Additionally, other innovative methods have been introduced to further refine DDPMs sampling, such as reverse SDEs (Song et al., 2020b) with unique coefficients, “corrector” steps, and probability flow ODEs (Liu et al., 2022b). As the exploration of efficient sampling in diffusion models continues, our research stands on the shoulders of these pioneering works, aiming to push the boundaries of what’s achievable in generative models.

We propose the Skip-Step Diffusion Models, a diffusion model that adapts better to accelerated sampling algorithms by simply adding an additional leap step loss during the training process. We demonstrate how to incorporate leap step loss into the original loss function during training. Our results qualitatively and quantitatively show that under the same accelerated sampling algorithms, S²-DMs significantly improve the sample quality of image generation. Our method successfully explores adding leap step information during training, allowing the model to reach a symmetrical state during sampling, thereby enhancing sample quality.

Our findings pave a new direction for future diffusion model research. By investigating the asymmetry between the training of diffusion models and accelerated sampling algorithms, we can ensure that the trained models are better suited to accelerated sampling algorithms. This ensures that even with algorithms using fewer sampling steps, the model can still generate high-quality outputs. Perhaps this offers an effective solution to the challenges diffusion models face in balancing sampling speed and sample quality. In the future, we will continue to explore how to better integrate leap-step information to ensure greater consistency between the trained model and accelerated sampling algorithms, thereby improving sample quality. Additionally, we plan to investigate how to incorporate leap-step information into ODEs, while also exploring possibilities in non-continuous spaces.