Efficient Diffusion Models via Time Step Optimization with Consistent Training and Inference Constraints

Diffusion probabilistic models (DPMs)’ sampling process is often inefficient, requiring hundreds to thousands of iterative steps to accurately approximate the diffusion trajectory. This inefficiency limits their practical applicability. Although recent advances in sampling efficiency—such as numerical solvers for diffusion ordinary differential equations (ODEs)—have made progress, significant challenges remain: training-free numerical solvers suffer from the suboptimality of manually designed timestep selection rules and the inherent inconsistency between the forward diffusion process (typically involving thousands of steps) and the reverse denoising process (usually limited to tens of steps). Since timestep selection is inherently a discrete problem and cannot be optimized via gradients, we propose an innovative approach—reparameterizing the timestep scheduling through probabilistic masking, thereby enabling gradient-based optimization of sampling timesteps. To circumvent backpropagation, we employ policy gradient methods. Furthermore, to address the inconsistency between forward diffusion (training) and reverse denoising (inference), we extend this framework into a bilevel optimization paradigm: the inner loop performs additional lightweight training on the model at specific timesteps determined by the outer mask to align forward and reverse processes, while the outer loop optimizes the timestep distribution via probabilistic masking and policy gradient based on generation quality. Under mild assumptions, we theoretically analyze the convergence of the proposed algorithm. Extensive experiments across diverse datasets and samplers demonstrate that this framework effectively enhances sampling efficiency and generation quality while maintaining compatibility with various DPM architectures and advanced ODE solvers.

Diffusion models are powerful tools for generating images and other data, but they often need many repeated steps to produce high-quality results. This makes them slow and costly to use in practice. Existing methods can reduce the number of steps, but they usually rely on hand-designed schedules that may not be optimal. In addition, these models are typically trained with many small steps, while they are often used with far fewer steps during generation, creating a mismatch that can hurt quality. In this paper, we propose a method that automatically learns which generation steps are most useful. Instead of manually choosing when the model should perform each step, our approach learns a probability-based schedule that favors the most important steps. We also lightly adapt the model to the selected steps, so that the way it is used during generation better matches how it behaves after training. This helps the model generate good results with fewer steps. We provide theoretical analysis showing that the proposed learning procedure can converge under reasonable conditions. Experiments on multiple datasets and generation methods show that our approach improves both speed and output quality. Importantly, the method can be used with different diffusion model designs and existing fast generation techniques, making it broadly applicable.