让扩散模型的求解器自己学着变好,只用极少神经网络计算就提升生成质量。
S4S: Solving for a Diffusion Model Solver
- 用教师模型指导学习新求解器,直接优化生成效果。
- 仅5次神经网络计算,CIFAR10上FID达3.73,提升1.5倍。
- 无需训练数据,可通用适配各类扩散模型和采样方案。
扩散模型通过从随机噪声出发,迭代求解反向时间常微分方程(ODE)生成样本。由于每步需昂贵的神经函数评估(NFE),研究者致力于在不修改模型的前提下,仅用少量NFE近似求解。然而在极低NFE条件下,传统求解器无法追踪真实ODE演化。本文提出新方法S4S(Solving for the Solver),通过学习一个优秀求解器来提升生成质量,其通过匹配强教师求解器的输出进行优化。我们在六种预训练扩散模型上验证了S4S,涵盖像素空间与潜在空间、条件与非条件采样。所有场景下,S4S均优于传统求解器。该方法轻量、无数据依赖,可黑盒集成至任意离散化调度或架构中。进一步提出S4S-Alt,同时优化求解器与离散化调度。利用完整设计空间,在5次NFE下,CIFAR10上实现FID 3.73,MS-COCO上为13.26,相较此前免训练方法提升1.5倍。
原文摘要 · Abstract (English)
Diffusion models (DMs) create samples from a data distribution by starting from random noise and iteratively solving a reverse-time ordinary differential equation (ODE). Because each step in the iterative solution requires an expensive neural function evaluation (NFE), there has been significant interest in approximately solving these diffusion ODEs with only a few NFEs without modifying the underlying model. However, in the few NFE regime, we observe that tracking the true ODE evolution is fundamentally impossible using traditional ODE solvers. In this work, we propose a new method that learns a good solver for the DM, which we call Solving for the Solver (S4S). S4S directly optimizes a solver to obtain good generation quality by learning to match the output of a strong teacher solver. We evaluate S4S on six different pre-trained DMs, including pixel-space and latent-space DMs for both conditional and unconditional sampling. In all settings, S4S uniformly improves the sample quality relative to traditional ODE solvers. Moreover, our method is lightweight, data-free, and can be plugged in black-box on top of any discretization schedule or architecture to improve performance. Building on top of this, we also propose S4S-Alt, which optimizes both the solver and the discretization schedule. By exploiting the full design space of DM solvers, with 5 NFEs, we achieve an FID of 3.73 on CIFAR10 and 13.26 on MS-COCO, representing a $1.5\times$ improvement over previous training-free ODE methods.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。