Source-linked AI summary

Think-Then-Generate: Reasoning-Aware Text-to-Image Diffusion with LLM Encoders

Siqi Kou, Jiachun Jin, Zetong Zhou, Ye Ma, Yugang Wang, Quan Chen, Peng Jiang, Xiao Yang, Jun Zhu, Kai Yu, Zhijie Deng

arXiv:2601.10332v1cs.CV

TL;DR

Existing T2I models often use LLM encoders only as feature extractors, handling literal instructions but struggling with conceptual instructions. Think-then-generate fine-tunes an LLM to reason about and rewrite raw prompts, then co-optimizes it with the diffusion transformer using Dual-GRPO. On the T2I task, our method scores 0.79 on WISE, 30% above pretrained Qwen-Image and on par with GPT-4o.

  • Problem

    Existing T2I models often use LLM encoders only as feature extractors, handling literal instructions but struggling with conceptual instructions.

    Source: §1 · p. 1
  • Method

    Think-then-generate fine-tunes an LLM to reason about and rewrite raw prompts, then co-optimizes it with the diffusion transformer using Dual-GRPO.

    Source: §1 · p. 2, §3.1 · p. 3
  • Results

    On the T2I task, our method scores 0.79 on WISE, 30% above pretrained Qwen-Image and on par with GPT-4o.

    Source: §1 · p. 2
  • Takeaways & Limitations

    The authors conclude that T2G excels in conceptual image editing and enables knowledge-intensive visual generation.

    Source: §6 · p. 9

Abstract

from arXiv · show

Recent progress in text-to-image (T2I) diffusion models (DMs) has enabled high-quality visual synthesis from diverse textual prompts. Yet, most existing T2I DMs, even those equipped with large language model (LLM)-based text encoders, remain text-pixel mappers -- they employ LLMs merely as text encoders, without leveraging their inherent reasoning capabilities to infer what should be visually depicted given the textual prompt. To move beyond such literal generation, we propose the think-then-generate (T2G) paradigm, where the LLM-based text encoder is encouraged to reason about and rewrite raw user prompts; the states of the rewritten prompts then serve as diffusion conditioning. To achieve this, we first activate the think-then-rewrite pattern of the LLM encoder with a lightweight supervised fine-tuning process. Subsequently, the LLM encoder and diffusion backbone are co-optimized to ensure faithful reasoning about the context and accurate rendering of the semantics via Dual-GRPO. In particular, the text encoder is reinforced using image-grounded rewards to infer and recall world knowledge, while the diffusion backbone is pushed to produce semantically consistent and visually coherent images. Experiments show substantial improvements in factual consistency, semantic alignment, and visual realism across reasoning-based image generation and editing benchmarks, achieving 0.79 on WISE score, nearly on par with GPT-4. Our results constitute a promising step toward next-generation unified models with reasoning, expression, and demonstration capacities.

1. Introduction

Existing text-to-image models often use LLMs as frozen text encoders and struggle with conceptual instructions. Think-then-generate trains the encoder to reason and rewrite prompts, then co-optimizes it with the diffusion backbone for image generation.

Motivation

  • Models trained on descriptive image-caption pairs with frozen LLM encoders handle literal attributes but struggle with conceptual instructions, acting as text-pixel mappers.The proposed approach instead aims to use the encoder’s world knowledge to interpret conceptual tasks and enrich prompts for diffusion generation.Source: §1 · pp. 1–2

Method

  • Think-then-generate has the LLM reason over and rewrite raw prompts, then feeds embeddings of the rewritten prompts to the diffusion transformer as generation conditioning.Supervised fine-tuning on prompts enriched with chain-of-thought reasoning and rewrites teaches the encoder this think-then-rewrite pattern.Source: §1 · p. 2
  • Dual-GRPO co-optimizes the LLM encoder and diffusion transformer end to end, using image-based rewards for semantic alignment and conceptual understanding in the encoder and realism and aesthetics in the transformer.The joint optimization also adapts the transformer to the encoder’s evolved representation space.Source: §1 · p. 2

