多模態可控圖片生成統一模型來了，模型參數、推理代碼全部開源

作者：機器之心 2023-06-06 14:09:32

來自 Salesforce AI、東北大學、斯坦福大學的研究者提出了 MOE-style Adapter 和 Task-aware HyperNet 來實現 UniControl 中的多模態條件生成能力。UniControl 在九個不同的 C2I 任務上進行訓練，展示了強大的視覺生成能力和 zero-shot 泛化能力。

論文地址：https://arxiv.org/abs/2305.11147
代碼地址：https://github.com/salesforce/UniControl
項目主頁：https://shorturl.at/lmMX6

引言：Stable Diffusion 表現出了強大的視覺生成能力。然而，它們在生成具有空間、結構或幾何控制的圖像方面常常表現不足。ControlNet [1] 和 T2I-adpater [2] 等工作實現針對不同模態的可控圖片生成，但能夠在單一統一的模型中適應各種視覺條件，仍然是一個未解決的挑戰。UniControl 在單一的框架內合并了各種可控的條件到圖像（C2I）任務。為了使 UniControl 有能力處理多樣的視覺條件，作者引入了一個任務感知的 HyperNet 來調節下游的條件擴散模型，使其能夠同時適應不同的 C2I 任務。UniControl 在九個不同的 C2I 任務上進行訓練，展示了強大的視覺生成能力和 zero-shot 泛化能力。作者已開源模型參數和推理代碼，數據集和訓練代碼也將盡快開源，歡迎大家交流使用。

圖 1: UniControl 模型由多個預訓練任務和 zero-shot 任務組成

動機：現有的可控圖片生成模型都是針對單一的模態進行設計，然而 Taskonomy [3] 等工作證明不同的視覺模態之間共享特征和信息，因此本文認為統一的多模態模型具有巨大的潛力。

解決：本文提出了 MOE-style Adapter 和 Task-aware HyperNet 來實現 UniControl 中的多模態條件生成能力。并且作者建立了一個新的數據集 MultiGen-20M，包含 9 大任務，超過兩千萬個 image-condition-prompt 三元組，圖片尺寸≥512。

優點: 1) 更緊湊的模型 (1.4B #params, 5.78GB checkpoint)，更少的參數實現多個 tasks。2) 更強大的視覺生成能力和控制的準確性。3) 在從未見過的模態上的 zero-shot 泛化能力。

1.介紹

生成式基礎模型正在改變人工智能在自然語言處理、計算機視覺、音頻處理和機器人控制等領域的交互方式。在自然語言處理中，像 InstructGPT 或 GPT-4 這樣的生成式基礎模型在各種任務上都表現優異，這種多任務處理能力是最吸引人的特性之一。此外，它們還可以進行 zero-shot 或 few-shot 的學習來處理未見過的任務。

然而，在視覺領域的生成模型中，這種多任務處理能力并不突出。雖然文本描述提供了一種靈活的方式來控制生成的圖像的內容，但它們在提供像素級的空間、結構或幾何控制方面往往不足。最近熱門研究例如 ControlNet，T2I-adapter 可以增強 Stable Diffusion Model (SDM) 來實現精準的控制。然而，與可以由 CLIP 這樣的統一模塊處理的語言提示不同，每個 ControlNet 模型只能處理其訓練過的特定模態。

為了克服先前工作的限制，本文提出了 UniControl，一個能同時處理語言和各種視覺條件的統一擴散模型。UniControl 的統一設計可以享受到提高訓練和推理效率以及增強可控生成的優點。另一方面，UniControl 從不同視覺條件之間的固有聯系中獲益，來增強每個條件的生成效果。

UniControl 的統一可控生成能力依賴于兩個部分，一個是 "MOE-style Adapter"，另一個是 "Task-aware HyperNet"。MOE-style Adapter 有 70K 左右的參數，可以從各種模態中學習低級特征圖，Task-aware HyperNet 可以將任務指令作為自然語言提示輸入，并輸出任務 embedding 嵌入下游的網絡中，來調制下游模型的參數來適應不同模態的輸入。

該研究對 UniControl 進行預訓練，以獲得多任務和 zero-shot 學習的能力，包括五個類別的九個不同任務：邊緣 (Canny, HED, Sketch)，區域映射 (Segmentation, Object Bound Box)，骨架 (Human Skeleton)，幾何圖 (Depth, Normal Surface) 和圖片編輯 (Image Outpainting)。然后，該研究在 NVIDIA A100 硬件上訓練 UniControl 超過 5000 個 GPU 小時 (當前新模型仍在繼續訓練)。并且 UniControl 展現出了對新任務的 zero-shot 適應能力。

該研究的貢獻可以概括如下：

該研究提出了 UniControl，一個能處理各種視覺條件的統一模型 (1.4B #params, 5.78GB checkpoint)，用于可控的視覺生成。
該研究收集了一個新的多條件視覺生成數據集，包含超過 2000 萬個圖像 - 文本 - 條件三元組，涵蓋五個類別的九個不同任務。
該研究進行了實驗，證明了統一模型 UniControl 由于學習了不同視覺條件之間的內在關系，超過了每個單任務的受控圖像生成。
UniControl 表現出了以 zero-shot 方式適應未見過的任務的能力，展現了其在開放環境中廣泛使用的可能性和潛力。

2. 模型設計

圖 2: 模型結構。為了適應多個任務，該研究設計了 MOE-style Adapter，每個任務大約有 70K 個參數，以及一個任務感知 Task-aware HyperNet（約 12M 參數）來調制 7 個零卷積層。這個結構允許在一個單一的模型中實現多任務功能，既保證了多任務的多樣性，也保留了底層的參數共享。相比于等效的堆疊的單任務模型（每個模型大約有 1.4B 參數），顯著地減少了模型的大小。

UniControl 模型設計確保了兩個性質：

1) 克服來自不同模態的低級特征之間的不對齊。這有助于 UniControl 從所有任務中學習必要的和獨特的信息。例如，當模型將分割圖作為視覺條件時，可能會忽略 3D 信息。

2) 能夠跨任務學習元知識。這使得模型能夠理解任務之間的共享知識以及它們之間的差異。

