作者：Siyu Jiao等

解讀：AI生成未來

亮點(diǎn)直擊

首次提出思考驅(qū)動(dòng)的視覺生成框架： ThinkGen 是第一個(gè)顯式利用 MLLM 的思維鏈（CoT）推理來處理各種生成場景的思考驅(qū)動(dòng)視覺生成框架。這解決了現(xiàn)有方法在泛化性和適應(yīng)性方面的局限性，因?yàn)樗鼈兺ǔａ槍μ囟▓鼍霸O(shè)計(jì) CoT 機(jī)制。

解耦架構(gòu)： ThinkGen 采用解耦架構(gòu)，將預(yù)訓(xùn)練的 MLLM 和 Diffusion Transformer (DiT) 分開。其中，MLLM 負(fù)責(zé)根據(jù)用戶意圖生成定制指令，DiT 則根據(jù)這些指令生成高質(zhì)量圖像。這種設(shè)計(jì)克服了現(xiàn)有框架缺乏高級推理能力的問題。

視覺生成指令提煉 (VGI-refine) 模塊： 為了解決 CoT 推理過程中冗余信息的問題，提出了 VGI-refine 模塊。該模塊從 MLLM 的推理鏈中提取簡潔的指令信息，并將其與可學(xué)習(xí)的 Prepadding States 連接起來，從而實(shí)現(xiàn) MLLM 表示分布的自適應(yīng)調(diào)整，更好地與 DiT 的要求對齊。

可分離 GRPO-based 訓(xùn)練范式 (SepGRPO)： 提出了一種名為 SepGRPO 的可分離強(qiáng)化學(xué)習(xí)訓(xùn)練范式，它在 MLLM 和 DiT 模塊之間交替進(jìn)行強(qiáng)化學(xué)習(xí)。這種靈活的設(shè)計(jì)支持在不同數(shù)據(jù)集上進(jìn)行聯(lián)合訓(xùn)練，從而促進(jìn)了在廣泛生成場景中有效的 CoT 推理。

在多生成場景中實(shí)現(xiàn) SOTA 性能： 廣泛的實(shí)驗(yàn)證明，ThinkGen 在多個(gè)生成基準(zhǔn)測試中實(shí)現(xiàn)了穩(wěn)健的、最先進(jìn)的性能，尤其是在推理密集型任務(wù)中表現(xiàn)出色。

總結(jié)速覽

解決的問題

現(xiàn)有的多模態(tài)大語言模型（MLLMs）在理解任務(wù)中展現(xiàn)了思維鏈（CoT）推理的有效性，但其在生成任務(wù)中的擴(kuò)展仍然處于初級階段。

現(xiàn)有針對生成任務(wù)的 CoT 機(jī)制通常是針對特定場景定制的，這限制了其泛化能力和適應(yīng)性，導(dǎo)致在更廣泛的任務(wù)中性能下降，并且通常需要手動(dòng)干預(yù)來激活不同生成任務(wù)的 CoT 推理。

當(dāng)前框架普遍缺乏高級推理能力。

提出的方案

本文提出了 ThinkGen，這是一個(gè)通用且思考驅(qū)動(dòng)的視覺生成框架，旨在顯式地利用 MLLM 的 CoT 推理能力，以解決各種生成場景中的復(fù)雜任務(wù)。

通過解耦 MLLM 和 Diffusion Transformer (DiT) 的架構(gòu)，實(shí)現(xiàn)了在生成之前制定高質(zhì)量計(jì)劃的能力。

應(yīng)用的技術(shù)

解耦架構(gòu)： 框架包含一個(gè)預(yù)訓(xùn)練的 MLLM（用于生成定制指令）和一個(gè) Diffusion Transformer (DiT)（用于生成高質(zhì)量圖像）。

思維鏈（CoT）推理： 顯式地利用 MLLM 的 CoT 推理能力，通過生成明確的中間步驟來系統(tǒng)地解決復(fù)雜任務(wù)。

視覺生成指令提煉 (VGI-refine)： 引入 VGI-refine 模塊，從 MLLM 的自回歸 CoT 輸出中篩選掉冗余信息，并結(jié)合可學(xué)習(xí)的 Prepadding States 進(jìn)行對齊。

可分離 GRPO-based 訓(xùn)練范式 (SepGRPO)： 一種交替在 MLLM 和 DiT 模塊之間進(jìn)行強(qiáng)化學(xué)習(xí)的訓(xùn)練策略，旨在鼓勵(lì) MLLM 生成與 DiT 偏好對齊的指令，并使 DiT 基于這些指令生成高質(zhì)量圖像。

達(dá)到的效果

