Diffusion Probe: Generated Image Result Prediction Using CNN Probes

TL;DR

Diffusion Probe introduces a CNN-based framework that predicts final image quality early in the diffusion process by analyzing cross-attention maps, enabling more efficient and targeted image generation workflows.

Contribution

It presents a novel, model-agnostic method leveraging internal attention maps to accurately forecast image quality early in the diffusion process, reducing computational costs.

Findings

Achieves PCC > 0.7 in early quality prediction across models and settings.

Attains AUC-ROC > 0.9 for quality classification.

Enables early decision-making to improve efficiency and output quality.

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.