摘要 

擴散 Transformer（DiT）作為一種前景廣闊的視覺生成擴散模型，展現出了令人矚目的性能，但同時也帶來了巨大的計算開銷。有趣的是，對預訓練 DiT 模型的分析表明，全局自注意力往往存在冗余，主要捕捉局部模式，這凸顯了尋找更高效替代方案的潛力。在本文中，我們重新審視卷積，將其作為構建高效且富有表現力的擴散模型的替代構建塊。然而，直接用卷積替換自注意力通常會導致性能下降。我們的研究發現，這種性能差距歸因于卷積網絡相比 Transformer 存在更高的通道冗余。為解決這一問題，我們引入了一種緊湊通道注意力機制，該機制能夠促進更多樣化通道的激活，從而增強特征多樣性。由此產生了擴散卷積網絡（DiCo），這是一系列完全由標準卷積網絡模塊構建的擴散模型，具有強大的生成性能和顯著的效率提升。在類條件 ImageNet 基準測試中，DiCo 在圖像質量和生成速度方面均優于以往的擴散模型。值得注意的是，DiCo-XL 在 256×256 分辨率下的 FID 為 2.05，在 512×512 分辨率下為 2.53，相比 DiT-XL/2 分別實現了 2.7 倍和 3.1 倍的加速。此外，我們最大的模型 DiCo-H，參數規模達 10 億，在 ImageNet 256×256 上的 FID 達到 1.90，且在訓練過程中無需任何額外監督。 

1、引言 

擴散模型 [68,70,27,69,71] 在生成式學習領域引發了變革性的進步，在合成高度逼真的視覺內容方面展現出了卓越的能力。其多功能性和有效性使其在廣泛的實際應用中得到了廣泛采用，包括文本到圖像生成 [60,63,61]、圖像編輯 [53,40,7]、圖像修復 [39,1,2]、視頻生成 [30,82,4] 和 3D 內容創建 [58,81,78]。 

早期的擴散模型（如 ADM [11] 和 Stable Diffusion [61]）主要采用混合 U-Net [62] 架構，該架構將卷積層與自注意力相結合。最近，Transformer [77] 已成為一種更強大且可擴展的主干架構 [56,3]，促使人們轉向完全基于 Transformer 的設計。因此，擴散 Transformer（DiTs）正逐漸取代傳統的 U-Nets，如領先的擴散模型 Stable Diffusion 3 [16]、FLUX [44] 和 Sora [6] 所示。然而，自注意力的二次計算復雜度帶來了巨大挑戰，尤其是在高分辨率圖像合成中。最近的研究 [90,73,19,57,84] 探索了更高效的替代方案，聚焦于線性復雜度的類 RNN 架構，如 Mamba [20] 和門控線性注意力 [85]。盡管這些模型提高了效率，但其因果設計與視覺生成的雙向特性 [23,49] 本質上存在沖突，限制了它們的有效性。此外，如圖 3 所示，即使采用高度優化的 CUDA 實現，在高分辨率設置下，它們相比傳統 DiTs 的運行時優勢仍然有限。這引發了我們的關鍵問題：是否有可能設計一種硬件高效的擴散主干架構，同時保留如 DiTs 般強大的生成能力？ 

為解決這一問題，我們首先分析 DiTs 生成能力背后的特性。在視覺識別任務中，視覺 Transformer [15] 的成功通常歸因于自注意力捕捉長距離依賴的能力 [36,18,24]。然而，在生成任務中，我們觀察到不同的動態。如圖 4 所示，對于預訓練的類條件（DiT-XL/2 [56]）和文本到圖像（PixArt-α[9] 和 FLUX [44]）DiT 模型，當用錨定標記查詢時，注意力主要集中在附近的空間標記上，很大程度上忽略了遠處的標記。這一發現表明，計算全局注意力對于生成任務可能是冗余的，強調了局部空間建模的重要性。與需要長距離交互進行全局語義推理的識別任務不同，生成任務似乎更注重細粒度的紋理和局部結構的保真度。這些觀察揭示了 DiTs 中注意力的固有局部性，并促使我們尋求更高效的架構。 

