PaperlayerView exact revision on arXiv

Source-linked AI summary

Reflect-DiT: Inference-Time Scaling for Text-to-Image Diffusion Transformers via In-Context Reflection

Shufan Li, Konstantinos Kallidromitis, Akash Gokul, Arsh Koneru, Yusuke Kato, Kazuki Kozuka, Aditya Grover

arXiv:2503.12271v1cs.CV

TL;DR

Text-to-image inference-time scaling has largely relied on expensive best-of-N sampling, motivating a method that can refine generations from prior feedback. Reflect-DiT adds in-context reflection to Diffusion Transformers, achieving 0.81 on GenEval with 20 samples and a +0.19 gain over SANA-1.0-1.6B.

  • Problem

    Best-of-N inference improves text-to-image generation but requires generating many samples and selecting among them, creating an efficiency gap.

    Source: §1 · pp. 1–2

  • Method

    Reflect-DiT uses a VLM judge and a Diffusion Transformer to refine images from past generations and natural-language feedback.

    Source: §1 · p. 2, §3.3 · p. 4

  • Results

    0.81 on GenEval with 20 samples establishes a new state of the art, while Reflect-DiT gains +0.19 over SANA-1.0-1.6B.

    Source: §1 · p. 2, §4.3 · pp. 5–6

  • Takeaways & Limitations

    Reflect-DiT provides a more effective and efficient alternative to naive best-of-N sampling within the reported evaluation setting.

    Source: §1 · p. 2, §6 · p. 8

    The method inherits base-model and VLM-judge limitations, including hallucinated feedback and difficulty recognizing small objects.

    Source: §6 · p. 9, B. Limitations · p. 14

Abstract

from arXiv · show

The predominant approach to advancing text-to-image generation has been training-time scaling, where larger models are trained on more data using greater computational resources. While effective, this approach is computationally expensive, leading to growing interest in inference-time scaling to improve performance. Currently, inference-time scaling for text-to-image diffusion models is largely limited to best-of-N sampling, where multiple images are generated per prompt and a selection model chooses the best output. Inspired by the recent success of reasoning models like DeepSeek-R1 in the language domain, we introduce an alternative to naive best-of-N sampling by equipping text-to-image Diffusion Transformers with in-context reflection capabilities. We propose Reflect-DiT, a method that enables Diffusion Transformers to refine their generations using in-context examples of previously generated images alongside textual feedback describing necessary improvements. Instead of passively relying on random sampling and hoping for a better result in a future generation, Reflect-DiT explicitly tailors its generations to address specific aspects requiring enhancement. Experimental results demonstrate that Reflect-DiT improves performance on the GenEval benchmark (+0.19) using SANA-1.0-1.6B as a base model. Additionally, it achieves a new state-of-the-art score of 0.81 on GenEval while generating only 20 samples per prompt, surpassing the previous best score of 0.80, which was obtained using a significantly larger model (SANA-1.5-4.8B) with 2048 samples under the best-of-N approach.

1. Introduction

Reflect-DiT addresses the inefficiency of training-time scaling and naive best-of-N inference by enabling Diffusion Transformers to refine generations using past images and textual feedback. It achieves stronger GenEval performance with substantially fewer samples than prior best-of-N scaling.

  • Inference-time scaling offers an alternative to expensive training-time scaling by adding computation during generation.Existing approaches generate many samples and select the best output, or use additional reasoning and verification before producing an answer.Source: §1 · pp. 1–2
  • Text-to-image inference-time scaling has primarily relied on best-of-N sampling, including SANA-1.5’s 2048-sample GenEval result.Best-of-N uses random samples and a reward or judge model to select an output.Source: §1 · p. 2
  • Reflect-DiT equips DiTs to evaluate past generations, identify prompt misalignments, and refine later images using natural-language feedback.The method targets errors such as object count and spatial positioning rather than relying solely on random future samples.Source: §1 · p. 2
  • The framework encodes past images and feedback, processes them with a Context Transformer, and supplies the resulting conditional embeddings to the DiT.A fixed context length limits the number of past generations, with stochastic subset selection when necessary.Source: §1 · p. 2
  • Reflect-DiT improves GenEval by +0.19 over SANA-1.0-1.6B and reaches 0.81 with only 20 samples per prompt.This surpasses the previous 0.80 score achieved by SANA-1.5-4.8B with 2048 samples.Source: §1 · p. 2

