A Unified Sampling Framework for Solver Searching of Diffusion Probabilistic Models

W2: The USF isn’t really something new. It is just a “dumb” aggregation of a bunch of design options. It is only based on empirical evidence — but I wonder if it can be theoretically justified.

A: We would like to share our perspectives on the contribution of USF.

• Our work indeed draws inspiration from existing diffusion sampler methods. However, these methods are presented in very different ways, even when they are working on the same solver component and proposing similar strategies. The USF framework is the first to comprehensively summarize all components of diffusion samplers. It can unify all existing exponential integral-based solvers like DEIS (in the $\lambda$ domain), DPM-Solver, DPM-Solver++, and UniPC. In our view, a unified framework can provide guidance for future implementation and research.
• For instance, a unified framework can make the extension and combination of existing strategies much more straightforward. In fact, thanks to the USF framework, our work managed to propose and implement quite a few new strategies, like low-order derivative estimation, more types of scaling magnitude, and searchable time steps. In addition, USF also offers flexibility in incorporating other strategies like arbitrary prediction types (See the results in Tab.19 for a simple attempt at interpolation between data and noise prediction), free choice of corrector, and using arbitrary off-the-shelf numerical derivative estimators.
• In addition to the newly proposed single-step strategies mentioned above, USF also advocate the utilization of various strategies at different timesteps.
• To the best of our knowledge: (1) No existing work has systematically discussed the relationship between existing samplers. (2) The new strategies we proposed have not been used. (3) The idea of using different solver strategies at different timesteps is not fully explored either.

For theoretical justification, we provide a discussion of the accuracy order of solution at App I to ensure the convergence of USF. It is worth mentioning that the performance of a solver is not only decided by accuracy order but also influenced by the coefficient of the high order error term $\mathcal{O}(h^{p})$. However, this coefficient changes with the trajectory value and function evaluation value, making it difficult to provide theoretical analysis other than the convergence order. This fact is an important motivation for us to propose searching optimal solver schedules.

We're not sure we understand the concern on "theoretical justification" correctly. Besides convergence order analysis, what types of theoretical justification do the reviewer think are helpful for this work? We'll appreciate further suggestions on this point.

Q1, Q2 & Q3: it is not clear to me which part of the paper is being claimed to be the primary contribution — the search algorithm S3 itself, or the optimal solver schedules and corresponding FID numbers ? Is the paper supposed to be about the “search algorithm S3” or an empirical AutoML report ?

A: Our primary contribution lies both in the framework USF and the search method S3. We have discussed the novelty and contribution of USF in the above response. And here, we'd like to expand on the contribution of the S3 search method.

As we've discussed in our paper, the idea of multi-stage predictor-based search is not new in the AutoML literature, e.g., [1]. What we want to claim as our contribution is that we propose an actionable and relatively efficient search method. The effectiveness of S3 has been validated by our extensive experiments, showing the possibility of sampling with very few NFEs without retraining the diffusion U-Net. And the search cost of S3 is moderate as we have analyzed in the public comment and App F.5 in the revision. This means S3 can be practically used for solver optimization for new tasks.

Another on-the-side contribution is that our discovered schedules can be used directly, and empirical knowledge based on the search results can help guide future solver design.

According to your concern, we have adjusted the summarization of contributions in the revision as follows.

1. (Primary) Design of the USF, which unifies existing solver methods and provides flexible extension space.
2. (Primary) Practical search method for solver schedule optimization, whose effectiveness is validated by extensive experiments.
3. The knowledge analysis based on the search results and the searched solvers that can be directly used.

[1] GATES: A Generic Graph-based Neural Architecture Encoding Scheme for Predictor-based NAS, Ning et al., ECCV 2020

Q4: What exactly the paper “concludes” at the end the day — did you find any “pattern” in the optimal solver schedules that may motivate further work? I don’t see the exact configuration of the optimal solver schedules (found by S3) — are they written anywhere? You can certainly report the ones with low NFE. This is important for verifying them later on by others.

