論文鏈接：https://arxiv.org/pdf/2404.01143.pdf

項目鏈接：??https://github.com/mit-han-lab/efficientvit??

今天和大家一起學習條件感知神經網絡（CAN），這是一種向圖像生成模型添加控制的新方法。與先前的條件控制方法并行，CAN通過動態調整神經網絡的權重來控制圖像生成過程。這是通過引入一個條件感知權重生成模塊來實現的，該模塊根據輸入條件生成卷積/線性層的條件權重。

本文在ImageNet上進行了類別條件圖像生成以及在COCO上進行了文本到圖像的生成測試。CAN始終為擴散Transformer模型（包括DiT和UViT）提供顯著改進。特別是，結合EfficientViT（CaT）的CAN在ImageNet 512×512上實現了2.78的FID，超過了DiT-XL/2，同時每個采樣步驟需要的MAC數量減少了52倍。

部分成果展示

介紹

大規模圖像和視頻生成模型已經展示了在合成逼真圖像和視頻方面的驚人能力。要將這些模型轉化為人類的有效工具，一個關鍵步驟是添加控制。本文不希望模型隨機生成數據樣本，而是希望生成模型能夠按照本文的指示生成圖像（例如，類別標簽、文本、姿態）。

為了實現這一目標，已經進行了大量研究。例如，在GANs中，一個廣泛采用的解決方案是使用自適應歸一化，根據輸入條件動態調整和偏移中間特征圖。此外，另一種廣泛采用的技術是使用交叉注意力或自注意力將條件特征與圖像特征融合。盡管在使用的操作上有所不同，但這些方法分享著相同的基礎機制，即通過特征空間操作添加控制。與此同時，神經網絡權重（卷積/線性層）對不同條件保持不變。

本文旨在回答以下問題：

• 可以通過操縱權重來控制圖像生成模型嗎？
• 受控圖像生成模型能否從這種新的條件控制方法中獲益？

為此，本文引入了條件感知神經網絡（CAN），這是一種基于權重空間操作的新的條件控制方法。與普通神經網絡不同，CAN引入了一個額外的權重生成模塊（如下圖2）。該模塊的輸入是條件embedding，包括用戶指令（例如，類別標簽）和擴散模型的時間步長。模塊的輸出是用于調整卷積/線性層的靜態權重的條件權重。本文進行了大量消融研究實驗，探究了CAN在diffusion transformers上的實際應用。本文的研究揭示了CAN的兩個關鍵見解。

首先，與使所有層都具有條件感知能力相比，精心選擇一小部分模塊具有條件感知能力（下圖3）對于效率和性能（下表1）都是有益的。其次，本文發現，直接生成條件權重比自適應地合并一組基本靜態層以進行條件控制（下圖4）要更有效。

本文在兩種代表性的擴散transformer模型上評估了CAN，包括DiT和UViT。CAN為所有這些擴散transformer模型實現了顯著的性能提升，同時增加的計算成本可以忽略不計（下圖7）。本文還發現，CAN單獨為圖像生成模型提供了有效的條件控制，其FID較低，CLIP分數較高，優于先前的條件控制方法（下表3）。除了將CAN應用于現有的擴散transformer模型外，本文還通過將CAN與EfficientViT結合，構建了一個名為CaT的新型擴散transformer模型家族（下圖6）。本文總結本文的貢獻如下：

• 本文引入了一種用于控制圖像生成模型的新機制。據本文所知，本文的工作是第一個展示權重操作在條件控制中有效性的工作。
• 本文提出了Condition-Aware Neural Network（CAN），這是一種用于控制圖像生成的新條件控制方法。本文還提供了設計見解，使CAN在實踐中可用。
• 本文的CAN在圖像生成模型上提高性能，在很大程度上優于先前的條件控制方法。此外，CAN還可以有助于圖像生成模型的部署。在ImageNet 512×512上實現更好的FID，本文的CAN模型在每個采樣步驟中需要的MAC數量比DiT-XL/2少52倍（上圖1），為擴散模型在邊緣設備上的應用鋪平了道路。

方法

Condition-Aware Neural Network

