從OmniGen的數據、模型設計與訓練策略中深度了解多模態統一生成模型 原創
?從OmniGen的訓練數據、模型設計與訓練策略中深度了解多模態統一生成模型。
今天詳細講講多模態大模型OmniGen~之前有總結過一篇多模態大模型的研究趨勢，
感興趣的小伙伴也可以具體看一下這篇介紹多模態大模型的文。?
OmniGen就屬于統一視覺模型這個方向中，將多個特定任務轉化為具備通用能力的圖像生成統一模型；本篇來具體介紹OmniGen的細節，以及未來前瞻技術研究可能值得借鑒的部分。
下面是一個快捷目錄。
一、OmniGen簡介
二、X2I數據集
三、OmniGen模型設計
四、OmniGen訓練策略
五、一些啟發
一、OmniGen簡介
1. 簡介
OmniGen，是用于統一圖像生成的新擴散模型，并且不需要額外的模塊（如ControlNet或IP-Adapter）來處理多樣化的控制條件。
2. 特點
- 統一性：從文本到圖像，內在統一支持多種圖像編輯、主題驅動的生成和視覺條件下的生成，另外把邊緣檢測和人體姿態識別也轉換成了圖像生成任務。
- 簡潔性：架構高度簡化，可以通過指令完成復雜任務，除去了額外文本編碼器，無需額外的預處理步驟。
- 知識遷移：OmniGen有效地在不同的任務之間轉移知識，管理未見過的任務和領域。
另外，此模型也存在一些推理能力和鏈式思維機制的潛在應用，但這塊不夠成熟。
二、X2I數據集
OmniGen的訓練數據集整個構造思路是非常值得借鑒的，個人感覺這塊甚至比模型本身更有學習意義，所以放在前面說，整個X2I數據集包含大約1億張圖像。
1. 文本到圖像
輸入為純文本，輸出為圖像。
1）開源數據
數據來源主要是一些開源數據集和合成數據：
開源數據包括Recap-DataComp（5600萬張圖像的子集）、SAM-LLaVA 、ShareGPT4V 、LAION-Aesthetic （400萬張圖像的子集）、ALLaVA-4V 、DOCCI、DenseFusion 和 JourneyDB。
2）合成數據
合成數據主要使用InternVL2 為內部數據和LAION-Aesthetic創建合成注釋。
3）數據使用技巧
在訓練初期，主要使用“文本到圖像”數據集來學習廣泛的圖像-文本匹配關系和多樣化知識；經過第三階段后，需要利用內部收集的1600萬張高質量圖像來提升生成圖像的美學質量。
2. 多模態到圖像
1）常見混合模態提示
任務類型包括圖像編輯（SEED-Data-Edit 、MagicBrush 和 InstructPix2Pix ）、人體運動（Something-Something ）、虛擬試穿（HR-VITON 和 FashionTryon）和風格遷移（stylebooth ）
數據構造形式：輸入為任意交錯的文本和圖像，輸出為符合文本和圖像條件的新圖像。
另外使用MultiGen數據集來學習利用額外的視覺條件進行細粒度空間控制，選擇六種代表性視覺條件：Canny、HED、Depth、Skeleton、Bounding Box 和分割
數據構造形式：輸入為文本提示和特定的視覺條件（如分割圖和人體姿態圖）；輸出為符合文本和圖像條件的新圖像。
2）主題驅動的圖像生成
數據任務與來源也包括兩大類，一個大規模基礎數據集（GRIT-Entity數據集）和一個高質量高級數據集（Web Images數據集）。
- GRIT-Entity數據集
基于GRIT數據集中圖像物體名稱標注，應用Grounding DINO模型進行文本到邊界框的對齊。基于邊界框使用SAM分割裁剪后的圖像，獲得物體掩碼。進一步使用MS-Diffusion模型重新繪制物體圖像，提升數據質量。構造了600萬對數據。
- Web Images數據集
基于構造的GRIT-Entity數據集由于直接提取原始圖像中的輸入數據，導致模型陷入簡單的復制粘貼模式。因此Web Images數據集使用知名人物的自然圖像構建。
- 首先，從Datacomp數據集中抽取了2000萬個Alt-text條目，并使用spaCy進行命名實體識別；
- 選擇了出現頻率最高的名字，并使用GPT-4o篩選出真實且知名的人物，最終得到2000個名字；
- 進一步擴展這2000個名字，包括密切相關的個體，得到了約10000個名字對。然后從搜索引擎中抓取這些個體和名字對的圖像；
- 使用交叉驗證策略，使用InternVL過濾單人和群體圖像，保留的單人和群體圖像隨后被標注了詳細的服飾和動作信息。
構建了包含533,000對圖像的數據集。
3. 計算機視覺任務
輸入是源圖像和自然語言表達，輸出是突出顯示相應對象的圖像。
這里的目的是為了增強模型的圖像生成能力；把從傳統計算機視覺任務中獲得的知識轉移到圖像生成任務中，從而實現更復雜的圖像生成能力。
- 對于低級視覺任務：
低級任務包括低光圖像增強、去雨、去模糊、修復、外推和著色，其中注釋本身就是圖像；并從GPT-4o生成的指令中隨機抽取文本指令添加進去。
- 對于高級視覺任務：
將所有注釋表示為圖像。
使用LAION作為源圖像，并選取注釋作為目標來構建圖像對（如源圖像及其人體姿態映射）。注釋內容包括人體姿態、深度映射、Canny和分割。
此外引入圖像分割的數據集，包括RefCOCO 、ADE20k 和ReasonSeg。
4. 小樣本到圖像
隨機選擇了一些示例，并將原始輸入與這些示例結合形成新的輸入。
三、OmniGen模型設計
1. 設計原則
1）通用性：接受任何形式的圖像和文本輸入以用于各種任務；
2） 簡潔性，避免過于復雜的結構設計和眾多的額外組件。
2. 網絡架構
OmniGen的框架主要是VAE + 預訓練大規模transformer。
輸入文本被標記化為token，輸入圖像通過VAE轉換為嵌入。OmniGen可以接受自由形式的多模態提示，并通過修正流方法生成圖像。
1）VAE
這里用的架構是SDXL，在訓練中會先凍結。
Stable Diffusion XL是一個二階段的級聯擴散模型，包括Base模型和Refiner模型：
- Base模型：主要工作和Stable Diffusion一致，具備文生圖，圖生圖，圖像inpainting等能力。（由U-Net，VAE，CLIP Text Encoder（兩個）三個模塊組成）
- Refiner模型：對Base模型生成的圖像Latent特征進行精細化，其本質上是在做圖生圖的工作。（由U-Net，VAE，CLIP Text Encoder（一個）三個模塊組成）
SDXL和之前的版本也是基于latent diffusion架構，對于latent diffusion，首先會采用一個auto encoder模型來圖像壓縮為latent，然后擴散模型用來生成latent，生成的latent可以通過auto encoder的decoder來重建出圖像。
2）Transformer
文中主要是基于Phi-3初始化。
Phi-3是微軟研究院推出的一系列小型語言模型（SLM），包括phi-3-mini、phi-3-small和phi-3-medium三個不同規模的版本，分別適合不同的應用場景和需求。具體模型結構可以見下圖。
圖引用自“2024最新大模型詳查之：微軟Phi-3系列”(https://zhuanlan.zhihu.com/p/695257957)
3. 輸入格式
模型的輸入是自由形式的多模態交錯文本和圖像。
1）文本：
利用Phi-3的分詞器處理文本，無需任何修改。
2）圖像：
- 特征提取：使用帶有簡單線性層的VAE來提取潛在表示。然后，通過線性嵌入潛在空間中的每個patch，將它們展平為一系列視覺token。
- 對不同寬高比圖像的處理：對輸入的視覺token應用標準的基于頻率的位置嵌入，并使用與SD3 相同的方法處理具有不同寬高比的圖像。
3）文本和圖像拼接
在將每個圖像序列插入到文本令牌序列之前，用兩個特殊令牌：“<img>”和“<\img>”封裝，另外還在輸入序列的末尾添加了時間步嵌入。
4. 注意力機制
注意力機制這里，把圖像應該作為一個整體來建模。
相對于LLM中的注意力機制，論文中修改了LLM中常見的因果注意力機制，對序列中的每個元素應用因果注意力，但在每個圖像序列內部應用雙向注意力。
目的是每個patch關注同一圖像內的其他patch，同時確保每個圖像只能關注之前出現過的其他圖像或文本序列。
5. 推理
推理過程中，隨機采樣一個高斯噪聲，然后應用流匹配方法來預測目標速度，通過多次迭代獲得最終的潛在表示。
最后，使用VAE將潛在表示解碼為預測圖像。默認的推理步數設置為50。
另外，由于注意力機制的存在，OmniGen可以通過使用kv緩存像LLMs一樣加速推理；存儲輸入的之前時間步和當前時間步的鍵值狀態在GPU上來計算注意力機制，而不需要冗余計算。?
四、OmniGen訓練策略
1. 訓練目標
使用修正流來優化模型的參數，流匹配通過在線性插值噪聲和數據之間進行前向過程。在步驟t， xt 定義為
其中 x 是原始數據，?～N(0,1) 是高斯噪聲。模型被訓練以直接回歸給定的噪聲數據 xt 、時間步 t 和條件信息 c 的目標速度。具體而言，目標是最小化均方誤差損失：
這里， vθ 表示模型預測的目標速度，θ 是模型參數。通過最小化上述損失函數，模型能夠學習如何從噪聲中恢復原始數據，進而實現圖像生成任務。
這種修正流方法在處理復雜的多模態數據時提供了一種有效的方式來訓練擴散模型，不僅提高了模型的效率，而且通過減少訓練時間和資源消耗，使模型更加易于優化。
對于圖像編輯任務，目標是在保持其他區域不變的同時修改輸入圖像的特定區域。
OmniGen會在圖像發生變化的區域放大損失，來避免直接將輸入圖像復制為輸出的情況，具體是根據輸入圖像 x′和目標圖像 x的潛在表示來計算每個區域的損失權重：
2. 訓練流程
在訓練過程中需要逐漸增加圖像分辨率。低分辨率在數據效率方面表現良好，而高分辨率可以提高生成圖像的美學質量。
五、一些啟發
1. 多模態模型中多任務數據非常重要
論文對比了跟SD3的效果，
- OmniGen模型只有38億參數，而SD3模型總共有127億參數（超過三倍）；
- OmniGen模型無額外文本編碼器，而SD3模型采用Enc-Dec架構外，還使用了額外的大參數量Encoder來編碼文本；
- OmniGen訓練數據只有1億張圖像數據，SD3使用了超過10億張圖像。
這確實說明了多任務數據X2I在提升文本到圖像能力方面的作用，以及多任務學習使模型能夠展現出新興能力。
2. 多模態統一模型已經具備In-context learning能力
當提供一個示例后，OmniGen能夠做出準確的預測，這表明上下文學習可以增強模型在不同領域的泛化能力。因此也可以嘗試用思維鏈來探索多模態統一大模型了。
3. 現階段的多模態大模型還比較脆弱，針對這些點還有很大的探索和發論文的空間
主要體現在下面的點：
- 對文本提示非常敏感：描述的文本也有技巧，越詳細圖像生成質量也會越高
- 對序列長度敏感：目前只能處理短文本段落，受限于算力資源，長圖片序列的處理和生成也有問題
- 圖像生成的幻覺問題：比如生成細節不可控，生成了不合常識的圖片等
- 泛化性較差：無法處理沒有見過的圖像類型
其實多模態大模型如何通過數據集、模型設計和訓練策略來實現統一任務指令對齊這塊還有很大的研究空間，OmniGen確實是個很值得參考和學習的例子；還有其他的一些可以發論文的思路大家也可以一起討論~
參考文獻：
[1] OmniGen: Unified Image Generation（https://export.arxiv.org/pdf/2409.11340
[2] https://zhuanlan.zhihu.com/p/721471869
[3]https://blog.csdn.net/weixin_47748259/article/details/135541372
?
本文轉載自公眾號瓦力算法學研所，作者：喜歡瓦力的卷卷