2. Related Works

Related work develops inference-time scaling for language and image generation, while prior image self-correction and controllable-generation methods use more constrained mechanisms than Reflect-DiT. Reflect-DiT instead applies iterative reflection directly within the diffusion-generation process.

2.1. Inference-time scaling of LLMs

  • Inference-time language methods use repeated sampling, adaptive search, self-verification, and self-correction to improve results beyond training-time scaling.Prior work reports that adaptive search and iterative refinement can outperform naive Best-of-N under fixed compute.Source: §2.1 · p. 2

2.2. Inference-time scaling of Text-to-Image Diffusion Models

  • Recent text-to-image diffusion work explores inference-time scaling strategies, with random search plus Best-of-N selection reported as best; SANA-1.5 achieved a new GenEval state of the art using N=2048.Source: §2.2 · pp. 2–3

2.3. Self-Correction for Text-to-Image Generation

  • Earlier self-correction systems use LLMs or VLMs to control predefined image operations rather than letting diffusion models learn from natural feedback.Their performance depends on successful execution of specialized modules and heuristics.Source: §2.3 · p. 3

2.4. Controllable Text-to-Image Generation

  • Controllable-generation methods expand text conditioning with learnable embeddings or visual signals, but do not use those signals to learn from iterative feedback.Prior instruction-based methods mainly target image editing rather than iterative inference-time refinement.Source: §2.4 · p. 3

3. Method

Reflect-DiT performs iterative verification and reflection: a VLM critiques generated images, while a DiT uses selected past image-feedback pairs to generate refinements. Its architecture encodes and fuses multimodal context before injecting it into DiT cross-attention.

3.1. Overall Framework

  • Reflect-DiT loops between image generation and VLM feedback until feedback is empty or the maximum iteration count is reached.When context exceeds length K, the method randomly samples K past image-feedback pairs before generating the next image.Source: §3.1 · p. 4, §3.3 · p. 4

3.2. VLM Feedback

  • The VLM judge is trained on synthetic prompt-image data to provide concise feedback about object presence, counts, positions, and attributes.Structured templates express errors such as missing objects or incorrect counts.Source: §3.2 · p. 4

3.3. Diffusion Transformer

  • The DiT incorporates past generated images and text feedback as additional cross-attention context while using SANA-1.0-1.6B as the base model.SANA’s linear attention and smaller size support inference-time scaling with many samples.Source: §3.3 · p. 4
  • Past images are encoded as vision embeddings and feedback as text embeddings, then concatenated and processed by a two-layer Context Transformer.The resulting context is appended to standard prompt embeddings and passed into DiT cross-attention.Source: §3.3 · pp. 4–5
  • Table 1 reports that Reflect-DiT achieves 0.81 on GenEval with 20 samples and gains +0.19 over SANA-1.0-1.6B.The table compares Reflect-DiT with larger or more computationally demanding models.Source: Table 1 · p. 5
  • Training pairs a good image with bad images and corresponding feedback, using the flow-matching objective for DiT finetuning.The reflection context consists of in-context image-feedback pairs, and logit-normal weighting follows SANA.Source: §3.3 · p. 5

4. Experiments

Experiments evaluate Reflect-DiT on GenEval, qualitative trajectories, finetuning baselines, human preferences, and design ablations. Across these settings, reflection improves quality and efficiency relative to best-of-N and competing finetuning methods.

4.1.1. Dataset

  • 20 images per prompt are generated from 6,000 filtered templates, producing a 78.5k image-feedback-pair dataset for training.The VLM judge and Diffusion Transformer are trained on this synthetic dataset; image and text encoders remain frozen during training.Source: §4.1.1 · p. 5, §4.1.2 · p. 5

4.2. Sampling

  • Reflect-DiT uses 20 sampling steps, permits at most 20 images per prompt, and normally conditions on three previous feedback contexts.Baselines include best-of-N, SFT, and Diffusion-DPO, with best-of-N added to finetuning methods to equalize test-time compute.Source: §4.2 · p. 5, §4.2.1 · p. 5