圖像生成過程可以看作是從源域（噪聲或嘈雜圖像）到目標域（真實圖像）的映射。對于受控圖像生成，不同條件（例如，貓圖像的數據分布與城堡圖像的數據分布）給定時，目標數據分布是不同的。此外，擴散模型在不同時間步的輸入數據分布也是不同的。盡管存在這些差異，先前的模型對所有情況使用相同的靜態卷積/線性層，由于不同子任務之間的負遷移，限制了整體性能。為了緩解這個問題，一個可能的解決方案是為每個子任務都有一個專家模型。然而，由于巨大的成本，這種方法在實際使用中是不可行的。本文的條件感知神經網絡（CAN）通過使神經網絡能夠根據給定條件動態調整其權重來解決這個問題，而不是顯式地擁有專家模型。

上圖2展示了CAN的一般思想。與常規神經網絡的關鍵區別在于CAN具有額外的條件權重生成模塊。該模塊以條件embeddingc作為輸入，并輸出條件權重。除了條件權重外，每一層都有靜態權重W。在訓練和推理過程中，通過對權重值進行求和，將和W融合成單個核調用。這相當于在輸入圖像特征上獨立應用和W，然后將它們的輸出相加。

實用設計

哪些模塊應該是條件感知的？ 從理論上講，可以使神經網絡中的所有層都是條件感知的。然而，在實踐中，這可能不是一個好的設計。首先，從性能的角度來看，具有太多條件感知層可能會使模型的優化變得困難。其次，從效率的角度來看，雖然為所有層生成條件權重的計算開銷可以忽略不計，但會產生顯著的參數開銷。例如，假設將條件embedding的維度表示為d（例如，384、512、1024等），將模型的靜態參數大小表示為＃params。使用單個線性層將條件embedding映射到條件權重需要＃params×d個參數，這在實際應用中是不切實際的。在這項工作中，本文仔細選擇了一組模塊應用CAN來解決這個問題。

上圖3提供了將CAN應用于擴散transformer的概述。深度卷積的參數大小比常規卷積小得多，使其成為一個低成本的條件感知候選模塊。因此，本文在FFN中間添加了一個深度卷積，遵循先前的設計。本文使用UViT-S/2在ImageNet 256×256上進行消融研究實驗，以確定要成為條件感知模塊的一組模塊。所有模型，包括基線模型，都具有相同的架構。唯一的區別是條件感知模塊的集合不同。

在上表1中總結了結果。在消融研究實驗中，有以下觀察：

• 使模塊成為條件感知并不總是會提高性能。例如，使用靜態頭部比使用條件感知頭部的FID更低，CLIP分數更高（上表1中第2行與第4行對比）。
• 使深度卷積層、block embedding層和輸出投影層成為條件感知可以顯著提高性能。它將FID從28.32提高到8.82，將CLIP分數從30.09提高到31.74。

基于這些結果，本文選擇了這種CAN的設計。具體細節如上圖3所示。對于深度卷積層和block embedding層，本文為每個層使用單獨的條件權重生成模塊，因為它們的參數大小較小。相反，對于輸出投影層，本文使用一個共享的條件權重生成模塊，因為它們的參數大小較大。由于不同的輸出投影層具有不同的靜態權重，因此本文仍然對不同的輸出投影層使用不同的權重。

CAN與自適應核選擇。與直接生成條件權重不同，另一種可能的方法是維護一組基本卷積核，并動態生成縮放參數來組合這些基本核。這種方法的參數開銷小于CAN。然而，這種自適應核選擇策略無法達到CAN的性能（上圖4）。這表明，單純的動態參數化并不是獲得更好性能的關鍵；更好的條件感知適應能力是關鍵。

實現。 由于條件感知層在不同樣本下具有不同的權重，本文不能進行批處理訓練和推理。相反，本文必須獨立地為每個樣本運行核調用，如下圖5（左）所示。這將顯著減慢GPU上的訓練過程。為了解決這個問題，本文采用了一種高效的CAN實現方法（下圖5右）。其核心思想是將所有卷積核調用融合成一個分組卷積，其中#Groups是批處理大小B。本文在運行分組卷積之前進行批到通道轉換以保留功能。操作后，本文添加通道到批次轉換，將特征圖轉換為原始格式。

理論上，使用這種高效的實現方式，與運行靜態模型相比，訓練開銷將可以忽略不計。在實踐中，由于NVIDIA GPU對常規卷積的支持要比對分組卷積好得多，本文仍然觀察到30%-40%的訓練開銷。這個問題可以通過編寫定制的CUDA核來解決。本文將其留給未來的工作。