在多個(gè)生成基準(zhǔn)測試中實(shí)現(xiàn)了穩(wěn)健的、最先進(jìn)的性能。ThinkGen 在采用 CoT 推理時(shí)，在廣泛的生成場景中取得了卓越的性能（如下圖1所示）。實(shí)現(xiàn)了在多種生成場景下的有效 CoT 推理，從而增強(qiáng)了泛化能力。

架構(gòu)方法

ThinkGen 采用解耦架構(gòu)，包含一個(gè)預(yù)訓(xùn)練的 MLLM 和一個(gè) Diffusion Transformer (DiT)。MLLM 負(fù)責(zé)根據(jù)用戶意圖生成定制指令，而 DiT 則根據(jù)這些指令生成高質(zhì)量圖像。這種解耦設(shè)計(jì)確保了每個(gè)組件的最佳性能，同時(shí)保持了系統(tǒng)的可擴(kuò)展性和模塊化，如下圖3所示。

多模態(tài)大型語言模型 (MLLM)

ThinkGen 利用 MLLM 處理視覺和文本輸入，并通過自回歸生成進(jìn)行 CoT 推理。MLLM 使用 Qwen3-VL-8B-Think 初始化。對于圖像生成任務(wù)，本文設(shè)計(jì)了一個(gè)專門的系統(tǒng)提示（[SYS]）來促使 MLLM 理解用戶意圖并提供適當(dāng)?shù)闹貙懼噶睢ｋS后，從  標(biāo)記之后生成的隱藏狀態(tài)的最后兩層被提取出來，作為 DiT 的條件輸入。經(jīng)驗(yàn)結(jié)果表明，使用最后兩層隱藏狀態(tài)對視覺生成有顯著益處。

Diffusion Transformer (DiT)

ThinkGen 采用標(biāo)準(zhǔn)的 DiT 架構(gòu)，并使用 OmniGen2-DiT-4B初始化。MLLM 的輸出作為 DiT 的條件文本輸入。在圖像編輯任務(wù)中，額外的參考圖像通過 VAE處理并作為條件視覺輸入。視覺和文本輸入與噪聲潛在特征連接，實(shí)現(xiàn)跨模態(tài)的聯(lián)合注意力。本文采用一個(gè)簡單的線性層作為連接器，以對齊來自多個(gè)條件輸入的特征。

VGI-refine

為了解決 MLLM 自回歸思維鏈（CoT）輸出中存在的冗余信息，本文引入了視覺生成指令提煉（VGI-refine）模塊，包含兩個(gè)步驟。首先，從 MLLM 生成的文本標(biāo)記中提取  特殊標(biāo)記之后的指令標(biāo)記，從而分離出用于下游圖像生成的必要 CoT 結(jié)果。其次，將 K 個(gè)可學(xué)習(xí)的 Prepadding States 連接到提取的指令標(biāo)記上。這種連接調(diào)節(jié)了輸出隱藏狀態(tài)的數(shù)據(jù)分布，尤其對于短指令（例如，“生成一只狗”或“移除貓”）特別有益。最終得到的精煉指令狀態(tài)作為條件輸入提供給 DiT。

訓(xùn)練策略

ThinkGen 的訓(xùn)練分為五個(gè)不同的階段，如下圖4所示。首先，對 DiT 進(jìn)行監(jiān)督預(yù)訓(xùn)練（階段 1-3）以確保高質(zhì)量的圖像生成。隨后，引入了一種稱為 SepGRPO 的可分離 MLLM 和 DiT 強(qiáng)化學(xué)習(xí)方法（階段 4-5）。通過 SepGRPO 訓(xùn)練，MLLM 學(xué)習(xí)生成與 DiT 偏好最佳對齊的描述或編輯指令，而 DiT 則在此基礎(chǔ)上進(jìn)一步優(yōu)化以生成更優(yōu)質(zhì)的圖像。

監(jiān)督預(yù)訓(xùn)練

監(jiān)督預(yù)訓(xùn)練階段（階段 1-3）旨在將 DiT 與 MLLM 對齊，同時(shí)提高圖像生成質(zhì)量。本文采用 Rectified Flow訓(xùn)練范式，通過最小化 Flow Matching 目標(biāo)直接回歸速度場 。

其中  表示目標(biāo)速度場。

輸入格式: 在預(yù)訓(xùn)練階段，為了避免昂貴的重寫每個(gè)標(biāo)題或編輯指令的成本，本文構(gòu)建了偽 CoT 模板來模擬 MLLM 的 CoT 過程。具體來說，  內(nèi)的內(nèi)容留空，并簡單地重復(fù)原始標(biāo)題或編輯指令作為答案。

階段1 對齊: 在此階段，本文引入 K 個(gè)可學(xué)習(xí)的預(yù)填充狀態(tài)，并通過僅訓(xùn)練線性連接器來將 DiT 與 MLLM 對齊，同時(shí)保持 MLLM 和 DiT 凍結(jié)。

