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 表現出了更高的效率。