Non-uniform Timestep Sampling: Towards Faster Diffusion Model Training

Authors: tianyi  zheng
cong  geng
peng-tao  jiang
hao  zhang
jinwei  chen
jia  wang
ben  wan
bo  li
Venue: acmmm_org_ACMMM_2024_Conference
Published: 2024-03-24 09:24:26

Abstract: Diffusion models have garnered significant success in generative tasks, emerging as the predominant model in this domain. Despite their success, the substantial computational resources required for training diffusion models restrict their practical applications. In this paper, we resort to the optimal transport theory to accelerate the training of diffusion models, providing an in-depth analysis of the forward diffusion process. It shows that the upper bound on the Wasserstein distance of the distribution between any two timesteps in the diffusion process is an exponential decrease of the initial distance by a factor of times. This finding suggests that the state distribution of the diffusion model has a non-uniform rate of change at different points in time, thus highlighting the different importance of the diffusion timestep. To this end, we propose a novel non-uniform timestep sampling method based on the Bernoulli distribution, which favors more frequent sampling in significant timestep intervals. The key idea is to make the model focus on timesteps with larger differences, thus accelerating the training of the diffusion model. Experiments on benchmark datasets reveal that the proposed method significantly reduces the computational overhead while improving the quality of the generated images.

Main paper: 
 introduction: 
 In recent times, generative models, particularly diffusion-based generative models [6, 25, 46, 57] , have garnered significant attention for their notable achievements in computer vision [4, 35, 37, 51] , natural language processing [3, 47, 54] , temporal data modeling [1, 5, 12, 44] and AI Security [13, 14, 36] . The diffusion model comprises two primary stages: forward diffusion and reverse diffusion. In the forward diffusion stage, the objective is to facilitate the transformation of the original distribution into a Standard Gaussian distribution. On the other hand, the reverse diffusion stage aims to transition from a Standard Gaussian distribution back to the data distribution, thereby accomplishing the generation of new samples.Despite the significant achievements of diffusion models as evidenced by several studies [25, 41, 48] , the considerable computational demands required for their training restrict their utility in a broad spectrum of applications. ADM [10] presents an enhanced adaptive group normalization framework that more effectively integrates temporal data within residual blocks, thereby accelerating the training process of diffusion models. EDM [28] designs advanced noise schedules to improve training efficiency. Meanwhile, many studies propose the various re-weighted loss functions, effectively enhancing the training efficiency of diffusion models. For instance, the P2-weight method [7] focuses on images with medium signalto-noise (SNR) ratios, believed to contain rich semantic information. To leverage this, P2-weight introduces a re-weighted loss function specifically targeting these medium SNR ratios. This approach aims to enhance the model's ability to learn visual concepts, thereby boosting training efficiency. Moreover, Min-SNR [23] and ANT [20] discover that the optimal weight gradients for different noise distributions conflict, leading to a notable slowdown in the training of diffusion models. Therefore, Min-SNR proposes a clamped SNRweighted loss function to reduce this conflict and accelerate the training process. ANT [20] proposes an Uncertainty Weighting strategy [30] to faster training the diffusion model. Similarly, the Debias [56] methods also design a re-weighted loss function based on SNR to reduce the bias of different noise distributions and speed up the training. While these methods somewhat expedite the training process of the diffusion model, they all employ uniform timestep sampling methods in the training stage. Our subsequent theoretical analysis suggests that uniform timestep sampling methods may be sub-optimal for the training of diffusion models, primarily because the variability across different noise distributions in the forward diffusion is not uniform.In this paper, we introduce an approach that leverages optimal transport (OT) theory to examine the forward diffusion process of the diffusion model. OT theory provides enlightenment to the changes in distance between different distributions in the Wasserstein space. More importantly, the forward process of the diffusion model can be modeled as a gradient flow in the Wasserstein space, which provides a good understanding of the change in distribution of the diffusion process. As depicted in Figure 1a , there is a notable exponential diminution in the Wasserstein distances concomitant with the advancement of the forward diffusion process. This analysis confirms that the upper bound of the Wasserstein distance between any two specified moments exhibits an exponential reduction relative to their initial disparity. Figure 1b presents empirical validation of this conclusion. Based on this foundation, we propose a novel non-uniform timestep sampling method named Bernoulli Distribution-Based Sampling for training diffusion models. We call our method Denoising Diffusion Probabilistic Models with Bernoulli Sampling (DDPM-BS). This design concentrates on diffusion stages with distinct distribution differences, enhancing training efficiency and improving the generative performance of diffusion models. Furthermore, our approach is compatible with existing enhancement methods (e.g., P2-weight, EDM, etc.), allowing for seamless integration that can leverage the advantages of various designs to further elevate training efficiency and generative quality. Our contributions can be summarized as follows: Diffusion models are proposed by [48] and improved by [17, 25, 41] .Recently, the ADM [10] can generate higher-quality images than Generative Adversarial Networks (GANs) [21, 22, 52] . EDM [28] enhances training efficiency and sample quality through advanced noise schedules and network architectures. However, training a diffusion model like ADM and EDM needs substantial computational cost, thereby constraining the application of these models. To reduce the computational cost, the P2-Weight [7] method designs the re-weighted loss functions to speed up the training. Meanwhile, Min-SNR [23] and ANT [20] analyze the training process of diffusion models through a multi-task learning perspective. Therefore, they propose different re-weighted loss functions to speed up the training. Furthermore, E-TSDM [55] find that Lipschitz singularities pose a threat to the stability of the training. Therefore, E-TSDM shares the timestep with large Lipschitz constants to reduce the instabilities in the training. However, the uniform timestep sampling approach adopted by these methods during the training stage may not be optimal, potentially impacting training efficiency and the quality of the generated samples. Furthermore, these enhanced designs are orthogonal to our timestep sampling method, indicating that they can be integrated with our approach without conflict. 
