Bespoke Non-Stationary Solvers for Fast Sampling of Diffusion and Flow Models

Neta Shaul, Uriel Singer, Ricky T. Q. Chen, Matthew Le, Ali Thabet, Albert Pumarola, Yaron Lipman
Proceedings of the 41st International Conference on Machine Learning, PMLR 235:44603-44627, 2024.

Abstract

This paper introduces Bespoke Non-Stationary (BNS) Solvers, a solver distillation approach to improve sample efficiency of Diffusion and Flow models. BNS solvers are based on a family of non-stationary solvers that provably subsumes existing numerical ODE solvers and consequently demonstrate considerable improvement in sample approximation (PSNR) over these baselines. Compared to model distillation, BNS solvers benefit from a tiny parameter space ($<$200 parameters), fast optimization (two orders of magnitude faster), maintain diversity of samples, and in contrast to previous solver distillation approaches nearly close the gap from standard distillation methods such as Progressive Distillation in the low-medium NFE regime. For example, BNS solver achieves 45 PSNR / 1.76 FID using 16 NFE in class-conditional ImageNet-64. We experimented with BNS solvers for conditional image generation, text-to-image generation, and text-2-audio generation showing significant improvement in sample approximation (PSNR) in all.

Cite this Paper

BibTeX
@InProceedings{pmlr-v235-shaul24a, title = {Bespoke Non-Stationary Solvers for Fast Sampling of Diffusion and Flow Models}, author = {Shaul, Neta and Singer, Uriel and Chen, Ricky T. Q. and Le, Matthew and Thabet, Ali and Pumarola, Albert and Lipman, Yaron}, booktitle = {Proceedings of the 41st International Conference on Machine Learning}, pages = {44603--44627}, year = {2024}, editor = {Salakhutdinov, Ruslan and Kolter, Zico and Heller, Katherine and Weller, Adrian and Oliver, Nuria and Scarlett, Jonathan and Berkenkamp, Felix}, volume = {235}, series = {Proceedings of Machine Learning Research}, month = {21--27 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v235/main/assets/shaul24a/shaul24a.pdf}, url = {https://proceedings.mlr.press/v235/shaul24a.html}, abstract = {This paper introduces Bespoke Non-Stationary (BNS) Solvers, a solver distillation approach to improve sample efficiency of Diffusion and Flow models. BNS solvers are based on a family of non-stationary solvers that provably subsumes existing numerical ODE solvers and consequently demonstrate considerable improvement in sample approximation (PSNR) over these baselines. Compared to model distillation, BNS solvers benefit from a tiny parameter space ($<$200 parameters), fast optimization (two orders of magnitude faster), maintain diversity of samples, and in contrast to previous solver distillation approaches nearly close the gap from standard distillation methods such as Progressive Distillation in the low-medium NFE regime. For example, BNS solver achieves 45 PSNR / 1.76 FID using 16 NFE in class-conditional ImageNet-64. We experimented with BNS solvers for conditional image generation, text-to-image generation, and text-2-audio generation showing significant improvement in sample approximation (PSNR) in all.} }
Endnote
%0 Conference Paper %T Bespoke Non-Stationary Solvers for Fast Sampling of Diffusion and Flow Models %A Neta Shaul %A Uriel Singer %A Ricky T. Q. Chen %A Matthew Le %A Ali Thabet %A Albert Pumarola %A Yaron Lipman %B Proceedings of the 41st International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2024 %E Ruslan Salakhutdinov %E Zico Kolter %E Katherine Heller %E Adrian Weller %E Nuria Oliver %E Jonathan Scarlett %E Felix Berkenkamp %F pmlr-v235-shaul24a %I PMLR %P 44603--44627 %U https://proceedings.mlr.press/v235/shaul24a.html %V 235 %X This paper introduces Bespoke Non-Stationary (BNS) Solvers, a solver distillation approach to improve sample efficiency of Diffusion and Flow models. BNS solvers are based on a family of non-stationary solvers that provably subsumes existing numerical ODE solvers and consequently demonstrate considerable improvement in sample approximation (PSNR) over these baselines. Compared to model distillation, BNS solvers benefit from a tiny parameter space ($<$200 parameters), fast optimization (two orders of magnitude faster), maintain diversity of samples, and in contrast to previous solver distillation approaches nearly close the gap from standard distillation methods such as Progressive Distillation in the low-medium NFE regime. For example, BNS solver achieves 45 PSNR / 1.76 FID using 16 NFE in class-conditional ImageNet-64. We experimented with BNS solvers for conditional image generation, text-to-image generation, and text-2-audio generation showing significant improvement in sample approximation (PSNR) in all.
APA
Shaul, N., Singer, U., Chen, R.T.Q., Le, M., Thabet, A., Pumarola, A. & Lipman, Y.. (2024). Bespoke Non-Stationary Solvers for Fast Sampling of Diffusion and Flow Models. Proceedings of the 41st International Conference on Machine Learning, in Proceedings of Machine Learning Research 235:44603-44627 Available from https://proceedings.mlr.press/v235/shaul24a.html.

Related Material