Improving Text-to-Image Generation with Input-Side Inference-Time Scaling

TL;DR

Rather than retraining or fine-tuning complex text-to-image (T2I) models, this paper introduces an input-side inference-time scaling framework that uses an LLM to rewrite user prompts before generation. Trained purely via iterative Direct Preference Optimization (DPO) without any supervised fine-tuning (SFT) data, the rewriter consistently enhances image fidelity, compositional alignment, and aesthetics across both diffusion and autoregressive image generators.

---

What Problem Does This Paper Address?

Real-world user prompts for T2I generation are frequently short, vague, or underspecified (e.g., "a cat disney cartoon style"). Consequently, even state-of-the-art generators like FLUX or Stable Diffusion produce images with poor composition, missing objects, or inconsistent aesthetics.

Prior remedies have notable friction points:

1. In-Context Learning (ICL) and manual prompt engineering offer modest gains and rely heavily on manual prompt crafting.
2. Supervised Fine-Tuning (SFT) on paired prompts requires costly data curation, limits the LLM's exploratory capacity, and overfits the rewriter to specific model-caption distributions.
3. Output-side inference scaling (e.g., best-of-$N$ candidate sampling or iterative diffusion latent search) is computationally expensive and difficult to standardize across fundamentally different backbones (diffusion vs. autoregressive).

The authors ask: Can we treat any downstream T2I generator as a frozen black box and scale generation quality solely through automated input prompt optimization?

---

Methods

The authors train an LLM prompt rewriter purely through reinforcement learning via Iterative DPO without SFT:

User Prompt x 
   ──> [LLM Rewriter] ──> n Candidate Prompts {y_1, ..., y_n}
   ──> [Frozen T2I Model] ──> n Synthesized Images
   ──> [MLLM-as-a-Judge (Qwen2.5-VL-72B)] ──> Pairwise Tournament Voting
   ──> Identify Best (y_w) & Worst (y_l) ──> Update Rewriter via DPO Loss

• Iterative DPO (No SFT): In each round, the policy proposes $n$ candidate rewrites for a prompt. A frozen T2I model generates the corresponding images. Images are compared pairwise by a multimodal judge to identify the best ($y_w$) and worst ($y_l$) rewrites, which are then used in the standard DPO objective.
• Decomposed Multimodal Reward Design: Instead of relying on noisy pointwise scalar scores (which make algorithms like GRPO fragile in image generation), the authors use pairwise comparisons with Qwen2.5-VL-72B-Instruct across four distinct criteria:
• Image Quality ($r_{\text{Quality}}$): Evaluates anatomical plausibility, geometry, physical consistency (reflections, shadows), and artifacts.
• General Alignment ($r_{\text{General-Alignment}}$): First decomposes the prompt into atomic VQA questions (e.g., "Are there four apples?", "Is there a table?"), then scores adherence.
• Physical Alignment ($r_{\text{Physical-Alignment}}$): Assesses object counts, spatial layouts (left/right, front/behind), and attribute binding.
• Aesthetics ($r_{\text{Aesthetics}}$): Measures visual composition, styling, and artistic merit.
• Rewriter Flavors: By altering reward combinations, they train a General Rewriter (faithful to instructions and quality) and an Aesthetics Rewriter (prioritizing visual appeal).
• Parameter-Efficient Training: Using LoRA ($\text{rank}=64$) achieves performance on par with full-model fine-tuning while drastically reducing memory footprint.

---

Datasets

• Training Data: A text-only corpus of ~60k prompts combined from the Pick-a-Pic v2 user prompt training set and the T2I-CompBench++ training split. No paired ground-truth images or refined captions were needed.
• Evaluation Benchmarks:
• Text-Image Alignment: GenEval, T2I-CompBench++ test split, and TIFA-Benchmark.
• Image Quality & Aesthetics: MS-COCO-30K (FID), Pick-a-Pic v2 test split (500 prompts evaluated via GPT-4o pairwise win rate against DALL·E 3), PartiPrompts, and a LAION aesthetic predictor.

---

Key Insights and Ideas

• Iterative DPO beats SFT for prompt expansion: SFT restricts the LLM's generative variety, whereas RL-driven exploration allows the rewriter to discover prompt modifiers that best guide downstream image generation.
• The Alignment vs. Aesthetics Tension: Adding an explicit aesthetics reward drastically boosts visual appeal (aesthetics win rate against DALL·E 3 jumps from 0.476 to 0.818 on FLUX.1-schnell), but hurts alignment (drops from 0.561 to 0.424). Qualitative checks show the aesthetic rewriter frequently hallucinates decorative backgrounds, extra objects, and ornate lighting that dilute the original instruction.
• Strong Backbone Transferability: A rewriter trained with one image generator (e.g., SD-3.5) transfers cleanly to others (e.g., FLUX.1-schnell) without retraining. The learned prompt structure reflects general spatial and descriptive conventions common to modern vision-language datasets.
• Model Scalability & the Reasoning Quirk: Performance scales monotonically with the rewriter LLM's parameter count (from 3B to 70B, with Llama-3-70B performing best). However, reasoning-tuned models (e.g., DeepSeek-R1-Distill) underperform standard instruction-tuned models—explicit chain-of-thought fine-tuning seems to restrict creative and descriptive prompt expansion.

---

Key Results

• GenEval Alignment: Using Llama-3-70B as the General Rewriter, FLUX.1-dev improves from 0.70 to 0.79 overall, surpassing prior prompt-optimization approaches such as RePrompt (0.76) and PARM++ (0.70).
• Image Quality (MS-COCO FID): FID scores drop across all tested generators (FLUX.1-dev: 24.38 $\to$ 19.57; FLUX.1-schnell: 20.57 $\to$ 17.76; SD-3.5-medium: 17.81 $\to$ 17.13; JanusPro: 19.28 $\to$ 16.71).
• GPT-4o Win Rates vs. DALL·E 3: On Pick-a-Pic v2 with FLUX.1-dev, our General Rewriter achieves a 0.575 alignment win rate (up from 0.360 baseline), while the Aesthetics Rewriter hits 0.800 on aesthetics (up from 0.350).

---

Limitations and Future Work

• Diminishing Returns Across Rounds: As DPO rounds exceed 6 (one complete pass over the 60k training dataset), prompt lengths inflate significantly, causing quality and aesthetics to plateau or degrade due to over-optimization.
• Inherent Metric Trade-offs: The pipeline cannot maximize aesthetics and strict semantic faithfulness simultaneously in a single policy, requiring users to pick between specialized model variants.
• Anomalies on Specific Architectures: On JanusPro, prompt rewriting failed to improve GenEval scores despite improving T2I-CompBench++ and TIFA, hinting that certain models might be overfitted to specific benchmark prompt distributions.