Make U-Nets Great Again！北大&華為提出擴散架構U-DiT，六分之一算力即可超越DiT

机器之心發表於2024-11-15
AI架構
圖片
AIxiv專欄是機器之心釋出學術、技術內容的欄目。過去數年，機器之心AIxiv專欄接收報導了2000多篇內容，覆蓋全球各大高校與企業的頂級實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或者聯絡報導。投稿郵箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com

Sora 的釋出讓廣大研究者及開發者深刻認識到基於 Transformer 架構擴散模型的巨大潛力。作為這一類的代表性工作，DiT 模型拋棄了傳統的 U-Net 擴散架構，轉而使用直筒型去噪模型。鑑於直筒型 DiT 在隱空間生成任務上效果出眾，後續的一些工作如 PixArt、SD3 等等也都不約而同地使用了直筒型架構。

然而令人感到不解的是，U-Net 結構是之前最常用的擴散架構，在影像空間和隱空間的生成效果均表現不俗；可以說 U-Net 的 inductive bias 在擴散任務上已被廣泛證實是有效的。因此，北大和華為的研究者們產生了一個疑問：能否重新拾起 U-Net，將 U-Net 架構和 Transformer 有機結合，使擴散模型效果更上一層樓？帶著這個問題，他們提出了基於 U-Net 的 DiT 架構 U-DiT。
圖片

  • 論文標題：U-DiTs: Downsample Tokens in U-Shaped Diffusion Transformers

  • 論文地址：https://arxiv.org/pdf/2405.02730

  • GitHub 地址：https://github.com/YuchuanTian/U-DiT

從一個小實驗談開去

首先，研究者開展了一個小實驗，在實驗中嘗試著將 U-Net 和 DiT 模組簡單結合。然而，如表 1 所示，在相似的算力比較下，U-Net 的 DiT（DiT-UNet）僅僅比原始的 DiT 有略微的提升。
圖片
在圖 3 中，作者們展示了從原始的直筒 DiT 模型一步步演化到 U-DiT 模型的過程。

根據先前的工作，在擴散中 U-Net 的主幹結構特徵圖主要為低頻訊號。由於全域性自注意力運算機制需要消耗大量算力，在 U-Net 的主幹自注意力架構中可能存在冗餘。這時作者注意到，簡單的下采樣可以自然地濾除噪聲較多的高頻，強調資訊充沛的低頻。既然如此，是否可以透過下采樣來消除對特徵圖自注意力中的冗餘？

Token 下采樣後的自注意力

由此，作者提出了下采樣自注意力機制。在自注意力之前，首先需將特徵圖進行 2 倍下采樣。為避免重要資訊的損失，生成了四個維度完全相同的下采樣圖，以確保下采樣前後的特徵總維度相同。隨後，在四個特徵圖上使用共用的 QKV 對映，並分別獨立進行自注意力運算。最後，將四個 2 倍下采樣的特徵圖重新融為一個完整特徵圖。和傳統的全域性自注意力相比，下采樣自注意力可以使得自注意力所需算力降低 3/4。

令人驚訝的是，儘管加入下采樣操作之後能夠顯著模型降低所需算力，但是卻反而能獲得比原來更好的效果（表 1）。
圖片
U-DiT：全面超越 DiT

根據此發現，作者提出了基於下采樣自注意力機制的 U 型擴散模型 U-DiT。對標 DiT 系列模型的算力，作者提出了三個 U-DiT 模型版本（S/B/L）。在完全相同的訓練超參設定下，U-DiT 在 ImageNet 生成任務上取得了令人驚訝的生成效果。其中，U-DiT-L 在 400K 訓練迭代下的表現比直筒型 DiT-XL 模型高約 10 FID，U-DiT-S/B 模型比同級直筒型 DiT 模型高約 30 FID；U-DiT-B 模型只需 DiT-XL/2 六分之一的算力便可達到更好的效果（表 2、圖 1）。
圖片
圖片
在有條件生成任務（表 3）和大圖（512*512）生成任務（表 5）上，U-DiT 模型相比於 DiT 模型的優勢同樣非常明顯。
圖片
圖片
研究者們還進一步延長了訓練的迭代次數，發現 U-DiT-L 在 600K 迭代時便能優於 DiT 在 7M 迭代時的無條件生成效果（表 4、圖 2）。
圖片
圖片
U-DiT 模型的生成效果非常出眾，在 1M 次迭代下的有條件生成效果已經非常真實。
圖片
論文已被 NeurIPS 2024 接收，更多內容，請參考原論文。

相關文章