S4S: Solving for a Fast Diffusion Model Solver

Thank you so much for the time taken to review our work and your helpful feedback! Please find our responses below:

However, the focus is largely on improvements in the few-NFE regime for training-free methods, and it remains somewhat unclear how the student solver matches the output of a teacher solver (which is supposed to run for many NFEs) as claimed in the conclusion.

You raise an important point about our objective of matching teacher outputs. I'd like to clarify this aspect of our approach. While our training objective involves matching teacher solver outputs, this is primarily a means to an end rather than the ultimate goal. Our core innovation is recognizing that in the few-NFE regime, traditional error control approaches break down, so we instead directly optimize for what matters most: high-quality final samples. We use teacher matching as a proxy for quality for several reasons, including:

• It provides a clear optimization target that avoids the need for dataset access, unlike distillation methods.
• It learns to reproduce the outcomes of many-NFE solvers without being constrained to follow their intermediate steps.
• The relaxed objective (Equation 7) further enhances this by allowing flexibility in finding solutions that prioritize final quality over exact trajectory matching.

Empirically, our results demonstrate that this approach succeeds not just at matching teacher outputs but often exceeding them at lower NFE counts. For instance, on CIFAR-10, S4S-Alt with 7 NFEs (FID 2.52) outperforms the teacher with 20 NFEs (FID 2.87). This indicates our method is learning a fundamentally more efficient sampling strategy, not merely approximating the teacher's steps. The effectiveness of our global error minimization approach suggests that, while traditional solvers accumulate errors across many small steps, our learned solvers take optimal large steps that directly target high-quality outputs. This represents a paradigm shift in diffusion model sampling that could inspire further research into learning-based ODE solvers beyond diffusion models. We'll emphasize this distinction more clearly in the final version of the paper to ensure the conceptual contribution is properly understood.

Since S4S is designed for few-NFE, while the teacher solvers are specifically for many-NFE. It would be insightful to examine how far are the student solvers away from ideal performance (many-NFE for teacher solvers).

You raise an excellent point about comparing student solvers to teacher performance. We did analyze this gap in our experiments, though we may not have emphasized it enough in the main paper due to space constraints. In our full results tables (H.4), we provide FID scores for a range of NFEs up to 10, while in H.5 we mention the FID scores of the various teacher solvers. For example, on CIFAR-10, our S4S-Alt achieves 2.18 FID at 10 NFEs, while the teacher solver (UniPC with logSNR) achieves 2.03 at 20 NFEs. For FFHQ, we see a similar pattern: S4S-Alt with 10 NFEs achieves 2.91 FID, while the teacher at 20 NFEs achieves 2.62.

The convergence of quality vs. NFE (at how many NFEs do the student solvers cease to improve) may be a useful reference of quality-speed tradeoff

This is an insightful suggestion. While we didn't include full convergence curves in the paper, our results in Tables 13-18 show the quality-speed tradeoff across 3-10 NFEs. We observe that:

• For most datasets, improvements diminish beyond 8 NFEs, with minimal gains between 8-10 NFEs.
• S4S-Alt shows more rapid convergence than traditional solvers, achieving near-optimal quality at just 5-6 NFEs on most datasets.
• Different datasets show varying convergence patterns: CIFAR-10 and AFHQ-v2 converge more quickly (plateau around 6-7 NFEs), while more complex datasets like MS-COCO and LSUN-Bedroom continue improving up to 10 NFEs.

Thank you for your constructive feedback, we address your questions and concerns, starting with two of the weaknesses you mentioned.

Lack of direct experimental comparisons in the main context

In the linked image, we provide direct comparisons against the recent works [2,3] that focus on coefficient distillation. Compared to these works, on the same discretization schedule (Time EDM), S4S achieves superior performance in the low-NFE regime across several key datasets. This helps reinforce the strength of S4S, even when only optimizing over the solver coefficients. Here, * refers to FID values that were estimated from Figures if no results table was present, and the "Rel. Fig." column details which Figure the value was taken from.

Regarding BNS, it is difficult to provide a wholesale, fair comparison to the original paper for two main reasons: (1) the authors do not release their code for learning BNS, and (2) the model checkpoints used for BNS are not publicly released. As a result, we cannot (1) port an exact implementation of BNS over onto the pretrained DMs that we used, or (2) port over our own implementation onto the pretrained (flow) models that BNS used.