在這項工作中，我們重新審視卷積神經網絡（ConvNets），并提出擴散卷積網絡（DiCo），這是一種專為擴散模型設計的簡單而高效的卷積主干架構。與自注意力相比，卷積操作對硬件更友好，在大規模和資源受限的部署中具有顯著優勢。雖然用卷積替代自注意力可大幅提高效率，但通常會導致性能下降。如圖 5 所示，這種直接替換會引入明顯的通道冗余，許多通道在生成過程中保持不活躍狀態。我們假設，這種性能下降源于自注意力相比卷積具有更強的動態內容依賴表示能力。為解決這一限制，我們引入了緊湊通道注意力（CCA）機制，該機制通過輕量級線性投影動態激活信息豐富的通道。作為一種逐通道的全局建模方法，CCA 在保持低計算開銷的同時，增強了模型的表示能力和特征多樣性。與依賴大型昂貴內核的現代識別卷積網絡 [12,21] 不同，DiCo 采用了完全基于高效 1×1 逐點卷積和 3×3 深度卷積的簡化設計。盡管架構簡單，DiCo 仍實現了強大的生成性能。 

如圖 2 和圖 3 所示，DiCo 模型在 ImageNet 256×256 和 512×512 基準測試中均優于最先進的擴散模型。值得注意的是，我們的 DiCo-XL 模型在 256×256 和 512×512 分辨率下分別實現了令人印象深刻的 FID 分數 2.05 和 2.53。除了性能提升外，DiCo 模型相比基于注意力 [77]、基于 Mamba [20] 和基于線性注意力 [38] 的擴散模型還展現出顯著的效率優勢。具體而言，在 256×256 分辨率下，DiCo-XL 的 Gflops 減少了 26.4%，速度比 DiT-XL/2 [56] 快 2.7 倍。在 512×512 分辨率下，DiCo-XL 的運行速度分別比基于 Mamba 的 DiM-H [73] 和 DiS-H/2 [19] 模型快 7.8 倍和 6.7 倍。我們最大的模型 DiCo-H 擁有 10 億參數，進一步將 ImageNet 256×256 上的 FID 降低至 1.90。這些結果共同凸顯了 DiCo 在基于擴散的生成建模中的強大潛力。 

總體而言，這項工作的主要貢獻可總結如下： 

?我們分析了預訓練 DiT 模型，揭示了其全局注意力機制中存在的顯著冗余和局部性。這些發現可能會啟發研究人員開發更高效的策略來構建高性能擴散模型。 

?我們提出了 DiCo，這是一種簡單、高效且強大的擴散模型卷積主干架構。通過引入緊湊通道注意力，DiCo 在不犧牲效率的情況下顯著提高了表示能力和特征多樣性。 

?我們在 ImageNet 256×256 和 512×512 基準測試上進行了廣泛實驗。DiCo 在生成質量和速度方面均優于現有擴散模型。據我們所知，這是第一項證明精心設計的全卷積主干架構可以在基于擴散的生成建模中實現最先進性能的工作。 

2、相關工作 

2.1 擴散模型架構 

早期的擴散模型通常采用 U-Net [62] 作為基礎架構 [11,28,61]。最近，越來越多的研究探索將視覺 Transformer（ViTs）[15] 作為擴散模型的替代主干架構，取得了顯著成果 [56,3,52,88,59,48]。值得注意的是，DiT [56] 展示了基于 Transformer 架構的卓越性能，在 ImageNet 生成任務中實現了最先進的性能。然而，ViTs 固有的二次計算復雜度在長序列建模的效率方面帶來了巨大挑戰。為緩解這一問題，最近的研究探索了使用具有線性復雜度的類 RNN 架構，如 Mamba [20] 和線性注意力 [38]，作為擴散模型的主干架構 [19,90,73,84,57]。DiS [19] 和 DiM [73] 采用 Mamba 來降低計算開銷，而 DiG [90] 則利用門控線性注意力 [85] 來實現具有改進效率的競爭性能。在這項工作中，我們重新審視卷積網絡作為擴散模型的主干架構。我們表明，通過適當設計，純卷積架構可以實現最先進的生成性能，為 DiTs 提供了一種高效且強大的替代方案。 