Results

  • 0.79 on WISE surpasses pretrained Qwen-Image by 30% and is on par with GPT-4o, while substantially outperforming Bagel and Janus-Pro-7B.The method also scores 92.2 on T2I-ReasonBench, surpassing Gemini-2.0.Source: §1 · p. 2
  • For image editing, qualitative results show more faithful, fine-grained, instruction-aligned modifications, while challenging human-activity schematics demonstrate knowledge grounding and visual plausibility.The paper also reports improved aesthetic quality in these scenarios.Source: §1 · p. 2

2. Preliminary: Group Relative Policy Optimization

GRPO uses group-wise reward normalization to estimate relative advantages without a value model, making it suited to large generative models. The section describes its outcome-reward setup for LLMs and Flow-GRPO’s stochastic reformulation for flow-matching models.

2. Preliminary: Group Relative Policy Optimization

  • GRPO [30] extends policy-gradient methods such as PPO [28] by comparing normalized rewards within groups of similar prompts to calculate relative advantages, eliminating the need for a value model.This design is described as more efficient and easier to scale to LLMs and diffusion or flow-matching models [16] [22] [36].Source: §2 · pp. 2–3

GRPO for LLMs

  • In typical LLM GRPO, a reward model scores only the final generated text, and every token in the rollout receives the same reward.This setup implicitly assumes each token contributes equally to the final outcome.Source: §2 · p. 3

Flow-GRPO for Flow Matching Models

  • Flow-GRPO [22] converts deterministic flow-ODE sampling into a stochastic differential equation and discretizes it with Euler–Maruyama, enabling diverse trajectories for advantage estimation and exploration.The resulting transition kernels are tractable Gaussian distributions, allowing direct application of the GRPO policy update to flow-matching models.Source: §2 · p. 3

3. Method

The method activates an LLM encoder to reason about and rewrite prompts before image conditioning, then jointly trains the encoder and diffusion model with image-grounded, stage-specific rewards.

3.2. Dual-GRPO for LLM-DiT Composite Models

  • Dual-GRPO jointly optimizes the LLM and DiT through J reasoning traces and K images per trace, using image-based rewards, group-relative advantages, PPO clipping, and KL regularization.The LLM advantage aggregates rewards across images generated from the same reasoning output.Source: §3 · p. 3, §3.2 · p. 4

3.1. Reasoning-aware Behavior Activation

  • Gemini-2.5 supplies raw-prompt→CoT→refined-prompt examples for SFT; embeddings totally overlap before and after, preserving the DiT’s capacity for stable, reasonable outputs.The refined-prompt embedding conditions the DiT during image generation.Source: §3.1 · pp. 3–4, Figure 3 · p. 3

3.3. The Reward Function and Scheduler

  • Stage-specific rewards give LLM reasoning average semantic consistency across K images, while diffusion receives weighted aesthetic, physical-consistency, and semantic-consistency scores; schedulers vary stage emphasis during training [31].The diffusion-stage reward scores the final generated image.Source: §3.3 · pp. 4–5

4. Experiments

Experiments evaluate T2G on text-to-image generation and image editing, with quantitative and qualitative results supporting its effectiveness.

4.2.1. Supervised Fine-tuning

  • The T2I SFT stage uses 7,000 curated prompts requiring world knowledge and reasoning, with Gemini-2.5-generated reasoning and refined prompts as targets for Qwen2.5-VL fine-tuning.The resulting training teaches the model to reproduce the think-then-rewrite pattern on new prompts.Source: §4.2.1 · pp. 5–6

4.2.2. Quantitative Results

  • 0.79 on WISE and 68.3 accuracy on T2I-ReasonBench; the post-trained models lead open-source T2I models and match or surpass closed alternatives.The comparison covers reasoning-focused generation benchmarks and includes diffusion, unified multimodal, and proprietary models.Source: Table 1 · p. 6, §4.2.2 · pp. 6–7, Table 2 · p. 7