A: Thanks for your reminder! We have already demonstrate about some of our searched solver schedules in Sec G.5. More general patterns have been discussed in our response to W1.2. Additionally, we upload our code and these schedules to https://anonymous.4open.science/r/USF-anonymous-code-56E6 for reproduction. Due to the time limitation, we didn't have enough time to clean our code and systematically arrange more searched schedules. We will rearrange and open source our code together with all searched solver schedules in the future.

(Continued from the Previous Page)

Then we analyze the pattern of the searched schedules.

• Timestep: Current solvers use "logSNR uniform" for low-resolution datasets and "time uniform" for high-resolution datasets. (1) We find that for low-resolution datasets, more time points should be placed when $t$ is small. (2) However, our observations reveal that the default "logSNR uniform" excessively emphasizes very small and large timesteps. According to our observation, we suggest putting more points at $0.2<t<0.5$ rather than putting too much at $t<0.05$. (3) For high-resolution datasets, we recommend a slightly smaller step size at $0.35<t<0.75$ on top of the uniform timestep schedule.
• Prediction Type: (1) We find that the data prediction model outperforms the noise prediction model by a large margin on low-resolution datasets, especially when the budget is very low (i.e., 4 or 5). (2) Current solvers commonly apply the data prediction model on all datasets. But on large resolution datasets, we find noise prediction often outperforms data prediction, opposite to current methods. (3) When sampling with guidance in pixel space, noise prediction with dynamic thresholding can also outperform data prediction under very low NFE budgets. (4) Moreover, we find that noise prediction is not suitable when using a high Taylor order.
• Derivative estimation methods: Existing solvers use equal order for Taylor expansion and derivative estimation. USF proposes to decouple the Taylor order and derivative estimation order and apply a lower order for the latter. We find that a low-order derivative estimation is preferred for high-order Taylor expansion when the step size is very large. Oppositely, full-order estimation is more likely to outperform low-order estimation when the step size is not so large.
• Sample-wise Consistency: We find that the performance of solver strategies remains consistent among data samples. We use a very small number of images (1k in our main experiments, 500 and 250 in our ablation studies) to evaluate the solver schedule during the search, while using the standard setting for the final evaluation (i.e., 10k for ImageNet-256 and MS-COCO, 50k for other datasets). The performance ranking of solver strategies is highly consistent between evaluating with a few images and many images. Therefore, we can directly apply the solver strategies discovered with a few images to more samples.

We update Sec 5.3 of the revision to discuss these observations.

W1.1: While I agree that figuring out good FID numbers for existing benchmarks is pretty good, it is not usually preffered to lock the contribution to very specific models/dataset/solver. If a new dataset/model/solver comes out in future, one need to perform the entire search again.

A: Thanks for raising this valuable concern. For new models, our experimental results in Tab.21 show that the searched schedules of S3 have the potential to be transferred across models.

As for new tasks, we would like to highlight the efficiency of our search method, S3 (please refer to the analysis of search cost in the global response). The efficiency of S3 makes it feasible to employ S3 for new tasks.

Specifically, on low-resolution datasets like CIFAR-10, S3 is about 100 $\times$ more efficient than training-based methods like Consistency Models [1]. And on high-resolution datasets like LSUN-Bedroom, S3 is around 700 $\times$ more efficient. Compared with a concurrent work of efficient diffusion sampler, DPM-Solver-v3 [2], the search cost of S3 is 4 $\times$ smaller and can achieve superior performance in the meantime on CIFAR-10.

Moreover, our experiments find that the search cost of S3 can be easily decreased by evaluating fewer schedules with fewer images without significant performance dropping, according to the results in Tab.3 and Tab.16.

[1] Consistency Model, Song et al., ICML 2023

[2] DPM-Solver-v3: Improved Diffusion ODE Solver with Empirical Model Statistics, Zheng et al., NeurIPS 2023

W1.2: The major weakness I see is that the paper can be called “just a search” without any solid conclusion.

A: We conduct detailed analysis based on the search results and give some knowledge that can offer guidance for designing solver schedules in the future.

