Source-linked AI summary
From Reflection to Perfection: Scaling Inference-Time Optimization for Text-to-Image Diffusion Models via Reflection Tuning
Le Zhuo, Liangbing Zhao, Sayak Paul, Yue Liao, Renrui Zhang, Yi Xin, Peng Gao, Mohamed Elhoseiny, Hongsheng Li
TL;DR
Text-to-image diffusion models still struggle with complex scenes and fine-grained details, motivating more effective use of existing models within fixed inference budgets. ReflectionFlow adds noise-, prompt-, and reflection-level scaling, supported by the 1-million-triplet GenRef dataset and reflection tuning. Experiments report better performance than naive noise-level scaling, especially on challenging prompts.
Problem
Text-to-image diffusion models show inconsistent performance on complex scenes and fine-grained generation tasks, while improving them through training requires scaling compute, parameters, and data.
Method
ReflectionFlow performs inference-time iterative refinement across noise-level, prompt-level, and reflection-level scaling, using GenRef reflection triplets to tune a diffusion transformer.
Results
ReflectionFlow outperforms naive noise-level scaling under identical inference budgets and achieves particularly substantial gains on challenging prompts.
Takeaways & Limitations
Inference-time compute can be flexibly allocated through iterative reflection to improve image generation quality without additional large-scale training.
Abstract
from arXiv · showhide
Recent text-to-image diffusion models achieve impressive visual quality through extensive scaling of training data and model parameters, yet they often struggle with complex scenes and fine-grained details. Inspired by the self-reflection capabilities emergent in large language models, we propose ReflectionFlow, an inference-time framework enabling diffusion models to iteratively reflect upon and refine their outputs. ReflectionFlow introduces three complementary inference-time scaling axes: (1) noise-level scaling to optimize latent initialization; (2) prompt-level scaling for precise semantic guidance; and most notably, (3) reflection-level scaling, which explicitly provides actionable reflections to iteratively assess and correct previous generations. To facilitate reflection-level scaling, we construct GenRef, a large-scale dataset comprising 1 million triplets, each containing a reflection, a flawed image, and an enhanced image. Leveraging this dataset, we efficiently perform reflection tuning on state-of-the-art diffusion transformer, FLUX.1-dev, by jointly modeling multimodal inputs within a unified framework. Experimental results show that ReflectionFlow significantly outperforms naive noise-level scaling methods, offering a scalable and compute-efficient solution toward higher-quality image synthesis on challenging tasks.
1. Introduction
ReflectionFlow addresses inconsistent performance on complex image-generation tasks by shifting some scaling from training to inference-time iterative refinement. It combines three scaling axes with GenRef data and reflection tuning to improve outputs without additional large-scale training.
- Complex human poses, multi-object compositions, and difficult lighting remain challenging despite progress in high-resolution, photorealistic synthesis.
- ReflectionFlow scales inference through noise-level, prompt-level, and reflection-level optimization to iteratively assess and correct generated images.The three axes search for better noise initialization, more precise semantic guidance, and actionable corrections to previous outputs.
- ReflectionFlow enables diffusion models to use iterative self-refinement by jointly leveraging textual reflections and previously generated images.The approach is motivated by the relationship between self-refinement and image editing, while addressing the lack of dedicated reflection-guided datasets.
- GenRef contains 1 million reflection triplets across multiple domains, with an additional 227K chain-of-thought annotations.Its automatic construction pipeline uses verifiable objectives, ensemble reward models, and diverse rollout strategies.
- Compared with naive noise-level scaling under identical inference budgets, ReflectionFlow performs better and gains are particularly substantial on challenging prompts.Experiments also examine search-width versus reflection-depth trade-offs and show progressive correction toward superior outputs.
- The framework integrates three inference-time scaling axes and reflection tuning for diffusion transformers, substantially boosting performance through additional inference-time compute.
2. Related Work
Related work traces text-to-image diffusion from U-Net architectures and hand-designed schedules toward diffusion transformers and flow-based training. It also situates ReflectionFlow within efforts to improve pretrained diffusion models at inference time.
- Text-to-image diffusion architectures have shifted from prevalent U-Net models toward diffusion transformers.
- Training strategies have moved from complex hand-designed diffusion schedules toward simpler flow-based formulations and progressive multi-resolution training.These developments improved training efficiency and supported larger diffusion transformer models.
- Recent work increasingly focuses on enhancing pretrained diffusion models during inference rather than relying only on training-scale growth.
3. Method
The method reframes image self-refinement as conditional image editing, using GenRef reflection triplets to train a corrector that jointly processes prompts, reflections, and flawed images. ReflectionFlow then combines noise-, prompt-, and reflection-level inference scaling with adjustable search width and refinement depth.
- Iterative refinement: The corrector model refines an image conditioned on the original prompt, the previous image, and a textual reflection describing shortcomings and improvement directions.External reward models and multimodal language models assess previous outputs and produce instructive reflections.
- Reflection dataset: GenRef contains 1 million flawed-image, high-quality-image, and reflection triplets collected across multiple domains for text-guided image refinement.The dataset addresses the lack of open-source data specifically curated for refinement guided by textual feedback.
- Reflection dataset: GenRef is built through an automated pipeline using rule-based, reward-based, long-short prompt, and editing data sources.The pipeline targets common generator errors, substantial quality differences, and actionable textual corrections.
- Reflection dataset: A chain-of-thought annotation pipeline compares image pairs, identifies differences, judges preference, and produces concise correction instructions.After filtering and quality control, approximately 270K annotations yield 227K high-quality GenRef-CoT reflections.
- Reflection tuning: Reflection tuning adapts pretrained diffusion transformers without additional modules by jointly exchanging information across prompt and image modalities.The concatenated token sequence includes the original and reflection prompts plus flawed and target images, enabling unified attention.
- Test-time scaling: ReflectionFlow scales inference through noise level, prompt level, and reflection level while flexibly balancing search width N against reflection depth M.Search width controls parallel initial-noise chains, whereas reflection depth controls iterative refinement rounds.
4. Experiments
Experiments evaluate ReflectionFlow with different verifiers, inference budgets, and refinement strategies. Results show strong GenEval gains over simpler scaling methods, with deeper sequential refinement particularly effective on challenging prompts.
- 4.2. Main Results: ReflectionFlow reaches a GenEval score of 0.91 with all three scaling dimensions, improving from 0.67 for FLUX.1-dev and 0.85 with noise-level scaling alone.Prompt-level scaling raises the score to 0.87 before reflection-level scaling produces the largest additional gain.
- 4.2. Main Results: ReflectionFlow outperforms Reflect-DiT, with either the SANA verifier or the 16-sample setting exceeding Reflect-DiT’s 0.81 achieved with 20 samples per prompt.The comparison is attributed to the dataset, reflection-model training, and overall inference-time scaling framework.
- 4.3. Ablation Studies: With 32 samples per prompt, the SANA verifier achieves a GenEval score of 0.91, while the oracle upper bound reaches 0.98, indicating substantial verifier-dependent headroom.GPT-4o approaches its limit quickly, whereas the BT-loss verifier continues improving as the sample count increases.
- 4.3. Ablation Studies: ReflectionFlow’s performance rapidly improves as the budget increases from 1 to 4, then improves more slowly, reaching 0.91 with 32 samples per prompt.Noise Scaling and Noise & Prompt Scaling underperform ReflectionFlow across all evaluated budgets, which were capped at 32 samples.
- 4.3. Ablation Studies: Under a fixed budget of N × M = 16, sequential refinement achieves 0.78, outperforming parallel refinement at 0.74.Here, N is the search width and M is the reflection depth; combined strategies balance branching and iterative refinement.
- 4.4. Analysis and Discussion: ReflectionFlow raises correctness on hard prompts from 0.10 to 0.81, compared with increases from 0.55 to 0.85 for medium prompts and 0.95 to 0.97 for easy prompts.The results support deeper iterative reflection as especially beneficial for challenging tasks.
5. Conclusion
ReflectionFlow equips text-to-image diffusion models with iterative self-refinement through reflection-guided correction. It combines GenRef-based fine-tuning with inference-time scaling across noise, reflection, and prompt dimensions.
- ReflectionFlow equips text-to-image diffusion models with iterative self-refinement capabilities.
- GenRef contains one million triplets of flawed images, high-quality images, and textual reflections for reflection-guided refinement.
- The framework formulates self-refinement as generalized image editing and fine-tunes pretrained diffusion transformers without deviating from pretrained distributions.
- ReflectionFlow scales inference across noise, reflection, and prompt dimensions to balance computational efficiency and generation quality.
A. Dataset Preview
The GenRef dataset is previewed through samples from multiple curated subsets and a separate Qwen fine-tuning dataset. The examples pair prompts with reflections, while visual borders distinguish starting, final, correct, and incorrect images.
- Figures 7, 8, and 9 show GenRef samples divided according to the subsets used during data curation.Edit samples were sourced from the OmniEdit dataset.
- The dataset previews provide prompt-and-reflection pairs for each image.
- Red and green borders indicate the starting and final images in the GenRef previews.
- Figure 10 presents samples from the dataset used to fine-tune the Qwen model, with green and red borders denoting correct and incorrect images.Correctness was determined by closed-source APIs.
B. Qualitative Results
The qualitative-results section presents examples of ReflectionFlow's three scaling steps and includes a complex-reasoning case. These visualizations expose the framework's refinement process in more detail.
- Figure 11 presents qualitative results produced by the ReflectionFlow framework.
- The figure includes the detailed process of ReflectionFlow's three scaling steps.The presentation is intended to provide a more nuanced understanding of the framework.
- One listed example uses the prompt “a photo of a backpack right of a sandwich.”
- Figure 11 is identified as showing qualitative results involving complex reasoning.
C. Algorithm Process
The ReflectionFlow algorithm takes a prompt, generator, corrector, verifier, and scaling widths and depths, then iteratively generates, scores, and refines image sets. Its process begins with multiple noise-initialized candidates and repeats correction over refinement steps.
- The algorithm requires a prompt y, generator Gθ, corrector Cϕ, verifier V, scaling width N, and scaling depth M.It aims to return a high-quality image that best realizes user intent.
- The procedure initializes an empty image set and samples zj from N(0, I) for j = 1 to N.
- At each refinement step, the verifier scores the previous images before a new image set is formed.
- The algorithm repeats its refinement loops through the specified scaling depth M.
D. Prompts
The appendix provides the prompts used throughout the work and organizes verifier prompts by the image property being assessed. These properties include object identity, composition, counting, color, position, and color attribution.
- D. Prompts: The appendix states that it provides all prompts used throughout the work.These prompts were inspired by Figure 16 from [34].
- D.1.1. Single Object: The single-object verifier prompt targets images containing one object.
- D.1.2. Two Objects: The two-object verifier prompt targets images containing two objects.
- D.1.3. Counting: The counting verifier prompt targets image-based counting tasks.
- D.1.4. Colors: The colors verifier prompt targets image-based color tasks.
- D.1.5. Position: The position verifier prompt targets image-based position tasks.
- D.1.6. Color Attribution: The color-attribution verifier prompt targets associating colors with image elements.
D.2. Reflection Prompt
The appendix includes prompts for generating reflection instructions and refining prompts or images. These prompts use generated images and assessments to identify shortcomings and guide improvement.
- D.2. Reflection Prompt: The appendix includes a prompt for generating reflection instructions.
- D.3. Refine Prompt: The refinement prompt receives an original prompt, current prompt, generated images, reflections, and corresponding assessments.
- D.3. Refine Prompt: The refinement prompt asks for analysis of the current prompt’s strengths and drawbacks.
- D.3. Refine Prompt: The refinement prompt is intended to improve the overall quality of generated images.
- D.3. Refine Prompt: The refinement materials operationalize reflection as input for improving image-generation prompts.
- D.3. Refine Prompt: The appendix includes a prompt for refinement after reflection-based assessment.
D.4. Chain-of-Though Image Reflection Prompt
The chain-of-thought image reflection prompt compares two generated images against a prompt and produces an editing instruction for the inferior image.
- D.4. Chain-of-Though Image Reflection Prompt: The reflection assistant compares two generated images with respect to their alignment to the given prompt.
- D.4. Chain-of-Though Image Reflection Prompt: The assistant decides which image better matches the prompt and labels it left or right.
- D.4. Chain-of-Though Image Reflection Prompt: The analysis checks missing elements, incorrect attributes, spatial relationships, and unnecessary elements.
- D.4. Chain-of-Though Image Reflection Prompt: The prompt requires exactly one most important editing instruction for modifying the inferior image.
- D.4. Chain-of-Though Image Reflection Prompt: The appendix includes a prompt for generating chain-of-thought image reflection annotations.