擊敗SDXL登頂SOTA | CoMat:端到端的擴散模型微調策略
文章鏈接：https://arxiv.org/pdf/2404.03653
項目鏈接：??https://caraj7.github.io/comat/??
擴散模型在文本到圖像生成領域取得了巨大成功。然而，緩解文本提示和圖像之間的不對齊仍然具有挑戰性。不對齊的根因尚未得到廣泛調查。不對齊是由于token注意力激活不足引起的，進一步將這種現象歸因于擴散模型的訓練范式所導致的條件利用不足。
為了解決這個問題，本文提出了CoMat，這是一個端到端的擴散模型微調策略，具有圖像到文本概念匹配機制。利用圖像描述模型來衡量圖像到文本的對齊，并引導擴散模型重新訪問被忽視的token。此外還提出了一種新穎的屬性集中模塊來解決屬性綁定問題。在沒有任何圖像或人類偏好數據的情況下，僅使用了20K個文本提示來微調SDXL以獲得CoMat-SDXL。大量實驗證明，CoMat-SDXL在兩個文本到圖像對齊基準測試中明顯優于基線模型SDXL，并達到了最先進的性能水平。
介紹
最近，隨著擴散模型的引入，文本到圖像生成領域取得了可觀的進展。這些模型在基于文本提示創建高保真度和多樣化圖像方面表現出了顯著的性能。然而，對于這些模型來說，準確地與提示對齊仍然具有挑戰性，特別是對于復雜的提示。如下圖1所示，當前最先進的開源模型SDXL在生成實體或屬性方面存在失敗，例如，頂部行中提到的由蕾絲制成的羽毛和矮人。此外，它還未能理解提示中的關系。在圖1的中間行中，它錯誤地生成了一個維多利亞紳士和一塊上面有條河的被子。最近，各種工作提出了從語言學先驗或大語言模型（LLM）中引入外部知識以解決這個問題。
然而，對于不對齊問題的合理解釋仍然不足。為了進一步探索問題的原因，設計了一個試驗，研究文本token的交叉注意力激活值。如下圖2(a)所示的示例，禮服和文憑沒有出現在圖像中。與圖像中顯示的概念（即貓頭鷹和帽子）相比，這兩個token的激活值也處于較低水平。此外，我們還在圖2(b)中可視化了token激活的整體分布。
具體來說，模擬了訓練過程，使用預訓練的UNet對帶有文本標題的嘈雜圖像進行去噪。我們記錄了每個文本token的激活值，并在空間維度上取平均值。分布表明，在生成過程中激活仍然保持在較低水平。因此，確定了不對齊問題是由于對某些文本token的注意力不足引起的。這種行為最初源于文本到圖像擴散模型的訓練范式：給定文本條件c和配對圖像x，訓練過程旨在學習條件分布
。然而，文本條件僅僅是對去噪損失的附加信息。在學習文本中的每個概念時沒有明確的指導，擴散模型可能會輕易地忽略文本token的條件信息。
如前面圖2(a)所示，我們的方法大大增強了“禮服”和“文憑”的token激活，并使它們出現在圖像中。此外，圖2(b)表明我們的方法提高了整個分布中的激活。另外，由于描述模型在識別和區分屬性方面的不敏感性，發現屬性對齊仍然不理想。因此，引入了一個實體屬性集中模塊，其中屬性被強制在實體的區域內激活，以提高屬性對齊。最后，添加了一個保真度保持部分，以保持擴散模型的生成能力。值得注意的是，我們的訓練數據僅包括文本提示，不需要任何圖像文本對或人類偏好數據。此外，作為一種端到端方法，在推理過程中不引入額外的開銷，還展示了本文的方法與利用外部知識的方法是可組合的。
本文貢獻總結如下：
- 提出了CoMat，一種具有圖像到文本概念匹配機制的擴散模型微調策略。
- 引入了概念匹配模塊和實體屬性集中模塊來促進概念和屬性的生成。
- 與基線模型進行的廣泛定量和定性比較表明，所提出的方法具有更優越的文本到圖像生成能力。
相關工作
文本到圖像對齊
文本到圖像對齊是增強提示和生成圖像之間一致性的問題，涉及到存在性、屬性綁定、關系等多個方面。近期的方法主要通過三種方式來解決這個問題。
基于注意力的方法旨在修改或添加限制到UNet中的注意力模塊中的注意力圖。這種方法通常需要針對每個不對齊問題進行特定的設計。例如，Attend-and-Excite通過激勵每個對象的注意力分數來改善對象的存在性，而SynGen通過調節修飾語和實體之間的注意力圖的距離來增強屬性綁定。
基于規劃的方法首先從用戶的輸入或大語言模型（LLM）的生成中獲得圖像布局，然后在給定布局的條件下生成對齊的圖像。此外，一些工作提出使用其他視覺專家模型，如grounded-sam、多模態LLM或圖像編輯模型，進一步完善圖像。盡管這種集成將復合提示分解為單個對象，但它并不能解決下游擴散模型的不準確性，并且仍然存在不正確的屬性綁定問題。此外，在推理過程中會產生不可忽略的成本。
此外，一些工作旨在利用來自圖像理解模型的反饋來增強對齊。[21, 46]通過對由VQA模型選擇的對齊良好的生成圖像進行微調擴散模型，以策略性地偏置生成分布。其他工作提出以在線方式優化擴散模型。對于通用獎勵，[4,13]引入了RL微調。而對于可微分獎勵，[11, 55, 57]提出通過去噪過程直接反向傳播獎勵函數梯度。我們的概念匹配模塊可以被視為直接利用描述生成器作為可微分獎勵模型。
與本文的工作類似，[14]提出對生成的圖像添加描述，并優化生成的描述與文本提示之間的一致性。雖然圖像描述模型也參與其中，但它們未能提供詳細的指導。生成的描述可能會遺漏關鍵概念，并且可能會添加不必要的特征，這兩者都導致了優化目標的次優性。
圖像描述模型
這里的圖像描述模型是指那些在各種視覺和語言任務上進行了預訓練（例如，圖像文本匹配、（masked）語言建模），然后在圖像描述任務上進行了微調。已經提出了各種模型架構[25,26,49,50,60]。BLIP采用了融合編碼器架構，而GIT采用了統一的Transformer架構。最近，多模型大語言模型已經蓬勃發展。例如，LLaVA利用LLM作為文本解碼器，并取得了令人印象深刻的結果。
準備工作
方法
本文方法的整體框架如圖4所示，由三個模塊組成：概念匹配、屬性集中和保真度保持。將通過圖像描述模型介紹了圖像到文本概念匹配機制。然后詳細介紹了用于促進屬性綁定的屬性集中模塊。隨后，介紹了如何保留擴散模型的生成能力。最后，將三個部分結合起來進行聯合學習。
概念匹配
導致不對齊問題的根本原因在于對上下文信息的不完全利用。因此，即使所有的文本條件都呈現出來，擴散模型也很少關注某些token，因此在生成的圖像中缺少相應的概念。為了解決這個問題，我們的關鍵見解是在生成的圖像上添加監督以檢測缺失的概念。
通過利用圖像描述模型的圖像理解能力來實現這一目標，該模型可以根據給定的文本提示準確識別生成的圖像中不存在的概念。通過描述模型的監督，擴散模型被迫重新審視文本token，以搜索被忽視的條件信息，并且會賦予先前被忽視的文本概念更大的重要性，以實現更好的文本-圖像對齊。
實際上，來自描述模型的分數可以被視為微調擴散模型的差分獎勵。為了通過整個迭代去噪過程進行梯度更新，遵循 [55] 的做法，對去噪網絡 εθ 進行微調，通過簡單地停止去噪網絡輸入的梯度來確保訓練的有效性和效率。此外，值得注意的是，圖像中的概念涉及到一個廣泛的領域，我們的概念匹配模塊可以緩解各種不對齊問題，如對象存在和復雜關系。
屬性聚焦
已經有報道稱，將屬性與圖像聯系起來是文本到圖像擴散模型中的一個挑戰性問題。正如圖5所示的例子所基于的SDXL模型，單詞“紅色”和“藍色”的注意力大部分被激活在背景中，與其相應的對象幾乎沒有對齊。我們的概念匹配模塊可以在一定程度上緩解這個問題。然而，受限于說明模型對對象屬性的不敏感，性能改進是有限的。
其中 ⊙ 表示元素級乘法。像素級別的注意力損失進一步強制區域內的每個像素僅關注對象token，使用二元交叉熵損失函數：
其中 |A| 表示注意力圖中的像素數量。與 [52] 不同，提示中的某些對象可能由于不對齊而未出現在生成的圖像中。在這種情況下，像素級別的注意力損失仍然有效。當mask完全為零時，表示當前圖像中沒有任何像素應該關注缺失的對象token。此外，考慮到計算成本，僅在在線模型的圖像生成過程中的 r 個隨機選擇的時間步驟上計算上述兩個損失。
保真度保持
由于當前的微調過程完全由圖像描述模型和屬性與實體之間的先驗知識驅動，擴散模型可能會迅速過擬合獎勵，失去其原始能力，并生成退化的圖像，如下圖6所示。
目標是微調在線模型，以最小化這個對抗損失，同時訓練鑒別器以最大化它。
聯合學習
在這里，將描述模型損失、屬性集中損失和對抗損失結合起來，建立我們的在線擴散模型的訓練目標，具體如下:
其中 α、β 和 λ 是用于平衡不同損失項的縮放因子。
實驗
實驗設置
基礎模型設置。主要在 SDXL 上實現了所有實驗中的方法，它是最先進的開源文本到圖像模型。此外，在某些實驗中，還在 Stable Diffusion v1.5 (SD1.5) 上評估了我們的方法，以進行更全面的比較。對于描述模型，選擇了在 COCO 圖像描述數據上微調的 BLIP。至于保真度保持中的鑒別器，直接采用了 SD1.5 的預訓練 UNet。
數據集。由于傳遞給擴散模型的提示需要足夠具有挑戰性，以導致概念缺失，我們直接利用了幾個文本到圖像對齊基準上提供的訓練數據或文本提示。具體而言，訓練數據包括 T2I-CompBench提供的訓練集、HRS-Bench中的所有數據以及從 ABC-6K隨機選擇的 5,000 個提示。總體而言，這些約為 20,000 個文本提示。請注意，訓練集的構成可以根據目標改進的能力自由調整。
訓練細節。在本文的方法中，將 LoRA 層注入到在線訓練模型和鑒別器的 UNet 中，并保持所有其他組件凍結。對于 SDXL 和 SD1.5，在 8 個 NVIDIA A100 GPU 上訓練了 2,000 次迭代。對 SDXL 使用了batch大小為6，對 SD1.5 使用了batch大小為 4。從其他開放詞匯分割模型中選擇了 Grounded-SAM。DDPM 采樣器使用了 50 步來生成圖像，用于在線訓練模型和原始模型。特別地，遵循 [55]，僅在這 50 步中的 5 步中啟用梯度，其中屬性集中模塊也會被操作。此外，為了加快訓練速度，使用訓練提示提前生成并保存預訓練模型的生成潛在編碼，這些編碼稍后在微調過程中輸入鑒別器。
基準。在兩個基準上評估我們的方法：
- T2I-CompBench是一個用于組合式文本到圖像生成的基準。它包括來自 3 個類別（屬性綁定、對象關系和復雜構圖）和 6 個子類別（顏色綁定、形狀綁定、紋理綁定、空間關系、非空間關系和復雜構圖）的 6,000 個組合式文本提示。它采用了 BLIP-VQA 模型、CLIP 和 UniDet 來自動評估生成結果。每個子類別包含 700 個訓練提示和 300 個測試提示。
- TIFA是一個用于評估文本到圖像忠實度的基準。它使用預先生成的問題-答案對和 VQA 模型來評估生成結果。該基準包含了 4,000 個多樣化的文本提示和跨越 12 個類別的 25,000 個問題。
遵循這兩個基準的默認評估設置。
定量結果
將我們的方法與基線模型進行比較：SD1.5 和 SDXL，以及兩種最先進的開源文本到圖像模型：PixArt-α 和 Playground-v2。PixArt-α 采用了 transformer 架構，并利用了由大視覺語言模型自動標注的密集偽描述來輔助文本-圖像對齊學習。Playground-v2 的結構與 SDXL 類似，但在生成的圖像上更受歡迎。
T2I-CompBench。評估結果如下表1所示。
值得注意的是，由于評估代碼的演變，我們無法復現一些相關工作中報告的結果。我們展示的所有結果都基于 GitHub 上最新發布的代碼。與我們的基線模型相比，在所有六個子類別中，觀察到顯著的增益。具體而言，SD1.5 在顏色、形狀和紋理屬性上分別增加了 0.2983、0.1262 和 0.2004。在對象關系和復雜推理方面，CoMat-SD1.5 也獲得了顯著改進，空間關系的改進超過了 70%。使用我們的方法，SD1.5 甚至可以在與 PixArt-α 和 Playground-v2 相比中獲得更好或可比的結果。
當將我們的方法應用于更大的基礎模型 SDXL 時，仍然可以看到巨大的改進。CoMat-SDXL 在屬性綁定、空間關系和復雜構圖方面表現最佳。發現我們的方法在非空間關系中未能獲得最佳結果。推測這是由于我們文本提示中的訓練分布，其中大多數提示是描述性短語，旨在混淆擴散模型。在非空間關系中的提示類型可能僅占很小的一部分。CoMat-SDXL的架構流程如下圖8：
TIFA。 在下表2中展示了 TIFA 的結果。我們的 CoMat-SDXL 達到了最佳表現，比 SDXL 提高了 1.8 分。此外，CoMat 對 SD1.5 的提升達到了 7.3 分，遠遠超過了 PixArt-α。
定性結果
下圖3呈現了CoMat-SDXL與其他最先進的擴散模型的并排比較。觀察到這些模型與CoMat-SDXL相比，表現出較低的條件利用能力。圖3中的提示都包含與現實現象矛盾的概念。所有三個比較的模型都堅持原始偏見，并選擇忽略不現實的內容（例如，從茶壺中傾瀉出的瀑布、透明的小提琴、機器人企鵝和液態黃金的瀑布），這導致了不對齊。然而，通過訓練以與提示中的條件保持忠實對齊，CoMat-SDXL遵循了不現實的條件并提供了對齊良好的圖像。
更多 SDXL 和 CoMat-SDXL 的比較：
消融研究
在這里，旨在評估我們框架中每個組件和模型設置的重要性，并嘗試回答三個問題：
1）兩個主要模塊，即概念匹配和屬性集中，是否都是必要且有效的？
2）保真性保護模塊是否必要，以及如何選擇鑒別器？
3）如何選擇用于圖像描述模型的基礎模型？
概念匹配和屬性集中。在下表5中，展示了旨在確定概念匹配和屬性集中模塊有效性的T2I-CompBench結果。發現概念匹配模塊為基線模型帶來了主要的增益。此外，屬性聚焦模塊在T2I-CompBench的所有六個子類別中進一步改善了性能。
不同的鑒別器。 在這里消融了鑒別器的選擇。從COCO驗證集中隨機抽取了10K個文本-圖像對。計算了它的FID分數，以定量評估生成圖像的逼真程度。如下表3所示，第二行顯示了沒有鑒別器的情況，即沒有保真感知模塊，模型的生成能力大大惡化，FID分數從16.69增加到19.02。
還在下圖6中可視化了生成的圖像。如圖所示，沒有保真性保護，擴散模型生成的信封和天鵝形狀不良。這是因為擴散模型只嘗試欺騙描述，失去了其原始的生成能力。
此外，受到[45]的啟發，還嘗試使用預訓練的DINO來區分原始模型和在線訓練模型在圖像空間生成的圖像。然而發現DINO無法提供有效的指導，并嚴重干擾了訓練過程，導致FID分數甚至高于不應用鑒別器的情況。使用預訓練的UNet效果最好。FID分數與原始模型相同，在生成的圖像中沒有明顯的退化。
不同的圖像描述模型。 在下表4中展示了使用不同圖像描述模型的T2I-CompBench結果，其中底部行對應于沒有概念匹配優化的基線擴散模型。
發現所有三個描述模型都可以通過我們的框架提升擴散模型的性能，其中BLIP實現了最佳性能。特別要注意的是，通用多媒體大語言模型LLaVA無法捕捉與其他兩個在圖像描述任務上訓練的描述模型相比的性能。
限制
如何有效地將多模態大語言模型（MLLMs）納入我們提出的方法中，以改進文本到圖像擴散模型，仍然未經深入探討。鑒于它們在圖像-文本理解方面的SOTA水平，我們將專注于利用MLLMs實現更細粒度的對齊和生成保真度。此外，還觀察到CoMat在適應3D領域方面的潛力，促進文本到3D生成，實現更強的對齊。
結論
本文提出了CoMat，一種配備圖像到文本概念匹配的端到端擴散模型微調策略。利用圖像描述模型感知圖像中缺失的概念，并引導擴散模型審查文本token，找出被忽視的條件信息。這種概念匹配機制顯著增強了擴散模型的條件利用率。此外，還介紹了屬性集中模塊，進一步促進屬性綁定。我們的方法只需要文本提示進行訓練，無需任何圖像或人工標注的數據。通過大量實驗證明，CoMat在多個方面遠遠優于其基線模型，甚至超過了商業產品。希望我們的工作能夠啟發未來對齊問題及其解決方案的研究。
本文轉自 AI生成未來，作者：Dongzhi Jiang等