To try to compare against BNS, we note that it has two significant differences from our work:

• In BNS, the order of the learned solver is equal to the number of NFEs used; that is, it assigns a coefficient to all of the previous denoising steps.
• BNS jointly learns the solver coefficients and the time discretization steps.
• BNS uses PSNR as the global objective function.

In our submission, we tried to explore the first two components of this tradeoff in the third row of Table 6, although we should have made the connection to BNS more clear. To further extend this analysis, we explicitly re-implemented the approach from BNS into our own setup to directly compute this comparison. We include this in the table below for CIFAR-10 and FFHQ.

Pathologies/errors from the teacher trajectory distilled into the student

We agree that more experiments are needed in the revision. We reference a visualization from recent work [4] that visualizes trajectories from few-step sampling at this link. At the beginning of the boomerang shape, there is greater variance in the trajectories. Therefore, training a model to explicitly match these trajectories can result in distilling these pathological behaviors directly into the student solver. To further assess this claim, we conducted several experiments to understand the consequences of using this as an objective. In the attached image, we see that training on a dataset that includes these "pathological" trajectories teaches the student solver to mimic them. In particular, conducting a full FID evaluation of a solver trained in this manner reveals worse FID performance.

Concerns about novelty

To best characterize our novelty, we compare with two types of learning diffusion solvers. First, unlike $\Pi$A or D-DDIM, which mimic the ODE trajectory, our approach directly matches the output of the teacher solver. While this view is shared at a high-level by LD3 in how discretization points should be selected, at its core, LD3 still views sampling from a diffusion model as solving an ODE and instead focuses on selecting discretization points that minimize the global error. In contrast, S4S argues that in the very few-NFE regime, an ODE-centric view is limiting.

Second, S4S and S4S-Alt are similar in spirit to BNS but differ in key methodological choices. Through a careful exploration of the diffusion solver design space, S4S-Alt greatly outperforms BNS. This improvement stems from four key decisions:

• Using local rather than global information—learning coefficients for only the three most recent denoising points.
• Employing an alternating objective to learn coefficients and discretization points.
• Relaxing the objective for more effective learning.
• Exploring a wider array of distance functions.

While these choices may seem minor, they significantly enhance performance. BNS, in contrast, attempts to generalize across solvers, optimizing a complex objective over a large coefficient space, which makes it vastly slower (requiring orders of magnitude more compute) and leads to a more difficult optimization landscape. Our ablation study confirms that jointly optimizing solver coefficients and schedules underperforms compared to S4S-Alt.

LPIPS losses

Thank you for this question! Although our results achieve stronger FID scores when used with the LPIPS loss, our overall results still hold when using an alternative distance function in our objective. Below, we include a table that characterizes our performance using non-LPIPS losses for both S4S and S4S-Alt. Broadly speaking, using an alternative loss decreases FID scores, though not to such an extent that it weakens our overall approach.

[4] https://arxiv.org/pdf/2409.19681 SFD

Thank you for your time reviewing our work, and for recommending acceptance! We hope we answer your outstanding questions below.

How sensitive is performance to the radius $r$? Does the $r\propto m^{-5/2}$ heuristic hold across varying $m$?

In practice, the heuristic $r \propto m^{-5/2}$ works reasonably well "out of the box" across the experimental settings we looked at. Sensitivity emerges in two different directions though on both sides of the parameter scale. With very few number of parameters, choosing the radius to be too small makes the optimization problem more difficult. On the other hand, with a larger number of parameters, allowing for too large of a radius enables more overfitting. While $m^{-5/2}$ is a good heuristic that achieves baseline strong performance in both settings, we would expect improvements to come by carefully tuning this parameter based on the experimental settings.

To help clarify this, we ran some additional experiments on CIFAR-10 and ImageNet-256 to characterize this dependence; our results are in the table below.

As shown in the table, the performance is generally robust within a reasonable range around our heuristic, with the optimal value typically falling between $0.5m^{-5/2}$ and $1.0m^{-5/2}$. The sensitivity increases with larger model sizes, supporting our conclusion that a well-calibrated radius is more important as parameter count increases to prevent overfitting.

