PaperlayerView exact revision on arXiv

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 text-to-image diffusion models often use LLMs as passive encoders, limiting them to literal instructions rather than conceptual visual reasoning. This paper introduces think-then-generate, combining prompt rewriting with supervised fine-tuning and Dual-GRPO co-optimization, and reports stronger reasoning-based generation and editing, including 0.79 WISE performance nearly on par with GPT-4o.

  • Problem

    Existing text-to-image diffusion models typically use frozen LLM encoders as feature extractors, so they handle literal descriptions but struggle with conceptual instructions requiring reasoning and world knowledge.

  • Method

    Think-then-generate fine-tunes an LLM to reason over and rewrite prompts, then uses Dual-GRPO to jointly optimize image-grounded reasoning and diffusion rendering with stage-specific rewards.

  • Results

    The approach improves factual consistency, semantic alignment, and visual realism across reasoning-based generation and editing benchmarks, including 0.79 on WISE, nearly matching GPT-4o.

  • Takeaways & Limitations

    The results support T2G as a promising route for open models to perform knowledge-intensive visual generation and conceptual image editing.

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 T2I diffusion models often use LLMs only as frozen text encoders, limiting them to literal instructions rather than conceptual visual reasoning. The proposed think-then-generate paradigm trains the encoder to reason and rewrite prompts, then jointly optimizes it with the diffusion model for image-grounded generation.

  • Motivation: Existing models often function as text–pixel mappers because frozen LLM encoders are trained primarily on descriptive image–caption pairs.This supports literal attributes such as colors and textures but not conceptual instructions requiring world knowledge.
  • Method: Think-then-generate trains an LLM encoder to reason over and rewrite raw prompts, using rewritten-prompt embeddings to condition the diffusion transformer.The method first uses supervised fine-tuning to acquire the think-then-rewrite pattern, then applies Dual-GRPO to co-optimize the encoder and DiT.
  • Method: Dual-GRPO uses image-based rewards tailored to each component: semantic alignment and conceptual understanding for the encoder, and realism and aesthetics for the DiT.The rewritten prompt bridges text reasoning and image synthesis during end-to-end optimization.
  • Results: 0.79 WISE score is achieved on T2I with Qwen-Image, surpassing pretrained Qwen-Image by 30% and performing on par with GPT-4o.The method also reaches 92.2 on T2I-ReasonBench and substantially outperforms other open-source models.
  • Results: The approach produces more faithful, fine-grained, and instruction-aligned image edits, while improving knowledge grounding, visual plausibility, and aesthetics in schematic activity illustrations.These results are reported for Qwen-Image-Edit and more challenging real-world scenarios.

2. Preliminary: Group Relative Policy

GRPO is a policy-gradient method that normalizes rewards within groups to estimate relative advantages without a value model. Its extensions address language-model and flow-matching optimization, including stochastic trajectory generation for deterministic flow sampling.

  • GRPO: GRPO introduces group-wise relative reward normalization to handle high-variance training signals in large generative models.It is presented as a variant of policy-gradient methods such as PPO.
  • GRPO: GRPO groups trajectories with similar prompts or conditions and computes normalized relative advantages directly from their rewards.Unlike actor–critic methods, it does not learn a value model for advantage estimation.
  • Objective: The GRPO objective combines clipped policy updates with KL-divergence regularization against a reference policy.The likelihood ratio is defined relative to the old policy, while β controls KL regularization.
  • LLM Application: For LLMs, outcome-based rewards assigned uniformly across rollout tokens implicitly assume that every token contributes equally to the final result.The language policy generates tokens conditioned on previous tokens and the user prompt.
  • Flow-GRPO: Flow-GRPO converts deterministic flow-ODE sampling into a stochastic differential equation and discretizes it with Euler–Maruyama to enable trajectory exploration.The resulting transition kernels provide the stochasticity needed for GRPO advantage estimation in flow-matching models.
  • Flow-GRPO: The transition kernels are reduced to tractable Gaussian distributions, enabling direct GRPO policy updates for flow-matching models.This makes the optimization applicable despite the original deterministic sampling dynamics.

3. Method

