Reflect-DiT: Inferentie-tijd schaling voor tekst-naar-beeld diffusie Transformers via In-Context Reflectie

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

March 15, 2025

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

cs.AI

Samenvatting

De overheersende aanpak om tekst-naar-beeldgeneratie te verbeteren, is het schalen tijdens de training, waarbij grotere modellen worden getraind met meer data en grotere rekenkracht. Hoewel effectief, is deze aanpak rekenkundig kostbaar, wat heeft geleid tot een groeiende interesse in schalen tijdens de inferentie om de prestaties te verbeteren. Momenteel is schalen tijdens de inferentie voor tekst-naar-beeld diffusiemodellen grotendeels beperkt tot best-of-N sampling, waarbij meerdere afbeeldingen per prompt worden gegenereerd en een selectiemodel de beste uitvoer kiest. Geïnspireerd door het recente succes van redeneermodellen zoals DeepSeek-R1 in het taaldomein, introduceren we een alternatief voor naïeve best-of-N sampling door tekst-naar-beeld Diffusion Transformers uit te rusten met in-context reflectiecapaciteiten. We stellen Reflect-DiT voor, een methode die Diffusion Transformers in staat stelt hun generaties te verfijnen met behulp van in-context voorbeelden van eerder gegenereerde afbeeldingen samen met tekstuele feedback die de benodigde verbeteringen beschrijft. In plaats van passief te vertrouwen op willekeurige sampling en te hopen op een beter resultaat in een toekomstige generatie, past Reflect-DiT zijn generaties expliciet aan om specifieke aspecten die verbetering vereisen aan te pakken. Experimentele resultaten tonen aan dat Reflect-DiT de prestaties op de GenEval-benchmark verbetert (+0.19) met SANA-1.0-1.6B als basismodel. Bovendien behaalt het een nieuwe state-of-the-art score van 0.81 op GenEval terwijl slechts 20 samples per prompt worden gegenereerd, wat de vorige beste score van 0.80 overtreft, die werd behaald met een aanzienlijk groter model (SANA-1.5-4.8B) met 2048 samples onder de best-of-N aanpak.

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.