From Reflection to Perfection: Scaling Inference-Time Optimization for Text-to-Image Diffusion Models via Reflection Tuning

Motivation

• Diffusion models remain inconsistent on complex poses, multi-object compositions, and lighting, while existing inference optimizations emphasize parallel noise search rather than sequential self-correction.The paper proposes exploiting pretrained models at inference instead of continually scaling training compute, parameters, and data; it identifies a lack of reflection-refinement datasets as another obstacle.Source: §1 · pp. 1–2

ReflectionFlow

• ReflectionFlow scales inference through better noise initialization, optimized prompts for semantic guidance, and explicit reflections that assess and correct prior generations.The three axes let diffusion models use additional inference-time computation to improve image quality iteratively.Source: §1 · p. 2

GenRef

• GenRef provides 1 million reflection triplets, while GenRef-CoT adds 227K progressive annotations from GPT-4o and Gemini 2.0.Its automated pipeline draws on four data sources, verifiable objectives, ensemble reward models, and diverse rollout strategies to produce actionable reflections and higher-quality target images.Source: §1 · p. 2

Reflection tuning

• Reflection tuning adapts diffusion transformers such as FLUX.1-dev by jointly attending to the original prompt, reflection prompt, flawed image, and high-quality image in one sequence.This unified multimodal strategy requires no additional modules.Source: §1 · p. 2

Experiments

• ReflectionFlow outperforms naive noise-level scaling under identical inference budgets, improves with better verifiers, and yields particularly substantial gains on challenging prompts.The authors also report flexible trade-offs between search width and reflection depth, and qualitative convergence through iterative correction.Source: §1 · p. 2

Text-to-Image Diffusion Models

• Diffusion research has moved from U-Net models [44] toward DiTs [37], from hand-designed schedules [19] [50] to flow formulations [30] [33], and toward larger datasets and models [12] [41] [59] [67].Multi-resolution progressive training is cited as a technique that enhances training efficiency.Source: §2 · pp. 2–3

Diffusion Inference-Time Enhancement

• Prior diffusion inference methods target initialization [2] [34] [40] [66], sampling or denoising and inversion [3] [47] [64] [49], or prompt augmentation [28] [59]; ReflectionFlow adds sequential generation and refinement to parallel scaling.The paper characterizes existing methods as focusing on parallel single-pass generation and positions its approach as unifying parallel and sequential inference-time scaling.Source: §2 · p. 3

Scaling Inference-Time Compute

• LLM inference scaling studies verifier-guided search [13] [62], elicited self-reflection [23] [35] [46], and SFT or RL for reflective improvement [11] [17] [25] [42] [56]; prior work explored autoregressive image generation [16].Meta’s framework [48] unifies these directions and examines tradeoffs between pretraining scale and inference-time computation, inspiring this approach.Source: §2 · p. 3

3.1. Problem Formulation

• The framework initializes an image from random noise, then conditions each correction on the original prompt, previous image, and a textual reflection.External reward models or multimodal language models assess the previous image and generate reflections; the correction task is framed as generalized image editing.Source: §3.1 · p. 3

3.2. Reflection Dataset

• GenRef comprises 1 million flawed-image, high-quality-image, and reflection triplets gathered through four data pipelines.The pipelines cover rule-based and reward-based samples, long-short prompt pairs, and editing data; reflections should identify shortcomings and offer actionable refinement guidance.Source: §3.2 · pp. 3–5, Figure 2 · p. 4
• A CoT annotation pipeline and trained verifier improve reflection and pair quality, yielding 227K high-quality GenRef-CoT reflections and filtering lower-quality samples.Closed-source APIs analyze image pairs, state preferences, and produce corrective reflections; the verifier is trained with Bradley–Terry pairwise comparisons using preference annotations.Source: §3.2 · pp. 5–6

3.3. Reflection Tuning