4.2.3. Qualitative Results

  • GRPO penalizes incorrect reasoning rollouts, such as mislabeling dumplings as Dragon Boat Festival food, and the resulting images are text-consistent, coherent, and aesthetically pleasing.The qualitative examples demonstrate improved inference from raw user prompts.Source: §4.2.3 · p. 7

4.3.1. Supervised Fine-tuning

  • For image-editing SFT, the authors augment UniREdit-Data-100K with Gemini-2.5 refined prompts and use 5,000 raw prompts for Dual-GRPO training.The data follows the same training format as T2I generation.Source: §4.3.1 · pp. 7–8

4.3.2. Quantitative Results

  • Post-training gives QwenImage-Edit strong image-editing performance, surpassing Gemini-2.5-Flash-Image on UniREdit and substantially improving over the SFT-only baseline.The evaluation uses UniREditBench and RISEBench, which assess reasoning-aware editing quality.Source: Table 3 · p. 8, §4.3.2 · p. 8

4.3.3. Qualitative Results

  • T2G correctly interprets conceptual editing instructions, depicting sun-exposed ice cream as melting while producing coherent, aesthetically pleasing edits with high consistency to the original image.Vanilla Qwen-Image instead places the reference image under sunlight without depicting the intended melting.Source: Figure 5 · p. 7, §4.3.3 · p. 8

4.4. Ablation Studies

  • The SFT ablation shows that teaching the T2G pattern outperforms the no-SFT variant by aligning chain-of-thought reasoning with refined prompts and giving the DiT more coherent conditioning.This alignment supports semantically grounded prompts rather than shallow pattern matching or implicit reasoning.Source: §4.4 · p. 8

5. Related Work

Prior diffusion-RL methods largely leave text encoders static, whereas this framework jointly optimizes the LLM encoder and diffusion decoder. Related multimodal generation work adds reasoning within unified models, while this framework places text-based CoT reasoning in the LLM encoder.

5.1. RL for Diffusion Models

  • Prior diffusion-RL methods include DDPO [3], Flow-GRPO [22], DanceGRPO [36], DiffusionNFT [39], and Diffusion-SDPO [12], but optimize only the decoder.DDPO discretizes reverse sampling; Flow-GRPO converts ODEs to SDEs for exploration, DanceGRPO uses group-wise relative optimization, and DiffusionNFT optimizes the forward process.Source: §5.1 · pp. 8–9
  • Dual-GRPO jointly optimizes the LLM encoder and diffusion decoder through a composite policy, unlike prior approaches that treat the text encoder as static.This addresses the limitation of encoder-freezing approaches for knowledge-intensive tasks.Source: §5.1 · pp. 8–9

5.2. Multimodal Models for Image Generation

  • Unified multimodal models such as HunyuanImage [6] and BAGEL [10] combine understanding and generation, while T2I-R1 [18] and Uni-CoT add multimodal reasoning.These models remain biased toward literal generation, which this framework addresses by incorporating text-based CoT reasoning within the LLM encoder.Source: §5.2 · p. 9

6. Conclusion

The think-then-generate framework uses LLM reasoning and world knowledge to improve text-to-image generation. It combines supervised fine-tuning with Dual-GRPO and outperforms open-source baselines, surpasses Gemini-2.0, and excels in conceptual image editing.

6. Conclusion

  • The approach outperforms open-source baselines, surpasses Gemini-2.0, and excels in conceptual image editing for knowledge-intensive visual generation.The conclusion reports these outcomes without specifying numerical results.Source: §6 · p. 9
  • Think-then-generate addresses literal text-to-pixel mapping by turning LLMs from passive encoders into active reasoning agents that use reasoning and world knowledge.Source: §6 · p. 9
  • The framework pairs lightweight supervised fine-tuning on raw prompts with Dual-GRPO, jointly optimizing LLMs for image-grounded semantic consistency and DiT for visual realism.Source: §6 · p. 9

with LLM Encoders

The supplied passage identifies Table 5 as comparing reward schedulers on T2I-ReasonBench but provides no comparative results.

