作者：Siyu Jiao等

解讀：AI生成未來

亮點直擊

首次提出思考驅(qū)動的視覺生成框架： ThinkGen 是第一個顯式利用 MLLM 的思維鏈（CoT）推理來處理各種生成場景的思考驅(qū)動視覺生成框架。這解決了現(xiàn)有方法在泛化性和適應(yīng)性方面的局限性，因為它們通常針對特定場景設(shè)計 CoT 機制。

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

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

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

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

總結(jié)速覽

解決的問題

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

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

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

提出的方案

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

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

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

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

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

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

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

達(dá)到的效果

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

架構(gòu)方法

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

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

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

Diffusion Transformer (DiT)

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

VGI-refine

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

訓(xùn)練策略

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

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

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

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

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

階段1 對齊: 在此階段，本文引入 K 個可學(xué)習(xí)的預(yù)填充狀態(tài)，并通過僅訓(xùn)練線性連接器來將 DiT 與 MLLM 對齊，同時保持 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)建了一個 0.7M 的高質(zhì)量子集，以增強 DiT 的指令遵循能力和圖像美學(xué)。

SepGRPO

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

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

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

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

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

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

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

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

推理生成

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

推理編輯

如下表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)出改進，并在許多知名模型中取得了最佳結(jié)果。這些結(jié)果表明 ThinkGen 具有強大的指令遵循和文本渲染能力。

圖像編輯

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

消融研究

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

僅訓(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）進一步增強了圖像細(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）進一步提高了圖像生成質(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 極大地增強了 ThinkGen 的推理能力（WISE: 0.55 → 0.74）。

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

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

SepGRPO 過程分析

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

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

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

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

結(jié)論

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

參考文獻

[1] ThinkGen: Generalized Thinking for Visual Generation

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