本文目錄

1 CAN：條件感知的擴散模型
(來自 MIT HAN LAB, Song Han 團隊)
1 CAN 論文解讀
1.1 可控圖像生成模型：從控制特征到控制權(quán)重
1.2 哪一層設(shè)計成條件感知？
1.3 與自適應(yīng)核選擇的對比
1.4 CAN 方法具體實現(xiàn)
1.5 實驗設(shè)置和評價指標(biāo)
1.6 消融實驗結(jié)果
1.7 與 SOTA 模型對比

太長不看版

本文提出的 CAN 模型 (Condition-Aware Neural Network) 是一種對圖像生成模型添加控制的方法。CAN 可以通過動態(tài)操縱神經(jīng)網(wǎng)絡(luò)的權(quán)重來控制圖像生成過程。如圖1所示，具體的方法是通過一個條件感知權(quán)重生成模塊，這個模塊的輸入是條件 (比如類別標(biāo)簽，時間步)，作用是為卷積/線性層生成權(quán)重。作者在 ImageNet 圖像生成任務(wù)以及 COCO 文生圖任務(wù)上面測試了 CAN 方法。CAN 始終為擴散 Transformer 模型提供顯著的改進，比如 DiT 和 UViT。

本文做了什么工作

1. 引入了一個控制圖像生成模型的新機制：通過操控權(quán)重來控制條件生成模型。
2. 提出了條件感知神經(jīng)網(wǎng)絡(luò)，一種用于條件圖像生成的控制方法。
3. CAN 可以用來提升圖像生成模型的性能，大大優(yōu)于之前的方法，而且對部署很有幫助。比如在 ImageNet 512×512 的圖像生成任務(wù)，CAN 方法的 FID 比 DiT-XL/2 更小，且每個采樣步驟的 MACs 少 52 倍，為邊緣設(shè)備上的擴散模型應(yīng)用提供支持。

效果圖如下圖2所示。

1 CAN：條件感知的擴散模型

論文名稱：Condition-Aware Neural Network for Controlled Image Generation (CVPR 2024)

論文地址：http://arxiv.org/pdf/2404.01143.pdf

1.1 可控圖像生成模型：從控制特征到控制權(quán)重

大規(guī)模圖像和視頻生成模型[1][2][3]在合成逼真圖像和視頻方面表現(xiàn)出了驚人的能力。為了將這些模型轉(zhuǎn)換為人類的生產(chǎn)工具，關(guān)鍵步驟是添加控制。我們希望生成模型遵循我們的指令 (例如類別標(biāo)簽、文本、姿勢等等[4])，而不是讓模型隨機生成數(shù)據(jù)樣本。

一些先前的工作通過添加 Cross-Attention[5]或者 Self-Attention[6]將條件特征與圖像特征融合。雖然使用的操作不同，但這些方法的共同之處是：通過特征空間操作來添加控制。同時，對于不同的條件 (Condition)，神經(jīng)網(wǎng)絡(luò)權(quán)重 (卷積/線性層) 保持不變。

這項工作旨在回答以下問題：

1. 是否可以通過操控圖像生成模型的權(quán)重來控制圖像的生成過程？
2. 圖像生成模型是否能夠受益于這種新的控制方法？

本文提出一種條件感知神經(jīng)網(wǎng)絡(luò) (Condition-Aware Neural Network, CAN)，一種基于權(quán)重空間操作的新條件控制方法。CAN 引入了一個權(quán)重生成模塊來產(chǎn)生權(quán)重，這個模塊的輸入是條件的嵌入，比如用戶指令 (類別標(biāo)簽) 和擴散模型的時間步。模塊的輸出是卷積層/線性層的權(quán)重。

作者通過消融實驗調(diào)研了 CAN 對于擴散模型的實際作用。首先，作者發(fā)現(xiàn)，仔細(xì)選擇一部分模塊使其權(quán)重是條件感知的，而不是將所有的模塊都變成條件感知的，這樣做更有利于性能和效率的權(quán)衡。其次，作者發(fā)現(xiàn)根據(jù)條件直接生成權(quán)重比自適應(yīng)地合并靜態(tài)的權(quán)重更加有效。