4.3. Main results

  • 0.81 overall GenEval score with at most 20 samples establishes Reflect-DiT as the highest-scoring method, including models with substantially more parameters.It improves the SANA-1.0-1.6B baseline by +0.19 and surpasses the previous 0.80 score from SANA-1.5-4.8B using N=2048.Source: §4.3 · pp. 5–6

4.4. Qualitative Examples

  • Reflect-DiT progressively corrects spatial relations, object counts, object identity, and shapes in generation trajectories using feedback.Examples include satisfying “a cup left of an umbrella,” counting butterflies and dandelions, and refining a crown-shaped splash that best-of-N fails to find.Source: Figure 3 · p. 6, §4.4 · p. 6

4.5. Comparison with Finetuning

  • 4 samples suffice for Reflect-DiT to outperform SFT and Diffusion-DPO baselines using 20 samples, saving 80% of the compute budget.Across varying inference budgets, Reflect-DiT consistently outperforms both finetuning methods.Source: Figure 4 · p. 6, §4.5 · p. 7, Table 2 · p. 7

4.6. Human Evaluation

  • 53.3% versus 32.3% of human selections favor Reflect-DiT over best-of-N on 100 PartiPrompts prompts.Evaluators compared paired images with a maximum of 20 samples per prompt for each method.Source: §4.6 · p. 7, Figure 5 · p. 7

4.7.1. Number of Feedback

  • Three in-context feedback pairs achieve the best ablation result, while increasing Context Transformer layers improves performance.These studies evaluate the number of feedback pairs and Transformer layers as separate design choices.Source: §4.7.1 · p. 7, §4.7.2 · p. 7, Table 3 · p. 7

4.7.3. Number of Image Embeddings

  • Using more tokens to represent past generations improves performance in the image-embedding ablation.The default image-embedding down-sampling size is 8 × 8.Source: §4.7.3 · pp. 7–8, Table 5 · p. 8

5. Discussion

Reflect-DiT differs from instruction-guided image editing by requiring only good images that avoid problems in bad images, while using multi-round feedback to progressively refine generations.

5.1. Connection with Instruction-Guided Image Editing

  • Reflect-DiT requires only good images that avoid defects in bad images, rather than paired input-edited images that preserve visual consistency.Source: §5.1 · p. 8
  • Multi-round feedback progressively refines generations, with images converging toward correctness across iterations even when individual defects are not fixed immediately.This ability emerges despite randomly sampled in-context feedback during training.Source: §5.1 · p. 8

6. Conclusion

Reflect-DiT uses reflection on past generations and natural-language feedback to improve text-to-image generation, outperforming naive best-of-N and larger models with higher sampling budgets. Its main limitation is inherited bias and error from both the base model and VLM judge.

6. Conclusion

  • Reflect-DiT outperforms naive best-of-N and establishes a new state-of-the-art on GenEval despite using larger models and substantially higher sampling budgets as comparators.Source: §6 · p. 8
  • Reflect-DiT inherits biases and shortcomings from both the base text-to-image model and its VLM judge, which may hallucinate feedback and struggle with small objects.Source: §6 · p. 9

A.1. Generalizability of Reflect-DiT

The appendix reports Reflect-DiT's generalization across harder prompts, comparable throughput to best-of-N, simpler scaling than tool-based agents, stable supervised training, and advantages over concurrent approaches.

A.1. Generalizability of Reflect-DiT

  • Reflect-DiT performs better than the base model and best-of-N across DPG-Bench, with larger gains on the Hard-246 and Hard-56 subsets.The full benchmark contains many easy prompts, while the hard subsets select prompts where SANA-1.0-1.6B scores below 0.8 or 0.5.Source: §A.1 · p. 13

A.2. Inference Speed

  • 11.32 samples per minute for Reflect-DiT versus 10.12 for best-of-N indicates similar throughput, while N=20 latency is 118.57 versus 105.98 seconds.Reflect-DiT's extra overhead comes from context encoding and cross-attention; 75.5% of time is image generation and 24.5% VLM inference.Source: §A.2 · p. 13

A.3. Connection with Self-Correcting T2I Agent

  • Reflect-DiT avoids the proprietary APIs, multiple external tools, and execution overhead required by self-correcting T2I agents.Its refinement instead uses a smaller VLM judge to produce concise natural-language feedback for each image.Source: §A.3 · p. 13