The method trains an LLM encoder to reason over raw prompts and rewrite them before diffusion, then jointly optimizes the encoder and DiT with stage-specific image-grounded rewards.

  • T2G and SFT: The supervised fine-tuning dataset teaches the LLM to perform chain-of-thought reasoning and produce descriptive refined prompts for diffusion conditioning.Raw prompts are processed with Gemini-2.5 to construct training examples.
  • T2G and SFT: SFT activates prompt rewriting while preserving the embedding distribution sufficiently to maintain compatibility with the DiT.The before-and-after SFT embeddings overlap in the t-SNE visualization, and the paper associates this with stable visual outputs.
  • Dual-GRPO: Dual-GRPO models the LLM and DiT as a composite policy whose rollout first generates reasoning tokens and then diffusion predictions conditioned on the rewritten representation.The objective jointly optimizes both components using separate reward functions for the LLM and DiT.
  • Dual-GRPO: Hierarchical sampling generates J reasoning sequences and K images per rewritten prompt, enabling group-relative advantages for both the LLM and DiT.The rollout uses PPO clipping and KL-divergence regularization in the resulting Dual-GRPO objective.
  • Reward function and scheduler: Stage-specific rewards address the two-stage pipeline: LLM reasoning uses averaged semantic consistency, while diffusion uses aesthetic, physical-consistency, and semantic-consistency scores.Reward schedulers control the relative emphasis assigned to the LLM and diffusion stages during training.

4. Experiments

Experiments evaluate T2G across reasoning-intensive text-to-image and image-editing benchmarks, comparing post-training stages and competing models. The results show that Dual-GRPO improves reasoning-grounded generation and editing beyond literal text–pixel mapping.

  • Experimental Setup: The evaluation covers T2I and image-editing tasks with quantitative and qualitative comparisons against diffusion, unified multimodal, and proprietary models.The benchmark suite includes WISE, T2I-ReasonBench, UniREditBench, and RISEBench.
  • Supervised Fine-tuning: A 7,000-sample dataset uses Gemini-2.5 chain-of-thought reasoning and refined prompts to fine-tune Qwen2.5-VL for think-then-rewrite behavior.The prompts target cases requiring world knowledge and reasoning for semantically coherent images.
  • Quantitative Results: 0.79 WISE score is achieved after SFT and Dual-GRPO, while zero-shot chain-of-thought prompting raises Qwen-Image only from 0.61 to 0.65.The optimized model also reaches 68.3 accuracy on T2I-ReasonBench and leads open-source T2I models.
  • Quantitative Results: Dual-GRPO jointly optimizes the LLM and DiT with image-grounded rewards, addressing the mismatch between reasoning outputs and renderable visual semantics.SFT alone leaves the CoT process unaware of the DiT, which can produce tokens the image generator cannot render reasonably.
  • Image Editing: Post-trained QwenImage-Edit surpasses Gemini-2.5-Flash-Image on UniREdit and improves substantially over the SFT-only baseline.The authors attribute the improvement in CoT reasoning to image-based rewards in Dual-GRPO.
  • Qualitative Results: The T2G model infers melting ice cream in conceptual editing, whereas vanilla Qwen-Image merely places the reference image under sunlight.The resulting edits are described as coherent, aesthetically pleasing, and consistent with the original image.

5. Related Work

Prior work applies reinforcement learning and unified multimodal modeling to generative systems, but typically leaves the text encoder as a static feature extractor. The paper instead emphasizes joint optimization of language reasoning and diffusion generation.

  • Diffusion Reinforcement Learning: Earlier diffusion reinforcement-learning methods address optimization challenges through reverse-process policy gradients, forward-process objectives, or preference learning.These approaches face issues including solver restrictions or likelihood-estimation difficulties.
  • Diffusion Reinforcement Learning: Existing preference-optimization approaches still treat the text encoder as a static feature extractor and optimize only the decoder.The paper positions Dual-GRPO as differing by jointly optimizing the LLM encoder and diffusion decoder.
  • Unified Multimodal Architectures: Unified multimodal models integrate vision-language understanding and generation, with examples including HunyuanImage and BAGEL.HunyuanImage uses an MoE architecture exceeding 80B parameters, while BAGEL is trained on interleaved text-image-video data.

6. Conclusion