為了提供這些屬性，模型引入了兩個新穎的模塊：MOE-style Adapter 和 Task-aware HyperNet。

MOE-style Adapter 是一組卷積模塊，每個 Adapter 對應一個單獨的模態，靈感來自專家混合模型（MOE），用作 UniControl 捕獲各種低級視覺條件的特征。此適配器模塊具有約 70K 的參數，計算效率極高。此后視覺特征將被送入統一的網絡中處理。

Task-aware HyperNet 則是通過任務指令條件對 ControlNet 的零卷積模塊進行調節。HyperNet 首先將任務指令投影為 task embedding，然后研究者將 task embedding 注入到 ControlNet 的零卷積層中。在這里 task embedding 和零卷積層的卷積核矩陣尺寸是對應的。類似 StyleGAN [4]，該研究直接將兩者相乘來調制卷積參數，調制后的卷積參數作為最終的卷積參數。因此每個 task 的調制后零卷積參數是不一樣的，這里保證了模型對于每個模態的適應能力，除此之外，所有的權重是共享的。

3. 模型訓練

不同于 SDM 或 ControlNet，這些模型的圖像生成條件是單一的語言提示，或如 canny 這樣的單一類型的視覺條件。UniControl 需要處理來自不同任務的各種視覺條件，以及語言提示。因此 UniControl 的輸入包含四部分: noise, text prompt, visual condition, task instruction。其中 task instruction 可以自然的根據 visual condition 的模態得到。

有了這樣生成的訓練配對，該研究采用 DDPM [5] 對模型進行訓練。

4. 實驗結果

圖 6: 測試集視覺對比結果。測試數據來自于 MSCOCO [6] 和 Laion [7]

與官方或該研究復現的 ControlNet 對比結果如圖 6 所示，更多結果請參考論文。

5.Zero-shot Tasks 泛化

模型在以下兩個場景中測試 zero-shot 能力:

混合任務泛化：該研究考慮兩種不同的視覺條件作為 UniControl 的輸入，一個是分割圖和人類骨骼的混合，并在文本提示中添加特定關鍵詞 "背景" 和 "前景"。此外，該研究將混合任務指令重寫為結合的兩個任務的指令混合，例如 "分割圖和人類骨骼到圖像"。

新任務泛化：UniControl 需要在新的未見過的視覺條件上生成可控制的圖像。為了實現這一點，基于未見過的和見過的預訓練任務之間的關系估計任務權重至關重要。任務權重可以通過手動分配或計算嵌入空間中的任務指令的相似度得分來估計。MOE-style Adapter 可以與估計的任務權重線性組裝，以從新的未見過的視覺條件中提取淺層特征。

可視化的結果如圖 7 所示，更多結果請參考論文。