360AI推出DiT架構下“省錢版”ControlNet,引數量驟減85%效能達到SOTA!

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
—  —
學術投稿請於工作日發郵件到:
標題註明【投稿】,告訴我們:
你是誰,從哪來,投稿內容
附上論文/專案主頁連結,以及聯絡方式哦
我們會(儘量)及時回覆你
一鍵關注 👇 點亮星標
科技前沿進展每日見
一鍵三連「點贊」「轉發」「小心心」
歡迎在評論區留下你的想法!

相關文章