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

TL;DR

The paper tackles the limitation of T2I diffusion models that rely on static LLM encoders by introducing Think-Then-Generate (T2G), where an LLM encoder reasons about prompts and rewrites them before conditioning the diffusion model. It presents a two-part solution: a lightweight supervised fine-tuning to instill chain-of-thought prompting and a Dual-GRPO training regime that jointly optimizes the LLM and the DiT with image-grounded rewards. Results across reasoning-intensive generation and editing benchmarks show substantial gains in semantic alignment, factual consistency, and visual realism, with performance approaching GPT-4o on open benchmarks. This work demonstrates a promising step toward unified models capable of reasoning, expression, and demonstration in visual generation tasks.

Abstract

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.

Figures (10)

• Figure 1: An overview of our think-then-generate method. Beyond using LLM as a frozen text encoder, we train it to think and refine the raw user prompts guided by the reward of output images.
• Figure 2: Comparison of T2I models on conceptual visual generation. Our Qwen-Image under think-then-generate pipeline produces semantically aligned and visually coherent results, correctly interpreting user intent given prompt "Holiday celebrating the birth of Jesus Christ." (e.g., generating a warm Christmas celebration rather than literally depicting Jesus), whereas vanilla Qwen-Image behaves like a simple text–pixel mapper and often fails to capture conceptual meanings.
• Figure 3: t-SNE visualization of the embedding distributions before and after SFT. The distributions remain highly consistent, indicating that our SFT procedure preserves the embedding space structure and thus maintains compatibility with the DiT, enabling it to render stable and coherent visual outputs.
• Figure 4: Dual-GRPO training trajectories. (a) Tree-structured rollout for a given user prompt $q$: the LLM encoder samples $J$ reasoning traces, each rewritten prompt conditions the DiT to generate $K$ images. Image-grounded rewards are aggregated to compute group-relative advantages for updating both the LLM and the DiT. (b) Evolution of alignment and style scores during training, demonstrating how DiT training improves both semantic alignment and visual quality over time.
• Figure 5: Comparison of T2I models on conceptual image editing. Vanilla Qwen-Image fails to interpret instructions (e.g., showing an ice cream under sunlight instead of melting), behaving as a text–pixel mapper. Our model correctly infers intended semantics, producing coherent, aesthetically pleasing edits with high consistency to the original image, outperforming unified models like Emu2 and Bagel.