Diffusion or flow-based models are powerful generative paradigms that are notoriously hard to sample as samples are defined as solutions to high-dimensional Ordinary or Stochastic Differential Equations (ODEs/SDEs) which require a large Number of Function Evaluations (NFE) to approximate well. Existing methods to alleviate the costly sampling process include model distillation and designing dedicated ODE solvers. However, distillation is costly to train and sometimes can deteriorate quality, while dedicated solvers still require relatively large NFE to produce high quality samples. In this paper we introduce "Bespoke solvers", a novel framework for constructing custom ODE solvers tailored to the ODE of a given pre-trained flow model. Our approach optimizes an order consistent and parameter-efficient solver (e.g., with 80 learnable parameters), is trained for roughly 1% of the GPU time required for training the pre-trained model, and significantly improves approximation and generation quality compared to dedicated solvers. For example, a Bespoke solver for a CIFAR10 model produces samples with Fr\'echet Inception Distance (FID) of 2.73 with 10 NFE, and gets to 1% of the Ground Truth (GT) FID (2.59) for this model with only 20 NFE. On the more challenging ImageNet-64$\times$64, Bespoke samples at 2.2 FID with 10 NFE, and gets within 2% of GT FID (1.71) with 20 NFE.
翻译:扩散模型或基于流的模型是强大的生成范式,但由于样本被定义为高维常微分方程或随机微分方程(ODE/SDE)的解,且需要大量函数评估(NFE)才能获得良好近似,因此采样尤为困难。现有缓解高成本采样过程的方法包括模型蒸馏和设计专用ODE求解器。然而,蒸馏训练成本高且有时会降低生成质量,而专用求解器仍需要相对较多的NFE才能生成高质量样本。本文提出"定制化求解器"(Bespoke solvers)——一种为给定预训练流模型的ODE构建自定义ODE求解器的新框架。我们的方法优化了阶数一致且参数高效的求解器(例如仅含80个可学习参数),训练所需GPU时间仅为预训练模型的约1%,相比专用求解器显著提升了近似与生成质量。例如,针对CIFAR10模型的定制化求解器在10次NFE下生成样本的弗雷歇初始距离(FID)为2.73,仅需20次NFE即可达到该模型真实分布FID(2.59)的1%。在更具挑战性的ImageNet-64×64数据集上,该求解器在10次NFE下实现2.2 FID,20次NFE即可达到真实分布FID(1.71)的2%以内。