RepControlNet:一种新的扩散模型推理加速技术 !

2024-09-10 20:52:41 浏览数 (2)

随着扩散模型的广泛应用,推理资源的高昂成本成为其普适应用的重要瓶颈。可控制生成,例如ControlNet,是扩散模型研究的关键方向之一,与推理加速和模型压缩相关的研究更是重要且迫切。 为了解决这个问题,本文提出了一种模态重参化方法RepControlNet,可以在不增加计算的情况下实现扩散模型的可控制生成。 在训练过程中,RepControlNet使用 Adapter 将模态信息调制到特征空间,将原扩散模型的卷积和MLP可学习层作为模态网络进行复制,并根据原始权重和系数初始化这些权重。训练过程只优化模态网络的参数。 在推理过程中,模态网络中原扩散模型的权重进行重参化,这可以与或甚至超越使用额外参数和计算量的方法,如ControlNet,在增加参数数量的同时不增加参数数量。 作者在SD1.5和SDXL上进行了大量实验,实验结果显示了所提出RepControlNet的有效性和效率。

Introduction

随着扩散模型应用的逐渐普及和用户群体的指数增长,对扩散模型的压缩[1]和推理加速[2]的研究变得越来越重要。这些努力旨在在实际应用中节省推理时间和能耗,从而减少碳排放,为环境保护做出无声的贡献。由SnapFusion[1]代表的模型压缩工作通过减少模型参数的数量并使用降香来提高推理速度,但同时牺牲了一些生成质量。扩散模型的推理加速目标通常是要实现步骤降香[2],从而减少推理步骤的数量。以前的方法主要集中在优化扩散模型本身。在受多模态信息控制的生成领域,关于加速的研究几乎没有。

受控扩散模型的生成一直是一个研究热点,ControlNet[19]被广泛用作即插即用的模块。然而,每个ControlNet都会使推理过程的计算负荷增加约一半,影响效率和计算成本。为了解决这个问题,受到RepVGG[3]和ControlNet[19]的启发,作者提出了一种模态重参化方法,RepControlNet。在训练过程中,RepControlNet冻结原始扩散模型,并复制其卷积和线性层来学习输入模态控制信息。在推理时,复制的卷积和线性层与原始扩散模型重参化,实现基于条件的受控生成,而不会增加超过扩散模型本身的额外计算负荷。在Stable Diffusion 1.5 (SD1.5)[8]和Stable Diffusion XL (SDXL)[9]上的实验结果表明了RepControlNet的有效性,表明它在与ControlNet几乎无损失的情况下,提高了计算负荷的一半。

总的来说,作者的贡献如下:

  • 在扩散模型的受控生成领域,作者提出了一种模态重参化方法,RepControlNet,实现了在不增加计算量的条件下基于条件的受控生成。
  • 在SD1.5和SDXL上,作者验证了RepControlNet方法在类似canny,depth,语义图和骨架等模态下的模态重参化条件控制生成功能。

Related Work

Controllable Diffusion Models

稳定扩散 [8] 是图像生成领域的一个基础模型,可以提供高质量的输出。然而,其可控性有限,限制了其在需要特定输出的场景中的应用。为了解决这个问题,ControlNet [19] 引入了可插拔模块,通过添加额外的控制信号来指导生成过程。尽管 ControlNet 在可控制生成方面取得了显著的结果,但其计算负担大幅增加,使其效率降低。 T2I-Adapter [20] 进一步扩展了这个想法,引入了文本到图像的可控性,允许基于文本输入进行更精确的生成。然而,T2I-Adapter 也需要额外的参数,这可能在其资源受限的环境中成为一个限制。与此同时,IP-Adapter [21] 专注于推理加速和模型压缩,使用 Adapter 模块来调节输入信息。尽管 IP-Adapter 可以提高速度和效率,但它仍然需要额外的参数。 InstantID [22] 提供实例级可控性,使得在类别内部生成特定实例成为可能。然而,InstantID 也面临计算复杂性增加的问题。