optimal transport.: 
 The Wasserstein distance in optimal transport theory is a distance function defined between probability distributions on a given metric space. The 2-Wasserstein distance between two probability measures 𝜇, 𝜈 is:EQUATION.Here, Π(𝜇, 𝜈) represents the set of all joint distributions (couplings) on R 𝑑 × R 𝑑 that have 𝜇 and 𝜈 as their respective marginals. The optimal transport theory and stochastic differential equations (SDE) are closely related [15] . Therefore, the Wasserstein distance has been utilized to explain the diffusion-based generative model [18, 19, 31, 33] in several works. Moreover, DPM-OT [34] accelerates inference speed during generation by solving a semidiscrete OT problem. Even though these studies provide valuable theoretical insights, they have yet to effectively translate these findings into tangible improvements in the training of diffusion models. In this paper, we not only provide an in-depth theoretical analysis of the diffusion model, but also propose an improved strategy for training, which greatly improves the training speed and generation quality of the diffusion model. 
method: 
 In this section, we provide an introduction to the background of the diffusion model and the optimal transport theory in Section 3.1. Subsequently, in Section 3.2, we explore the analysis of the diffusion process using the Wasserstein gradient flow and experimental W 2 ( 1 t , 2 t ) W 2 ( 1 t , 3 t ) W 2 ( 1 t , 4 t ) (b)The Wasserstein distance varies with the diffusion process. verification of the conclusion. Finally, in Section 3.3, we present our timestep sampling method based on theoretical insights.The aim of the forward diffusion process is to transform the complex image distribution 𝑞(𝑥 0 ) into simple normal Gaussian distribution N (0, 𝑰 ). In the forward diffusion process, we define the noise schedule 𝛽 𝑡 , Meanwhile, with 𝛼 𝑡 := 1 -𝛽 𝑡 and ᾱ𝑡 := 𝑡 𝑠=0 𝛼 𝑠 we assume in each time 𝑡, the probability density function of 𝑞(𝑥 𝑡 |𝑥 0 ) is𝑞 (𝑥 𝑡 | 𝑥 0 ) = N 𝑥 𝑡 ; √ ᾱ𝑡 𝑥 0 , (1 -ᾱ𝑡 ) 𝑰 𝑥 𝑡 = √ ᾱ𝑡 𝑥 0 + √ 1 -ᾱ𝑡 𝜖.(1) In the training stage, we train a model 𝜖 𝜃 (𝑥 𝑡 , 𝑡) to predict the added noise 𝜖 in each timestep.In the inference stage, we start from a known distribution, such as the standard Gaussian distribution, and reverse each step of the noise-corrupted latent 𝑥 𝑡 -1 from 𝑥 𝑡 . However, the distribution 𝑞(𝑥 𝑡 -1 |𝑥 𝑡 ) is based on the entire data distribution. Therefore, in DDPM, a neural network parameterized by 𝜃 is employed to estimate 𝑞(𝑥 𝑡 -1 |𝑥 𝑡 ) using the following equation:𝑝 𝜃 (𝒙 𝑡 -1 | 𝒙 𝑡 ) := N (𝒙 𝑡 -1 ; 𝜇 𝜃 (𝒙 𝑡 , 𝑡) , Σ 𝜃 (𝒙 𝑡 , 𝑡)) .Leveraging Bayes' rule, we can use the trained diffusion model to predict the previous denoising distribution step by step, employing the following equation:𝑥 𝑡 -1 = 1 √︁ 1 -𝛽 𝑡 𝑥 𝑡 - 𝛽 𝑡 √ 1 -𝛼 𝑡 𝜖 𝜃 (𝑥 𝑡 , 𝑡) + 𝜎 𝑡 𝑧.When we make time tend to be continuous [50] , the generalized forward diffusion process can be expressed as stochastic differential equations (SDEs) form: d𝑥 = 𝑓 (𝑥, 𝑡)d𝑡 + 𝑔(𝑡)d𝑤 .(EQUATIONThe SDE 2 in Euclidean space describes the change of probability distribution over time in the forward diffusion process, while each distribution in Euclidean space corresponds to an element in the Wasserstein space. In order to gain a better understanding of the different distributions in the forward diffusion process, we can shift our attention to the Wasserstein space. Indeed, the Wasserstein gradient flow elucidates the dynamics of probability density as it follows the steepest descent direction of the functional F (e.g. KL divergence), relative to the Riemannian metric induced by the 2-Wasserstein distance [2, 11] . The Wasserstein gradient can be expressed in the following form [2] , where F ′ (𝜇) represents the variance of F with respect to the measure 𝜇.∇ 𝑊 2 F = ∇F ′ (𝜇).The Wasserstein gradient of the functional KL(.∥𝜋) at 𝜇 is∇ 𝑊 2 KL(.∥𝜋) = ∇ log 𝜇 𝜋 .Then using the functional KL(.∥𝜋) and the continuity equation, we obtain the equation of the Wasserstein gradient flow as:𝜕𝜇 𝑡 𝜕𝑡 = ∇ • 𝜇 𝑡 ∇ log 𝜇 𝑡 𝜋 .Wasserstein gradient flow is a differential equation of probability measures. When we scale a time-varying 𝛽 𝑡 at time 𝑡, we can get the regular gradient flow [19] . The regular gradient flow satisfies the following continuity equation:EQUATIONIn the regular gradient flow, the velocity vector is defined as𝜈 𝑡 = -𝛽 𝑡 ∇ 𝑊 2 F .Then we can prove that the corresponding Fokker-Planck equation [45] of the SDE 2 is equivalent to the regular Wasserstein gradient flow [19] when the functional on the Wasserstein space is defined by F (𝜇) = KL(.∥𝜋). The corresponding Fokker-Planck equation [45] is:EQUATIONBased on this relationship, we have the following lemma.Lemma 1. Consider a regular Wasserstein gradient flow, as defined in Equation 3, initiating from a data distribution 𝜇 0 and converging to a normal Gaussian distribution 𝜇 𝑇 = N (0, 𝑰 ). With the selection of 𝑓 = 𝛽 𝑡 𝑥 and 𝑔 𝑡 = √︁ 2𝛽 𝑡 , the family of measures {𝜇 𝑡 } 𝑇 𝑡 =0 derived from the Fokker-Planck equation 4 is equivalent to the family of measures corresponding to this gradient flow.The detailed proof of Lemma 1 is in the Appendix Section 2.1. 𝜇 𝑡 is the noisy data distribution in the forward diffusion at time 𝑡. This leads us to the following remark. Proposition 1. When set the 𝑓 = 𝛽 𝑡 𝑥 and 𝑔 𝑡 = √︁ 2𝛽 𝑡 , the SDE 2 can be written as:EQUATION)This SDE describes the forward diffusion process of the diffusion model [25, 50] .Proposition 1 implies that the forward diffusion process of diffusion model [25, 50] can be equivalently thought as a regular Wasserstein gradient flow [19] . Meanwhile, at any time 𝑡, the constant 𝛽 𝑡 corresponding to different starting distributions is the same. Therefore, based on the solid mathematical framework of Wasserstein gradient flow [2, 26, 53] , We can analyze the forward diffusion process based on the properties of Wasserstein gradient flow in the Wasserstein space. Therefore, we have the following theorem.Theorem 1. Consider two distinct initial distributions 𝜇 0 and μ0 on the data manifold 𝑀, which is equipped with a reference measure 𝜈 = 𝑒 -𝑉 vol, satisfying 𝐻𝑒𝑠𝑠 𝜇 ≥ 𝐾. Let 𝜇 𝑡 and μ𝑡 represent the distributions at time 𝑡 in the forward diffusion process described in Proposition 1, originating from 𝜇 0 and μ0 respectively. For all 𝑡 > 0, the following inequality holdsEQUATIONThis theorem can be proved in terms of Riemannian geometric (in pure Otto's formalism), and the detailed proof procedure is in the Appendix Section 2.2. The determination of a precise value for 𝐾, denoting the lower bound of the curvature on the data manifold 𝑀 [15] , presents a significant challenge. Despite this, Theorem 1 elucidates critical theoretical insights on the forward diffusion mechanism inherent in the diffusion model. During the forward diffusion stage, the upper bound of Wasserstein distances between different initial distributions declines rapidly over time. This means that the different initial distributions become very similar at an exponential rate. Then we have the following proposition.Proposition 2. Based on Theorem 1, we observe that during the forward diffusion process, the upper bound of the 2-Wasserstein distance between different distributions exhibits an exponential decrease over time. For a sufficiently small 𝛿, when 𝑡 > ln 𝛿 |𝐾 | , the following inequality holds:EQUATIONThis proposition tells us that when diffusion has proceeded to a certain point, the 2-Wasserstein distance between the different distributions is already almost 0. The distributions at this point are virtually the same. Time intervals before this point can be considered significant intervals in the diffusion stage and this provides a theoretical basis for designing improved training methods for diffusion models.Then, we use the entropic regularization algorithm [9, 16] to empirically evaluate the Wasserstein distance between different distributions in the forward diffusion process. We select different images in the ImageNet dataset as a distribution. Then we use Equation 8 to empirically evaluate the 2-Wasserstein distance between different distributions by setting C(𝑥, 𝑦) = ∥𝑥 -𝑦 ∥ 2 . This algorithm can compute the 2-Wasserstein distance between different distributions efficiently.EQUATIONAs illustrated in Figure 1b , we observe an exponential decrease in the 2-Wasserstein distance among various distributions. Initially distinct due to their unique characteristics, these distributions become more similar as the forward diffusion process progresses. This finding challenges the efficacy of uniform timestep sampling in training diffusion models, suggesting that a more targeted timestep sampling approach increasing sampling density at timesteps with notable distributional shifts could better.In this section, we propose the Denoising Diffusion Probabilistic Models with Bernoulli Sampling (DDPM-BS) inspired by our analysis of optimal transport theory. Based on the Theorem 1 and Proposition 2, we can divide the timesteps of the diffusion process into significant timestep intervals. Within the significant timestep intervals, notable disparities are observed among the various distributions. Conversely, in the remaining timestep intervals, these disparities diminish, leading to smaller differences. Therefore, we introduce a non-uniform time sampling method that focuses more on significant distributional time intervals. This approach is designed to ensure the model is effectively trained to manage diffusion processes with notable distributional differences. Specifically, we use 𝑡 𝛿 to represent the threshold value of significant timestep intervals. Consequently, we partition the set of intervals 𝑡 1 , 𝑡 2 , ..., 𝑡 𝛿 , ..., 𝑡 𝑇 into two subsets: 𝑡 1 , ...𝑡 𝛿 and 𝑡 𝛿 , ..., 𝑡 𝑇 . Our strategy samples with probability 𝑝 in significant intervals and 1-𝑝 otherwise. This approach yields the following proposition. Proposition 3. During the training stage, let 𝑋 𝑛 be a random variable representing a sample at time 𝑡 𝑛 in the timestep series {𝑡 1 , 𝑡 2 , . . . , 𝑡 𝛿 , . . . , 𝑡 𝑇 }. Specifically, 𝑃𝑟 [𝑋 𝑛 = 1] denotes the probability of sampling within the significant intervals at time n. The probability of having exactly 𝑘 samples in the significant intervals, denoted by 𝑃𝑟 [𝑆 𝑛 = 𝑘], is given by the binomial formula 𝑛 𝑘 𝑝 𝑘 (1 -𝑝) 𝑛-𝑘 , consistent with a Bernoulli distribution.This proposition elucidates the probability density function (PDF) for sampling within significant intervals. Our non-uniform timestep sampling method, following the Bernoulli distribution, is thus named Bernoulli Distribution-Based Sampling. Moreover, we maintain uniform sampling in each sub-interval, both significant and non-significant. Therefore, the sampling weight of our method in each timestep is non-uniform and can be written as a Bernoulli distribution:𝑓 (𝑡) = 𝑝 for 𝑡 ∈ [0, 𝑡 𝛿 ] 1 -𝑝 for 𝑡 ∈ [𝑡 𝛿 , 𝑡 𝑇 ].The proposed Bernoulli Distribution-Based Sampling training algorithm for the diffusion model is presented in Algorithm 1. Here, U denotes a uniform distribution. It is important to note that accurately estimating 𝑡 𝛿 is often challenging. Further discussion about the choice of 𝑡 𝛿 is available in Section 4.4. In all our experiments, we use 𝑡 𝛿 /𝑇 = 0.8 to classify significant time sub-intervals. Then we present a comparison plot of various sampling methods in Figure 2 . It becomes evident that we can regulate the sampling frequency during the training stage within different intervals by adjusting the parameter 𝑝 for the Bernoulli distribution. Specifically, setting 𝑝 > 0.5 allows us to focus timestep sampling more intensively on the significant time sub-intervals. Therefore, the corresponding training loss of DDPM-BS can be written asL (𝜖 𝜃 ) := E 𝑡 ∼B U (0,𝑇 ),x 0 ∼𝑞 (x 0 ) ∥𝜖 𝜃 (𝛼 𝑡 x 0 + 𝜎 𝑡 𝜖, 𝑓 T (𝑡)) -𝜖 ∥ 2 2 (9)We represent our timestep sampling strategy with B U (0,𝑇 ), and denote 𝜖 as a random variable following the standard Gaussian distribution N (0, I). Furthermore, DDPM-BS only adjusts the timestep sampling strategy, indicating it is orthogonal to other improvement strategies. Therefore, it can easily be integrated with these strategies for enhanced performance of the diffusion model.Beyond our Bernoulli distribution-based non-uniform timestep sampling method, we quantitatively evaluate other distributionbased methods in Section 3.4 in the Appendix. Experimental results indicate that similarly focusing on significant time intervals in other non-uniform sampling methods also boosts training speed and enhances generate quality.Require: Upper bound 𝛿, Bernoulli parameter 𝑝.1: Estimating the significant interval 𝑡 𝛿 using Eq. 8 2: repeat 3:𝑥 0 ∼ 𝑞(𝑥 0 ) 4: Generating random numbers 𝑢 ∼ U[0, 1] 5: if 𝑢 < 𝑝 then 6:𝑡 ∼ U({1, . . . , 𝑡 𝛿 }), 𝝐 ∼ N (0, 𝑰 ) advanced than DDPM. Throughout the training stage, we maintain 𝑇 = 1, 000 for all experiments. We conduct comparative experiments using the CIFAR-10 (32×32), ImageNet (32×32), CelebA (64×64) [38] , Stanford Cars (128×128) [32] FFHQ (128×128) [29] , CelebAHQ (256×256) [27] and AFHQ-D (256×256) [8] datasets for unconditional image generation. More details about our experiment setup are shown in Section 4 in the Appendix.Evaluation. We adopt the widely-used Frechet Inception Distance score (FID) [24] and sFID [40] to evaluate the sample quality.When comparing against previous methods, we follow the previous work [25, 41] Comparison with previous state-of-the-art methods. We conduct a comparative analysis of diffusion models trained using our method against various previous state-of-the-art models [7, 23, 41, 56] on the CIFAR-10 (32 × 32), FFHQ (128 × 128) and AFHQ-D (256 × 256) datasets, exploring different sampling step settings 𝑇 . For a fair comparison, we employed the same model architecture and focused on the best-performing iterations of the different methods. Moreover, we provide the FID scores concerning the number of training iterations of each method in the Appendix Section 3.1. Additionally, we assess the impact of varying sampling steps, utilizing the respacing technique [10] for step reduction. The results presented in Table 1 , indicate that our method consistently outperforms others across multiple datasets and sampling step configurations. Notably, certain weighting strategies (e.g., Min-SNR [23] , Debias [56] ) do not perform as well as the baseline model under some sampling step scenarios in the low-resolution CIFAR-10 dataset, suggesting a lack of robustness. In contrast, our method demonstrated superior performance under all tested conditions and datasets. Integrating with other improvement method. Our non-uniform timestep sampling strategy is designed to not conflict with other training improvement methods for diffusion models. Because of this compatibility, we attempt to combine our method with previous methods, such as the Min-SNR [23] and P2-Weight [7] . We apply this combination in the training on the AFHQ-D dataset. Table 2 shows significant improvements in FID and sFID metrics. This suggests that combining our method with other training improvement algorithms can further boost model performance, highlighting our method's generalizability. Moreover, our ADM-BS also can be integrated with the Input Perturbation (IP) method [43] to mitigate the bias issue in the diffusion model. For detailed experimental details and results, please refer to 3.3 in the Appendix. Fine-Grained Data. While commonly used datasets such as Ima-geNet, FFHQ, and AFHQ-D are prevalent in generative modeling studies, they often lack the expressive power required to capture extremely fine-grained differences. These standard datasets primarily serve the purpose of broader category identification, rather than emphasizing fine-grained distinctions within categories. Given that the intricate details hold significance for generation applications, 4 reveal that ADM-BS achieved about 1.51× acceleration in attaining an equivalent FID score compared to the ADM. Furthermore, it is observed that ADM-BS consistently outperforms ADM in terms of generation quality.  