階段2 預(yù)訓(xùn)練: 在此階段，所有 DiT 參數(shù)都是可訓(xùn)練的。訓(xùn)練語料庫包含 60M 圖像樣本，包括文本到圖像、圖像編輯、文本渲染和上下文生成數(shù)據(jù)。

階段3 高質(zhì)量微調(diào): 在監(jiān)督微調(diào)階段，本文構(gòu)建了一個(gè) 0.7M 的高質(zhì)量子集，以增強(qiáng) DiT 的指令遵循能力和圖像美學(xué)。

SepGRPO

SepGRPO，一種強(qiáng)化學(xué)習(xí)訓(xùn)練策略，旨在鼓勵(lì) MLLM 生成與 DiT 偏好最佳對齊的標(biāo)題/編輯指令，同時(shí)使 DiT 能夠根據(jù)這些指令生成更高質(zhì)量的圖像。SepGRPO 解耦了文本和視覺的展開過程：首先，固定 DiT，通過聯(lián)合多任務(wù)訓(xùn)練將 GRPO 應(yīng)用于 MLLM；然后，固定 MLLM，將 GRPO 應(yīng)用于 DiT。

輸入格式: 在策略訓(xùn)練期間，本文設(shè)計(jì)了一個(gè)專門的 [SYS] 來促進(jìn)冷啟動(dòng)，允許 MLLM 探索 DiT 偏好的文本條件。具體來說，本文將 [SYS]、輸入樣本 [C] 和特殊  標(biāo)記連接起來作為 MLLM 的輸入。

階段4 MLLM-GRPO: 在此階段，將 GRPO 應(yīng)用于 MLLM，以鼓勵(lì)生成與 DiT 偏好對齊的重寫文本。本文在多個(gè)場景下優(yōu)化 MLLM，以增強(qiáng) CoT 推理的泛化能力。本文選擇了五個(gè)代表性生成場景：語義組合、推理生成、文本渲染、圖像編輯和反射。如下表1所示，針對每個(gè)場景，本文收集并整理了專用數(shù)據(jù)集，并設(shè)計(jì)了相應(yīng)的規(guī)則模型來指導(dǎo)優(yōu)化。

對于 MLLM 的每個(gè)輸入，本文從策略  執(zhí)行  次 rollout 以生成軌跡 ，然后 DiT 使用這些軌跡生成相應(yīng)的圖像。規(guī)則模型用于計(jì)算每個(gè)軌跡的獎(jiǎng)勵(lì) 。隨后，以組相對的方式計(jì)算第  個(gè)軌跡的優(yōu)勢 ：

然后通過優(yōu)化 GRPO 目標(biāo)來更新策略 ，該目標(biāo)是帶有 KL 散度正則化的裁剪替代函數(shù)：

其中  表示  和  輸出當(dāng)前 token 的概率比。在此過程中，DiT 和規(guī)則模型共同充當(dāng)獎(jiǎng)勵(lì)模型。

階段5 DiT-GRPO: 在此階段，本文應(yīng)用 FlowGRPO來增強(qiáng) DiT 的指令遵循能力。本文利用來自 Simple Scene 和 Text Rendering 場景的數(shù)據(jù)，以及它們相應(yīng)的獎(jiǎng)勵(lì)計(jì)算方法。實(shí)驗(yàn)結(jié)果總結(jié)

ThinkGen 在多種生成場景下進(jìn)行了評估，并與現(xiàn)有方法進(jìn)行了比較。結(jié)果表明，ThinkGen 在推理生成、推理編輯、文本到圖像生成和圖像編輯方面都取得了顯著的性能提升。

推理生成

如下表2所示，在 WISEBench 基準(zhǔn)測試中，ThinkGen 表現(xiàn)出優(yōu)于直接生成方法的顯著優(yōu)勢。通過利用 CoT 推理，ThinkGen 實(shí)現(xiàn)了 +21% (0.55 → 0.76) 的顯著改進(jìn)，并在 WISEBench 上建立了新的最先進(jìn)性能。

推理編輯

如下表3所示，在 RISEBench 上，ThinkGen 的 CoT 推理能力顯著超越了開源模型（3.6 → 13.0），并取得了與閉源模型 Gemini-2.0 相當(dāng)?shù)慕Y(jié)果。

文本到圖像生成

如下表4所示，ThinkGen 在 GenEval、DPG-Bench 和 CVTG 基準(zhǔn)測試中，通過 CoT 推理，始終在所有場景中表現(xiàn)出改進(jìn)，并在許多知名模型中取得了最佳結(jié)果。這些結(jié)果表明 ThinkGen 具有強(qiáng)大的指令遵循和文本渲染能力。