We first investigate the contributions of each component to the final results. Specifically, we fix other components and search only for one component. We run the evolutionary search for a few steps and report the results. Below are the search results on CIFAR-10 and LSUN-Bedroom ("LOE" stands for "low order estimation" of derivatives). More details can be found in App C in the revision.

The results show that: (1) There is room for improvement across all components; (2) The timestep schedule has the largest influence on the performance; (3) The order schedule is the second most crucial component. According to these observations, we know that (1) Future work might get higher performance gains by focusing on refining the timesteps and order schedule. (2) When applying S3, if the search budget is tight, one can conduct the search on a search space only encompassing a subset of components.

Dear reviewer qhfh,

Thanks once again for your valuable time and helpful suggestions. Did we address the concerns satisfactorily? If you have any further comments, please do not hesitate to tell us. We are more than willing to provide further clarifications.

W1: The caption of Table 3 should be improved, it’s hard to figure our the meaning.

A: Tab.3 shows the FID results on the MS-COCO dataset. "Ours" stands for our method with default setting. "Ours-500" and "Ours-250" stand for generating only 500 and 250 images when evaluating solver schedules, which are demonstrated to show the performance of our method with a tighter search cost.

Thanks for your questions. We have improved this caption in our revision for better understanding.

W2：The search budget should be analysed more since it is very important to implement to previous diffusion models.

A: Thank you for raising the need for more discussion on this valuable question. We discuss the search budget thoroughly in the public comment above. Please refer to it.

W3: Why not demonstrated more categories of generated images when NFE is less? A fuller qualitative analysis would add strength on the USF.

A: Thank you for your suggestions! We add more qualitative results in App J of our revision. Please refer to it.

W4: How about the reproducibility of USF? Can you share the code? The open source code will help researchers continue to improve slow sampling speed since USF has great performance.

A: Thank you for your advice! Due to the time limitation, we can only upload the unclean code and some of our searched solver schedules to https://anonymous.4open.science/r/USF-anonymous-code-56E6. We will surely open source our code and upload all searched solver schedules in the future after carefully rearranging them.

Dear Reviewer HfC4,

Thanks once again for your valuable time and helpful suggestions. Did we address the concerns satisfactorily? If you have any further comments, please do not hesitate to tell us. We are more than willing to provide further clarifications.

W1:The experiment part seems to lack exact data on the computational cost of the overall pipeline of proposed methods. The overall pipeline needs iterative sampling images under a given sampling configuration, evaluating the performance and training of the proposed predictor models.

A: Thank you for pointing out our oversight of inadequate discussion of this question. We have discussed the search overhead thoroughly in the global response above. Please refer to it.

W2:The content of section 4.2 is not well-organized and needs further explanation.

A: Thank you for your question. Sec 4.2 describes the workflow of our predictor-based multi-stage search algorithm, S3. We first initialize the population (a set of truly evaluated solver schedules) with several baseline schedules under all NFE budgets. Then we evolutionary sample a bunch of solver schedules from the search space and truly evaluate them to expand the initial population. Then we iteratively conduct the following 3 steps: 1. Training predictor with all schedules in the population; 2. Using the predictor to sample a new bunch of promising solver schedules from the search space; 3. Evaluating these new schedules to expand the population. After all iterations end, we finally choose the schedule with the best performance in the population as our search results. This method works because of the efficiency of evolutionary search and the acceleration brought by the predictor for exploring the search space.

We have reconstructed Sec 4.2 in the updated paper for better presentation.

Q:Does the predictor model need to be trained separately on different NFEs?

A: The predictor doesn't need to be trained on different NFEs. The predictor is trained on the whole solver schedule population, which contains schedules under all NFE budgets. We pad all schedules to the same length for the predictor. The predictor is used to give a predicted score of any solver schedule with any NFE in the range.

Dear Reviewer VK39,

Thanks once again for your valuable time and helpful suggestions. Did we address the concerns satisfactorily? If you have any further comments, please do not hesitate to tell us. We are more than willing to provide further clarifications.

W3 & Q2: This paper may lack the results when the NFE is larger than 10. The gap between the baselines and the proposed method may decrease with larger NFE.

