Reflect-DiT : Mise à l'échelle au moment de l'inférence pour les transformateurs de diffusion texte-image via la réflexion en contexte

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

Résumé

L'approche prédominante pour faire progresser la génération texte-image a été la mise à l'échelle lors de l'entraînement, où des modèles plus volumineux sont entraînés sur davantage de données en utilisant des ressources computationnelles accrues. Bien qu'efficace, cette approche est coûteuse en termes de calcul, suscitant un intérêt croissant pour la mise à l'échelle lors de l'inférence afin d'améliorer les performances. Actuellement, la mise à l'échelle lors de l'inférence pour les modèles de diffusion texte-image se limite largement à l'échantillonnage best-of-N, où plusieurs images sont générées par prompt et un modèle de sélection choisit la meilleure sortie. Inspirés par le récent succès des modèles de raisonnement comme DeepSeek-R1 dans le domaine du langage, nous introduisons une alternative à l'échantillonnage best-of-N naïf en dotant les Transformers de Diffusion texte-image de capacités de réflexion en contexte. Nous proposons Reflect-DiT, une méthode qui permet aux Transformers de Diffusion d'affiner leurs générations en utilisant des exemples en contexte d'images précédemment générées ainsi que des retours textuels décrivant les améliorations nécessaires. Au lieu de se reposer passivement sur un échantillonnage aléatoire en espérant un meilleur résultat dans une génération future, Reflect-DiT adapte explicitement ses générations pour répondre à des aspects spécifiques nécessitant des améliorations. Les résultats expérimentaux montrent que Reflect-DiT améliore les performances sur le benchmark GenEval (+0,19) en utilisant SANA-1.0-1.6B comme modèle de base. De plus, il atteint un nouveau score record de 0,81 sur GenEval en générant seulement 20 échantillons par prompt, surpassant le précédent meilleur score de 0,80, obtenu avec un modèle significativement plus volumineux (SANA-1.5-4.8B) et 2048 échantillons dans le cadre de l'approche 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.