Paper: A Unified Sampling Framework for Solver Searching of Diffusion Probabilistic Models (ICLR-2024)

Reviewer comment (Novelty):
The major weakness I see is that the paper can be called “just a search” without any solid conclusion. While I agree that figuring out good FID numbers for existing benchmarks is pretty good, it is not usually preferred to lock the contribution to very specific models/dataset/solver. If a new dataset/model/solver comes out in future, one needs to perform the entire search again.

Author response:
**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.

|  |  | CIFAR10 |  |  | Bedroom |  |
| --- | --- | --- | --- | --- | --- | --- |
| NFE | 5 | 7 | 10 | 5 | 7 | 10 |
| Baseline | 23.44 | 6.47 | 3.90 | 33.92 | 25.94 | 23.65 |
| Timestep | **10.06** | **5.36** | **2.88** | **27.23** | **22.99** | **18.64** |
| Order | 17.76 | 5.86 | 3.90 | 31.89 | 25.41 | 22.38 |
| Prediction | 23.44 | 6.47 | 3.90 | 32.43 | 24.84 | 23.25 |
| Corrector | 22.58 | 6.23 | 3.90 | 33.40 | 25.60 | 23.27 |
| Scaling | 21.67 | 5.56 | 3.90 | 32.84 | 25.60 | 23.25 |
| LOE | 23.44 | 5.91 | 3.84 | 33.92 | 25.62 | 23.52 |

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.

Reviewer comment (Novelty):
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.

Author response:
**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.

Reviewer comment (Novelty):
In high-level, 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?

Author response:
**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

Reviewer comment (Presentation):
The caption of Table 3 should be improved, it’s hard to figure our the meaning.

Author response:
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.

Reviewer comment (Experiments):
The search budget should be analysed more since it is very important to implement to previous diffusion models.

Author response:
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.

Reviewer comment (Writing):
In Equation (1), dt and dw_t should be \mathrm{d} t and \mathrm{d}w_t. Similarly, in other parts of this paper, the authors sometimes use dt and sometimes use \mathrm{d} t. The notation should be unified.

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

Reviewer comment (Theory):
In Equation (5), the \mathrm{d} t is missing in the right hand side of the equation.

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

Reviewer comment (Novelty):
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++?

Author response:
**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++.

Reviewer comment (Experiments):
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.

Author response:
**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.

Reviewer comment (Experiments):
Does the predictor model need to be trained separately on different NFEs?

Author response:
**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.

Reviewer comment (Theory):
This paper illustrates the phenomenon that the suitable strategies vary among timesteps but has no further explanation on why this phenomenon happens. The proposed method paper may lack motivation.

Author response:
**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.

Reviewer comment (Experiments):
On empirical results, 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.

Author response:
**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.

|  |  | CIFAR-10 |  |  | CelebA |  |
| --- | --- | --- | --- | --- | --- | --- |
| NFE | 12 | 15 | 20 | 12 | 15 | 20 |
| Baseline-W | 5.31 | 4.52 | 3.54 | 6.12 | 4.20 | 3.56 |
| Baseline-B | 3.67 | 3.03 | 2.80 | 3.82 | 2.72 | 2.70 |
| Ours | **2.65** | **2.41** | **2.41** | **2.32** | **2.06** | **2.06** |

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).

Reviewer comment (Theory):
Can you provide an intuitive explanation of the phenomenon observed in the paper?

Author response:
**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.

Reviewer comment (Writing):
Can you give a brief explanation of the meaning and effect of multi-stage in Algorithm 2?

Author response:
**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.

Reviewer comment (Experiments):
Since the USF can include the method DEIS in the domain $t$, why do you not include this method as a baseline in experiments?

Author response:
**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

|  |  | CIFAR-10 |  |  |  | CelebA  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| NFE | 5 | 7 | 10 | 15 | 5 | 7 | 10 | 20 |
| DEIS | 15.37 | - | 4.17 | 3.37 | 25.07 | - | 6.95 | 3.41 |
| Ours | **7.65** | **3.91** | **3.09** | **2.41** | **5.17** | **3.80** | **2.73** | **2.06** |

Reviewer comment (Writing):
There are some typos in Appendix D. In analyzing the current solvers, the index of $\tilde{x}_{i-1}$ may be incorrect.

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

Reviewer comment (Theory):
They should be $t_{i-1}$ in DPM-Solver-2, DPM-Solver++(2S) and DPM-Solver++(2M).

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

Reviewer comment (Presentation):
When using the Assumption in the Appendix, it may not be Thm A.1, A.2 etc.

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