360AI推出DiT架构下“省钱版”ControlNet, 参数量骤减85%性能达到SOTA！

今天小编分享的科学经验：360AI推出DiT架构下“省钱版”ControlNet, 参数量骤减85%性能达到SOTA！，欢迎阅读。

现有的可控 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 表现出了更高的效率。