Diffusion Sampling Correction via Approximately 10 Parameters

We sincerely appreciate the reviewer's recognition of our work and meticulous review. Please find below our responses to all the questions. We would greatly appreciate it if you could consider increasing the score if you are satisfied with our response.

Abbreviation: CaE (Claims And Evidence), EDoA (Experimental Designs Or Analyses)

CaE and EDoA2: The study lacks a comparison of training efficiency between PAS and other methods.

A: We provide a comparison of PAS with the previously well-performing AMED and GITS methods, while also introducing the high-cost PD method as a contrast. The comparative experiments, based on sample quality and training cost on CIFAR10, are shown in the table below:

FID↓(A100 hours↓)

As shown in the table above, PAS demonstrates higher training efficiency compared to AMED, and better sample quality than GITS with negligible additional cost. In contrast, PD, which yields higher sample quality, incurs a significantly larger training cost. Notably, PAS is theoretically orthogonal to PD, AMED and GITS, and can further improve their sampling efficiency.

Theoretical Claims: Lacking theoretical justification for sampling trajectories in a low-dimensional subspace.

A: Inspired by [WV23] and [WV24], we have added more theoretical analysis to the paper (see Reviewer pbLK, Expand theory). In brief, since analyzing neural networks directly is inherently intractable, [WV23] and [WV24] approximate diffusion trajectories using Gaussian score structures, providing both theoretical and empirical insights.

Furthermore, [WV23] offers a theoretical analysis in Sec. 4.2, suggesting that diffusion trajectories resemble 2D rotations on the plane spanned by the initial noise ($x_T$) and the final sample ($x_0$). This offers a theoretical justification for why sampling trajectories lie in a low-dimensional subspace. It also explains why trajectories corresponding to different initial noises lie in different subspaces (see our Fig. 2b).

EDoA1: Add a comparison with other trajectory-based methods (e.g., AMED, GITS) combined with other SOTA solvers (e.g., iPNDM).

A: We did not include additional comparisons between PAS and other trajectory-based methods (e.g., AMED, GITS) primarily because the integration of PAS with DDIM already demonstrates significant improvements in sampling efficiency compared to AMED and GITS, thereby validating the effectiveness of PAS. Furthermore, PAS is theoretically orthogonal to methods like AMED and GITS, and can potentially be combined with them to further enhance sampling efficiency. Therefore, we believe that exploring how to better combine PAS with trajectory-based methods (e.g., AMED, GITS) is a more meaningful direction, which will be considered in our future work.

Questions For Authors: Have you experimented with alternative dimensionality reduction techniques in place of PCA?

A: We have experimented with both standard PCA (torch.svd) and low-rank or sparse matrix PCA (torch.pca_lowrank). Notably, PCA_lowrank (used in our paper) computes approximately 1.58× faster than standard PCA, with some loss in accuracy. We provide a comparative experiment on the CIFAR10, as shown in the following table:

PAS+DDIM, FID↓

From the table, we can observe that PCA_lowrank and standard PCA yield comparable results. Upon further analysis, we believe this is because only around 4 principal components are sufficient to span the complete sampling trajectory space (which consists of 1000 vectors). Since extracting 4 principal components from trajectory space is relatively straightforward, the choice of dimensionality reduction technique is not a critical factor in the effectiveness of the proposed PAS method.

[R1] Salimans T et al. Progressive distillation for fast sampling of diffusion models. ICLR 2022.

[R2] Zhou, Zhenyu et al. Fast ode-based sampling for diffusion models in around 5 steps. CVPR 2024.

[R3] Defang Chen et al. On the Trajectory Regularity of ODE-based Diffusion Sampling. ICML 2024.

We sincerely appreciate the reviewer's efforts and valuable review. Below are our responses to all questions. We kindly hope you could consider increasing the score if you are satisfied.

Experimental Designs Or Analyses: Direct performance comparisons with training-based methods are lacking, highlighting PAS's advantage.