The paper addresses literal text-to-pixel mapping by activating LLM reasoning through T2G and jointly optimizing reasoning and visual rendering. Experiments report gains over open-source baselines and Gemini-2.0, including conceptual editing improvements.

  • Conclusion: T2G turns LLM-based text encoders into active reasoning agents that rewrite raw prompts before diffusion generation.The framework combines lightweight supervised fine-tuning with Dual-GRPO for image-grounded semantic consistency and visual realism.
  • Conclusion: The approach outperforms open-source baselines, surpasses Gemini-2.0, and excels in knowledge-intensive visual generation and conceptual image editing.These results support the paper’s stated direction toward unified models combining reasoning and generation.

with LLM Encoders

The table compares reward-scheduler designs on T2I-ReasonBench, with scores organized by scheduler configuration.

  • Reward Schedulers: Table 5 compares different reward schedulers on T2I-ReasonBench.Use the table to compare how scheduler choices affect benchmark performance.

1. Ablation on Different Reward Schedulers

The study compares balanced and staged reward-weighting schedulers for Dual-GRPO training. Jointly updating the LLM and DiT with balanced weights performs better than staged optimization on T2I-ReasonBench.

  • The balanced scheduler sets β1(τ) = β2(τ) = 0.5, updating the LLM and DiT jointly at every iteration.The staged alternative updates only the LLM early and only the DiT later.
  • Balanced scheduling consistently outperforms staged scheduling on T2I-ReasonBench.The authors hypothesize that joint optimization coordinates prompt refinement with visual rendering more tightly.

2. Detailed Results on Image Editing Task

The image-editing evaluation compares the proposed think-then-generate training with existing editing models using benchmark scores, training dynamics, and qualitative examples. The reported results indicate stronger editing accuracy, visual quality, semantic alignment, and coherence.

  • Think-then-generate training significantly outperforms vanilla Qwen-Image-Edit and surpasses Seedream-4.0 on image editing.The authors attribute the gains to incorporating reasoning into the generation process.
  • Training improves semantic alignment and visual quality over time under the new scheduler.Figure 6 visualizes score evolution and corresponding DiT samples during training.
  • User-study scores compare the model with GPT-4o and Qwen-Image in challenging real-world T2I scenarios.The figure caption reports performance close to GPT-4o and clearly above Qwen-Image.
  • The human-evaluation interface presents Bagel-think, the proposed model, vanilla Qwen-Image, and GPT-4o as images 1–4.The caption states that the proposed model gives the most comprehensive math deduction and reaches the correct final answer.
  • RISEBench reports detailed image-editing scores, with bold values marking the highest scores among open-source models.The table caption specifies the comparison convention but does not provide the numeric values here.

3. User Study

A user study evaluates T2I outputs in challenging real-world scenarios using 46 prompts and rankings across semantic, visual, and aesthetic criteria. The proposed model performs close to GPT-4o and above Qwen-Image, with particularly strong math-related reasoning examples.

  • The study uses 46 prompts and compares the model with vanilla Qwen-Image, Bagel-think, and GPT-4o.Users rank outputs by semantic alignment and conceptual understanding, visual realism and coherence, and aesthetic quality.
  • Each model receives prompt-level reward si = 5 − ri, where ri ∈ {1, 2, 3, 4} is its rank among four models.Average scores are obtained by summing rewards and averaging across prompts.
  • The model performs closely to GPT-4o and clearly surpasses Qwen-Image in challenging real-world T2I scenarios.The study reports the comparison through aggregated prompt-level rankings.
  • For math-related problems, the model produces the most comprehensive reasoning and uniquely reaches the correct final answer.This result is reported as evidence of stronger deductive capability and practical effectiveness in the evaluated tasks.

4. More Demos

Additional T2I and image-editing demonstrations illustrate the model’s broader instruction understanding and output quality. The reported examples emphasize conceptual alignment, appearance consistency, plausibility, and coherence.

  • Additional demonstrations report stronger conceptual instruction understanding and alignment than existing baselines.The examples cover both T2I generation and image editing.
  • The demonstrations show more robust appearance consistency, visual plausibility, and coherence across generated content.Figure 10 contains more T2I examples, while Figure 9 contains additional editing examples.
  • A sample prompt asks the model to depict what a scene will look like after a year in a high-school classroom.
Loading 2601.10332v1…