拳打開源SOTA腳踢商業閉源的LI-DiT是怎樣煉成的？ 精華
文章地址：https://arxiv.org/pdf/2406.11831
今天和大家一起學習的這個工作展示的效果非常好，對提示的理解能力達到了新的高度。
僅基于解碼器的 Transformer 的大語言模型（LLMs）與 CLIP 和 T5 系列模型相比，已經展示出卓越的文本理解能力。然而，在文本到圖像擴散模型中利用當前先進的大語言模型的范例仍有待探索。本文觀察到一個不尋常的現象：直接使用大語言模型作為提示編碼器會顯著降低圖像生成中遵循提示的能力。本文發現了這個問題背后的兩個主要障礙：一個是大語言模型中下一token預測訓練與擴散模型中對有區別性的提示特征的要求之間的不一致；另一個是僅解碼器架構引入的內在位置偏差。為了解決這個問題，本文提出了一個新穎的框架來充分利用大語言模型的能力。通過精心設計的使用指南，有效地增強了用于提示編碼的文本表示能力，并消除了其內在的位置偏差。這使得能夠靈活地將最先進的大語言模型集成到文本到圖像生成模型中。
此外，本文還提供了一種將多個大語言模型融合到框架中的有效方式。 考慮到transformer架構所展示出的出色性能和擴展能力，本文進一步基于該框架設計了一個注入大語言模型的擴散Transformer（LI-DiT）。本文進行了廣泛的實驗，以在模型大小和數據大小方面驗證 LI-DiT。得益于大語言模型的固有能力和本文的創新設計，LI-DiT 的提示理解性能輕松超越了最先進的開源模型以及包括 Stable Diffusion 3、DALL-E 3 和 Midjourney V6 在內的主流閉源商業模型。強大的 LI-DiT-10B 將在進一步優化和安全檢查后可用。
介紹
擴散概率模型在高質量圖像合成方面帶來了顯著的改進。在諸如 CLIP 文本編碼器和 T5 系列等強大的提示編碼器的協助下，DALL-E 3 和 Stable Diffusion 3極大地增強了文本到圖像擴散模型中的提示理解能力。受 GPT 成功的鼓舞，一系列僅解碼器的大語言模型（LLM）出現了，并且與 CLIP 和 T5 系列模型相比展示出了卓越的文本理解能力，例如 LLaMA。然而，在擴散模型中有效利用這些強大的 LLM 的方法仍有待探索。
為了更好地理解 LLM 在擴散模型中的固有特性，本文首先使用基于transformer的擴散模型（DiT）進行實驗，并在 T2I-CompBench 基準上進行評估。遵循 DiT 和 PixArt-α 的設計，通過交叉注意力層將來自 LLM 最后一層的文本條件信息注入到擴散Transformer中。如下圖 2 所示，盡管 LLaMA3-8B 表現出更強的語言理解能力，但它在圖像到文本對齊基準上仍然未能趕上較小模型 T5-XL 的性能。同時，較大的變體 T5-XXL 相對于 T5-XL 取得了顯著的優勢。LLM 在文本理解和邏輯推理方面的強大能力在這種情況下尚未得到展示。基于這種異常現象，本文旨在探索 LLM 在提示編碼中的作用。
本文首先分析類T5的編解碼器模型和類GPT 的僅解碼器模型在優化目標和模型架構上的差異。masked語言模型優化和編解碼器架構設計賦予了 T5 編碼器內在的有效信息理解能力。然而，僅解碼器的大語言模型的優化目標側重于根據訓練數據分布預測具有最高概率的下一個token。
如下圖 4 所示，預訓練的大語言模型對給定的圖像提示提供了無意義的延續。這意味著大語言模型不關注給定圖像描述中的基本元素，并且大語言模型提取的文本表示不適合總結給定圖像的語義信息，導致與擴散模型的需求不一致。同時，本文發現大語言模型在理解提示后半部分提到的對象或屬性時通常會導致錯誤或遺漏。
這一觀察通過定量評估得到了進一步驗證。本文將此問題歸因于僅解碼器的大語言模型的因果注意力機制。在因果注意力層中，每個token只能關注自身和其他先前的token，而無法捕獲后面token的信息。這種結構性的信息不平衡挑戰了擴散模型理解復雜提示的能力。因此，這種不一致和位置偏差極大地阻礙了大語言模型成為擴散模型有效的文本編碼器。
為了解決這些問題，本文提出了一個新穎的框架，即融入大語言模型的擴散器（LLM-infused Diffuser），以充分利用強大的大語言模型來促進擴散模型在文本理解及后續方面的表現。首先，本文在提示之前明確插入一個指令，以減輕信息不一致的情況。基于大語言模型的指令遵循能力，本文利用人類指令來鼓勵語言模型關注與圖像生成相關的概念，包括對象、屬性和空間關系。此外，本文提出了一個語言token精煉器來解決位置偏差問題。這樣的設計通過雙向注意力機制促進了有效的全局表示建模。最后，協作精煉器合并并精煉來自多個大語言模型的文本表示，以進一步提升文本理解能力。這些有針對性的設計提供了一種在擴散模型中利用大語言模型能力的有效方式。
本文的融入大語言模型的擴散器可以輕松且靈活地整合到擴散模型中。考慮到Transformer架構的出色性能和擴展能力，本文進一步設計了一個融入大語言模型的擴散Transformer（LI-DiT）。本文進行了廣泛的實驗，以在不同的模型大小和數據大小上驗證 LI-DiT。得益于大語言模型的固有能力和本文的創新設計，LI-DiT 的提示理解性能輕松超越了最先進的開源模型以及包括 Stable Diffusion 3、DALL-E 3 和 Midjourney V6 在內的主流閉源商業模型。如上圖 1 所示，本文展示了一些由 LI-DiT-10B 生成的隨機抽樣案例。
使用語言模型進行提示編碼
如上一節所述，本文觀察到僅解碼器的大語言模型和編解碼器模型之間有兩個差異：優化目標和模型架構。具體來說，僅解碼器的大語言模型通常使用下一個token預測任務進行優化，而編解碼器模型則通過掩蔽語言建模任務進行訓練。此外，在僅解碼器的大語言模型中，序列中的前一個token不能關注后一個token，而在編解碼器模型中，序列中的每個token都可以相互關注。基于這些觀察結果，本文進行了精心的實驗，以研究這些差異如何影響大語言模型的提示編碼能力。
探索保留提示信息的能力
在 T5 模型的預訓練期間，輸入序列用mask進行格式化，模型通過預測被mask的內容從大量語言數據中學習。在這個過程中，編碼器負責從當前token序列中的所有token中提取信息。然而，僅解碼器的語言模型更側重于預測未來信息，而不是表示當前文本表示，這與擴散模型的使用不一致。為了更好地理解語言模型如何編碼提示的特征，本文將一個圖像提示輸入到 LLaMA3-8B 和 T5-XXL 中，以分析它們的輸出。如上 4 所示，T5-XXL 的輸出是輸入提示的重復，而 LLaMA3-8B 生成了一個不相關的擴展。這種現象進一步驗證了本文的假設。因此，盡管大語言模型具有更強的文本理解和推理能力，但這種限制損害了它們編碼提示的能力。
僅解碼器大語言模型的位置偏差
本文構建了一個基準來評估圖像提示中不同位置的所有形容詞-名詞組合的圖像-文本對齊情況。遵循傳統的文本到圖像生成基準，提取所有形容詞-名詞組合，并獲得它們在每個圖像提示中的相對位置。這些形容詞-名詞組合可以很容易地轉換為問題。然后，本文將生成的圖像和問題輸入到一個視覺問答模型中以獲得其對齊分數。請參考補充材料以獲取關于構建測試集的更多細節。如下圖 3 所示，本文計算每個形容詞-名詞組合在提示內的平均對齊分數和相對位置。可以觀察到帶有 T5 編碼器的擴散模型對位置變化表現出很強的穩健性，而帶有僅解碼器大語言模型的模型在后面的位置表現不佳。這種固有的位置偏差顯著損害了僅解碼器大語言模型的提示編碼能力。
LLM-infused Diffuser
整合大語言模型和擴散模型
為了彌合預訓練優化和提示編碼之間的差距，本文利用大語言模型的指令跟隨能力，以鼓勵它關注給定標題中的圖像內容。此外，本文還提出了精煉器模塊來減輕大語言模型文本embedding的固有位置偏差。通過結合這些設計，開發了一個名為“ LLM-infused Diffuser”的框架，它可以靈活地注入當前最先進的大語言模型，以釋放其強大的文本理解能力。
如下圖 5 所示，LLM-infused Diffuser的流程包括四個部分：
- 在圖像提示之前插入系統提示和指令，以鼓勵大語言模型關注圖像內容并突出其屬性。
- 帶有指令的圖像提示可以由多個凍結的大語言模型分別編碼。
- 采用不同的語言token精煉器模塊來消除這些大語言模型的文本embedding的位置偏差。
- 借助協作精煉器，來自大語言模型的文本特征被協同精煉，從而產生更穩健的表示。
輸入提示。受到大語言模型強大的指令跟隨能力的啟發，本文旨在利用這種能力迫使大語言模型關注提示中的關鍵圖像內容，并促進文本表示與文本到圖像合成任務之間的對齊。具體來說，本文提議在傳統的圖像描述之前插入自定義指令。這樣的指令提示大語言模型關注關鍵圖像內容，例如圖像中物體的屬性以及物體之間的空間關系。在本文的實驗中，采用了一個簡單的指令：通過詳細描述物體的顏色、形狀、大小、紋理、數量、文字和空間關系來描述圖像。如上圖 4 所示，如果不提供明確的指令，大語言模型往往會生成與圖像上下文無關的內容。當向大語言模型提供指令和圖像提示時，它將遵循指令專注于與圖像相關的概念，以詳細描述圖像，并根據給定的提示提供對齊的表示。大語言模型的輸出embedding將由后續的精煉器模塊進一步處理。
語言token精煉器。在大語言模型的因果注意力層中，當前token只能關注到之前的token，因此這極大地損害了全局文本表示建模。例如，文本token序列中的最后一個token只能被其自身關注。為了減輕僅解碼器大語言模型的這種位置偏差，本文插入一個語言token精煉器模塊來精煉每個大語言模型有偏差的輸出表示。如上圖 5 所示，每個精煉器模塊包含一堆Transformer塊，它由一個自注意力層、一個前饋層（FFN）和一個自適應門控模塊組成。對于自注意力層，直接丟棄大語言模型的因果mask來執行全注意力，這使得后面的token的表示可以被前面的token關注。每層的輸出特征由自適應門控網絡控制，其權重初始化為零以獲得更好的訓練穩定性。具體來說，本文首先對大語言模型表示進行平均池化，然后將池化后的表示與時間步 t 的embedding通過元素級求和進行合并。門控網絡將這種時間步感知和上下文感知的表示作為輸入來執行精確的信息注入。精煉器的最終輸出表示將共同被饋送到協作精煉器中進行增強。
協作精煉器。為了進一步提高文本理解能力，本文采用多個大語言模型和語言token精煉器進行提示編碼，并通過所提出的協作精煉器協同精煉這些表示。來自多個語言token精煉器的表示由多個平行分支分別處理，并且一個分支中的每個塊都包含一個交叉注意力和前饋神經網絡層。此外，本文使用一種調制機制，根據時間步和文本上下文來調節協作精煉器的每一層。這種調制采用與語言token精煉器中上述門控網絡相同的輸入。該模塊中的分支通過多個平行的交叉注意力層連接，在這里文本表示可以進行協同精煉。具體來說，交叉注意力層將當前分支的特征作為查詢，將其他分支的特征作為鍵和值來精煉當前特征。最后截斷輸出的token序列，丟棄指令token，并通過連接混合這兩種表示。這種混合并精煉后的表示可以靈活地集成到擴散模型中，以提供有區分度的文本條件信息。
LLM-infused Diffuser Transformer
本文提出的注入大語言模型的 Transformer可以靈活地集成到當前的擴散模型中。考慮到擴散 Transformer 顯著的擴展能力，本文開發了一個名為注入大語言模型的擴散 Transformer（LI-DiT）的擴散模型。
遵循 DiT 的范例，LI-DiT 將來自變分自動編碼器（VAE）潛在空間的有噪表示作為輸入，并將空間輸入轉換為token序列。LI-DiT 的每個Transformer塊包含一個自注意力層、一個交叉注意力層、一個前饋神經網絡層和調制模塊。交叉注意力層可以將注入大語言模型的擴散器提取的文本條件信息注入token序列中。調制模塊接收時間步embedding和文本表示以提供額外的條件信息。與之前工作中的二維位置embedding設計不同，本文采用基于卷積的位置embedding。在擴散Transformer中的補丁化層之后，直接采用一個 ResBlock 作為位置embedding模塊。卷積算子的平移不變性可以有效地為Transformer算子引入位置信息。因此，LI-DiT 可以支持任意分辨率的圖像生成，而無需額外的設計修改。
大規模的 Transformer 模型通常會遭受不穩定的梯度和數值精度問題，導致在訓練過程中損失發散。為了解決訓練不穩定的問題，本文納入了在大規模視覺或語言模型訓練中采用的若干策略。首先，在自注意力層和交叉注意力層中都引入了 QK 歸一化。RMSNorm 層將在點積計算注意力分數之前對查詢和關鍵token進行歸一化。這樣的操作確保了注意力分數的數值穩定性，并避免了來自分布外值的不穩定梯度。此外，考慮到 bfloat16 更廣泛的數值表示范圍，本文最終采用了 bfloat16 混合精度訓練策略。
實驗
實現細節
模型架構。本文的實驗默認在較小的模型 LI-DiT-1B 上進行。采用 LLaMA3-8B 和 Qwen1.5-7B并結合多模態指令微調 作為 LI-DiT-1B 和 LI-DiT-10B 的雙文本編碼器。對于消融研究基線，本文只保留 LLaMA3-8B 以降低訓練成本。本文在語言token精煉器中采用 2 個塊，在協作精煉器中采用 1 個塊。在實驗中， 本文將來自倒數第三個Transfomer 塊的文本embedding作為每個大語言模型的輸出。關于 LI-DiT-1B 和 LI-DiT-10B 的詳細架構，請參考補充材料。
訓練數據。所有的探索和消融實驗都在 ImageNet 數據集和 CC12M 數據集的一個子集上進行訓練。本文為 ImageNet 的每個樣本分配“一張{class}的照片”的文本提示，并從 CC12M 中隨機選擇 130 萬圖像-文本對。遵循之前的工作，本文混合了原始描述和由CogVLM生成的合成描述。將 LI-DiT 與其他領先的模型進行比較時，本文使用了一個具有十億級圖像-文本對的大規模訓練數據集，包括 LAION-5B和其他包含英語和中文的內部數據集，這使 LI-DiT 具有雙語理解能力。遵循Stable Diffusion，當來自 LAION 的圖像-文本對的美學評分低于 4.7 時，本文將其移除。低分辨率圖像和包括網址和標簽的低質量提示也被移除。具體來說，本文僅從這個大規模數據集中抽樣 3000 萬圖像-文本對來訓練 LI-DiT-1B，并使用所有十億級對來訓練 LI-DiT-10B。
訓練細節。遵循潛在擴散模型（LDM）的范例，本文利用一個變分自編碼器（VAE）編碼器將圖像表示投影到潛在空間。本文訓練一個具有 8 倍下采樣率和 16 個通道的變分自編碼器以獲得更好的圖像生成效果。本文不使用任何數據增強策略。遵循 RAPHEL中的多尺度訓練，根據圖像的縱橫比將圖像分組。只有具有相似縱橫比的圖像才會構成一個批次。對于在 300 萬圖像-文本對上進行的消融實驗，本文以 256 的分辨率用 256 的batch大小和 的學習率訓練模型 30 萬次迭代。對于 LI-DiT-1B 的訓練，本文將批次大小增加到 2048 并將迭代次數增加到 50 萬。在訓練 LI-DiT-10B 時，批次大小是 4096，迭代次數超過 100 萬。本文在訓練期間直接采用 512 的分辨率，然后用高質量數據將其微調至 1024 分辨率以進一步提高美學質量。
評估指標。對于定量評估，本文主要考慮 T2I-CompBench、DPG-Bench和 GenEval 基準。本文還引入了人類評估以更好地理解藝術和美學質量。請注意，消融研究中的“T2I-平均”是指 T2I-CompBench 屬性指標的平均得分。
性能比較
定量評估。在定量評估中，本文專注于生成圖像與輸入提示之間的一致性。如下表 1 所示，本文選擇 T2I-CompBench、DPG-Bench 和 GenEval 基準來評估 LI-DiT-1B 和 LI-DiT-10B 的生成能力。T2I-CompBench 和 GenEval 基準由簡短提示組成，側重于組合評估。DPG-Bench 是用復雜密集的提示構建的。與像 SDXL 和 PixArt-α這樣的開源學術作品相比，LI-DiT-1B 在所有基準上都大幅優于它們。本文還將 LI-DiT-10B 與 DALL-E 3 和Stable Diffusion 3（8B）這兩個主流的閉源商業模型進行比較。這一顯著的改進進一步驗證了本文的大語言模型融合擴散器的有效性。
人類評估。定量評估指標不能直接衡量藝術和美學質量。遵循先前的工作，本文也進行人類評估，以令人信服地將 LI-DiT-10B 與Stable Diffusion 3、DALL-E 3 和Midjourney V6 進行比較。本文的評估數據集包含 200 個具有不同風格和場景的提示。來自 LI-DiT-10B 的圖像和來自競爭對手的圖像將構成一個評估對。人類評估者將從圖像質量和圖像-文本對齊的角度比較圖像對。下圖 6 中的結果表明，LI-DiT-10B 在圖像-文本對齊和圖像質量方面都可以超過 DALLE-3 和Stable Diffusion 3。與最受歡迎的商業模型Midjourney V6 相比，LI-DiT-10B 在圖像-文本對齊方面表現出領先的能力，同時具有類似的圖像-文本質量表現。在下圖 7 中，本文展示了一些隨機抽樣的案例以進行清晰比較。此外還在圖 8 和圖 9 中提供了一些高質量圖像。
消融研究
組件級消融研究。如下表 2 所示，本文進行組件級消融研究。本文采用預訓練的 LLaMA3-8B 的 DiT 作為基線設置。首先，觀察到在將指令引入輸入提示或把語言token精煉器合并到基線后性能有一致的提升。當同時利用這兩種設計時，在兩個基準上的圖像-文本對齊性能繼續提高。此外，本文引入一個額外強大的大語言模型，Qwen1.5-7B，并進行多模態微調以驗證協同精煉器的有效性。大語言模型融合策略進一步增強了擴散模型對提示的理解能力。這些結果清楚地驗證了每個提出組件的有效性。
因果Mask的影響。在這個實驗中本文研究因果Mask對提示編碼的影響。如下表 3 所示，在大語言模型之后插入具有完全注意力的語言token精煉器顯著提高了性能。然而，這個精煉器未能提高具有雙向注意力的 T5 編碼器的性能。如果將大語言模型的因果Mask引入到精煉器中，在 LLaMA3-8B 和 Qwen1.5-7B 中都會出現嚴重的性能下降。這些結果表明因果Mask是損害大語言模型提示編碼能力的一個核心因素，而本文提出的精煉器可以消除這種位置偏差。
指令的效果。為了驗證指令的有效性，本文在下表 4 中進行了一項消融實驗。首先，本文發現提示指令對于使用沒有指令微調的基礎 LLaMA3-8B 模型未能帶來增益。如果將基礎模型改為多模態指令微調變體，對齊分數可以顯著提高。由于指令微調帶來的強大遵循指令能力，插入指令可以進一步提升性能。這個結果表明多模態指令微調數據有助于大語言模型更好地描述圖像并突出圖像內的關鍵元素。此外，指令能夠鼓勵大語言模型關注給定提示中的圖像內容。
語言token精煉器設計。如下表 5 所示，本文對語言token精煉器的設計進行實驗。首先，將本文的模型與精煉器中具有不同塊數的其他變體進行比較。觀察到當精煉器中的塊數增加時性能有一致的提升。然而，當語言token精煉器中有 2 個塊時這種提升并不顯著。因此，本文在token精煉器中采用 2 個塊來實現復雜性和性能之間的最佳平衡。此外，本文還消融了精煉器中門控網絡的效果。當移除門控網絡時，在兩個基準上的性能都下降了。這表明時間和文本上下文的條件信息有助于更好的圖像-文本對齊。
協同精煉器的效果。如下表 6 所示，本文觀察到具有簡單融合技術的模型可以優于具有單個大語言模型的其他對應模型。此外，協同精煉器可以在這種拼接融合的基礎上進一步提升性能。這樣的結果表明一種有效的表示融合方法可以進一步增強大語言模型的能力。
結論
針對在采用僅解碼器的大語言模型來編碼提示時在文本到圖像生成任務中表現不佳的情況，本文探索了大語言模型在擴散模型提示編碼中的作用。通過實驗和分析，本文確定了限制僅解碼器的大語言模型作為擴散模型的有效文本編碼器的核心因素，即下一個token預測訓練與擴散模型中對判別性提示特征的要求之間的不一致，以及僅解碼器架構引入的內在位置偏差。為了處理這些問題，提出了一個新穎的框架來充分利用大語言模型的能力。本文還基于該框架進一步設計了一個注入大語言模型的擴散 Transformer（LI-DiT）。LI-DiT 超越了最先進的開源模型以及包括 Stable Diffusion 3、DALLE-3 和 Midjourney V6 在內的主流閉源商業模型。
限制和潛在的社會負面影響
由于計算資源有限，本文對具有 70 億參數的大語言模型進行實驗。在未來的工作中，本文將在具有 130 億或 700 億參數的更大的大語言模型中進一步驗證注入大語言模型的擴散的有效性。潛在的負面社會影響是圖像可能包含有誤導性或虛假信息。本文將在數據處理方面進行廣泛的努力來處理這個問題。
本文轉自 AI生成未來 ，作者：Bingqi Ma等
