穩(wěn)定提升！CAN和Diffusion Transformer才是最佳拍檔(MIT&清華&英偉達(dá))
論文鏈接：https://arxiv.org/pdf/2404.01143.pdf
項目鏈接：??https://github.com/mit-han-lab/efficientvit??
今天和大家一起學(xué)習(xí)條件感知神經(jīng)網(wǎng)絡(luò)（CAN），這是一種向圖像生成模型添加控制的新方法。與先前的條件控制方法并行，CAN通過動態(tài)調(diào)整神經(jīng)網(wǎng)絡(luò)的權(quán)重來控制圖像生成過程。這是通過引入一個條件感知權(quán)重生成模塊來實現(xiàn)的，該模塊根據(jù)輸入條件生成卷積/線性層的條件權(quán)重。
本文在ImageNet上進(jìn)行了類別條件圖像生成以及在COCO上進(jìn)行了文本到圖像的生成測試。CAN始終為擴散Transformer模型（包括DiT和UViT）提供顯著改進(jìn)。特別是，結(jié)合EfficientViT（CaT）的CAN在ImageNet 512×512上實現(xiàn)了2.78的FID，超過了DiT-XL/2，同時每個采樣步驟需要的MAC數(shù)量減少了52倍。
部分成果展示
介紹
大規(guī)模圖像和視頻生成模型已經(jīng)展示了在合成逼真圖像和視頻方面的驚人能力。要將這些模型轉(zhuǎn)化為人類的有效工具，一個關(guān)鍵步驟是添加控制。本文不希望模型隨機生成數(shù)據(jù)樣本，而是希望生成模型能夠按照本文的指示生成圖像（例如，類別標(biāo)簽、文本、姿態(tài)）。
為了實現(xiàn)這一目標(biāo)，已經(jīng)進(jìn)行了大量研究。例如，在GANs中，一個廣泛采用的解決方案是使用自適應(yīng)歸一化，根據(jù)輸入條件動態(tài)調(diào)整和偏移中間特征圖。此外，另一種廣泛采用的技術(shù)是使用交叉注意力或自注意力將條件特征與圖像特征融合。盡管在使用的操作上有所不同，但這些方法分享著相同的基礎(chǔ)機制，即通過特征空間操作添加控制。與此同時，神經(jīng)網(wǎng)絡(luò)權(quán)重（卷積/線性層）對不同條件保持不變。
本文旨在回答以下問題：
- 可以通過操縱權(quán)重來控制圖像生成模型嗎？
- 受控圖像生成模型能否從這種新的條件控制方法中獲益？
為此，本文引入了條件感知神經(jīng)網(wǎng)絡(luò)（CAN），這是一種基于權(quán)重空間操作的新的條件控制方法。與普通神經(jīng)網(wǎng)絡(luò)不同，CAN引入了一個額外的權(quán)重生成模塊（如下圖2）。該模塊的輸入是條件embedding，包括用戶指令（例如，類別標(biāo)簽）和擴散模型的時間步長。模塊的輸出是用于調(diào)整卷積/線性層的靜態(tài)權(quán)重的條件權(quán)重。本文進(jìn)行了大量消融研究實驗，探究了CAN在diffusion transformers上的實際應(yīng)用。本文的研究揭示了CAN的兩個關(guān)鍵見解。
首先，與使所有層都具有條件感知能力相比，精心選擇一小部分模塊具有條件感知能力（下圖3）對于效率和性能（下表1）都是有益的。其次，本文發(fā)現(xiàn)，直接生成條件權(quán)重比自適應(yīng)地合并一組基本靜態(tài)層以進(jìn)行條件控制（下圖4）要更有效。
本文在兩種代表性的擴散transformer模型上評估了CAN，包括DiT和UViT。CAN為所有這些擴散transformer模型實現(xiàn)了顯著的性能提升，同時增加的計算成本可以忽略不計（下圖7）。本文還發(fā)現(xiàn)，CAN單獨為圖像生成模型提供了有效的條件控制，其FID較低，CLIP分?jǐn)?shù)較高，優(yōu)于先前的條件控制方法（下表3）。除了將CAN應(yīng)用于現(xiàn)有的擴散transformer模型外，本文還通過將CAN與EfficientViT結(jié)合，構(gòu)建了一個名為CaT的新型擴散transformer模型家族（下圖6）。本文總結(jié)本文的貢獻(xiàn)如下：
- 本文引入了一種用于控制圖像生成模型的新機制。據(jù)本文所知，本文的工作是第一個展示權(quán)重操作在條件控制中有效性的工作。
- 本文提出了Condition-Aware Neural Network（CAN），這是一種用于控制圖像生成的新條件控制方法。本文還提供了設(shè)計見解，使CAN在實踐中可用。
- 本文的CAN在圖像生成模型上提高性能，在很大程度上優(yōu)于先前的條件控制方法。此外，CAN還可以有助于圖像生成模型的部署。在ImageNet 512×512上實現(xiàn)更好的FID，本文的CAN模型在每個采樣步驟中需要的MAC數(shù)量比DiT-XL/2少52倍（上圖1），為擴散模型在邊緣設(shè)備上的應(yīng)用鋪平了道路。
方法
Condition-Aware Neural Network
圖像生成過程可以看作是從源域（噪聲或嘈雜圖像）到目標(biāo)域（真實圖像）的映射。對于受控圖像生成，不同條件（例如，貓圖像的數(shù)據(jù)分布與城堡圖像的數(shù)據(jù)分布）給定時，目標(biāo)數(shù)據(jù)分布是不同的。此外，擴散模型在不同時間步的輸入數(shù)據(jù)分布也是不同的。盡管存在這些差異，先前的模型對所有情況使用相同的靜態(tài)卷積/線性層，由于不同子任務(wù)之間的負(fù)遷移，限制了整體性能。為了緩解這個問題，一個可能的解決方案是為每個子任務(wù)都有一個專家模型。然而，由于巨大的成本，這種方法在實際使用中是不可行的。本文的條件感知神經(jīng)網(wǎng)絡(luò)（CAN）通過使神經(jīng)網(wǎng)絡(luò)能夠根據(jù)給定條件動態(tài)調(diào)整其權(quán)重來解決這個問題，而不是顯式地?fù)碛袑＜夷Ｐ汀?/p>
上圖2展示了CAN的一般思想。與常規(guī)神經(jīng)網(wǎng)絡(luò)的關(guān)鍵區(qū)別在于CAN具有額外的條件權(quán)重生成模塊。該模塊以條件embeddingc作為輸入，并輸出條件權(quán)重。除了條件權(quán)重外，每一層都有靜態(tài)權(quán)重W。在訓(xùn)練和推理過程中，通過對權(quán)重值進(jìn)行求和，將和W融合成單個核調(diào)用。這相當(dāng)于在輸入圖像特征上獨立應(yīng)用和W，然后將它們的輸出相加。
實用設(shè)計
哪些模塊應(yīng)該是條件感知的？ 從理論上講，可以使神經(jīng)網(wǎng)絡(luò)中的所有層都是條件感知的。然而，在實踐中，這可能不是一個好的設(shè)計。首先，從性能的角度來看，具有太多條件感知層可能會使模型的優(yōu)化變得困難。其次，從效率的角度來看，雖然為所有層生成條件權(quán)重的計算開銷可以忽略不計，但會產(chǎn)生顯著的參數(shù)開銷。例如，假設(shè)將條件embedding的維度表示為d（例如，384、512、1024等），將模型的靜態(tài)參數(shù)大小表示為＃params。使用單個線性層將條件embedding映射到條件權(quán)重需要＃params×d個參數(shù)，這在實際應(yīng)用中是不切實際的。在這項工作中，本文仔細(xì)選擇了一組模塊應(yīng)用CAN來解決這個問題。
上圖3提供了將CAN應(yīng)用于擴散transformer的概述。深度卷積的參數(shù)大小比常規(guī)卷積小得多，使其成為一個低成本的條件感知候選模塊。因此，本文在FFN中間添加了一個深度卷積，遵循先前的設(shè)計。本文使用UViT-S/2在ImageNet 256×256上進(jìn)行消融研究實驗，以確定要成為條件感知模塊的一組模塊。所有模型，包括基線模型，都具有相同的架構(gòu)。唯一的區(qū)別是條件感知模塊的集合不同。
在上表1中總結(jié)了結(jié)果。在消融研究實驗中，有以下觀察：
- 使模塊成為條件感知并不總是會提高性能。例如，使用靜態(tài)頭部比使用條件感知頭部的FID更低，CLIP分?jǐn)?shù)更高（上表1中第2行與第4行對比）。
- 使深度卷積層、block embedding層和輸出投影層成為條件感知可以顯著提高性能。它將FID從28.32提高到8.82，將CLIP分?jǐn)?shù)從30.09提高到31.74。
基于這些結(jié)果，本文選擇了這種CAN的設(shè)計。具體細(xì)節(jié)如上圖3所示。對于深度卷積層和block embedding層，本文為每個層使用單獨的條件權(quán)重生成模塊，因為它們的參數(shù)大小較小。相反，對于輸出投影層，本文使用一個共享的條件權(quán)重生成模塊，因為它們的參數(shù)大小較大。由于不同的輸出投影層具有不同的靜態(tài)權(quán)重，因此本文仍然對不同的輸出投影層使用不同的權(quán)重。
CAN與自適應(yīng)核選擇。與直接生成條件權(quán)重不同，另一種可能的方法是維護(hù)一組基本卷積核，并動態(tài)生成縮放參數(shù)來組合這些基本核。這種方法的參數(shù)開銷小于CAN。然而，這種自適應(yīng)核選擇策略無法達(dá)到CAN的性能（上圖4）。這表明，單純的動態(tài)參數(shù)化并不是獲得更好性能的關(guān)鍵；更好的條件感知適應(yīng)能力是關(guān)鍵。
實現(xiàn)。 由于條件感知層在不同樣本下具有不同的權(quán)重，本文不能進(jìn)行批處理訓(xùn)練和推理。相反，本文必須獨立地為每個樣本運行核調(diào)用，如下圖5（左）所示。這將顯著減慢GPU上的訓(xùn)練過程。為了解決這個問題，本文采用了一種高效的CAN實現(xiàn)方法（下圖5右）。其核心思想是將所有卷積核調(diào)用融合成一個分組卷積，其中#Groups是批處理大小B。本文在運行分組卷積之前進(jìn)行批到通道轉(zhuǎn)換以保留功能。操作后，本文添加通道到批次轉(zhuǎn)換，將特征圖轉(zhuǎn)換為原始格式。
理論上，使用這種高效的實現(xiàn)方式，與運行靜態(tài)模型相比，訓(xùn)練開銷將可以忽略不計。在實踐中，由于NVIDIA GPU對常規(guī)卷積的支持要比對分組卷積好得多，本文仍然觀察到30%-40%的訓(xùn)練開銷。這個問題可以通過編寫定制的CUDA核來解決。本文將其留給未來的工作。
實驗
setups
數(shù)據(jù)集。 由于資源限制，本文使用ImageNet數(shù)據(jù)集進(jìn)行類別條件的圖像生成實驗，并使用COCO進(jìn)行文本到圖像生成實驗。對于大規(guī)模文本到圖像生成實驗，本文將其留給未來的工作。
評估指標(biāo)。 按照常見做法，本文使用FID作為圖像質(zhì)量的評估指標(biāo)。此外，本文使用CLIP分?jǐn)?shù)作為可控性的指標(biāo)。本文使用公開的CLIP ViT-B/32來衡量CLIP分?jǐn)?shù)，遵循的做法。文本prompt是根據(jù)CLIP的zero-shot圖像分類設(shè)置構(gòu)建的。
實現(xiàn)細(xì)節(jié)。 本文將CAN應(yīng)用于最近的擴散transformer模型，包括DiT和UViT。本文遵循官方論文或GitHub存儲庫中建議的訓(xùn)練設(shè)置。除非另有說明，默認(rèn)情況下，所有模型均使用無分類器引導(dǎo)。基線模型的架構(gòu)與CAN模型相同，都在FFN層中使用深度卷積。本文使用Pytorch實現(xiàn)本文的模型，并使用A6000 GPU進(jìn)行訓(xùn)練。訓(xùn)練過程中使用自動混合精度。除了將CAN應(yīng)用于現(xiàn)有模型外，本文還通過將CAN與EfficientViT結(jié)合，構(gòu)建了一個名為CaT的新型擴散transformer家族。CaT的宏觀架構(gòu)如上圖6所示。
消融研究
本文對所有模型進(jìn)行了80個epoch的訓(xùn)練，批量大小為1024（大約10萬次迭代），用于消融研究實驗，除非另有明確說明。所有模型使用DPM-Solver進(jìn)行圖像采樣，步數(shù)為50。
CAN的有效性。 上圖7總結(jié)了CAN在各種UViT和DiT變體上的結(jié)果。對于所有變體，CAN顯著提高了圖像質(zhì)量和可控性，超過了基線。此外，這些改進(jìn)帶來的計算成本增加可以忽略不計。因此，CAN通過以更低成本的模型實現(xiàn)相同的FID和CLIP分?jǐn)?shù)，也提高了效率。
下圖8比較了UViT-S/2和DiT-S/2上CAN和基線模型的訓(xùn)練曲線。本文可以看到，即使對這兩個模型進(jìn)行更長時間的訓(xùn)練，絕對改善仍然顯著。這表明改善不是由于更快的收斂而來，而是通過添加CAN提高了模型的性能上限。
分析。 對于擴散模型，條件embedding包含類別標(biāo)簽和時間步長。為了分析哪一個對條件權(quán)重生成過程更重要，本文使用UViT-S/2進(jìn)行消融研究實驗，并在下表2中總結(jié)結(jié)果。本文發(fā)現(xiàn)：
- 在權(quán)重生成過程中，類別標(biāo)簽信息比時間步信息更重要。僅添加類別標(biāo)簽比僅添加時間步長可以使FID降低5.15，CLIP分?jǐn)?shù)提高0.33。
- 包括類別標(biāo)簽和時間步長在條件embedding中可以獲得最佳結(jié)果。因此，在接下來的實驗中，本文堅持采用這種設(shè)計。
?
與先前條件控制方法的比較在之前的實驗中，本文保留了原來的條件控制添加 CAN 時，DiT（自適應(yīng)歸一化）和 UViT（以條件為標(biāo)記的注意力）方法保持不變。為了看看CAN是否可以單獨工作以及CAN與以前的條件控制方法的比較，進(jìn)行實驗并提供上表3中的結(jié)果，有以下發(fā)現(xiàn)
- CAN單獨可以作為一種有效的條件控制方法。例如，CAN單獨在DiT-S/2上比自適應(yīng)歸一化實現(xiàn)了13.00更好的FID和0.97更高的CLIP分?jǐn)?shù)。此外，CAN單獨在UViT-S/2上比注意力（條件作為標(biāo)記）實現(xiàn)了19.53更低的FID和1.66更高的CLIP分?jǐn)?shù)。
- CAN可以與其他條件控制方法結(jié)合以獲得更好的結(jié)果。例如，將CAN與自適應(yīng)歸一化結(jié)合，可以為DiT-S/2提供最佳結(jié)果。
- 對于UViT模型，將CAN與注意力（條件作為tokens）結(jié)合會略微降低性能。因此，在接下來的實驗中，本文將在UViT模型上使用單獨的CAN。
與最先進(jìn)模型的比較
將本文的最終模型與ImageNet和COCO上的其他擴散模型進(jìn)行比較。結(jié)果總結(jié)在下表4和下表6中。對于CaT模型，本文使用UniPC來采樣圖像，以減少采樣步驟的數(shù)量。
在ImageNet 256×256上進(jìn)行類別條件生成。 如上表4（底部）所示，使用無分類器引導(dǎo)（cfg），本文的CaT-B0在ImageNet 256×256上實現(xiàn)了2.09的FID，優(yōu)于DiT-XL/2和UViT-H/2。更重要的是，本文的CaT-B0比這些模型要高效得多：比DiT-XL/2少了9.9×的MAC，比UViT-H/2少了11.1×的MAC。在沒有分類器引導(dǎo)的情況下，本文的CaT-B0在所有比較的模型中也實現(xiàn)了最低的FID（8.81 vs. 9.62 vs. 10.56）。
ImageNet 512×512 上的類條件生成 在更具挑戰(zhàn)性的512×512圖像生成任務(wù)上，本文觀察到CAN的優(yōu)點變得更加顯著。例如，本文的CAN（UViT-S-Deep/4）可以與UViT-H的性能相匹配（4.04 vs. 4.05），而每個擴散步驟只需要UViT-H計算成本的12%。此外，本文的CaT-L0在ImageNet 512×512上實現(xiàn)了2.78的FID，優(yōu)于每個擴散步驟需要52×更高計算成本的DiT-XL/2（3.04 FID）。此外，通過略微擴大模型，本文的CaT-L1將FID從2.78提高到2.48。
除了計算成本比較外，下表5還比較了CaT-L0和DiT-XL/2在NVIDIA Jetson AGX Orin上的性能。延遲是使用TensorRT，fp16測量的。在ImageNet 512×512上提供更好的FID，結(jié)合訓(xùn)練無關(guān)的快速采樣方法（UniPC），CaT-L0在Orin上比DiT-XL/2快229倍。通過將CaT與基于訓(xùn)練的少步驟方法結(jié)合起來，有可能進(jìn)一步推動效率的前沿，展示在邊緣設(shè)備上實現(xiàn)實時擴散模型應(yīng)用的潛力。
除了定量結(jié)果外，上圖9展示了由CAN模型隨機生成的圖像樣本，展示了本文模型在生成高質(zhì)量圖像方面的能力。
在COCO 256×256上的文本到圖像
相關(guān)工作
受控圖像生成。 受控圖像生成要求模型將條件信息納入計算過程中，以生成相關(guān)的圖像。社區(qū)中已經(jīng)開發(fā)了各種技術(shù)用于受控圖像生成。一個典型的例子是自適應(yīng)歸一化，它從條件信息中回歸出尺度和偏移參數(shù)，并應(yīng)用特征級別的仿射變換來影響輸出。除了自適應(yīng)歸一化，另一個典型的方法是將條件信息視為標(biāo)記，并使用交叉注意力或自注意力來融合條件信息。ControlNet是另一種代表性技術(shù)，它使用特征級別的加法來為預(yù)訓(xùn)練的文本到圖像擴散模型添加額外的控制。與這些技術(shù)并行，本工作探索了另一種機制，用于將條件控制添加到圖像生成模型中，即使神經(jīng)網(wǎng)絡(luò)層（conv/linear）的權(quán)重具有條件感知性。
動態(tài)神經(jīng)網(wǎng)絡(luò)。 本文的工作可以被視為一種新型的動態(tài)神經(jīng)網(wǎng)絡(luò)。除了在本文中探討的添加條件控制外，動態(tài)調(diào)整神經(jīng)網(wǎng)絡(luò)可以應(yīng)用于許多深度學(xué)習(xí)應(yīng)用。例如，CondConv提出根據(jù)輸入圖像特征動態(tài)組合一組基礎(chǔ)卷積核，以增加模型容量。類似地，專家混合技術(shù)使用門控網(wǎng)絡(luò)動態(tài)地將輸入路由到不同的專家。為了實現(xiàn)高效部署，一次性全網(wǎng)絡(luò)和可調(diào)整神經(jīng)網(wǎng)絡(luò)根據(jù)給定的效率約束動態(tài)調(diào)整神經(jīng)網(wǎng)絡(luò)架構(gòu)，以實現(xiàn)效率和準(zhǔn)確性之間更好的權(quán)衡。
權(quán)重生成網(wǎng)絡(luò)。 本文的條件權(quán)重生成模塊可以被視為一種專門設(shè)計用于向生成模型添加條件控制的新型權(quán)重生成網(wǎng)絡(luò)。在其他情況下，一些先前的工作利用權(quán)重生成網(wǎng)絡(luò)。例如，[33]提出使用一個小網(wǎng)絡(luò)為一個更大的網(wǎng)絡(luò)生成權(quán)重。這些權(quán)重對于數(shù)據(jù)集中的每個示例都是相同的，以提高參數(shù)效率。此外，權(quán)重生成網(wǎng)絡(luò)已被應(yīng)用于神經(jīng)結(jié)構(gòu)搜索，以預(yù)測給定其架構(gòu)的神經(jīng)網(wǎng)絡(luò)的權(quán)重，以降低神經(jīng)結(jié)構(gòu)搜索的訓(xùn)練和搜索成本。
高效深度學(xué)習(xí)計算。 本文的工作也與高效深度學(xué)習(xí)計算聯(lián)系在一起，旨在提高深度學(xué)習(xí)模型的效率，使它們適合在硬件上部署。最先進(jìn)的圖像生成模型具有巨大的計算和內(nèi)存成本，這使得在資源受限的邊緣設(shè)備上部署它們并保持高質(zhì)量具有挑戰(zhàn)性。本文的工作可以通過提供相同性能但使用更少的擴散步驟和更低成本模型來提高受控生成模型的效率。對于未來的工作，本文將探索將本文的工作與高效深度學(xué)習(xí)計算技術(shù)結(jié)合起來，以進(jìn)一步提高效率。
結(jié)論
這項工作研究了通過操縱權(quán)重向圖像生成模型添加控制。本文引入了一種新的條件控制方法，稱為條件感知神經(jīng)網(wǎng)絡(luò)（CAN），并提供了高效實用的設(shè)計，使CAN在實踐中可用。本文進(jìn)行了廣泛的實驗，使用ImageNet進(jìn)行類別條件生成和使用COCO進(jìn)行文本到圖像生成，以評估CAN的有效性。CAN相對于先前的條件控制方法實現(xiàn)了一致且顯著的改進(jìn)。本文還通過將CAN與EfficientViT結(jié)合，構(gòu)建了一個新的擴散transformer模型系列。對于未來的工作，本文將把CAN應(yīng)用于更具挑戰(zhàn)性的任務(wù)，如大規(guī)模文本到圖像生成、視頻生成等。
本文轉(zhuǎn)自 AI生成未來 ，作者：Han Cai等
