On Optimizing Timesteps of an EDM Based Diffusion Sampling Procedure

Abstract

Recently, a technique named improved integration approximation (IIA) has been introduced in [1], aiming to enhance the performance of the EDM [2] sampling procedure for small timesteps. The basic idea of IIA-EDM is to optimize the step-sizes in front of the gradient vectors per timestep when computing the next diffusion state via minimum mean squared error (MMSE). In this paper, we further improve the sampling performance of IIA-EDM by optimizing the timesteps. Our main contribution is to utilize the residual errors of the MMSE in IIA-EDM as the objective function and then apply the genetic algorithm to minimize the residual error w.r.t. the timesteps. We refer to the resulting sampling procedure as OTS-IIA-EDM, where OTS stands for optimized timesteps. With the OTS-IIA-EDM, we obtain substantial performance enhancements on different datasets (CIFAR-10, FFHQ, AFHQv2, ImageNet) in comparison to both IIA-EDM and EDM. Notably, when utilizing 6 timesteps, OTS-IIA-EDM demonstrates a noteworthy FID improvement, ranging from 10.8 to 18.4 percent compared to IIA-EDM, and a more substantial boost of 40.1 to 69.8 percent to EDM.