[論文レビュー] Diffusion Probe: Generated Image Result Prediction Using CNN Probes

Benlei Cui, Bukun Huang+8 more|arXiv (Cornell University)|Feb 27, 2026

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モデルにおける最終画像品質との関連を解明する。
• nascentなアテンション統計を最終品質スコアへマッピングする軽量なCNNベースのプローブを開発する。
• モデルに依存しない適用性を実証し、実践的ワークフローにおける効率向上を検証する。

提案手法

• 与えられたプロンプトに対して、初期のデノイジングステップで中間モデルブロックからクロスアテンションマップを抽出する。
• attentionマップと timestep 埋め込みを用いてスカラー品質スコアへマッピングする軽量プローブ E_theta を、実測指標に対する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は diverseモデルと初期デノイズステップにおいて高い予測精度（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