實驗

setups

數據集。 由于資源限制，本文使用ImageNet數據集進行類別條件的圖像生成實驗，并使用COCO進行文本到圖像生成實驗。對于大規模文本到圖像生成實驗，本文將其留給未來的工作。

評估指標。 按照常見做法，本文使用FID作為圖像質量的評估指標。此外，本文使用CLIP分數作為可控性的指標。本文使用公開的CLIP ViT-B/32來衡量CLIP分數，遵循的做法。文本prompt是根據CLIP的zero-shot圖像分類設置構建的。

實現細節。 本文將CAN應用于最近的擴散transformer模型，包括DiT和UViT。本文遵循官方論文或GitHub存儲庫中建議的訓練設置。除非另有說明，默認情況下，所有模型均使用無分類器引導。基線模型的架構與CAN模型相同，都在FFN層中使用深度卷積。本文使用Pytorch實現本文的模型，并使用A6000 GPU進行訓練。訓練過程中使用自動混合精度。除了將CAN應用于現有模型外，本文還通過將CAN與EfficientViT結合，構建了一個名為CaT的新型擴散transformer家族。CaT的宏觀架構如上圖6所示。

消融研究

本文對所有模型進行了80個epoch的訓練，批量大小為1024（大約10萬次迭代），用于消融研究實驗，除非另有明確說明。所有模型使用DPM-Solver進行圖像采樣，步數為50。

CAN的有效性。 上圖7總結了CAN在各種UViT和DiT變體上的結果。對于所有變體，CAN顯著提高了圖像質量和可控性，超過了基線。此外，這些改進帶來的計算成本增加可以忽略不計。因此，CAN通過以更低成本的模型實現相同的FID和CLIP分數，也提高了效率。

下圖8比較了UViT-S/2和DiT-S/2上CAN和基線模型的訓練曲線。本文可以看到，即使對這兩個模型進行更長時間的訓練，絕對改善仍然顯著。這表明改善不是由于更快的收斂而來，而是通過添加CAN提高了模型的性能上限。

分析。 對于擴散模型，條件embedding包含類別標簽和時間步長。為了分析哪一個對條件權重生成過程更重要，本文使用UViT-S/2進行消融研究實驗，并在下表2中總結結果。本文發現：

• 在權重生成過程中，類別標簽信息比時間步信息更重要。僅添加類別標簽比僅添加時間步長可以使FID降低5.15，CLIP分數提高0.33。
• 包括類別標簽和時間步長在條件embedding中可以獲得最佳結果。因此，在接下來的實驗中，本文堅持采用這種設計。

?

與先前條件控制方法的比較在之前的實驗中，本文保留了原來的條件控制添加 CAN 時，DiT（自適應歸一化）和 UViT（以條件為標記的注意力）方法保持不變。為了看看CAN是否可以單獨工作以及CAN與以前的條件控制方法的比較，進行實驗并提供上表3中的結果，有以下發現

• CAN單獨可以作為一種有效的條件控制方法。例如，CAN單獨在DiT-S/2上比自適應歸一化實現了13.00更好的FID和0.97更高的CLIP分數。此外，CAN單獨在UViT-S/2上比注意力（條件作為標記）實現了19.53更低的FID和1.66更高的CLIP分數。
• CAN可以與其他條件控制方法結合以獲得更好的結果。例如，將CAN與自適應歸一化結合，可以為DiT-S/2提供最佳結果。
• 對于UViT模型，將CAN與注意力（條件作為tokens）結合會略微降低性能。因此，在接下來的實驗中，本文將在UViT模型上使用單獨的CAN。

與最先進模型的比較

將本文的最終模型與ImageNet和COCO上的其他擴散模型進行比較。結果總結在下表4和下表6中。對于CaT模型，本文使用UniPC來采樣圖像，以減少采樣步驟的數量。

在ImageNet 256×256上進行類別條件生成。 如上表4（底部）所示，使用無分類器引導（cfg），本文的CaT-B0在ImageNet 256×256上實現了2.09的FID，優于DiT-XL/2和UViT-H/2。更重要的是，本文的CaT-B0比這些模型要高效得多：比DiT-XL/2少了9.9×的MAC，比UViT-H/2少了11.1×的MAC。在沒有分類器引導的情況下，本文的CaT-B0在所有比較的模型中也實現了最低的FID（8.81 vs. 9.62 vs. 10.56）。

