Linear Combination of Saved Checkpoints Makes Consistency and Diffusion Models Better

Fuente: arXiv
Salvato in:
Dettagli Bibliografici
Autori principali: Liu, Enshu, Zhu, Junyi, Lin, Zinan, Ning, Xuefei, Wang, Shuaiqi, Blaschko, Matthew B., Yekhanin, Sergey, Yan, Shengen, Dai, Guohao, Yang, Huazhong, Wang, Yu
Natura: Preprint
Pubblicazione: 2024
Soggetti:
Accesso online:
Tags: Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
_version_ 1866912246174580736
author Liu, Enshu
Zhu, Junyi
Lin, Zinan
Ning, Xuefei
Wang, Shuaiqi
Blaschko, Matthew B.
Yekhanin, Sergey
Yan, Shengen
Dai, Guohao
Yang, Huazhong
Wang, Yu
author_facet Liu, Enshu
Zhu, Junyi
Lin, Zinan
Ning, Xuefei
Wang, Shuaiqi
Blaschko, Matthew B.
Yekhanin, Sergey
Yan, Shengen
Dai, Guohao
Yang, Huazhong
Wang, Yu
contents 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 https://github.com/imagination-research/LCSC.
format Preprint
id arxiv_https___arxiv_org_abs_2404_02241
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Linear Combination of Saved Checkpoints Makes Consistency and Diffusion Models Better
Liu, Enshu
Zhu, Junyi
Lin, Zinan
Ning, Xuefei
Wang, Shuaiqi
Blaschko, Matthew B.
Yekhanin, Sergey
Yan, Shengen
Dai, Guohao
Yang, Huazhong
Wang, Yu
Computer Vision and Pattern Recognition
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 https://github.com/imagination-research/LCSC.
title Linear Combination of Saved Checkpoints Makes Consistency and Diffusion Models Better
topic Computer Vision and Pattern Recognition
url https://arxiv.org/abs/2404.02241