圖像編輯

如下表5所示，在 ImgEdit 上，ThinkGen 與一系列開源模型相比，顯示出顯著優(yōu)越的指標(biāo)，取得了與 GPT-4o 相當(dāng)?shù)男阅堋?/p>

消融研究

訓(xùn)練階段消融: 如下表6所示，逐步應(yīng)用每個(gè)訓(xùn)練階段對 ThinkGen 的性能都有貢獻(xiàn)。

僅訓(xùn)練連接器（Stage1）導(dǎo)致文本渲染性能不佳（CVTG: 0.28），表明 MLLM 和 DiT 之間缺乏足夠的細(xì)粒度對齊。大規(guī)模預(yù)訓(xùn)練（Stage2）顯著改善了圖像質(zhì)量，GenEval 增加了 10%，WISE 增加了 9%，CVTG 增加了 35%。高質(zhì)量微調(diào)（Stage3）進(jìn)一步增強(qiáng)了圖像細(xì)節(jié)，CVTG 提高了 +12.0%。將 GRPO 應(yīng)用于 MLLM（Stage4）略微影響了 GenEval (-0.01) 和 WISE (-0.01) 上的圖像生成，但結(jié)合 CoT 顯著提升了推理和生成能力（WISE: 0.55 → 0.76）。DiT-GRPO（Stage5）進(jìn)一步提高了圖像生成質(zhì)量，尤其是在細(xì)粒度文本渲染任務(wù)中（CVTG: 0.79 → 0.84）。

Prepadding States: 如下表7所示，預(yù)填充狀態(tài)顯著改善了短提示基準(zhǔn)測試的性能（GenEval: 0.64→0.78，WISEBench: 0.37→0.46，CVTG: 0.24→0.28，ImgEdit: 3.46→3.93）。

訓(xùn)練策略: 如表8所示，直接將 SFT 應(yīng)用于具有推理數(shù)據(jù)的 DiT 并未改善推理基準(zhǔn)測試的性能。然而，使用 MLLM-GRPO 訓(xùn)練 MLLM 極大地增強(qiáng)了 ThinkGen 的推理能力（WISE: 0.55 → 0.74）。

VGI-refine 中的提取策略: 如下表13所示，僅使用  標(biāo)記后面的隱藏狀態(tài)（CUT）比使用所有隱藏狀態(tài)（ALL）在所有基準(zhǔn)測試中都產(chǎn)生了持續(xù)的改進(jìn)，尤其是在短提示生成任務(wù)中。這表明截?cái)?nbsp;pre- 隱藏狀態(tài)可以有效消除冗余信息，從而提高圖像生成質(zhì)量。

連接器設(shè)計(jì): 如下表12所示，線性層連接器優(yōu)于 MLP 或 Transformer 連接器。

SepGRPO 過程分析

如圖5所示，通過可視化 SepGRPO 的中間過程，觀察到以下關(guān)鍵現(xiàn)象：

CoT 長度增加: 平均 CoT 長度逐漸增長，表明模型在訓(xùn)練期間發(fā)展出更復(fù)雜的推理能力。

統(tǒng)一獎(jiǎng)勵(lì)增長: 隨著訓(xùn)練的進(jìn)行，多任務(wù)獎(jiǎng)勵(lì)穩(wěn)步增加，表明 ThinkGen 學(xué)習(xí)在不同場景中自適應(yīng)地思考。

圖像質(zhì)量改進(jìn): 50、300 和 700 步的圖像可視化顯示出圖像生成質(zhì)量的明顯提升趨勢，生成的圖像展現(xiàn)出更豐富的細(xì)節(jié)和更高的保真度。

結(jié)論

ThinkGen，一個(gè)新穎的思考驅(qū)動(dòng)框架，能夠自動(dòng)地將思維鏈（CoT）推理應(yīng)用于多樣化的生成任務(wù)。本文的方法采用解耦的 MLLM-DiT 架構(gòu)，并通過 SepGRPO 進(jìn)行訓(xùn)練，使其能夠在生成之前制定高質(zhì)量的計(jì)劃。廣泛的實(shí)驗(yàn)證明，ThinkGen 在推理密集型任務(wù)上取得了顯著的改進(jìn)。本工作代表著構(gòu)建更智能、更通用、無縫整合推理與創(chuàng)造的生成模型邁出了關(guān)鍵一步。

參考文獻(xiàn)

[1] ThinkGen: Generalized Thinking for Visual Generation

       原文標(biāo)題 : 北交&字節(jié)最新開源ThinkGen：首次顯式利用多模態(tài)CoT處理生成任務(wù)，多項(xiàng)任務(wù)性能SOTA