潜在一致性模型:[2310.04378] Latent Consistency Models: Synthesizing High-Resolution Images with Few-Step Inference (arxiv.org)
原文:Paper page - Latent Consistency Models: Synthesizing High-Resolution Images with Few-Step Inference (huggingface.co)
简介:LCM 只需 4,000 个训练步骤(约 32 个 A100 GPU/小时)即可从任何预训练的稳定扩散 (SD) 中提取出来,只需 2~4 个步骤甚至一步即可生成高质量的 768 x 768 分辨率图像,从而显着加速文本转换 -图像生成。
潜在一致性模型
介绍
潜在扩散模型(Latent Diffusion models, ldm)在高分辨率图像合成方面取得了显著的成果。然而,迭代采样过程计算量大,导致生成速度慢。受一致性模型的启发,我们提出了潜在一致性模型(Latent Consistency Models, lcm),能够在任何预训练的ldm上以最小的步骤进行快速推理,包括稳定扩散。
原理:将引导反向扩散过程视为求解增强概率流ODE (PF-ODE), lcm设计用于直接预测潜在空间中此类ODE的解,从而减少了多次迭代的需要,并允许快速,高保真采样。有效地从预训练的无分类器引导扩散模型中提取,高质量的768×768 2 ~ 4步LCM仅需32 A100 GPU小时即可进行训练。此外,引入了潜在一致性微调(LCF),这是一种针对自定义图像数据集微调LCF的新方法。
一致性模型(CMs):作为一种新型生成模型显示出巨大的潜力,可以在保持生成质量的同时加快采样速度。一致性模型采用一致性映射,直接将ODE轨迹中的任意点映射到原点,实现快速一步生成。可以通过提取预训练的扩散模型或作为独立的生成模型进行训练。
原理
潜在空间中的一致性蒸馏
在诸如稳定扩散(Stable Diffusion, SD)(Rombach et al, 2022)等大规模扩散模型中,利用图像的潜在空间有效地提高了图像生成质量并减少了计算负载。在SD中,首先训练一个自编码器(E, D)来将高维图像数据压缩为低维潜在向量 𝑧=𝐸(𝑥)z=E(x),然后解码以重建图像 𝑥ˆ=𝐷(𝑧)xˆ=D(z)。在潜在空间中训练扩散模型与基于像素的模型相比,大大降低了计算成本并加快了推理过程;潜在扩散模型(LDMs)使得在笔记本电脑的GPU上生成高分辨率图像成为可能。
对于潜在一致性模型(LCMs),我们利用潜在空间的一致性蒸馏优势,与一致性模型(CMs)(Song et al, 2023)中使用的像素空间形成对比。这种方法被称为潜在一致性蒸馏(LCD),应用于预训练的SD,允许在1至4步内合成高分辨率的768×768图像。我们专注于条件生成。回顾一下逆扩散过程的PF-ODE:
其中𝑧𝑡是图像潜在变量,𝜖𝜃(𝑧𝑡,𝑐,𝑡) 是噪声预测模型,𝑐 是给定的条件(例如文本)。通过从 𝑇 到 0 解决 PF-ODE 可以抽取样本。为了执行潜在一致性蒸馏(LCD),我们引入一致性函数 𝑓𝜃:(𝑧𝑡,𝑐,𝑡)→𝑧0,直接预测 𝑡=0 时 PF-ODE 的解(公式8)。通过噪声预测模型 𝜖^𝜃参数化 𝑓𝜃,如下所示:
其中 𝑐skip(0)=1,𝑐out(0)=0,且 𝜖^𝜃(𝑧,𝑐,𝑡) 是噪声预测模型,其初始参数与教师扩散模型相同。假设有一个高效的ODE求解器 Ψ(𝑧𝑡,𝑡,𝑠,𝑐),用于近似积分公式8的右侧,从时间 𝑡 到 𝑠。在实际操作中,可以使用DDIM,DPM-Solver或DPM-Solver++ 作为 Ψ(⋅,⋅,⋅,⋅)。
只在训练/蒸馏中使用这些求解器,而不是在推理中。潜在一致性模型(LCM)旨在通过最小化一致性蒸馏损失来预测PF-ODE的解:
通过求解增强的PF-ODE进行单阶段引导蒸馏
无分类器引导(Classifier-free guidance, CFG)对于在稳定扩散(SD)中合成高质量的文本对齐图像至关重要,通常需要大于6的CFG比例 𝜔。因此,将CFG集成到蒸馏方法中变得不可或缺。之前的方法 Guided-Distill引入了一个两阶段蒸馏以支持从引导扩散模型中进行少步采样。然而,这种方法计算密集估计,2步推理至少需要45个A100 GPU天)。相比之下,潜在一致性模型(LCM)仅需要32个A100 GPU小时的训练时间来进行2步推理,如图1所示。此外,两阶段引导蒸馏可能导致累积误差,导致性能不佳。相反,LCM通过求解增强的PF-ODE采用高效的单阶段引导蒸馏。回顾在逆扩散过程中使用的CFG:
其中用条件噪声和无条件噪声的线性组合代替原有的噪声预测,ω称为引导标度。为了从引导逆向过程中采样,我们需要求解以下增广的PF-ODE(即,与ω相关的项增广):
为了有效地进行一级导向蒸馏,我们引入增广一致性函数fθ:(zt, ω, c, t)→z0来直接预测t = 0时增广PF-ODE (Eq. 13)的解。我们以与Eq. 9相同的方式参数化fθ,除了λ θ(z, c, t)被λ ϵθ(z, ω, c, t)取代,这是一个用与教师扩散模型相同的参数初始化的噪声预测模型,但还包含额外的可训练参数,用于ω的调节。一致性损失与Eq. 10相同,只是我们使用增广一致性函数fθ(zt, ω, c, t)。
跳过时间步加速蒸馏
离散扩散模型通常通过长时间步长计划 {𝑡𝑖}𝑖(也称为离散化计划或时间计划)训练噪声预测模型,以实现高质量的生成结果。例如,稳定扩散(SD)有一个长度为1000的时间计划。然而,直接将潜在一致性蒸馏(LCD)应用于具有如此长时间计划的SD可能会有问题。模型需要在所有1000个时间步长上进行采样,而一致性损失试图使LCM模型 𝑓𝜃(𝑧𝑡𝑛+1,𝑐,𝑡𝑛+1) 的预测与在相同轨迹上下一步 𝑓𝜃(𝑧𝑡𝑛,𝑐,𝑡𝑛) 的预测对齐。由于 𝑡𝑛 − 𝑡𝑛+1 很小,𝑧𝑡𝑛 和 𝑧𝑡𝑛+1(因此 𝑓𝜃(𝑧𝑡𝑛+1,𝑐,𝑡𝑛+1)和 𝑓𝜃(𝑧𝑡𝑛,𝑐,𝑡𝑛))已经彼此接近,导致一致性损失很小,因此收敛速度慢。
为了解决这个问题,我们引入了跳步方法(SKIPPING-STEP),大大缩短了时间计划的长度(从数千缩短到几十),以实现快速收敛,同时保持生成质量。
一致性模型(CMs)使用EDM连续时间计划,并使用欧拉或Heun求解器作为数值连续PF-ODE求解器。对于LCMs,为了适应稳定扩散中的离散时间计划,我们使用DDIM,DPM-Solver或DPM-Solver++作为ODE求解器。
现在,我们介绍潜在一致性蒸馏(LCD)中的跳步方法。与确保相邻时间步长 𝑡𝑛+1→𝑡𝑛 之间的一致性不同,LCMs旨在确保当前时间步长和相隔 𝑘 步的时间步长 𝑡𝑛+𝑘→𝑡𝑛之间的一致性。注意,设置 𝑘=1k=1 会恢复到中的原始计划,导致收敛速度慢,而非常大的 𝑘 可能会导致ODE求解器的大近似误差。在我们的主要实验中,我们设置 𝑘=20,将时间计划的长度从数千减少到几十。第5.2节的结果显示了不同 k 值的效果,并揭示跳步方法在加速LCD过程中的重要性。具体来说,公式14中的一致性蒸馏损失被修改为确保从 𝑡𝑛+𝑘 到 𝑡𝑛 的一致性:
上述推导类似于公式15。对于LCM,我们在此使用三种可能的ODE求解器:DDIM (Song et al, 2020a)、DPM-Solver (Lu et al, 2022a)、DPM-Solver++。