2.2 卷積網絡設計 

在過去的十年中，卷積神經網絡（ConvNets）在計算機視覺領域取得了顯著成功 [42,67,25,35,83]。為實際部署開發了許多輕量級卷積網絡 [33,65,32,13]。盡管 Transformer 已逐漸成為廣泛任務的主導架構，但其巨大的計算開銷仍然是一個重大挑戰。許多現代卷積網絡設計在保持高效率的同時實現了競爭性能。ConvNeXt [51] 探索了標準卷積網絡的現代化，并取得了優于基于 Transformer 模型的結果。RepLKNet [12] 研究了大內核卷積的使用，將內核大小擴展至 31×31。UniRepLKNet [14] 進一步將大內核卷積網絡推廣到音頻、點云、時間序列預測等領域。在這項工作中，我們探索了純卷積網絡在基于擴散的圖像生成中的潛力，并表明簡單高效的卷積網絡設計也可以實現優異性能。 

3、方法 

3.1 預備知識 

擴散公式 

我們首先回顧擴散模型的基本概念 [27,71]。擴散模型的特點是前向加噪過程，該過程將噪聲逐步注入數據樣本x_0。具體而言，前向過程可表示為： 

其中\overline{\alpha}_t是預定義的超參數。擴散模型的目標是學習反向過程：，其中神經網絡對過程的均值和協方差進行參數化。訓練涉及優化x_0對數似然的變分下界，該下界可簡化為： 

為簡化訓練，模型的預測均值\mu_\theta可重新參數化為噪聲預測器\epsilon_\theta。目標函數隨后簡化為預測噪聲與真實噪聲\epsilon_t之間的均方誤差：。遵循 DiT [56]，我們使用簡化損失L_{simple}訓練噪聲預測器\epsilon_\theta，而協方差\sum_\theta則使用完整損失進行優化。 

無分類器引導 

無分類器引導（CFG）[29] 是一種在條件擴散模型中提高樣本質量的有效方法。它通過將采樣過程引導至與給定條件c密切相關的輸出來實現這一點。具體而言，它修改預測噪聲以獲得高p(x|c)，如下所示： 

其中s\geq1控制引導強度，\epsilon_\theta(x_t,\emptyset)是在訓練期間通過隨機省略條件信息獲得的無條件預測。遵循先前的工作 [56,90]，我們采用該技術來提高生成樣本的質量。 

3.2 網絡架構 

當前，擴散模型主要分為三種架構類型：（1）無任何下采樣層的各向同性架構，如 DiT [56]；（2）具有長跳躍連接的各向同性架構，以 U-ViT [3] 為例；（3）U 形架構，如 U-DiT [76]。受圖像去噪中多尺度特征關鍵作用的啟發 [89,80]，我們采用 U 形設計來構建分層模型。我們還進行了廣泛的消融研究，以系統比較表 3 中這些不同架構選擇的性能。 

如圖 6（a）所示，DiCo 采用由堆疊 DiCo 塊組成的三階段 U 形架構。模型以 VAE 編碼器生成的空間表示z為輸入。對于大小為 256×256×3 的圖像，對應的z維度為 32×32×4。為處理該輸入，DiCo 應用 3×3 卷積，將z轉換為具有D通道的初始特征圖z_0。對于條件信息，即時間步t和類別標簽y，我們采用多層感知機（MLP）和嵌入層，分別作為時間步和標簽嵌入器。在 DiCo 的每個塊l中，特征圖z_{l-1}通過第l個 DiCo 塊生成輸出z_l。 