尽管这些方法在可控制生成方面取得了显著的进展,但它们通常在这些方面存在权衡,即计算负担和效率之间的权衡。作者提出的 RepControlNet 旨在通过在不增加计算复杂性的情况下实现可控生成来解决这些问题。

Accelerating Diffusion Models

在加速扩散模型的方面,SnapFusion [1] 专注于通过减少模型参数数量和应用蒸馏技术来压缩模型来加速推理速度。基于步进的蒸馏通过压缩推理步数来加速推理过程。步进蒸馏的核心思想是训练一个教师模型生成中间步进目标,并使用这些目标训练一个学生模型,从而减少推理所需的步骤数。例如,Consistency Models [5] 在训练期间引入一致性约束来减少推理步骤,同时保持生成质量。Progressive Distillation [6] 采用逐步蒸馏方法,逐步优化推理步骤以实现高效的推理加速。然而,这些方法通常要以牺牲生成质量为代价。

重参化技术在训练和推理阶段优化模型架构。在训练期间,使用复杂模型结构增强模型的表达能力,而在推理期间,将这些复杂结构转换为更简单的形式以减少计算复杂度。RepVGG [3] 和 RepViT [4] 是重参化技术的典型应用。RepVGG 将训练期间使用的复杂卷积层分为推理期间使用的简单卷积层,从而减少计算复杂度。RepViT 应用这种技术到视觉 Transformer (Vision Transformers),在保持模型性能的同时显著提高推理速度。然而,这些方法并未特别解决可控生成问题的附加复杂性。

尽管这些方法在加速扩散模型方面做出了显著贡献,但它们通常只关注模型本身,而没有考虑由可控生成引起的额外复杂性。受到 RepVGG 和 ControlNet 的启发,作者提出了一种模式重参化方法 RepControlNet,以解决 ControlNet 中增加的计算开销问题。在训练过程中,RepControlNet 冻结原始扩散模型并复制其卷积和线性层学习输入模态控制信息。在推理期间,将复制的卷积和线性层用原始扩散模型重参化,实现条件可控生成,而不增加超过扩散模型本身的额外计算负载。

Methods

Preliminaries

在控制噪声模型(ControlNet)中,控制元素 c 通过训练可学习的复制 Theta_{c} 注入到扩散模型中,其中原块 Theta 在训练过程中被冻结,并且使用零卷积。

y_{c}=F(x;Theta) Z(F(x Z(c;Theta_{z1});Theta_{c});Theta_{z2}) tag{1}

图1: RepControlNet 训练和重参化过程的流程图。

在稳定扩散中,控制噪声模型在 U-Net 的12个编码块和1个中间块中应用。训练过程中冻结的原编码块不计算梯度,这加速了训练过程并节省了GPU内存。然而,这通过增加参数数量和计算开销降低了推理的计算效率。目前,许多方法的重点是优化模型架构,以提高模型性能和推理效率。重参化是一种最有效的方法,它通过复杂的模型在训练过程中提高模型的能力,而在推理过程中应用使计算复杂性较少的重参化方法。 RepVGG [3] 是重参化的典型方法。 RepVGG 在训练过程中使用多分枝拓扑结构,并在推理过程中将其重新参数化为简单的平面结构,从而在推理速度和准确性之间实现平衡。受到 RepVGG 启发,作者提出了一种模态重参化方法 RepControlNet,在不增加扩散模型的额外计算复杂度的情况下,实现控制噪声模型所产生的等效效果。

Structure

类似于ControlNet,RepControlNet中的原块在训练过程中被冻结,而可训练的副本用于学习条件。与ControlNet不同,RepControlNet的可训练副本仅用于卷积和线性层,以便在推理过程中方便地重新参数化。具体来说,在RepControlNet中,所有卷积和线性层在训练过程中的预训练扩散模型中都已被冻结,然后这些层的新副本作为图1所示的那样进行训练。在RepControlNet中,特征嵌入同时输入到原始层和每个卷积和线性层中的副本,然后将两个输出的特征相加作为输出特征。在UNet的第一层中应用一种适配层,注入条件信息。条件信息通过将 Adapter 输出张量和原始输入张量(这两个张量具有相同的维数)相加来注入。

