機器之心專欄
機器之心編輯部
來自 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 所示，更多結果請參考論文。
圖 7: UniControl 在 Zero-shot tasks 上的可視化結果
6.總結
總的來說，UniControl 模型通過其控制的多樣性，為可控視覺生成提供了一個新的基礎模型。這種模型能夠為實作圖像生成任務的更高水準的自主性和人類控制能力提供可能。該研究期待和更多的研究者讨論和合作，以進一步推動這一領域的發展。
更多視覺效果
[1] Zhang, Lvmin, and Maneesh Agrawala. "Adding conditional control to text-to-image diffusion models." arXiv preprint arXiv:2302.05543 (2023).
[2] Mou, Chong, et al. "T2i-adapter: Learning adapters to dig out more controllable ability for text-to-image diffusion models." arXiv preprint arXiv:2302.08453 (2023).
[3] Zamir, Amir R., et al. "Taskonomy: Disentangling task transfer learning." Proceedings of the IEEE conference on computer vision and pattern recognition. 2018.
[4] Karras, Tero, Samuli Laine, and Timo Aila. "A style-based generator architecture for generative adversarial networks." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2019.
[5] Ho, Jonathan, Ajay Jain, and Pieter Abbeel. "Denoising diffusion probabilistic models." Advances in Neural Information Processing Systems 33 (2020): 6840-6851. APA
[6] Lin, Tsung-Yi, et al. "Microsoft coco: Common objects in context." Computer Vision–ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part V 13. Springer International Publishing, 2014.
[7] Schuhmann, Christoph, et al. "Laion-400m: Open dataset of clip-filtered 400 million image-text pairs." arXiv preprint arXiv:2111.02114 (2021).