通过深度维度的ODE加速扩散模型采样,降低延迟与内存占用。
ODE$_t$(ODE$_l$): Shortcutting the Time and the Length in Diffusion and Flow Models for Faster Sampling
- 用深度维度的ODE重写Transformer块,实现采样自由控制。
- 采样速度最高提升2倍,高质量生成时FID提升2.8点。
- 适用于追求高效采样的研究者,兼容多种求解器。
连续归一化流(CNFs)和扩散模型(DMs)从噪声分布生成高质量数据,但其采样过程需多次迭代求解高复杂度常微分方程(ODE)。现有方法主要通过减少离散时间步数提升效率。本文提出互补思路:在神经网络长度维度上控制质量-复杂度权衡。通过重连基于Transformer的架构,在深度方向求解内嵌离散ODE;训练中引入长度一致性正则项,使采样可任意选择时间步数与变换器块数。不同于以往方法,本工作在时间维度上对求解器无依赖,同时降低延迟与内存消耗。在CelebA-HQ与ImageNet上的实验表明,最高效采样模式下延迟减少最多2倍,应用于已有方法时高质量生成的FID提升最高达2.8点。代码与检查点已开源于github.com/gudovskiy/odelt。
原文摘要 · Abstract (English)
Continuous normalizing flows (CNFs) and diffusion models (DMs) generate high-quality data from a noise distribution. However, their sampling process demands multiple iterations to solve an ordinary differential equation (ODE) with high computational complexity. State-of-the-art methods focus on reducing the number of discrete time steps during sampling to improve efficiency. In this work, we explore a complementary direction in which the quality-complexity tradeoff can also be controlled in terms of the neural network length. We achieve this by rewiring the blocks in the transformer-based architecture to solve an inner discretized ODE w.r.t. its depth. Then, we apply a length consistency term during flow matching training, and as a result, the sampling can be performed with an arbitrary number of time steps and transformer blocks. Unlike others, our ODE$_t$(ODE$_l$) approach is solver-agnostic in time dimension and reduces both latency and, importantly, memory usage. CelebA-HQ and ImageNet generation experiments show a latency reduction of up to $2\times$ in the most efficient sampling mode, and FID improvement of up to $2.8$ points for high-quality sampling when applied to prior methods. We open-source our code and checkpoints at github.com/gudovskiy/odelt.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。