360AI推出DiT架构下"省钱版"ControlNet_风闻

RelaCtrl团队 投稿至 凹非寺

量子位 | 公众号 QbitAI

现有的可控Diffusion Transformer方法，虽然在推进文本到图像和视频生成方面取得了显著进展，但也带来了大量的参数和计算开销。

同时，由于这些方法未能有效考虑不同Transformer层之间控制信息的相关性，导致资源分配效率低下。

360人工智能研究院的研究团队提出了新一代控制相关性引导的高效可控生成框架——RelaCtrl。

该框架通过优化控制信号的集成方式，在Diffusion Transformer中实现了更加高效且资源优化的控制，从而有效解决了上述问题。

团队对常见的T2I Diffusion Transformer进行了实验验证，取得了显著成果。

论文一作为来自中国科学技术大学的二年级硕士生曹可，论文共同一作和项目负责人为来自360 AI Research的AIGC研究员马傲。

目前，相关的研究论文、项目主页和代码仓库已正式开放，并即将开源。

研究动机

当前，基于Diffusion Transformer（DiT）的可控生成方法（如PixArt-δ、OminiControl等）充分利用了其强大的架构和可扩展性，显著提升了生成结果的保真度和与给定文本描述的一致性。

然而，DiT的受控生成方法仍面临两个主要问题：

首先，这些方法引入了大量额外的参数和计算，导致训练和推理负担加重。

其次，不同网络层之间控制信息的相关性常常被忽视，进而造成计算资源分配效率低下。

在模型量化和剪枝领域的相关研究中，已发现具有层叠Transformer结构的LLM模型在不同层的影响力存在差异，因此可以设计针对不同层的细致化剪枝策略。

受此启发，研究人员通过在推理阶段跳过控制网络中的不同层，测量生成图像的质量和控制精度，从而得到“ControlNet相关性得分”，并以此来指导后续可控生成的条件注入。

如图所示，DiT中不同层对控制信息的相关性存在显著差异，这种相关性呈现先增加后减少的趋势，在前中心层表现出较高的相关性，而在较深层则较低。

现有方法未能充分考虑这种变化，而是对所有层应用统一的设置，导致参数和计算资源分配效率低下，特别是在冗余参数和相关性较低层中的计算。

因此，研究团队提出了一种基于相关性引导的高效可控生成框架，通过精确的参数分配和控制策略，结合轻量化设计的控制模块，从而有效减少模型参数并降低计算复杂度。

高效可控生成框架RelaCtrl

研究团队首先根据得出的ControlNet相关性得分排名，确定并选择了11个最关键的控制位置（按相关性从高到低排序）用于集成控制模块。

通过这种方法，团队实现了与使用了13个复制块的PixArt-δ相当的控制性能，同时将参数数量减少了约15%。

尽管这种方法有效地减少了模型的大小和计算开销，但控制模块的内部设计仍然存在一定的冗余。

MetaFormer将Transformers的有效性归因于两个关键组件：

通过自注意力机制实现的token混合器和通过前馈网络(FFN)层实现的通道混合器。

然而，研究表明，尽管FFN消耗了相当一部分模型参数，但它往往高度冗余。为此，研究团队引入了一个轻量级模块——相关性引导轻量级控制块（RGLC），该模块将token混合和通道混合统一为一个操作。具体来说，作者设计了一种新颖的二维shuffle mixer（TDSM），替代了原始PixArt Transformer块中的注意力层和FFN层，从而简化了架构并提高了效率。

该方法促进了token和通道维度之间的信息交互和建模，大幅度减少了复制块的参数数量和计算需求。

下面将详细介绍两个主要部分的实现：

DiT-ControlNet Relevance Prior

为了系统地评估DiT-ControlNet中各个层与生成质量和控制精度的相关性，研究团队训练了一个完全受控的PixArt-α网络，包含27个复制模块。

在推理过程中，研究者系统地跳过每个控制块层，并评估其对最终生成结果的影响。

为了进行定量评估，Fréchet Inception Distance(FID)被引入用于衡量图像生成的质量，Hausdorff Distance(HDD)则用于评估控制精度。

这些指标有助于分析跳过控制分支中各个层对整体性能的影响，并为每个控制块提供相关性分数。最终，基于这两个指标的组合，研究团队计算得到ControlNet的相关性分数，并从中选取了11个最关键的控制位置（按相关性从高到低排序）用于集成控制模块，如图中白色序号所示。

Two Dimensions Shuffle Mixer

为了减轻控制分支的计算负担，研究团队提出了对token进行分组计算的方法，并采用特定策略来增强token组之间的交互和建模能力。

具体来说，该操作首先进行随机通道选择，然后在token维度空间中对输入序列进行随机打乱，接着进行局部自注意力计算。

虽然后续的注意力机制局限于一个固定的组，但所涉及的token可能来自该组之外。

这一操作有效打乱了token之间的固有关系，并在一定程度上引入了其他局部窗口间的信息，从而打破了局部注意力通常施加的交互约束。

此外，为了确保划分操作不会影响恢复过程中latent code嵌入的语义信息，研究团队设计在自注意力计算后对token和通道维度应用逆恢复操作。

这种结合了shuffle和recovery的整体方法被称为二维shuffle混合器（TDSM），通过利用可逆变换对的能力，确保在自注意力计算期间能够保留信息，从而实现跨通道和token维度的高效非局部信息交互。

实验结果

作者在多种条件下的可控生成任务上进行了定性和定量评估，以验证RelaCtrl的效果。

定性比较

如图所示，RelaCtrl在视觉效果对比中表现出了极具竞争力的性能，能够在给定条件的约束下生成与真实图像高度一致的照片般逼真图像。即便不同条件下，也展现了生成图像的强大能力。

定量分析

研究团队在COCO验证集上对RelaCtrl与对比方法进行了定量评估。结果表明，RelaCtrl在控制指标、视觉质量和文本相似度上均实现了更优的表现，验证了其在各方面上的有效性。

计算复杂度分析

根据下表所示，基于PixArt-α的RelaCtrl仅带来了7.38%的参数增加和8.61%的计算复杂度增加，显著低于ControlNet方法的增量，后者的参数和复杂度均增加了近50%。在Flux.1-dev上的实验进一步验证了这一点，RelaCtrl在参数量和计算复杂度的增加上实现了均衡，同时在多个指标上显著优于ControlNet方法。这一优势同样体现在实际推理时间上，RelaCtrl表现出了更高的效率。

消融实验。

对社区模型的兼容能力：

研究团队使用通过Lora 微调的 PixArt权重进行推理。RelaCtrl可以有效地利用这些社区权重。下图展示了该模型在指定条件下生成的油漆、油画、古风和像素风格图像。

论文地址： http://arxiv.org/abs/2502.14377

项目主页： https://360cvgroup.github.io/RelaCtrl/

代码仓库： https://github.com/360CVGroup/RelaCtrl