CAN 可以單獨為圖像生成模型提供有效的條件控制，提供比以前條件控制方法更低的 FID 和更高的 CLIP 分?jǐn)?shù)。除了將 CAN 應(yīng)用于現(xiàn)有的擴散 Transformer 模型之外，作者還通過結(jié)合 CAN 和 EfficientViT 進一步構(gòu)建了一個名為 CaT 的新型擴散 Transformer 模型。

1.2 哪一層設(shè)計成條件感知？

理論上，可以使神經(jīng)網(wǎng)絡(luò)中的所有層做成條件感知的。但在實踐中，這不一定好。

首先，從性能的角度看，使用過多的條件感知層可能會使得優(yōu)化過程不穩(wěn)定，為模型優(yōu)化帶來挑戰(zhàn)。

其次, 從效率的角度來看, 雖然生成條件權(quán)重的計算開銷可以忽略不計, 但會帶來顯著的參數(shù)開銷。假設(shè)我們定義 Condition Embedding 為 (比如 384，512，1024 等等), 模型的靜態(tài)參數(shù)大小為 #params 。使用單個線性層將 Condition Embedding 映射到條件權(quán)重需要 #params 個參數(shù)。這對于現(xiàn)實世界的使用是不切實際的。因此在這項工作中，作者只選擇一部分模塊來應(yīng)用 CAN。

將 CAN 應(yīng)用與 Diffusion Transformer 的具體做法如圖3所示。Depth-Wise 卷積[7]的參數(shù)量比常規(guī)卷積小得多，將其作為條件感知的成本比較低。因此，作者按照[8]的設(shè)計，在 FFN 中間添加一個 Depth-Wise 卷積。作者使用 UViT-S/2 在 ImageNet 256×256 的圖像生成任務(wù)上進行了消融實驗研究，哪些模塊要使用條件感知方法。在消融實驗中，所有模型的架構(gòu)都相同，唯一的區(qū)別是條件感知的模塊不同。

消融實驗結(jié)果如圖4所示。作者給出了兩個觀察：

• 使用條件感知的模塊并不總是能夠提升性能，比如圖4第2行和第4行，使用靜態(tài)頭比使用條件感知頭得到更低的 FID 和更高的 CLIP 分?jǐn)?shù)。
• 將 Depth-Wise 卷積層、Patch Embedding 層和輸出投影層設(shè)置為條件感知，可以帶來，顯著的性能提升：將 FID 從 28.32 提高到 8.82，CLIP 分?jǐn)?shù)從 30.09 提高到 31.74。

基于這些結(jié)果，作者為 CAN 選擇了圖3這樣的設(shè)計。對于 Depth-Wise 卷積層和 Patch Embedding 層，作者為每個層使用單獨的條件權(quán)重生成模塊，因為它們的參數(shù)大小很小。對輸出投影層使用共享的條件權(quán)重生成模塊，因為它們的參數(shù)量很大。由于不同的輸出投影層本身具有不同的靜態(tài)權(quán)重，因此不同的輸出投影層還是具有不同的權(quán)重。

1.3 與自適應(yīng)核選擇的對比

自適應(yīng)核選擇 (Adaptive Kernel Selection[9][10])是另一種動態(tài)輸出神經(jīng)網(wǎng)絡(luò)參數(shù)的方法。Adaptive Kernel Selection 維護了一組基本的卷積核，然后動態(tài)地生成縮放參數(shù)來組合這些卷積核。這種方法的參數(shù)開銷小于 CAN。但是，這種自適應(yīng)核選擇策略的性能不如 CAN 方法，如圖5所示。這表明僅動態(tài)參數(shù)化并不是提高性能的關(guān)鍵，更好的條件感知適應(yīng)能力至關(guān)重要。

1.4 CAN 方法具體實現(xiàn)

由于條件感知層在給定不同樣本的情況下具有不同的權(quán)重, 因此不能進行批量化的訓(xùn)練和推理。因此, 必須針對每個樣本單獨運行內(nèi)核, 如圖6左側(cè)所示。為了解決這個問題, 作者提出了一種 CAN 的高效版本實現(xiàn)。核心思想是把所有的卷積核的調(diào)用封裝為一個分組卷積, 其中組數(shù) #Groups 為 Batch Size 。

在分組卷積之前, 作者進行了一步 batch-to-channel 的轉(zhuǎn)換, 把維度為 的特征轉(zhuǎn)換為維度為 的特征, 然后進行 Grouped Conv 操作。在結(jié)束之后, 再反過來通過 channel-to-batch 把特征變回原來的形狀。