A: Thanks for pointing out the need for larger NFE results. We report the FID result of baselines and our searched schedules below, where Baseline-B represents the baseline with the best setting and Baseline-W represents for the worst (see Sec 5 and Sec G.1 for details). We can see that our method still outperforms all baselines by a large margin at larger NFEs like 12, 15, and 20. Our method completely converges at NFE=15, much faster than existing solvers.

We would like to emphasize that when the NFE budget is adequate, the negative impact of sub-optimal empirical strategies diminishes. This fact can be verified by the decreasing gap between Baseline-W and Baseline-B. So, it is very reasonable that the gap between our method and baselines decreases. Considering this fact, we choose a more challenging and meaningful scenario to optimize, i.e., sampling under very few NFE budgets (4-10).

Q4: Since the USF can include the method DEIS in the domain $t$, why do you not include this method as a baseline in experiments?

A: For a more comprehensive comparison, we report the results of DEIS directly from their paper [3] and the results of our method below. We don't choose DEIS as a baseline in our paper because UniPC and DPM-Solver++ have already compared DEIS with their methods in their papers. DEIS fails to outperform DPM-Solver++ or UniPC in most cases.

• [3] FAST SAMPLING OF DIFFUSION MODELS WITH EXPONENTIAL INTEGRATOR, Zhang & Chen, ICLR 2023

Q3: Can you give a brief explanation of the meaning and effect of multi-stage in Algorithm 2?

A: Our main idea of S3 is to train a predictor to efficiently predict the performance of a solver schedule rather than evaluate its true performance by sampling a large number of images.

The meaning of "multi-stage": we iteratively conduct the following process with each iteration called a "stage": 1. Training predictor with all truly evaluated solver schedules in the population. 2. Using predictor to explore within the search space. 3. Truly evaluating the promising solver schedules selected by the predictor to expand the population. In Alg.2, the $for$ loop starting from line 3 represents the index of stage. The overall workflow is demonstrated in Fig.4.

The effect of multi-stage search: At the early stage of the search, the predictor is only required to explore in a limited sub-space. Therefore, only a small amount of schedule-performance data pair is needed to train the predictor, and the predictor can be quickly put into use. As the search process advances, the exploration range of the predictor increases. Thus, it needed to be updated with more data to ensure the accuracy of its prediction. Compared to the single-stage search (i.e., evaluate a large number of solver schedules in one go to train the predictor and perform the predictor-based search only once), the muli-stage method makes the sampling of solver schedule for true evaluation more effective since the predictor involves in this process.

Thanks for your question! For a better presentation of the search method, we have reconstructed Sec 4.2 in the revision.

Q: Typos

A: We appreciate you very much for pointing out these typos in our paper. We have corrected these mistakes in our revision.

W1, W2 & Q1: This paper may lack motivation, novelty and contributions in the sense that it does not provide a further understanding of why suitable solvers are different on different timesteps. Can you provide an intuitive explanation of the phenomenon observed in the paper?

A: Thank you for pointing out this question! Below we discuss the intuitions of the phenomenon "suitable strategies are different at different steps". We first lay out the high-level analyses, and then expand on each component.

• High level
  • The curvature of the ODE trajectory might vary across timesteps. When $t$ is small, the neural network can reconstruct $x_0$ from $x_t$ more precisely and the ODE trajectory points to the target data. Intuitively, the curvature in the small $t$ region is small. Conversely, when $t$ is not close to 0, the ODE trajectory does not necessarily orient to the final target point and could have a large curvature [1]. Therefore, the curvature of the ODE trajectory might change over time. Some other papers have discussed similar findings. [2] demonstrates the pattern of VP-SDE in Sec.3, from which we can see that $x_t$ changes slowly at large $t$ and changes rapidly when $t$ is small; [1] finds that the trajectory orients to the mean value of the distribution at large $t$ while points to the data when $t$ is small. Different curvature leads to varying truncation errors, making certain solver strategies more suitable for specific timesteps.
    • [1] Minimizing Trajectory Curvature of ODE-based Generative Models, Lee et al., ICML 2023
    • [2] Elucidating the Design Space of Diffusion-Based Generative Models, Karras et al., NeurIPS 2022
  • The fitting error of the neural network also varies among timesteps. The distributions of input for the neural network are different among timesteps. Therefore, the fitting tasks for the diffusion U-Net have unequal difficulty at different timesteps [3], leading to different distances between the predicted score and the ground truth score. Since the discretized solving error is related to the model's fitting error, it is natural for solver strategies to change across timesteps to adapt to different fitting accuracies.
    • [3] OMS-DPM: Optimize the Model Schedule of Diffusion Probabilistic Models, Liu et al., ICML 2023
  • The requirements of solving accuracy are different at different timesteps. At larger $t$, deviant values might be corrected to the right trajectories in subsequent solving processes. Conversely, for small $t$, a high solving error could directly cause a drop in sample quality. Therefore, different timesteps might need different solving strategies.
