Linear Combination of Saved Checkpoints Makes Consistency and Diffusion Models Better

Enshu Liu,Junyi Zhu,Zinan Lin,Xuefei Ning,Matthew B. Blaschko,Sergey Yekhanin,Shengen Yan,Guohao Dai,Huazhong Yang,Yu Wang

2024-04-08

Abstract:Diffusion Models (DM) and Consistency Models (CM) are two types of popular generative models with good generation quality on various tasks. When training DM and CM, intermediate weight checkpoints are not fully utilized and only the last converged checkpoint is used. In this work, we find that high-quality model weights often lie in a basin which cannot be reached by SGD but can be obtained by proper checkpoint averaging. Based on these observations, we propose LCSC, a simple but effective and efficient method to enhance the performance of DM and CM, by combining checkpoints along the training trajectory with coefficients deduced from evolutionary search. We demonstrate the value of LCSC through two use cases: $\textbf{(a) Reducing training cost.}$ With LCSC, we only need to train DM/CM with fewer number of iterations and/or lower batch sizes to obtain comparable sample quality with the fully trained model. For example, LCSC achieves considerable training speedups for CM (23$\times$ on CIFAR-10 and 15$\times$ on ImageNet-64). $\textbf{(b) Enhancing pre-trained models.}$ Assuming full training is already done, LCSC can further improve the generation quality or speed of the final converged models. For example, LCSC achieves better performance using 1 number of function evaluation (NFE) than the base model with 2 NFE on consistency distillation, and decreases the NFE of DM from 15 to 9 while maintaining the generation quality on CIFAR-10. Our code is available at

Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

This paper mainly explores the optimization problems of diffusion models (DM) and consistency models (CM) during the training process. The authors found that in the training process of DM and CM, the saved intermediate weight checkpoints are not fully utilized, and usually only the last converged checkpoint is used. The research found that high-quality model weights are often located in areas that SGD cannot reach, but can be obtained through proper linear combinations of checkpoints. Based on these observations, the paper proposes the LCSC (Linear Combination of Saved Checkpoints) method to determine the optimal weight combination through evolutionary search algorithms to improve the performance of DM and CM.

LCSC has two main applications:
1. Reduce training costs: By using LCSC, the model can be trained with fewer iterations or smaller batch sizes without sacrificing sample quality. For example, for consistency distillation (CD), a 23-fold training speed improvement can be achieved on CIFAR-10 and a 15-fold improvement on ImageNet-64.
2. Improve pre-trained models: Assuming training has been completed, LCSC can further improve the generation quality and speed of the final model. For example, in consistency training (CT), LCSC can reduce the number of function evaluations (NFE) of DM while maintaining quality.

In addition, the paper also points out that the commonly used exponential moving average (EMA) method may not be optimal in determining weight combinations, which emphasizes the importance of exploring better weight combination methods. LCSC does not require backpropagation and can be quickly optimized in low-dimensional space, making it suitable for metrics that are difficult to compute gradients and may be applicable to other tasks and neural networks.

In conclusion, the main contributions of the paper include analyzing the training process of DM and CM, proposing a simple and effective LCSC method to enhance model performance, and demonstrating the effectiveness of LCSC in reducing training costs and improving pre-trained models.