Reflect-DiT: Escalonamento em Tempo de Inferência para Transformers de Difusão Texto-para-Imagem via Reflexão em 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

Resumo

A abordagem predominante para avançar a geração de texto para imagem tem sido o escalonamento durante o treinamento, onde modelos maiores são treinados com mais dados utilizando recursos computacionais mais robustos. Embora eficaz, essa abordagem é computacionalmente custosa, levando a um crescente interesse no escalonamento durante a inferência para melhorar o desempenho. Atualmente, o escalonamento durante a inferência para modelos de difusão de texto para imagem é amplamente limitado à amostragem best-of-N, onde múltiplas imagens são geradas por prompt e um modelo de seleção escolhe a melhor saída. Inspirados pelo recente sucesso de modelos de raciocínio como o DeepSeek-R1 no domínio da linguagem, introduzimos uma alternativa à amostragem best-of-N ingênua ao equipar Diffusion Transformers (Transformadores de Difusão) para texto e imagem com capacidades de reflexão em contexto. Propomos o Reflect-DiT, um método que permite que Diffusion Transformers refinem suas gerações utilizando exemplos em contexto de imagens previamente geradas juntamente com feedback textual descrevendo melhorias necessárias. Em vez de depender passivamente da amostragem aleatória e esperar por um resultado melhor em uma geração futura, o Reflect-DiT ajusta explicitamente suas gerações para abordar aspectos específicos que requerem aprimoramento. Resultados experimentais demonstram que o Reflect-DiT melhora o desempenho no benchmark GenEval (+0,19) utilizando o SANA-1.0-1.6B como modelo base. Além disso, ele alcança um novo estado da arte com pontuação de 0,81 no GenEval enquanto gera apenas 20 amostras por prompt, superando a melhor pontuação anterior de 0,80, que foi obtida usando um modelo significativamente maior (SANA-1.5-4.8B) com 2048 amostras sob a abordagem best-of-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.