A.4. Connection with Reinforcement Learning (RL)

  • Reflect-DiT achieves state-of-the-art results with supervised fine-tuning, while reinforcement-learning objectives encountered training-stability issues.The authors leave further exploration of reinforcement-learning objectives to future work.Source: §A.4 · p. 14

A.5. Concurrent Works

  • Reflect-DiT outperforms SANA-1.5's best-of-2048 with only 20 samples and maintains constant memory, unlike concurrent approaches with solver or memory limitations.The cited comparison includes FK-Steering and an autoregressive chain-of-thought image-generation approach.Source: §A.5 · p. 14

B. Limitations

Reflect-DiT's judge is limited by prompt-alignment-focused training data, hallucinated feedback, and occasional inability of the diffusion model to correct feedback in one iteration.

B. Limitations

  • The VLM judge may provide poor feedback about image aesthetics because its training data primarily targets prompt alignment.Natural-language feedback datasets covering broader aspects are more difficult to collect.Source: B. Limitations · p. 14
  • VLM hallucinations can misidentify image content, while some diffusion-model corrections require multiple iterations or follow inconsistent trajectories.The authors report both gradual correction toward correctness and inconsistent progress in counting examples.Source: B. Limitations · p. 14

C.1. VLM Training

The method uses frozen vision and text encoders with trainable projection and context components to condition the DiT on image-feedback context. Its VLM training follows a conversation format, while the Context Transformer aligns modalities and associates feedback with images.

C.1. VLM Training

  • The VLM training data uses a conversation format based on SANA-1.5 but adapts it for the Qwen-2.5-VL 3B base model.The associated template is presented in Table 7.Source: §C.1 · pp. 14–15, Table 7 · p. 15
  • Feedback can introduce errors when lighting or reflected colors create ambiguity, while small or unusual objects may be overlooked by the feedback model.The examples distinguish subjective interpretations from more typical detection failures.Source: Figure 7 · p. 15

C.2.1. Vision-Encoder

  • SigLIP-Large encodes each image into a 24×24 feature map, which average pooling compresses to an 8×8 sequence of 64 embeddings.A two-layer GELU MLP projects the compressed features to the Context Transformer’s input dimension.Source: §C.2.1 · p. 15
  • The SigLIP vision encoder is frozen before training, while its projector is trained end-to-end with the rest of the DiT.An RMSNorm layer follows the projector to improve training stability.Source: §C.2.1 · p. 15

C.2.2. Text-Encoder

  • Gemma-2B provides text-feedback encoding and remains frozen, adding no parameters because SANA already uses it as the prompt encoder.The text encoder is part of the modality-specific conditioning pipeline.Source: §C.2.2 · p. 15

C.2.3. Context Transformer

  • The Context Transformer is a two-layer Transformer that aligns encoded features with the base DiT and associates feedback with its corresponding image.It uses self-attention, a feed-forward network, and rotary positional embeddings.Source: §C.2.3 · p. 15

C.2.4. Training

The training setup reports a CAME optimizer for DiT training and selects β = 2000 as the optimal Diffusion-DPO KL-divergence penalty.

C.2.4. Training

  • β = 2000 produces the optimal result among the tested Diffusion-DPO settings.The authors report training hyperparameters for Reflect-DiT, SFT, and Diffusion-DPO, using CAME for DiT training following SANA.Source: §C.2.4 · p. 15

C.3. Human Evaluation Details

Human evaluation compares Reflect-DiT(N=20) with best-of-20 using annotator judgments on PartiPrompts image pairs. Additional examples illustrate improvements in spatial placement, counting, and separating object identity from color.

C.3. Human Evaluation Details

  • Reflect-DiT(N=20) is compared with best-of-20 through three evaluations per image pair in a human study.The study uses Amazon Mechanical Turk and presents annotators with a dedicated interface.Source: §C.3 · p. 15, Figure 8 · p. 16
  • The human evaluation samples 100 PartiPrompts prompts, generates 100 image pairs, and collects 300 annotations.Each prompt contributes a corresponding pair for comparison.Source: §C.3 · p. 16

D. Additional Qualitative Examples

  • Qualitative examples show Reflect-DiT handling object placement, counting, complex multi-object prompts, and separating object identity from color attributes.One example concerns fulfilling a request such as “a black sandwich” without conflating color and object identity.Source: D. Additional Qualitative Examples · p. 16
Loading 2503.12271v1…