Reflect-DiT: Scalabilità al Tempo di Inferenza per Modelli di Diffusione Testo-Immagine tramite Riflessione In-Contesto

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

March 15, 2025

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

cs.AI

Abstract

L'approccio predominante per migliorare la generazione di immagini da testo è stato il ridimensionamento durante l'addestramento, in cui modelli più grandi vengono addestrati su più dati utilizzando maggiori risorse computazionali. Sebbene efficace, questo approccio è computazionalmente costoso, portando a un crescente interesse per il ridimensionamento durante l'inferenza per migliorare le prestazioni. Attualmente, il ridimensionamento durante l'inferenza per i modelli di diffusione testo-immagine è principalmente limitato al campionamento best-of-N, in cui vengono generate più immagini per prompt e un modello di selezione sceglie l'output migliore. Ispirati dal recente successo di modelli di ragionamento come DeepSeek-R1 nel dominio del linguaggio, introduciamo un'alternativa al semplice campionamento best-of-N dotando i Diffusion Transformer testo-immagine di capacità di riflessione in contesto. Proponiamo Reflect-DiT, un metodo che consente ai Diffusion Transformer di affinare le loro generazioni utilizzando esempi in contesto di immagini precedentemente generate insieme a feedback testuali che descrivono i miglioramenti necessari. Invece di fare affidamento passivamente sul campionamento casuale e sperare in un risultato migliore in una generazione futura, Reflect-DiT adatta esplicitamente le sue generazioni per affrontare aspetti specifici che richiedono miglioramenti. I risultati sperimentali dimostrano che Reflect-DiT migliora le prestazioni sul benchmark GenEval (+0.19) utilizzando SANA-1.0-1.6B come modello di base. Inoltre, raggiunge un nuovo punteggio record di 0.81 su GenEval generando solo 20 campioni per prompt, superando il precedente miglior punteggio di 0.80, ottenuto utilizzando un modello significativamente più grande (SANA-1.5-4.8B) con 2048 campioni con l'approccio 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.