ImageNet 512×512 上的類條件生成 在更具挑戰性的512×512圖像生成任務上，本文觀察到CAN的優點變得更加顯著。例如，本文的CAN（UViT-S-Deep/4）可以與UViT-H的性能相匹配（4.04 vs. 4.05），而每個擴散步驟只需要UViT-H計算成本的12%。此外，本文的CaT-L0在ImageNet 512×512上實現了2.78的FID，優于每個擴散步驟需要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，結合訓練無關的快速采樣方法（UniPC），CaT-L0在Orin上比DiT-XL/2快229倍。通過將CaT與基于訓練的少步驟方法結合起來，有可能進一步推動效率的前沿，展示在邊緣設備上實現實時擴散模型應用的潛力。

除了定量結果外，上圖9展示了由CAN模型隨機生成的圖像樣本，展示了本文模型在生成高質量圖像方面的能力。

在COCO 256×256上的文本到圖像 

相關工作

受控圖像生成。 受控圖像生成要求模型將條件信息納入計算過程中，以生成相關的圖像。社區中已經開發了各種技術用于受控圖像生成。一個典型的例子是自適應歸一化，它從條件信息中回歸出尺度和偏移參數，并應用特征級別的仿射變換來影響輸出。除了自適應歸一化，另一個典型的方法是將條件信息視為標記，并使用交叉注意力或自注意力來融合條件信息。ControlNet是另一種代表性技術，它使用特征級別的加法來為預訓練的文本到圖像擴散模型添加額外的控制。與這些技術并行，本工作探索了另一種機制，用于將條件控制添加到圖像生成模型中，即使神經網絡層（conv/linear）的權重具有條件感知性。

動態神經網絡。 本文的工作可以被視為一種新型的動態神經網絡。除了在本文中探討的添加條件控制外，動態調整神經網絡可以應用于許多深度學習應用。例如，CondConv提出根據輸入圖像特征動態組合一組基礎卷積核，以增加模型容量。類似地，專家混合技術使用門控網絡動態地將輸入路由到不同的專家。為了實現高效部署，一次性全網絡和可調整神經網絡根據給定的效率約束動態調整神經網絡架構，以實現效率和準確性之間更好的權衡。

權重生成網絡。 本文的條件權重生成模塊可以被視為一種專門設計用于向生成模型添加條件控制的新型權重生成網絡。在其他情況下，一些先前的工作利用權重生成網絡。例如，[33]提出使用一個小網絡為一個更大的網絡生成權重。這些權重對于數據集中的每個示例都是相同的，以提高參數效率。此外，權重生成網絡已被應用于神經結構搜索，以預測給定其架構的神經網絡的權重，以降低神經結構搜索的訓練和搜索成本。

高效深度學習計算。 本文的工作也與高效深度學習計算聯系在一起，旨在提高深度學習模型的效率，使它們適合在硬件上部署。最先進的圖像生成模型具有巨大的計算和內存成本，這使得在資源受限的邊緣設備上部署它們并保持高質量具有挑戰性。本文的工作可以通過提供相同性能但使用更少的擴散步驟和更低成本模型來提高受控生成模型的效率。對于未來的工作，本文將探索將本文的工作與高效深度學習計算技術結合起來，以進一步提高效率。

結論

這項工作研究了通過操縱權重向圖像生成模型添加控制。本文引入了一種新的條件控制方法，稱為條件感知神經網絡（CAN），并提供了高效實用的設計，使CAN在實踐中可用。本文進行了廣泛的實驗，使用ImageNet進行類別條件生成和使用COCO進行文本到圖像生成，以評估CAN的有效性。CAN相對于先前的條件控制方法實現了一致且顯著的改進。本文還通過將CAN與EfficientViT結合，構建了一個新的擴散transformer模型系列。對于未來的工作，本文將把CAN應用于更具挑戰性的任務，如大規模文本到圖像生成、視頻生成等。

本文轉自 AI生成未來 ，作者：Han Cai等

原文鏈接:??https://mp.weixin.qq.com/s/ROX6yyh2pfOg0pcz778GZw??