Diffusion Sampling with Momentum for Mitigating Divergence Artifacts

Despite the remarkable success of diffusion models in image generation, slow sampling remains a persistent issue. To accelerate the sampling process, prior studies have reformulated diffusion sampling as an ODE/SDE and introduced higher-order numerical methods. However, these methods often produce divergence artifacts, especially with a low number of sampling steps, which limits the achievable acceleration. In this paper, we investigate the potential causes of these artifacts and suggest that the small stability regions of these methods could be the principal cause. To address this issue, we propose two novel techniques. The first technique involves the incorporation of Heavy Ball (HB) momentum, a wellknown technique for improving optimization, into existing diffusion numerical methods to expand their stability regions. We also prove that the resulting methods have first-order convergence. The second technique, called Generalized Heavy Ball (GHVB), constructs a new high-order method that offers a variable trade-off between accuracy and artifact suppression. Experimental results show that our techniques are highly effective in reducing artifacts and improving image quality, surpassing state-of-the-art diffusion solvers on both pixel-based and latent-based diffusion models for low-step sampling. Our research provides novel insights into the design of numerical methods for future diffusion work. In this study, we investigate the potential causes of these artifacts and found that the narrow stability region of high-order numerical methods can cause solutions to diverge, resulting in divergence artifacts. To address this issue and enable low-step, artifact-free sampling, we propose two techniques. The first technique involves incorporating Polyak's Heavy Ball (HB) momentum [21], a well-known technique for improving optimization, into existing diffusion numerical methods. This approach effectively reduces divergence artifacts, but its accuracy only has first order of convergence. In this context, the accuracy measures how close the approximated, low-step solution is to the solution computed from a very high-step solver (e.g., 1,000-step DDIM). The second technique, called Generalized Heavy Ball (GHVB), is a new high-order numerical method that offers a variable tradeoff between accuracy and artifact suppression. Both techniques are training-free and incur negligible additional computational costs. Figure 1 demonstrates the superiority of both techniques in reducing divergence artifacts compared to previous diffusion sampling methods. Furthermore, our experiments show that our techniques are effective on both pixel-based and latent-based diffusion models. The paper is structured as follows. Section 2 covers background and related work on the diffusion sampling process in differential equation forms and stability region. Section 3 analyzes visual artifacts in diffusion sampling and establishes a connection to the stability region of the solver. Section 4 proposes a technique to apply momentum to existing numerical methods, as well as a technique that generalizes momentum to high-order numerical methods. Section 5 presents experiments and ablation studies. Finally, Section 6 concludes and discusses the implications and impacts of our work. Background This section first presents the theoretical foundation of diffusion sampling when modeled as an ordinary differential equation (ODE) and related numerical methods. Second, we discuss ODE forms for guided diffusion sampling and prior splitting numerical methods. Third, we cover the concept of stability region, which is our primary analysis tool.