Can S4S/S4S-Alt achieve competitive performance with 1-2 NFEs, or does it require a minimum step count?

Thank you for the interesting question! In practice, we found that going below 3 NFEs requires solving a very difficult problem that likely requires directly modifying the underlying score network. Concretely, 1 NFE generation entails directly generating an image from a noise latent; as such we completely lose any "degrees of freedom" over the time discretization; similarly, we may now only control ~2-4 coefficients, which is generally insufficient for producing high-quality outputs. While having 2 NFEs yields better performance, on traditional score network architectures, S4S and S4S-Alt still fall short of the performance seen in 1-2 step generation in training-based distillation methods. Nonetheless, this challenge persists for all other methods that don't modify the underlying score network.

To further characterize these results, we conducted a few more experiments. First, for CIFAR-10, we characterized S4S and S4S-Alt on 2 NFE generation for LD3 discretization.

As shown in the table, while S4S-Alt achieves significant improvements over traditional solvers at 2 NFEs, the FID scores are still substantially higher than at 3-4 NFEs. This supports our observation that there's a minimum effective step count (~3 NFEs) for maintaining reasonable image quality with training-free methods that don't modify the score network architecture.

On the other hand, we decided to replicate an experiment from the LD3 paper that we didn't have time to include in our results. Here, we examined results on InstaFlow, a flow network that is explicitly trained for high-quality few-step sampling. When the underlying model is explicitly trained to produce high-quality few-step samples, S4S correspondingly improves the image quality as well, on a scale competitive with the teacher solver (8 NFEs, Uniform Time Disc. = 14.16 FID).

Thank you for your time reviewing our work and for your recommendation of acceptance. Below, we hope to address the weaknesses you identified and the questions you raised.

With a larger parameter space, the method is more data-driven and less theory-grounded than traditional solvers. The coefficients require separate optimization under different NFE.

We agree that our approach is less grounded in theory relative to other diffusion model solvers. However, we think that this is more a feature of the particular problem setting -- very low-NFE sampling -- than of our specific approach. For instance, in many impressive theory-grounded diffusion model solvers, e.g. DPM-Solver++, the underlying assumptions that guarantee convergence in these solvers begin to break down in the low-NFE regime (~3-5 steps) as the step size for each solver step increases. Accordingly, we see significant degredation in performance in many of these theory-grounded solvers as the number of NFEs decrease; this decrease in performance can be seen our tables in Appendix H. S4S's data-centric approach avoids trying to make these strong assumptions in the low-NFE regime and as a result achieves stronger performance than its theoretically-grounded alternatives. Nonetheless, from a theoretical perspective, it's still not intuitive why S4S even works; in fact, it seems shocking that solvers like iPNDM or DPM-Solver++ get remotely reasonable performance in the low-NFE regime in the first place. We think that trying to characterize why these approaches even achieve a modicum of success with few NFEs is a very exciting direction of future work that we hope to provide answers to.

We also recognize that our data-centric approach necessitates S4S to learn a new solver for each number of NFEs. In practice, this is a similar limitation to alternative methods for learning diffusion model solvers (e.g. LD3 [1], BNS [2], $\Pi$A [3]). Creating a reusable method for crafting diffusion model sovlers without needing to retrain each time is similarly a direction for future work that we aim to address.

There are no intuitive visualizations of the learned coefficients and timestep discretizations.

Thank you for raising this point! Here, we provide some visualizations of the learned time-step discretizations for LSUN bedroom and for FFHQ. These visualizations indicate that the learned time-step discretizations are similar to the "best" heuristic time discretizations, i.e. the learned time discretization for latent diffusion models are similar to the uniform time discretization, while the learned time discretizations are similar to EDM / logSNR for pixel-space models.

Does the sampling trajectory of the learned solver still follow the teacher's, or is only the final sampler closer?

Thank you for asking this question! The answer to this question depends on the difficulty of the underlying task. Interestingly, in relatively simple domains (e.g. CIFAR-10), the trajectory of the learned solver still closely follows that of the teacher's, despite being trained to explicitly only match the output of the teacher solver. In contrast, however, on more complex domains (e.g. conditional generation in ImageNet-256 or MS-COCO text-to-image), the trajectories of the student solver can have notable differences from that of the teacher solver.