• Reflection tuning concatenates prompt, flawed-image, and refined-image tokens for unified multimodal attention, then trains on the refined target without specialized modules.To preserve alignment with the pretrained model, training drops inputs probabilistically and gradually shifts sampling from editing data toward the other three domains.Source: §3.3 · p. 6

3.4. Test-Time Scaling via Iterative Refinement

• ReflectionFlow combines noise-level, reflection-level, and prompt-level scaling to iteratively generate, assess, refine, and select images.It adjusts search width N and reflection depth M to balance performance and computational efficiency; an MLLM verifier ranks candidates and guides reflections and prompt updates between rounds.Source: §3.4 · pp. 6–7

4.1. Setup

• The experiments use FLUX.1-dev2 as the base generator, train the corrector on GenRef with rank-256 LoRA, and evaluate on GenEval’s 553 prompts, generating four 1024 × 1024 images per prompt.Images are generated with guidance scale 3.5 and 30 sampling steps; the main experiments use the SANA verifier and Qwen2.5-VL-7B for reflection.Source: §4.1 · pp. 7–8

4.2. Main Results

• GenEval scores rise from 0.67 for FLUX.1-dev to 0.85 with noise-level scaling, 0.87 with prompt-level scaling, and 0.91 with all three scaling levels.The authors attribute the gains to prompt refinement and verifier-provided textual reflections that let the corrector iteratively improve outputs.Source: §4.2 · p. 8
• ReflectionFlow with either SANA or its own verifier at 16 samples outperforms Reflect-DiT’s 0.81 score at 20 samples per prompt.The paper attributes this comparison to its dataset, reflection-model training, and inference-time scaling framework.Source: §4.2 · p. 8, Figure 5 · p. 8

4.3. Ablation Studies

• With GPT-4o verification, performance quickly approaches its limit at 32 samples; the trained verifier improves steadily, while SANA reaches 0.91 and the GenEval oracle upper bound is 0.98.The trained verifier has not reached its inference-time scaling limit, and results improve with increasingly capable verifiers, with little sign of saturation.Source: §4.3 · p. 9
• At a fixed budget of 16, sequential refinement (N1M16) scores 0.78, versus 0.74 for parallel refinement (N16M1); greater refinement depth generally outperforms wider branching.The authors describe the process as relatively unstable, requiring sequential iterations to identify and correct errors.Source: §4.3 · pp. 9–10, Table 2 · p. 9

4.4. Analysis and Discussion

• ReflectionFlow’s gains vary with prompt difficulty: correctness rises from 0.10 to 0.81 on hard prompts, 0.55 to 0.85 on medium prompts, and 0.95 to 0.97 on easy prompts.This suggests potential for allocating inference-time compute dynamically according to prompt difficulty.Source: §4.4 · p. 10
• Qualitative examples show iterative reflection correcting missed details and complex relations, while producing interpretable chains that address visual inconsistencies, stylistic mismatches, and compositional inaccuracies.The examples begin with incorrect generations and show progressively improved alignment with prompts and reflection instructions.Source: Figure 6 · p. 9, §4.4 · p. 10

• ReflectionFlow enables iterative self-refinement of text-to-image diffusion models by treating refinement as generalized image editing and efficiently fine-tuning pretrained diffusion transformers.The approach enhances performance without deviating from pretrained distributions.Source: §5 · p. 10
• GenRef comprises one million triplets pairing flawed images, high-quality images, and textual reflections to support image reflection tuning.Source: §5 · p. 10
• ReflectionFlow scales inference across noise, reflection, and prompt dimensions, allowing flexibility to balance computational efficiency and generation quality.Source: §5 · p. 10

A. Dataset Preview

