Diffusion Sampling Correction via Approximately 10 Parameters

TL;DR

Diffusion Probabilistic Models suffer slow sampling, and training-based distillation methods add substantial cost and can disrupt interpolation between modes. PAS introduces a PCA-driven, plug-and-play correction that learns a tiny set of coordinates in a low-dimensional sampling subspace to adjust sampling directions, combined with an adaptive search that targets high-curvature regions of the trajectory. The approach achieves significant improvements (e.g., CIFAR10 FID from 15.69 to 4.37 at $\text{NFE}=10$) with roughly $4$–$12$ parameters and minutes of training, while preserving the underlying ODE trajectory. PAS demonstrates strong, dataset- and solver-agnostic gains across unconditional and conditional pre-trained DPMs, offering a practical route to fast, high-quality diffusion sampling. The work also provides theoretical and empirical insights into why diffusion trajectories lie in low-dimensional subspaces and how consistent geometric structure across samples can be leveraged for efficient correction.

Abstract

While powerful for generation, Diffusion Probabilistic Models (DPMs) face slow sampling challenges, for which various distillation-based methods have been proposed. However, they typically require significant additional training costs and model parameter storage, limiting their practicality. In this work, we propose PCA-based Adaptive Search (PAS), which optimizes existing solvers for DPMs with minimal additional costs. Specifically, we first employ PCA to obtain a few basis vectors to span the high-dimensional sampling space, which enables us to learn just a set of coordinates to correct the sampling direction; furthermore, based on the observation that the cumulative truncation error exhibits an ``S"-shape, we design an adaptive search strategy that further enhances the sampling efficiency and reduces the number of stored parameters to approximately 10. Extensive experiments demonstrate that PAS can significantly enhance existing fast solvers in a plug-and-play manner with negligible costs. E.g., on CIFAR10, PAS optimizes DDIM's FID from 15.69 to 4.37 (NFE=10) using only 12 parameters and sub-minute training on a single A100 GPU. Code is available at https://github.com/onefly123/PAS.

Figures (16)

• Figure 1: PCA-based sampling correction. We first utilize PCA to obtain a few orthogonal unit vectors that span the space of the sampling trajectories, and then learn the coordinates to correct the sampling directions in regions of large curvature along the ground truth trajectory.
• Figure 2: We utilize PCA to analyze the sampling trajectories, illustrating the trend of cumulative percent variance as the number of principal components varies. The trajectories are obtained from 1k samples using the Euler solver EulerMaruyama in the EDM EDM pre-trained model with 100 NFE. (a) The average results of each trajectory $\{x_{T}, \left \{ d_{t_{i}} \right \} _{i=N}^{1} \}$. (b) The results of $K$ trajectories $\left \{ \{ x_{t_{i}}^{k} \} _{i=N}^{0} \right \}_{k=1}^{K}$ (FFHQ and ImageNet curves nearly overlap).
• Figure 3: The truncation errors are evaluated using the Euler solver both with and without the proposed PAS. We utilize the EDM EDM pre-trained model to sample 10k samples and compute the average $L_2$ distance of 10 NFE compared to the ground truth trajectory (100 NFE). (a) The "S"-shaped truncation error is produced by the Euler solver. (b) The truncation error is corrected using PAS. Notably, PAS adaptively corrects only the parts of the sampling trajectory with large curvature.
• Figure 4: Illustration of PCA-based Adaptive Search (PAS). We demonstrate the Euler solver with the proposed PAS method, where sampling directions are derived from the tangent direction of the ground truth trajectory. The details of the correction process for the sampling direction are presented in \ref{['fig:pca_process']}.
• Figure 5: Visualization results using DDIM with and without the proposed PAS. Left: Sampling results on Stable Diffusion v1.4 with a guidance scale of 7.5. Right: Sampling results on the CIFAR10, FFHQ 64$\times$64, ImageNet 64$\times$64, and LSUN Bedroom 256$\times$256 datasets.