A: We appreciate your suggestion! We did not include direct performance comparisons because high-cost training-based methods and PAS have their own application scenarios. Training-based methods can achieve one-step sampling but are expensive, making them difficult to apply to large models like Stable Diffusion. In contrast, PAS focuses on high-quality sampling under 10 NFE with negligible cost, making it a more attractive and practical solution. Furthermore, we added a comparison of parameter size and training time: Training-based one-step sampling methods require at least 35.7M parameters (DDPM, CIFAR10 model), while PAS uses only ~10. The table below summarizes training time↓ (A100 hours) on CIFAR10, using the estimation method and certain results from [R1].

Note: Training a CIFAR10 model from scratch takes ~200 A100 hours [R1].

W1: Theoretical explanation of why the sampling trajectories of all samples exhibit strongly consistent geometric characteristics.

A: Inspired by [WV23] and [WV24], we have added more theoretical analysis to the paper (Reviewer pbLK, Expand theory). Briefly, since analyzing neural networks directly is inherently intractable, [WV23] and [WV24] approximate diffusion trajectories using Gaussian score structures, providing both theoretical and empirical insights.

Specifically, [WV24] derives an analytical form of the EDM trajectory in Eq. 15(Gaussian structure):

$x_t=\mu+\frac{\sigma_t}{\sigma_T} x_T^\perp +\sum_{k=1}^r \psi (t,\lambda_k)c_k(T)u_k$,

which shows that $x_t$ is a linear combination of $x_T^\perp$ and $\mu, u_k$. Here, $\mu, u_k$ are the dataset-dependent mean and basis vectors that define the data manifold, while $x_T^\perp$ is an off-manifold component. Furthermore, given $x_T$, $c_k(T)$ is a constant, and the rate of change of the coefficients for the vectors $x_T^\perp, \mu$, and $u_k$ depends only on the dataset and the timestep $t$. For a fixed dataset, as $\sigma_t$ decreases during the sampling process, each sample is constrained by the same temporal decay scale, converging to the data manifold. This theory heuristically explains why all samples exhibit strongly consistent geometric characteristics.

W2 and Q2: Discuss the limitations of the reliance on PCA and the robustness of PAS to stochastic samplers.

A: The effectiveness of PAS relies on the accuracy of the basis derived from PCA, which may be influenced by noise from stochastic samplers. To further explore this, we added experiments using PAS on the stochastic sampler of EDM [R6] without the 2-order correction ([R6] Algorithm 2) on CIFAR10, with the following results:

FID↓

The table clearly shows that PAS is equally effective for the stochastic sampler, further highlighting the robustness of PAS. Nonetheless, in general, SDE has a higher quality upper bound, with 1000 NFE outperforming ODE. But for sampling with fewer steps, ODE performs better. Therefore, PAS with ODE would result in faster sampling speeds.

Q1: A comparison of sampling quality and computational trade-offs with training-based methods like PD[R2].

A: We have added some comparative results on the training cost and sample quality between PD and PAS, as shown in the following tables:

FID↓(A100 hours↓)

As shown in the tables above, PD requires ~146 A100 hours of training time for 8 NFE to achieve 2.57 FID, whereas PAS with iPNDM requires only 0.04 A100 hours with 8 parameters to achieve 2.84 FID at 10 NFE. Therefore, PD and PAS each have their own application scenarios. PAS achieves impressive performance with negligible cost, making it a more attractive and practical solution.

[R1] Zhou Z et al. Simple and fast distillation of diffusion models. NeurIPS 2024.

[R2] Salimans T et al. Progressive distillation for fast sampling of diffusion models. ICLR 2022.

[R3] Meng C et al. On distillation of guided diffusion models. CVPR 2023.

[R4] Song Y et al. Consistency models. ICML 2023.

[R5] Kim D et al. Consistency trajectory models: Learning probability flow ode trajectory of diffusion. arXiv:2310.02279.

[R6] Tero Karras et al. Elucidating the design space of diffusion-based generative models. NeurIPS 2022.

We sincerely appreciate the reviewer's detailed review and insightful suggestions!

Abbreviation: CaE (Claims And Evidence), MaEC (Methods And Evaluation Criteria), ERND (Essential References Not Discussed)

CaE1: The later trajectory may not be linear but instead takes smaller steps.

A: Your interpretation is reasonable. However, [KMTS22] and [WV23] only analyze the trajectory in 2D, which may lose some information. [R1] visualizes trajectory in 3D ([R1] Fig. 4) and offers a "boomerang" explanation (linear-nonlinear-linear) in Sec. 3, showing that the later part is indeed linear. We have added a detailed analysis in the revised version.