• GenRef preview images are grouped by curation subset, pair each image with a prompt and reflection, and use red and green borders for starting and final images.The “edit” samples come from OmniEdit [55].Source: A. Dataset Preview · p. 14
• The reward-based data samples pair prompts with reflections that request concrete changes to scenes, such as adding mist, rocks, or flowers.The examples include a floating island, a glass containing a miniature landscape, and a woman ice-skating.Source: Figure 7 · p. 14, A. Dataset Preview · p. 14
• The rule-based data samples show reflections specifying object or appearance edits, including replacing carpet-cleaning equipment and changing a blanket’s texture.The figure also includes a colored-pencil example.Source: Figure 8 · p. 14, A. Dataset Preview · p. 14
• Figure 10 shows chain-of-thought samples for Qwen fine-tuning, with analysis comparing two images and instructions describing corrective edits.The red and green borders identify incorrect and correct images, respectively, as judged by closed-source APIs.Source: A. Dataset Preview · p. 14, Figure 10 · p. 16
• The long-short prompt samples show reflections correcting specified portrait attributes, including clothing, lip color, jewelry, and superhero accessories.The figure presents three portrait examples with prompts and corresponding reflections.Source: A. Dataset Preview · p. 14, Figure 9 · p. 15
• The chain-of-thought examples use prompt-specific analyses to identify image mismatches and give instructions targeting composition, object form, or color and lighting.The examples cover a diner booth, Arabic calligraphy forming a coffee pot, and a man in a living room.Source: Figure 10 · p. 16

B. Qualitative Results

• Figure 11 presents qualitative results involving complex reasoning.Source: B. Qualitative Results · p. 17, Figure 11 · p. 17
• The qualitative results show the detailed process of ReflectionFlow’s three scaling steps.Source: B. Qualitative Results · p. 17
• The examples provide a more nuanced understanding of the ReflectionFlow framework.Source: B. Qualitative Results · p. 17

C. Algorithm Process

• The algorithm takes prompt y, generator Gθ, corrector Cϕ, verifier V, scaling width N, and scaling depth M, and aims to produce an image that best realizes user intent.Source: C. Algorithm Process · p. 18
• It initializes the image set as empty and samples noise z_j from N(0, I) for each of N width iterations.Source: C. Algorithm Process · p. 18
• Across M depth iterations, the verifier scores the previous images against the prompt and the next image set is initialized as empty.Source: C. Algorithm Process · p. 18

D.1.1. Single Object

• Figure 12 provides a verifier prompt for images containing a single object.Source: Figure 12 · p. 19

D.1.2. Two Objects

• Figure 13 provides a verifier prompt for images containing two objects.Source: Figure 13 · p. 20

D.1.3. Counting

• Figure 14 provides a verifier prompt for counting images.Source: Figure 14 · p. 21

D.1.4. Colors

• Figure 15 provides a verifier prompt for images involving colors.Source: Figure 15 · p. 22

D.1.5. Position

• Figure 16 provides a verifier prompt for images involving positions.Source: Figure 16 · p. 23

D.1.6. Color Attribution

• Figure 17 provides a verifier prompt for images involving color attribution.Source: Figure 17 · p. 24

D.2. Reflection Prompt

• Figure 18 presents the prompt for generating reflection instructions.Source: Figure 18 · p. 24

D.3. Refine Prompt

• The refinement prompt uses the original and current prompts, generated images, reflection, and multi-domain assessments as inputs.Source: §D.3 · p. 25
• It asks the model to assess the current prompt’s strengths and drawbacks based on the images and evaluations.Source: §D.3 · p. 25
• Refinement may address subject, scene, style, lighting, tone, mood, camera style, and composition.Source: §D.3 · p. 25
• The refined prompt must preserve the original description and must not contradict the reflection prompt.Source: §D.3 · p. 25
• The model is instructed to output only the refined prompt, without additional text.Source: §D.3 · p. 25

D.4. Chain-of-Though Image Reflection Prompt

• The assistant compares both images with the prompt, identifies the better match, and generates concise editing instructions to turn the inferior image into the superior one.The comparison considers missing elements, incorrect object attributes or relationships, and unnecessary elements.Source: §D.4 · p. 26
• The assistant must output “left” or “right” to identify the better-aligned image, then give exactly one concise, accurate, actionable imperative edit.The edit must not explicitly refer to the left or right image.Source: §D.4 · p. 26