• Specific to components: Here we discuss the analyses that motivate using different strategies for each component.
  • Step size: For low-resolution datasets, the processing of details is more important than large-resolution datasets. Therefore, the step size should be smaller near $t=0$ than near $t=T$ for low-resolution datasets. This pattern has been adopted empirically for low-resolution datasets in existing solvers like DDIM and DPM-Solver, which motivates us to use different step sizes and search for better schedules than empirical patterns.
  • Order: Existing high-order solvers (e.g., DPM-Solver, DPM-Solver++, and UniPC) empirically use low orders for the last several steps when the NFE budget is tight, indicating that the suitable orders for different timesteps may not be the same. We are inspired by this setting to use changeable orders at all timesteps for better performance.
  • Prediction Type: The output of the data prediction model changes rapidly when $t$ is large and changes slowly when $t$ is small since the data ratio of $x_t$ changes from 0 to 1 along timesteps. For the noise prediction model, the trend is the opposite. The change speed of model output has a close relationship with the stability of high-order derivatives. Therefore, the proper prediction types at different timesteps might be different.
  • Corrector: The ODE corrector can be viewed as a special ODE predictor with the involvement of the function evaluation $f_\theta(x_t, t)$ on the target point $t$. Thus, the above analysis also holds for corrector. We find that different corrector patterns are used by concurrent work [1] empirically to boost its results.
    • [1] DPM-Solver-v3: Improved Diffusion ODE Solver with Empirical Model Statistics, Zheng et al., NeurIPS 2023
  • The choices of different components are coupled with each other. For example, low Taylor order and low derivative estimation order are intuitively suitable for a large step size; prediction type affects the choice of derivative scaling method. In addition, the choice of solver strategy at one step is impacted by the solution accuracy of previous steps. Therefore, we apply different strategies for all components at different timesteps and search for the optimal schedule.

Dear Reviewer HHFa,

Thanks once again for your valuable time and helpful suggestions. Did we address the concerns satisfactorily? If you have any further comments, please do not hesitate to tell us. We are more than willing to provide further clarifications.

W & Q2: Writing Style

A: We sincerely appreciate you for pointing out our typos. We have corrected them in the revision.

Q1: Can you briefly tell me the differences between this paper and Lu et al's DPM-solver++, since both of them seem to use different strategies as multi-step solver? What kind of improvements have you done in your paper compared with DPM Solver++?

A: The main contribution of DPM-Solver++ is the introduction of two new strategies: 1. multi-step solver and 2. data prediction model. Except for the first and last few steps that use a low-order solver, DPM-Solver++ uses the same solver strategies along timesteps.

USF comprehensively summarizes tunable components of all diffusion solvers based on the exponential integral. The two new strategies of DPM-Solver++ are also incorporated in USF. See Tab.1 and App E.2 for details. Moreover, USF introduces different solver strategies among all timesteps, expanding the design space. With searched solver schedules, USF outperforms DPM-Solver++ by a large margin. In conclusion, DPM-Solver++ is a special case of USF, and USF is far more effective and flexible than DPM-Solver++.

Dear Reviewer 7SgP,

Thanks once again for your valuable time and helpful suggestions. Did we address the concerns satisfactorily? If you have any further comments, please do not hesitate to tell us. We are more than willing to provide further clarifications.