fast sampling: 
 The extensive inference process of diffusion models restricts their practical applications, prompting the development of various accelerated inference algorithms, such as DDIM [49] , Epsilon Scaling [42] , and DPM-Solver [39] . These methods have significantly enhanced the inference efficiency of diffusion models. Consequently, it is essential to evaluate the compatibility and performance of our model in conjunction with these accelerated approaches. As depicted in Table 3 , ADM-BS surpasses the baseline in terms of performance under rapid inference conditions. Additionally, our approach demonstrates minimal variation across different sampling steps, indicating enhanced stability and robustness of ADM-BS. Bernoulli distribution parameter 𝑝. As the parameter 𝑝 governs the probability of sampling at significant intervals, it directly influences the sampling frequency at these intervals during training. Therefore, we conduct ablation experiments to investigate its impact. The results are shown in Figure 5 . We observe that as the parameter 𝑝 increases, the model converges at an accelerated rate.When 𝑝 reaches 0.8, the model achieves an FID score of less than 3 within 250K training iterations, outperforming the final converged results of the ADM method. This phenomenon is attributed to ADM-BS, which places greater emphasis on intervals characterized by significant differences. This effectively expedites convergence and enables the model to pay more attention to distribution disparities, resulting in improved generative quality. Additionally, we note that setting 𝑝 below 0.5 leads to slower convergence compared to the ADM method. By the time the number of training iterations reaches 500K, the ADM has essentially converged, whereas ADM-BS has not reached full convergence at that training iteration. For more discussion about 𝑝, please refer to 3.2 in the Appendix.0.4 0.5 0.6 0.7 0.8 0.9 diffusion step (t /T) The upper bound 𝛿. The parameter 𝛿 plays a pivotal role in defining the range of significant intervals within a distribution. This determination is based the premise that when the Wasserstein distance between two distributions is less than 𝛿 times their initial values, the distributions are considered markedly similar. However, the computation of the Wasserstein distance is very difficult and complex, although we empirically validate the decreasing trend of the Wasserstein distance in Figure . 1b, there is a bias in this estimation and only serves primarily as a motivational figure. Intuitively, we think 𝛿 should be a relatively small value because this implies a more similar distribution. In our experiments, we substitute 𝛿 with various 𝑡 𝛿 /𝑇 ratios and evaluate the model's performance across these different ratios, as illustrated in Figure 6 . We observe that optimal performance consistently occurs at a 𝑡 𝛿 /𝑇 ratio of approximately 0.8 for various values of 𝑝. Based on this finding, we select 𝑡 𝛿 /𝑇 = 0.8 as the significant interval threshold for all subsequent experiments. Notably, the model's performance with any chosen 𝑡 𝛿 /𝑇 ratio consistently surpasses the baseline, underscoring the efficacy of our approach. Robustness to Noise Schedule. Since the convergence phenomenon remains unaffected by various noise schedules, we undertook a comparison involving different noise schedules (e.g., cosine and linear schedules) on the CIFAR-10 dataset, as outlined in Table 4 .The results highlight ADM-BS's robustness against these diverse noise schedules. Notably, ADM-BS not only surpasses the baseline in terms of generation quality but also enhances training speed by a factor of 1.67× across both noise schedules. For more comparisons of different noise distribution (e.g., EDM [28] ), please refer to Section 3.4 in the Appendix. The experimental results demonstrate that our method exhibits robustness across various noise schedules.Fast training. In Figure 7 , we present a comparison of samples generated by ADM and ADM-BS, trained across varying iterations (i.e., 50K, 100K, 200K, and 500K). ADM-BS can generate high-quality samples after only 100K iterations, achieving this at a much quicker rate than ADM. Moreover, the samples generated by ADM-BS exhibit a noticeably higher quality than those from ADM. This suggests that our method effectively improves the training efficiency. Unconditional generation. In Figure 8 , we compare the unconditional generation results of ADM and ADM-BS on the FFHQ and AFHQ-D datasets. It is observed that the samples generated by ADM exhibit color shift, whereas those generated by ADM-BS do not demonstrate color-shift issues. This result implies that ADM-BS generates higher-quality samples.Fine-grained dataset. We conducted visual comparisons using the fine-grained dataset Stanford Cars. The results in Figure 9 , reveal that the images produced via the ADM method manifest considerable deficiencies in detail, characterized by pronouncedADM ADM-BS  
 
 Appendix:  