[R1] Chen D et al. On the trajectory regularity of ode-based diffusion sampling. ICML2024.

CaE2: The degradation without adaptive search may be due to the obtained PCs not being good.

A: Notably, this paper has conducted a similar analysis in lines 377-384(Right): "The errors in the linear part are negligible; the correction does not further reduce errors and instead introduces noise", highlighting the necessity and effectiveness of the adaptive search.

Expand theory (MaEC1, Theoretical Claims, and ERND_M1): The study misses citing [WV23] and [WV24], which could theoretically motivate Fig. 2b and the correction of the more curved parts(Fig. 1).

A: We have carefully revisited [WV23] and [WV24], which approximate diffusion trajectories by Gaussian score structure and provide both theoretical and empirical insights. These works are indeed inspiring and insightful!

Specifically, [WV23] offers a theoretical analysis in Sec. 4.2, suggesting that diffusion trajectories resemble 2D rotations on the plane spanned by the initial noise ($x_T$) and the final sample ($x_0$). This provides a theoretical underpinning for our Fig. 2b, where trajectories corresponding to different initial noises lie in different subspaces.

Correction of the more curved parts: [WV24] derives an analytical form of the EDM trajectory in Eq. 15 (Gaussian structure):

$x_t=\mu+\frac{\sigma_t}{\sigma_T} x_T^\perp +\sum_{k=1}^r \psi (t,\lambda_k)c_k(T)u_k$,

which shows that $x_t$ is a linear combination of $x_T^\perp$ and $\mu,u_k$. PAS is designed to learn statistical patterns of the dataset that are independent of the initial noise $x_T$, i.e., the coefficients of $\mu,u_k$. As $\sigma_t \to 0$, the influence of $x_T^\perp$ on $x_t$ diminishes. So, in the low-noise regime (i.e., more curved parts), PAS batter captures the PCs of the data ($\mu,u_k$), leading to more accurate correction. This offers a theoretical explanation for why PAS tends to work better in the more curved regions (Fig. 1). We have added relevant discussions in Sec. 3 and the Related Work of the revised manuscript. The alignment between PAS and the conclusions of [WV23], [WV24] is indeed intriguing!

MaEC2: Discuss the time taken to obtain the sampling trajectory for training (similar to [WV24]).

A: We did not discuss the cost because other training-based methods also require sampling full or partial trajectory for training. In the revised version, we added a discussion on the time taken: 3.26m for CIFAR10 and 1.79h for Bedroom256 with 50 NFE and a 5k trajectory on an A100 GPU.

ERND_M2 and ERND_R1-3: Discuss the connections and differences between PAS and [WV23], [WV24], and include them as a benchmark.

A: The proposed PAS corrects sampling in the low-noise regime, while [WV24] accelerates sampling in the high-noise regime. So, combining [WV24] with PAS is indeed interesting. Theoretically, we can analytically warm up the PCs from the Gaussian structure and then correct the sampling starting from the teleported solution $x_{t'}$, further improving the FID in [WV24] Fig. 15. However, [WV24] is limited by the computational cost of estimating means and covariances on large pixel datasets. Furthermore, another key challenge is how to effectively combine the Gaussian (analytical but imprecise under low noise) and neural (precise but costly) structure to obtain more exact PCs with minimal cost. This will be a focus of our future work.

ERND_R4: The PAS is hard to apply to very high-res images.

A: No, the PAS can be easily applied to high-res images. This is because we use torch.pca_lowrank to decompose the trajectory matrix as $\mathbf{X}^{\prime}\in \mathbb{R}^{(N-i+2)\times D}$ (Eq. 13), where $i\leq N,N\leq 10$, so its dimension is low and the rank $r\ll D$. For Stable Diffusion(SD), $D$=64x64x4 from the latent space. In the table below, we compare the time for 1PCA and 1NFE, and the PCA time is negligible.

pre 128 samples time↓ (s)

Other Suggestions: Table1 caption needs clarification.

A: Thanks for pointing out the issues; we have revised them. Briefly, the list of numbers denotes all the time points requiring correction from adaptive search, with each element $i\in [N,1]$.