通过这种方式,作者将单分支模型转换为多分支模型,其效果与ControlNet相当。为了充分利用预训练扩散模型的强大容量,同时尽量减少多分支模型在初始训练过程中与原始模型之间的差距,多分支模型中的权重Theta_m被初始化为原始权重Theta的乘积w和比值w

Theta_m=wTheta, tag{2}

在作者的实验中,比值w设置为0.1。

Reparameter

RepControlNet的独特之处在于其使用重参化技术在推理过程中实现高效的有条件可控制生成。在推理过程中,RepControlNet通过重参化将训练过程中复杂结构简化为高效的推理结构。具体来说,卷积层和线性层的副本使用原始扩散模型的权重进行重新参数化。通过线性结合训练过程中学习的权重与原始权重,可以在不增加额外计算复杂度的情况下实现条件可控制生成。

具体来说,对于每个卷积层和线性层,RepVGG已经证明通过结合不同分支的权重,可以将多分支模型等价转换为单分支模型。重参化过程中的权重组合公式为

Theta' = alphaTheta betaTheta_m, tag{3}

其中Theta是原始扩散模型的权重,Theta_m是副本(模态)模型的权重,alphabeta是比超参数。

通过上述方法,作者在推理过程中将 RepControlNet 转化为具有与原始扩散模型相同参数数量的新模型,从而在不增加额外计算复杂度的情况下实现高效的条件控制生成。

Identity-Preserving Generation

自定义生成人像合成领域的一个热门主题。作者还可以训练 RepControlNet 来生成给定身份和姿态的真实图像。作者将 RepControlNet 与作者自行开发的肖像逼真度生成方法(目前尚未公开)相结合。在训练期间,图像的控制条件基于人体骨骼,同时作者使用交叉注意力引入身份信息。作者使用 Clip Vision [32] 和 ArcFace [33] 编码器来提取细粒度人脸特征。然后将这些特征重采样为具有 4×768 维度的脸部嵌入 F_{id},并且通过交叉注意力将 F_{id} 注入到 RepControlNet 中,具体操作如公式 4 所示:

CA=mathrm{CrossAttn}left(Q,F_{text}right) mathrm{CrossAttn}left(Q,F_{id} right), tag{4}

其中 CA 是扩散模型中的交叉注意力特征。Q 是 RepControlNet 的注意力 Query ,骨骼作为输入条件图像,F_{text} 是文本嵌入。

在推理期间,作者重新参数化 RepControlNet,而身份信息的引入不能重新参数化。然而,与 F_{text} 一样,F_{id} 只需提取一次,与时间步无关。此外,由于 F_{id} 的可用标记数量有限,推理的额外计算开销可以忽略不计。

Experiments

为了验证提出的RepControlNet的有效性和高效性,作者针对不同的图像生成条件进行了定性的实验,并将其与根据语义分割图的条件模型进行了比较。此外,作者还进行了定量的评估实验,包括条件保真度的FID和CLIP分数,以及性能指标FLOPs、MACs和参数量。

为了确保翻译内容的学术性和规范性,我在这里对部分难懂的词汇和句子结构进行了调整,同时保持原文所表达的思想和信息不变。如有不妥,还请指正。

Implementation details

作者在SD1.5 [8]和SDXL [9]上实现RepControlNet,以测试各种条件,包括Canny边缘 [7],深度图 [10],语义分割 [12, 13]。

在本节接下来的部分中,RepControlNet 指在无需重新参数化的情况下进行的训练中的RepControlNet,而RepControlNet-指的是在参数化推理过程中使用的RepControlNet。

Qualitative comparison

Canny Edge.

如图5所示:使用ADE分割来进行图像生成。说明:一栋大石头房子。一间卧室中的两张床。一个带有两个水槽和一个浴缸的浴室。人们在装满鱼的水箱前观看。