Decision: This paper introduces an approach to enhance the training efficiency of diffusion models by employing a non-uniform timestep sampling strategy based on the theory of optimal transport. While concerns, such as claims regarding non-uniform sampling given the choice of hyperparameters, were initially raised, after the rebuttal, all reviewers concur with the merits of the work. The ACs agree with this assessment and, therefore, recommend accepting the paper. Congratulations to the authors! Please incorporate these suggestions into the camera-ready version to enhance the paper's clarity and novelty and include new results from the rebuttal.

Reviews: review: paper_summary: paper_summary: The paper proposes a non-uniform timestep sampling strategy for diffusion models, which can reduce training time and enhance the quality of generated images. This approach is based on the observation that the Wasserstein distance between the distributions of two timesteps decreases exponentially over time.
strength_weakness: strength_weakness: 1) The observation that the Wasserstein distance between the distributions of two timesteps decreases exponentially is interesting and meaningful.
2) The theoretical part of the Wasserstein gradient flow is solid.
3) The proposed non-uniform timesteps indeed reduce the training time and improve the quality of the generated images.
limitations: limitations: 1) In the Ablation Study, it is shown that performance improves as $p$ increases, with $p=0.8$ yielding the best results. What about $p=0.9$ or $p=1.0$? Does the diffusion model need to be trained on timesteps greater than $0.8T$? Additionally, when $\frac{t_\delta}{T}=0.8$, the proposed method is the same as uniform timestep sampling. As $\frac{t_\delta}{T}=0.8$, the probability density function for time sampling becomes $f(t)=\frac{1}{T}$​ for all timesteps. What is the difference between the proposed method and uniform timestep sampling? Figure 2 is confusing.

