Optimal Stepsize for Diffusion Sampling

2025-03-27

AI Summary

Diffusion model sampling suffers from high computational overhead due to suboptimal discrete step-size scheduling, hindering practical deployment. This paper proposes an optimal step-size distillation framework grounded in dynamic programming—first formalizing step scheduling as a global discrete approximation error minimization problem with optimal substructure, and recursively learning the theoretically optimal schedule via error-controlled optimization. The method is architecture-, solver-, and noise-schedule-agnostic, ensuring strong generalizability and robustness. Evaluated on text-to-image generation, it accelerates sampling by 10× while preserving 99.4% of the GenEval score, striking a superior efficiency–quality trade-off. Its core innovation lies in formulating step-size optimization as an efficiently solvable dynamic programming problem and enabling cross-configuration transfer via knowledge distillation.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Computer Vision: Diffusion Models for Vision
• Planning, Routing, and Scheduling: Planning with Language Models

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines

Abstract

Diffusion models achieve remarkable generation quality but suffer from computational intensive sampling due to suboptimal step discretization. While existing works focus on optimizing denoising directions, we address the principled design of stepsize schedules. This paper proposes Optimal Stepsize Distillation, a dynamic programming framework that extracts theoretically optimal schedules by distilling knowledge from reference trajectories. By reformulating stepsize optimization as recursive error minimization, our method guarantees global discretization bounds through optimal substructure exploitation. Crucially, the distilled schedules demonstrate strong robustness across architectures, ODE solvers, and noise schedules. Experiments show 10x accelerated text-to-image generation while preserving 99.4% performance on GenEval. Our code is available at https://github.com/bebebe666/OptimalSteps.

Problem

• Optimizing stepsize schedules for diffusion sampling efficiency
• Reducing computational cost while preserving generation quality
• Ensuring robustness across architectures and ODE solvers

Innovation

• Dynamic programming for optimal stepsize schedules
• Recursive error minimization with global bounds
• Robust distilled schedules across architectures