Reflect-DiT: Escalado en Tiempo de Inferencia para Transformadores de Difusión de Texto a Imagen mediante Reflexión en Contexto

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

March 15, 2025

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

cs.AI

Resumen

El enfoque predominante para avanzar en la generación de texto a imagen ha sido el escalamiento durante el entrenamiento, donde modelos más grandes se entrenan con más datos utilizando mayores recursos computacionales. Aunque efectivo, este enfoque es computacionalmente costoso, lo que ha generado un creciente interés en el escalamiento durante la inferencia para mejorar el rendimiento. Actualmente, el escalamiento durante la inferencia para modelos de difusión de texto a imagen se limita principalmente al muestreo de mejor-de-N, donde se generan múltiples imágenes por prompt y un modelo de selección elige la mejor salida. Inspirados por el reciente éxito de modelos de razonamiento como DeepSeek-R1 en el dominio del lenguaje, introducimos una alternativa al muestreo ingenuo de mejor-de-N al equipar a los Transformadores de Difusión (DiT) con capacidades de reflexión en contexto. Proponemos Reflect-DiT, un método que permite a los Transformadores de Difusión refinar sus generaciones utilizando ejemplos en contexto de imágenes previamente generadas junto con retroalimentación textual que describe las mejoras necesarias. En lugar de depender pasivamente del muestreo aleatorio y esperar un mejor resultado en una generación futura, Reflect-DiT adapta explícitamente sus generaciones para abordar aspectos específicos que requieren mejora. Los resultados experimentales demuestran que Reflect-DiT mejora el rendimiento en el benchmark GenEval (+0.19) utilizando SANA-1.0-1.6B como modelo base. Además, alcanza un nuevo estado del arte con una puntuación de 0.81 en GenEval mientras genera solo 20 muestras por prompt, superando la mejor puntuación anterior de 0.80, que se obtuvo utilizando un modelo significativamente más grande (SANA-1.5-4.8B) con 2048 muestras bajo el enfoque de mejor-de-N.

English

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.