2) The connection between the analysis and the method appears weak. In Figure 1(b), the Wasserstein distance approaches 0 after 600 timesteps. However, as shown in Figure 6, the performance is best when $\frac{t_\delta}{T} = 0.8$. Could the author discuss this discrepancy?

3) The Wasserstein distance appears to depend on the noise schedule. How does the method perform on other types of diffusion models, such as VE-SDE[2]? Additionally, there is a lack of comparison with EDM[1].


[1] Karras, Tero, et al. "Elucidating the design space of diffusion-based generative models." Advances in Neural Information Processing Systems 35 (2022): 26565-26577.

[2] Song, Yang, et al. "Score-based generative modeling through stochastic differential equations." arXiv preprint arXiv:2011.13456 (2020).
paper_summary: paper_summary: This paper introduces an innovative approach to accelerate the training of diffusion models. The optimal transport theory is leveraged to analyze the forward diffusion process, revealing that the state distribution of the diffusion model changes non-uniformly over time. Based on this insight, a non-uniform timestep sampling method is proposed using a Bernoulli distribution to concentrate on significant timesteps, thereby speeding up training and improving image generation quality.
strength_weakness: strength_weakness: 1. **Theoretical Contribution.** The paper provides a rigorous theoretical analysis using optimal transport theory to understand the diffusion process, offering new insights into the non-uniform rate of change in the diffusion model's state distribution.
2. **Empirical Methodology.** The authors validate their theoretical findings through empirical experiments, demonstrating the practical effectiveness of their proposed non-uniform sampling method in reducing computational overhead and enhancing image quality.
3. **Novelty in Approach.** The non-uniform timestep sampling based on the Bernoulli distribution is a novel approach that contrasts with the uniform sampling methods traditionally used in training diffusion models.
4. **Potential Impact.** The proposed method is likely to be widely adopted by researchers and practitioners working on generative models due to its potential to significantly reduce training times and improve results without compromising quality.
limitations: limitations: 1. **Relevance to ACM MM.** The method proposed only focuses on unimodal diffusion models. Diffusion-based models are prominent for multimodal generation tasks. If more experiments with multimodal settings are provided for illustration, this paper would be more suitable for ACM MM.
2. **Inconsistent Terminology.** The method proposed is called DDPM-BS before the Experiment section but is suddenly changed to ADM-BS from subsection 4.2.
paper_summary: paper_summary: This paper introduces a non-uniform timestep sampling method leveraging optimal transport theory to accelerate diffusion model training while maintaining image quality.
strength_weakness: strength_weakness: The proposed non-uniform timestep sampling method is straightforward, easy to implement, and appears compatible with existing diffusion models.
limitations: limitations: 1. Limited novelty: While the method is easy to use, its core concept might be too basic and lack significant innovation compared to existing techniques. It could be seen more as a tuning approach than a groundbreaking advancement.
2. Limited comparison and integration: The paper focuses solely on comparing the method to loss reweighting techniques. It would benefit from a more comprehensive evaluation that includes comparisons and potential integration with other state-of-the-art methods like consistency models, Rectified Flow, and Score Distillation Sampling.
3. Training efficiency concerns: The claimed faster training speed compared to the baseline ADM needs further clarification. I wonder if the number of samples per training iteration remains the same for both ADM and ADM-BS. Without this control, the observed speedup in the number of iterations might not be solely attributable to the proposed method. A more accurate comparison of training efficiency should focus on the actual training time, instead of training iteration.
4. This is a minor: Some equations in the paper lack numbering.
paper_summary: paper_summary: This paper introduces a non-uniform timestep sampling method to improve the training efficiency of diffusion models. This paper employs optimal transport theory to examine the forward process of the diffusion model, and find that the upper bound of the Wasserstein distance between any two specified moments exhibits an exponential reduction relative to their initial disparity. Based on this finding, this paper proposes a Bernoulli Distribution-Based sampling method for training diffusion models. Experiments show that the proposed method can accelerate diffusion model training.
strength_weakness: strength_weakness: 1. Good empirical analysis combined with the optimal transport theory provides insights into the forward diffusion process. 
2. The proposed method is both theoretically validated and empirically successful.
3. Comprehensive experiment results confirm the approach is effective.
limitations: limitations: 1. Some theorems and propositions are hard to follow. Maybe better if authors can explain these concepts in a more easy-to-understand manner.
2. What are the differences between this work and other OT-based diffusion models (e.g., [1] [2] [3])? Is it correct that this work focuses on improving the efficiency of the diffusion sampling process while other OT-based diffusion models focus on improving the quality of generated images? It will be better if a short discussion can be included in the related works section.

References

[1] Gushchin N, Kolesov A, Korotin A, et al. Entropic neural optimal transport via diffusion processes. Advances in Neural Information Processing Systems, 2024, 36.

[2] Li Z, Li S, Wang Z, et al. DPM-OT: A New Diffusion Probabilistic Model Based on Optimal Transport. Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023: 22624-22633.

[3] Wan Z Y, Baptista R, Boral A, et al. Debias coarsely, sample conditionally: Statistical downscaling through optimal transport and probabilistic diffusion models. Advances in Neural Information Processing Systems, 2024, 36.