理論上，通過這種高效的實現(xiàn)，與運行靜態(tài)模型相比，額外的訓(xùn)練的開銷將可以忽略不計。在實踐中，由于 NVIDIA GPU 相比于分組卷積，對常規(guī)卷積的支持更友好，作者仍然觀察到 30%-40% 的訓(xùn)練開銷。這個問題可以通過編寫定制的 CUDA Kernel 來解決。

1.5 實驗設(shè)置和評價指標(biāo)

數(shù)據(jù)集： 由于資源限制，作者使用 ImageNet 數(shù)據(jù)集進行類條件圖像生成實驗，并使用 COCO 進行文本到圖像生成實驗。對于大規(guī)模的文本到圖像實驗[11]，作者將其留給未來的工作。

評價指標(biāo)： 按照常見的做法，作者使用 FID[12]作為圖像質(zhì)量的評估指標(biāo)。此外，作者使用 CLIP 分?jǐn)?shù)[13]作為可控性的指標(biāo)。使用公共 CLIP ViT-B/32[14]來測量 CLIP 分?jǐn)?shù)，遵循[15]的做法。text prompt 按照 CLIP 的 Zero-Shot 圖像分類設(shè)置來構(gòu)建。

具體實現(xiàn)： 作者將 CAN 應(yīng)用在了 DiT[16]和 UViT[17]模型中。所有模型都使用無分類器指導(dǎo) (Classifier-Free Guidance)，除非另有說明?；€模型的架構(gòu)與 CAN 模型的架構(gòu)相同，在 FFN 層中有 Depth-Wise Convolution。在訓(xùn)練期間使用自動混合精度。除了將 CAN 應(yīng)用于現(xiàn)有模型外，作者還通過將 CAN 和 EfficientViT 結(jié)合起來構(gòu)建了一個稱為 CaT 的新型 Diffusion Transformer 模型。宏觀的架構(gòu)如圖7所示。

1.6 消融實驗結(jié)果

除非另有說明，否則作者在消融實驗中訓(xùn)練 80 個 Epoch，Batch Size 為 1024。所有模型都使用 DPM-Solver[18] 和 50 步對圖像進行采樣。如下圖8所示為 UViT 和 DiT 模型應(yīng)用了 CAN 方法之后的結(jié)果。CAN 顯著地提高所有模型的生成圖像質(zhì)量和可控性，而且這些改進的計算成本開銷可以忽略不計。

下圖9比較了 CAN 方法在 UViT-S/2 和 DiT-S/2 上的訓(xùn)練曲線。可以看到，當(dāng)兩個模型的訓(xùn)練時間更長時，絕對的改進仍然顯著。這表明改進不是由于更快的收斂。相反，添加 CAN 可以提高模型的性能上限。

對于擴散模型，Condition Embedding 包含類別標(biāo)簽和時間步長。為了剖析哪個對條件權(quán)重生成過程更重要，作者使用 UViT-S/2 進行了消融實驗，并把結(jié)果總結(jié)在圖10中，可以發(fā)現(xiàn)：

• 類別標(biāo)簽信息比權(quán)重生成過程中的時間步信息更重要。 僅添加類別標(biāo)簽比單獨添加時間步得到更好的 FID 和 CLIP 分?jǐn)?shù)。
• 在 Condition Embedding 中包含類別標(biāo)簽和時間步長可以獲得最佳的結(jié)果。 因此，在接下來的實驗中，作者堅持這種設(shè)計。

與之前的條件控制方法對比

為了對比 CAN 和之前的條件控制方法，作者在下圖11中展示了實驗結(jié)果，并有以下發(fā)現(xiàn)：

• CAN 就已經(jīng)可以作為一種有效的條件控制方法。
• CAN 可以與其他條件控制方法相結(jié)合，以獲得更好的結(jié)果。
• 對于 UViT 模型，將 CAN 與注意力 (Condition 作為 tokens) 相結(jié)合會略微損害性能。因此，在接下來的實驗中作者僅在 UViT 模型上使用 CAN。

1.7 與 SOTA 模型對比

作者將本文 CaT 模型與其他方法在 ImageNet 圖像生成任務(wù)和 COCO 文生圖任務(wù)中進行了比較，實驗結(jié)果如圖12和15所示。對于 CaT 模型，作者使用了 UniPC[19]的技術(shù)加速采樣。