with LLM Encoders

  • Table 5 compares different reward schedulers on T2I-ReasonBench.The supplied caption does not include the comparison outcomes.Source: Table 5 · p. 11

1. Ablation on Different Reward Schedulers

The ablation finds that the balanced scheduler outperforms staged training on T2I-ReasonBench, comparing joint LLM-and-DiT updates with an LLM-then-DiT schedule.

1. Ablation on Different Reward Schedulers

  • The balanced scheduler uses constant weights β1(τ) = β2(τ) = 0.5, updating the LLM and DiT jointly at every iteration.The staged scheduler updates only the LLM early, with β1(τ) = 1 and β2(τ) = 0, then only the DiT later, with β1(τ) = 0 and β2(τ) = 1; Figure 6 visualizes reward evolution and samples before and after Dual-GRPO under this scheduler.Source: §1 · p. 11
  • The balanced scheduler consistently outperforms the staged scheduler on T2I-ReasonBench.The authors hypothesize that joint optimization better coordinates prompt refinement and visual rendering.Source: §1 · p. 11

2. Detailed Results on Image Editing Task

On the RISE benchmark, the think-then-generate model outperforms vanilla Qwen-Image-Edit and surpasses Seedream-4.0, with improvements in edited-image accuracy and visual quality.

2. Detailed Results on Image Editing Task

  • The reported results indicate that reasoning during generation improves both the accuracy and visual quality of edited images.Source: §2 · p. 11
  • The model significantly outperforms vanilla Qwen-Image-Edit on the RISE image-editing benchmark.The supplied passage reports the comparison but provides no specific score values.Source: §2 · p. 11, Table 6 · p. 12
  • The model also surpasses the strong proprietary image-editing model Seedream-4.0 on RISE.No numerical scores for this comparison are supplied.Source: §2 · p. 11, Table 6 · p. 12

3. User Study

In a user study on 46 challenging real-world prompts, the model performed closely to GPT-4o and surpassed Qwen-Image, while producing the most comprehensive reasoning on math problems and uniquely reaching the correct final answer.

3. User Study

  • Users ranked outputs for 46 prompts against vanilla Qwen-Image [34], Bagel-think [10], and GPT-4o [17] on semantic alignment, realism and coherence, and aesthetics.The preference-collection interface is shown in Figure 8.Source: §3 · p. 12
  • For each prompt, a model received reward 5 − r_i, where r_i ∈ {1, 2, 3, 4} was its rank among four models; scores were averaged across prompts.The average score is shown in Figure 7.Source: §3 · p. 12
  • Our model performed closely to GPT-4o [17] and clearly surpassed Qwen-Image [34] in challenging real-world T2I scenarios.The study aggregated prompt-level rankings across four models.Source: §3 · p. 12
  • On math-related problems, our model produced the most comprehensive reasoning and uniquely arrived at the correct final answer.The authors describe this as evidence of deductive capability and practical effectiveness in real-world tasks.Source: §3 · p. 12

4. More Demos

Additional text-to-image and image-editing demos illustrate the model’s stronger instruction understanding, appearance consistency, and visual plausibility and coherence relative to existing baselines.

4. More Demos

  • The demos show stronger conceptual instruction understanding and alignment than existing baselines across text-to-image and image editing.The examples appear in Figure 10 for text-to-image and Figure 9 for image editing.Source: §4 · p. 12, Figure 9 · p. 13, Figure 10 · p. 14
  • The demos show more robust appearance consistency across generated content than existing baselines.Examples include the image-editing demos in Figure 9 and text-to-image demos in Figure 10.Source: §4 · p. 12, Figure 9 · p. 13, Figure 10 · p. 14
  • The demos show improved overall visual plausibility and coherence compared to existing baselines.The paper provides additional image-editing examples in Figure 9 and text-to-image examples in Figure 10.Source: §4 · p. 12, Figure 9 · p. 13, Figure 10 · p. 14
Loading 2601.10332v1…