Few-Step Diffusion Sampling Through Instance-Aware Discretizations

TL;DR

This paper introduces an instance-aware discretization method for diffusion models that adapts timestep schedules based on input complexity, improving generation quality with minimal additional cost.

Contribution

It proposes a novel, learnable discretization framework that tailors timestep allocations to individual samples, surpassing traditional uniform schedules in diffusion-based generative models.

Findings

Consistently improves generation quality across various diffusion tasks.

Achieves these gains with marginal tuning and negligible inference overhead.

Abstract

Diffusion and flow matching models generate high-fidelity data by simulating paths defined by Ordinary or Stochastic Differential Equations (ODEs/SDEs), starting from a tractable prior distribution. The probability flow ODE formulation enables the use of advanced numerical solvers to accelerate sampling. Orthogonal yet vital to solver design is the discretization strategy. While early approaches employed handcrafted heuristics and recent methods adopt optimization-based techniques, most existing strategies enforce a globally shared timestep schedule across all samples. This uniform treatment fails to account for instance-specific complexity in the generative process, potentially limiting performance. Motivated by controlled experiments on synthetic data, which reveals the suboptimality of global schedules under instance-specific dynamics, we propose an instance-aware discretization framework. Our method learns to adapt timestep allocations based on input-dependent priors, extending gradient-based discretization search to the conditional generative setting. Empirical results across diverse settings, including synthetic data, pixel-space diffusion, latent-space images and video flow matching models, demonstrate that our method consistently improves generation quality with marginal tuning cost compared to training and negligible inference overhead.