ImageNet 256×256 類別條件圖像生成任務(wù)

使用無分類器指導(dǎo) (classifier-free guidance, cfg)，本文的 CaT-B0 在 ImageNet 上實現(xiàn)了 2.09 的 FID，超越了 DiT-XL/2 和 UViT-H/2。更重要的是，CaT-B0 比這些模型的計算效率要高得多：MAC 比 DiT-XL/2 少 9.9 倍，MAC 比 UViT-H/2 少 11.1 倍。在沒有無分類器指導(dǎo)的情況下，CaT-B0 在所有比較模型中也實現(xiàn)了最低的 FID。

ImageNet 512×512 類別條件圖像生成任務(wù)

在更具挑戰(zhàn)性的 512×512 圖像生成任務(wù)中，可以觀察到 CAN 的漲點變得更加顯著。例如，CAN (UViT-S-Deep/4) 可以匹配 UViT-H (4.04 vs. 4.05) 的性能，而每個擴散步驟只需要 12% 的 UViT-H 的計算成本。此外，CaT-L0 在 ImageNet 512×512 上得到 2.78 的 FID，優(yōu)于 DiT-XL/2 (3.04 FID)，DiT-XL/2 的每個擴散步驟需要 52× 高的計算成本。此外，通過縮放模型， CaT-L1 進一步將 FID 從 2.78 提高到 2.48。

除了計算成本比較之外，圖13也比較了在 NVIDIA Jetson AGX Orin 上的 CaT-L0 和 DiT-XL/2 的延時。延遲是用 TensorRT, fp16 測量的。CaT-L0 可以在 ImageNet 512×512 圖像生成任務(wù)中得到更好的 FID 結(jié)果，且在與快速采樣方法 UniPC 結(jié)合之后，在 Orin 上的運行比 DiT-XL/2 快 229 倍。

除了定量結(jié)果之外，下圖14為 CAN 模型隨機生成的圖像的樣本，證明了本文模型在生成高質(zhì)量圖像方面的能力。

COCO 256×256 文生圖任務(wù)

對于 COCO 文生圖實驗，作者遵循 UViT 中使用的相同設(shè)置。模型在 COCO 2014 訓(xùn)練集上從頭開始訓(xùn)練。在 UViT 之后，作者從 COCO 2014 驗證集中隨機抽取 30K 個文本提示來生成圖像，然后計算 FID。作者使用與 UViT 中相同的 CLIP 編碼器來編碼文本提示。

實驗結(jié)果如圖15所示，CaT-S0 實現(xiàn)了與 UViTS-Deep/2 相似的 FID 結(jié)果，同時計算成本要低得多 (19GMACs → 3GMACs)，證明了本文模型的泛化能力。這個實驗說明 CAN 方法不僅僅適用于圖像生成任務(wù)，也適用于文生圖任務(wù)。

參考

1. ^High-resolution image synthesis with latent diffusion models
2. ^Video generation models as world simulators
3. ^Stable video diffusion: Scaling latent video diffusion models to large datasets
4. ^Adding conditional control to text-to-image diffusion models
5. ^High-Resolution Image Synthesis with Latent Diffusion Models
6. ^All are Worth Words: a ViT Backbone for Score-based Diffusion Models
7. ^Xception: Deep Learning with Depthwise Separable Convolutions
8. ^EfficientViT: Lightweight Multi-Scale Attention for High-Resolution Dense Prediction
9. ^Scaling up GANs for Text-to-Image Synthesis
10. ^CondConv: Conditionally Parameterized Convolutions for Efficient Inference
11. ^Improving image captioning with better use of captions
12. ^GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash Equilibrium
13. ^CLIPScore: A Reference-free Evaluation Metric for Image Captioning
14. ^Learning Transferable Visual Models From Natural Language Supervision
15. ^Improving Image Captioning with Better Use of Captions
16. ^Scalable Diffusion Models with Transformers
17. ^All are Worth Words: a ViT Backbone for Score-based Diffusion Models
18. ^DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps
19. ^UniPC: A Unified Predictor-Corrector Framework for Fast Sampling of Diffusion Models

公眾號后臺回復(fù)“極市直播”獲取100+期極市技術(shù)直播回放+PPT