在每個階段內，編碼器和解碼器之間的跳躍連接促進了中間特征之間的高效信息流。 concatenation 后，應用 1×1 卷積以降低通道維度。為實現跨階段的多尺度處理，我們利用像素反洗牌操作進行下采樣，利用像素洗牌操作進行上采樣。最后，輸出特征z_L經過歸一化，并通過 3×3 卷積頭傳遞，以預測噪聲和協方差。 

3.3 DiCo 塊 

動機 

如圖 4 所示，DiT 模型中的自注意力計算（無論是類條件還是文本到圖像生成）均表現出明顯的局部結構和顯著冗余。這一觀察促使我們用硬件效率更高的操作取代 DiT 中的全局自注意力。自然的替代方案是卷積，其以高效建模局部模式而聞名。我們首先嘗試用 1×1 逐點卷積和 3×3 深度卷積的組合替代自注意力。 

然而，直接替換會導致生成性能下降。如圖 5 所示，與 DiT 相比，修改后的模型中許多通道保持不活躍狀態，表明存在顯著的通道冗余。我們假設，這種性能下降源于自注意力（動態且依賴內容）相比卷積（依賴靜態權重）具有更強的表示能力。為解決這一限制，我們引入緊湊通道注意力機制，以動態激活信息豐富的通道。我們在下面詳細描述完整設計。 

塊設計 

DiCo 的核心設計圍繞卷積模塊展開，如圖 6（c）所示。我們首先應用 1×1 卷積來聚合逐像素的跨通道信息，然后應用 3×3 深度卷積來捕捉逐通道的空間上下文。采用 GELU 激活進行非線性變換。為進一步解決通道冗余問題，我們引入緊湊通道注意力（CCA）機制，以激活更多信息豐富的通道。如圖 6（d）所示，CCA 首先通過跨空間維度的全局平均池化（GAP）聚合特征，然后應用可學習的 1×1 卷積，接著進行 sigmoid 激活以生成逐通道的注意力權重。一般來說，卷積模塊的整個過程可描述為： 

3.4 架構變體 

我們建立了四種模型變體 ——DiCo-S、DiCo-B、DiCo-L 和 DiCo-XL—— 其參數數量分別與 DiT-S/2、DiT-B/2、DiT-L/2 和 DiT-XL/2 一致。與對應的 DiT 模型相比，我們的 DiCo 模型在計算成本上顯著降低，Gflops 僅為 DiT 的 70.1% 至 74.6%。此外，為探索我們設計的潛力，我們將 DiCo 擴展至 10 億參數，得到 DiCo-H。這些模型的架構配置詳見附錄表 5。 

4、實驗 

4.1 實驗設置 

數據集和指標 

遵循先前的工作 [56,90, *75]，我們在類條件 ImageNet-1K [10] 生成基準上進行了 256×256 和 512×512 分辨率的實驗。我們使用 Fréchet Inception Distance（FID）[26] 作為評估模型性能的主要指標。此外，我們還報告了 Inception Score（IS）[64]、Precision 和 Recall [43] 作為次要指標。所有這些指標均使用 OpenAI 的 TensorFlow 評估工具包 [11] 計算。 

實現細節 

對于 DiCo-S/B/L/XL，我們采用與 DiT 完全相同的實驗設置。具體來說，我們使用恒定學習率1?—10^{-4}，不使用權重衰減，批量大小為 256。唯一應用的數據增強是隨機水平翻轉。在訓練過程中，我們對 DiCo 的權重進行指數移動平均（EMA），衰減率為 0.9999。使用預訓練的 VAE [61] 提取潛在特征。對于我們最大的模型 DiCo-H，我們遵循 U-ViT [3] 的訓練設置，將學習率提高到2?—10^{-4}，并將批量大小擴展到 1024 以加速訓練。更多細節見附錄 B 節。 

4.2 主要結果 

