Reflect-DiT: 텍스트-이미지 확산 트랜스포머를 위한 인퍼런스 시점 스케일링 - 인-컨텍스트 리플렉션 기반 접근

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

March 15, 2025

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

cs.AI

초록

텍스트-이미지 생성 기술을 발전시키기 위한 주된 접근 방식은 더 큰 모델을 더 많은 데이터로 훈련시키고 더 많은 계산 자원을 사용하는 훈련 시간 스케일링이었습니다. 이 방법은 효과적이지만 계산 비용이 많이 들어, 성능을 개선하기 위한 추론 시간 스케일링에 대한 관심이 점차 증가하고 있습니다. 현재, 텍스트-이미지 확산 모델에 대한 추론 시간 스케일링은 주로 best-of-N 샘플링에 국한되어 있으며, 이는 각 프롬프트에 대해 여러 이미지를 생성하고 선택 모델이 최적의 출력을 선택하는 방식입니다. 언어 영역에서 DeepSeek-R1과 같은 추론 모델의 최근 성공에 영감을 받아, 우리는 텍스트-이미지 Diffusion Transformer에 컨텍스트 내 반영 기능을 추가하여 단순한 best-of-N 샘플링에 대한 대안을 제시합니다. 우리는 Reflect-DiT라는 방법을 제안하며, 이는 Diffusion Transformer가 이전에 생성된 이미지의 컨텍스트 내 예시와 필요한 개선 사항을 설명하는 텍스트 피드백을 사용하여 생성물을 개선할 수 있도록 합니다. Reflect-DiT는 무작위 샘플링에 수동적으로 의존하고 미래의 생성에서 더 나은 결과를 기대하는 대신, 개선이 필요한 특정 측면을 명시적으로 해결하기 위해 생성물을 맞춤화합니다. 실험 결과는 Reflect-DiT가 SANA-1.0-1.6B를 기본 모델로 사용하여 GenEval 벤치마크에서 성능을 개선(+0.19)함을 보여줍니다. 또한, Reflect-DiT는 각 프롬프트당 20개의 샘플만 생성하면서 GenEval에서 0.81의 새로운 최고 점수를 달성하여, 이전의 최고 점수인 0.80(훨씬 더 큰 모델인 SANA-1.5-4.8B를 사용하여 best-of-N 접근 방식으로 2048개의 샘플을 생성한 결과)을 능가했습니다.

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.