Linear Combination of Saved Checkpoints Makes Consistency and Diffusion Models Better

TL;DR

This work tackles the underutilization of intermediate model checkpoints in diffusion and consistency models by introducing Linear Combination of Saved Checkpoints (LCSC), a gradient-free method that searches for optimal linear coefficients to merge saved weights along training. The authors provide theoretical analysis showing EMA is suboptimal for merging coefficients and demonstrate empirically that carefully merged checkpoints can reach high-quality basins inaccessible to standard SGD optimization, enabling substantial training-time savings and improved final generation quality. Across diffusion and consistency modeling, including LoRA-augmented and text-to-image scenarios, LCSC yields major speedups (up to ~23× in CM training and reduced NFE in DM) and enhances both pre-trained and converged models. The results underscore LCSC's practicality, generalizability across metrics and data, and potential for broader application beyond the studied diffusion tasks.

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 https://github.com/imagination-research/LCSC.

Figures (22)

• Figure 1: Comparison of LCSC and vanilla training. LCSC achieves more than 14$\times$ training speed up on CIFAR-10 with Consistency Distillation (CD) and more than 15$\times$ training speed up compared to official released model on ImageNet-64 with CD. LCSC can also enhance the final converged model significantly and achieves 1.7$\times$ inference speedup for DM.
• Figure 2: Comparison between the images generated by LCSC and vanilla training (on LCM-LoRA model lcmlcm_lora). The prompts for the left and right images are "A car that seems to be parked illegally behind a legally parked car" and "Two women waiting at a bench next to a street" respectively. The vanilla method produces images with unrealistic front wheels (left) and an unnatural sitting posture (right), whereas LCSC produces more realistic images.
• Figure 3: A schematic diagram of LCSC. Given a set of checkpoints from training (left), LCSC use evolutionary search to find the optimal linear combination coefficients (middle). LCSC can be applied on checkpoints from a training process with fewer training iterations or batch sizes and still gets similar performance (abbreviated as "Perf."), thus reducing training cost and enabling faster training. LCSC can also enhance the final model in terms of generation quality or speed (right).
• Figure 4: The metric landscape of DM and CM. Selected checkpoints $\theta_{n_0}$, $\theta_{n_1}$, and $\theta_{n_2}$ are aligned sequentially along the training trajectory, with $n_0<n_1<n_2$ denoting the progression in the number of training iterations. The origin point $(0,0)$ corresponds to the checkpoint $\theta_{n_1}$, while the X and Y axes quantify the differences between $\theta_{n_1}-\theta_{n_0}$ and $\theta_{n_2}-\theta_{n_0}$, respectively. A weight located at coordinate $(x,y)$ is formulated as $\theta_{(x,y)} = \theta_{n_0} + x(\theta_{n_1}-\theta_{n_0}) + y(\theta_{n_2}-\theta_{n_0})$. Additional visualizations are provided in \ref{['app:landscape']}.
• Figure 5: Visualization of weight combination coefficients obtained using LCSC and EMA.

Theorems & Definitions (4)

• Theorem 3.1: pmlr-v28-shamir13
• Theorem 3.2
• Theorem 3.3
• Lemma A.1