4.2.1 在 DiT 設置下的比較 

除了 DiT [56]，我們還選擇了最近的最先進擴散模型 DiG [90] 和 DiC [75] 作為基線，因為它們同樣遵循 DiT 的實驗設置。表 1 展示了在 ImageNet 256×256 上的比較結果。在不同模型規模下訓練 40 萬次迭代后，我們的 DiCo 在所有指標上始終實現最佳或次佳性能。此外，當使用無分類器引導（CFG）時，我們的 DiCo-XL 實現了最先進的 FID 2.05 和 IS 282.17。除了性能提升，DiCo 相比基線和基于 Mamba 的模型還展示了顯著的效率提升。 

表 2 展示了在 ImageNet 512×512 上的結果。在更高分辨率下，我們的模型在性能和效率上都有更大的提升。具體來說，DiCo-XL 實現了最先進的 FID 2.53 和 IS 275.74，同時將 Gflops 減少了 33.3%，并實現了比 DiT-XL/2 快 3.1 倍的加速。這些結果表明，我們的卷積架構在高分辨率圖像生成中仍然高效且有效。

4.2.2 模型擴展 

為了進一步探索我們模型的潛力，我們將其擴展到 10 億參數。如表 1 所示，與 DiCo-XL 相比，更大的 DiCo-H 模型在 FID 上進一步提升（1.90 vs. 2.05），展示了我們架構的強大可擴展性。我們在附錄 C 節和 D 節提供了額外的可擴展性分析和比較結果。更多生成樣本可在附錄 E 節和提交的補充材料中找到。 

4.3 消融研究 

對于消融研究，我們使用小規模模型并在 ImageNet 256×256 基準上評估性能。所有模型均訓練 40 萬次迭代。在本節中，DiT 中的自注意力未使用 FlashAttention-2 加速，以確保與其他高效注意力機制公平比較。我們分析了 DiCo 的整體架構和各個組件的貢獻，以更好地理解它們對模型性能的影響。 

4.3.1 架構消融 

我們評估了 DiCo 在各種架構設計下的性能，并與 DiT 進行了公平比較。如表 3 所示，DiCo 在所有結構中始終優于 DiT，同時還實現了顯著的效率提升。這些結果凸顯了 DiCo 作為 DiT 的強大高效替代方案的潛力。 

4.3.2 組件消融 

我們對 DiCo 進行了組件級分析，考察了激活函數、卷積核大小、緊湊通道注意力（CCA）和卷積模塊（CM）的影響。整體消融結果總結在表 4 中。增大卷積核尺寸會帶來進一步的性能提升，但以降低效率為代價，突顯了性能和計算成本之間的權衡。引入 CCA 導致 FID 提升 4.81 點。如圖 7 的特征可視化所示，CCA 顯著增強了特征多樣性，證明了其在提高模型表示能力方面的有效性。我們還將 CCA 與 SE 模塊 [34] 和通道自注意力 [89] 進行了比較；盡管簡單，CCA 仍實現了更優的性能和更高的效率。對于卷積模塊，我們將其與幾種先進的高效注意力機制（窗口注意力 [50]、聚焦線性注意力 [22]、代理注意力 [59]）進行了基準測試。結果表明，我們的 CM 在性能和計算效率上均更優。 

5、結論 

我們提出了一種新的擴散模型主干架構 —— 擴散卷積網絡（DiCo），作為擴散 Transformer（DiT）的有力替代方案。DiCo 用 1×1 逐點卷積和 3×3 深度卷積的組合取代了自注意力，并引入了緊湊通道注意力機制以減少通道冗余并增強特征多樣性。作為全卷積網絡，DiCo 在 ImageNet 256×256 和 512×512 基準上超越了最先進的擴散模型，同時實現了顯著的效率提升。我們期待進一步擴展 DiCo 并將其應用于更廣泛的生成任務，如文本到圖像生成。 

?

本文轉載自??AIRoobt?? ，作者：Yuang Ai1等