如图3所示:使用深度图来进行图像生成。说明:一个停车场,一个交通标志和一座建筑。一个停在红色砖墙附近的汽车。一座巨大的红砖建筑。

作者从LLAVA-558k[14]数据集中获取RGB图像和说明,使用Canny边缘检测器[7](带随机阈值)生成558K边缘图像-说明对。使用的基模型是SD1.5。图2显示了生成的图像和相应的输入边缘图像-说明对。

深度图。作者从DIODE数据集[15]的户外部分中获取RGB图像和深度图,使用BLIP2[16]为RGB图像生成说明,以获得16K深度图像-说明对。使用的基模型是SD1.5。图3显示了生成的图像和相应的输入深度图像-说明对。

语义分割。作者从COCO-Stuff数据集[17]中获取RGB图像、语义分割图和说明,以获得分割图像-说明对,包括118K对用于训练和5K对用于测试。使用的基模型是SD1.5。图4显示了生成的图像和相应的输入分割图像-说明对。

骨骼。为了验证RepControlNet的通用性,作者在方法部分的自开发保留身份生成算法中引入了它来进行实验。实验结果的比较如图7所示,其中训练过程将骨骼视为RepControlNet的条件。结果表明,它可以轻松有效地移植到其他算法方案中。

Quantitative comparison

为了定量评估条件可靠性,作者使用ADE20K数据集进行训练,并使用COCO-Stuff数据集进行测试,类似于ControlNet [19]。在进行全面的评估时,作者使用FID分数[34]和文本-图像CLIP分数[23]作为指标。具体来说,作者从COCO验证集中选择10000个随机的提示,并使用One-Former来生成这些样本的配对语义分割图。对于SD1.5基础模型,作者使用512×512的随机裁剪图像作为FID的目标集,而使用配对语义分割图生成的512×512图像作为FID的源集。对于SDXL基础模型,作者使用了768×768的图像。结果如表1所示。可以看出,作者的方法可以实现与使用额外参数的类似甚至更优越的性能,而没有增加参数的数量。

为了验证RepControlNet相对于ControlNet在减少参数和计算方面的优势,作者还计算了模型参数和每个模型在批处理大小为1且一步推理时间为基准。在这里,作者使用FLOPs(浮点运算)和MACs(乘法加法运算)作为计算速度的指标,如表2所示。可以看出,与SD基础模型相比,RepControlNet在推理过程中增加了极少的计算开销和模型参数,相比ControlNet具有显著优势。

作者进一步进行了保持身份的生成定量实验,以验证作者的RepControlNet的移植性。作者在Unsplash-50 [28]数据集上评估了各种模型在保持身份的生成能力,通过像CLIP分数、面部相似性(FaceSim)和FID等指标对他们的性能进行全面评估。作者与最先进的开源解决方案进行了比较,包括InstantID [22],PuLID [29]和PhotoMakerV2 [30]。为了确保公平的比较,作者遵循了他们各自的官方配置,通过增加默认提示来增强生成质量,并在生成过程中利用了Openpose ControlNet模型来约束姿势。在unsplash测试集中,如表3所示,作者的方法在所有指标上实现了最佳性能,并在FaceSim指标上取得了显着领先。这些结果全面验证了作者的提出的RepControlNet从客观指标的移植性。

Conclusions

作者提出了一种不需要额外计算开销的基于条件控制的生成方法 RepControlNet。作者在两种基础模型 SD1.5 和 SDXL 上,在四个模态(Canny、Depthmap、语义图和骨架)上进行了广泛的实验。

为了验证 RepControlNet 的普适性,作者将它集成到了作者自行开发的保持身份的生成方法中,将骨架作为控制条件进行训练。

实验结果,无论是主观评估还是客观度量,都验证了其便携性和效率。

普及扩散模型的应用还有一些距离要走,作者希望作者的工作可以给研究行人带来一些启发。

参考

[1].RepControlNet: ControlNet Reparameterization.

0 人点赞