[논문 리뷰] Diffusion Probe: Generated Image Result Prediction Using CNN Probes

Benlei Cui, Bukun Huang외 8명|arXiv (Cornell University)|2026. 02. 27.

Generative Adversarial Networks and Image Synthesis인용 수 0

한 줄 요약

Diffusion Probe는 확산 모델의 초기 교차 어텐션 맵을 사용하여 최종 이미지 품질을 예측하고, 전체 생성을 필요로 하지 않는 조기 품질 평가 및 효율적인 다운스트림 최적화를 가능하게 한다.

ABSTRACT

Text-to-image (T2I) diffusion models lack an efficient mechanism for early quality assessment, leading to costly trial-and-error in multi-generation scenarios such as prompt iteration, agent-based generation, and flow-grpo. We reveal a strong correlation between early diffusion cross-attention distributions and final image quality. Based on this finding, we introduce Diffusion Probe, a framework that leverages internal cross-attention maps as predictive signals. We design a lightweight predictor that maps statistical properties of early-stage cross-attention extracted from initial denoising steps to the final image's overall quality. This enables accurate forecasting of image quality across diverse evaluation metrics long before full synthesis is complete. We validate Diffusion Probe across a wide range of settings. On multiple T2I models, across early denoising windows, resolutions, and quality metrics, it achieves strong correlation (PCC > 0.7) and high classification performance (AUC-ROC > 0.9). Its reliability translates into practical gains. By enabling early quality-aware decisions in workflows such as prompt optimization, seed selection, and accelerated RL training, the probe supports more targeted sampling and avoids computation on low-potential generations. This reduces computational overhead while improving final output quality.Diffusion Probe is model-agnostic, efficient, and broadly applicable, offering a practical solution for improving T2I generation efficiency through early quality prediction.

연구 동기 및 목표

• 확산 기반 T2I 모델에서 초기 단계의 크로스-어텐션 패턴과 최종 이미지 품질 사이의 연결 고리를 밝힌다.
• 나선형(attention) 통계치를 최종 이미지 품질 점수로 매핑하는 경량 CNN 기반 프로브를 개발한다.
• 모델-독립적 적용 가능성을 입증하고, 실용적 워크플로에서의 효율성 향상을 검증한다.

제안 방법

• 주어진 프롬프트에 대한 초기 디노이징 단계에서 중간 모델 블록으로부터 크로스-어텐션 맵을 추출한다.
• 경량 프로브 E_theta를 훈련시켜 어텐션 맵과 타임스텝 임베딩을 이용해 지상 트루(metric)와의 MSE를 통해 스칼라 품질 점수를 매핑한다.
• 프로브를 예측기로 사용하여 전체 이미지 생성을 필요로 하지 않는 다운스트림 작업(프롬프트 최적화, 시드 선택, RL 학습)을 안내한다.
• 다양한 T2I 백본(SDXL, FLUX, Qwen-Image 등)에서 SRCC, KTC, PCC, AUC-ROC를 통해 프로브 정확도를 평가한다.
• 프로브를 다운스트림 작업에 적용하여 프롬프트 필터링, 시드 선택 또는 Flow-GRPO 학습의 보상 신호를 제공한다.

Figure 1 : Illustration of early cross-attention dispersion. Here, we present the prompt, the corresponding four cross-attention activation maps in the early denoising stage, and the final generated image. Compared to other tokens, the cross-attention activation maps of the “bird” token shows signif

실험 결과

연구 질문

• RQ1 초기 단계의 크로스-어텐션 분포가 서로 다른 T2I 모델에서 최종 이미지 품질을 예측할 수 있는가?
• RQ2 경량 프로브로 얼마나 이른 시점에서 품질을 신뢰성 있게 예측할 수 있는가?
• RQ3 모델-독립적 프로브가 전체 생성을 거치지 않고도 프롬프트 최적화, 시드 선택, RL 학습을 가능하게 하는가?

주요 결과

Base Model	Resolution	Step	SRCC	AUC-ROC	KTC	PCC
SDXL	1024×1024	1	0.49	0.53	0.35	0.48
SDXL	1024×1024	5	0.73	0.86	0.57	0.72
SDXL	1024×1024	10	0.76	0.89	0.61	0.75
SDXL	1024×1024	15	0.75	0.89	0.60	0.74
FLUX	1024×1024	1	0.52	0.62	0.38	0.50
FLUX	1024×1024	5	0.76	0.88	0.60	0.75
FLUX	1024×1024	10	0.79	0.91	0.64	0.78
FLUX	1024×1024	15	0.78	0.91	0.63	0.77
Qwen-Image	1024×1024	1	0.45	0.67	0.32	0.44
Qwen-Image	1024×1024	5	0.69	0.84	0.53	0.68
Qwen-Image	1024×1024	10	0.72	0.87	0.56	0.71
Qwen-Image	1024×1024	15	0.71	0.86	0.55	0.70

• Diffusion Probe는 다양한 모델 및 초기 디노이징 단계에서 높은 예측 정확도(SRCC, KTC, PCC)와 강력한 분류 성능(AUC-ROC)을 달성한다.
• FLUX에서 프롬프트의 예측 지표가 10단계 즈음에 피크에 도달한다( SRCC 0.79, AUC 0.91, PCC 0.78).
• 프로브는 SDXL 및 Qwen-Image에 일반화되어 높은 상관성(SRCC 약 0.72–0.76)과 AUC (>0.86)을 유지한다.
• 다운스트림 작업에서 프롬프트 최적화 및 시드 선택 지표를 개선하고 더 무거운 LLM 기반 방법에 비해 계산을 줄이면서 경쟁력 있는 성능을 보인다.
• Flow-GRPO에 프로브를 통합하면 더 높은 품질의 샘플로 배치를 풍부하게 만들어 RL 학습을 가속하고 수렴 안정성을 향상시킨다.

Figure 2 : Overview of the Diffusion Probe framework. Our framework takes as input the early-stage cross-attention feature maps (derived from the CrossAttn module at a probed timestep $t$ ) and the TimeStep Embedding . A lightweight network processes these inputs, ultimately outputting a quality sco