텍스트 이미지 확산 모델의 조기 품질 평가를 위한 ProbeSelect

초록

본 논문은 텍스트‑이미지 확산·플로우‑매칭 모델에서 이미지 생성 초기에 나타나는 구조적 신호를 활용해 최종 이미지 품질을 예측하는 ProbeSelect 모듈을 제안한다. 20 % 정도의 디노이징 단계만 수행한 뒤 중간 활성화를 입력으로 품질 점수를 추정하고, 저품질 시드를 조기에 종료함으로써 전체 샘플링 비용을 60 % 이상 절감하면서도 최종 이미지 품질을 향상시킨다.

상세 분석

ProbeSelect는 기존 텍스트‑이미지 생성 파이프라인을 변경하지 않고, 디노이저 내부의 중간 활성화(hₜ)를 가벼운 프로브 네트워크에 연결하는 플러그‑인 방식이다. 핵심 가정은 디노이징 초기 단계에서도 U‑Net의 특정 레이어가 물체 레이아웃, 공간 구도, 의미적 군집과 같은 거친 구조 정보를 안정적으로 보존한다는 점이다. 이러한 구조는 시간에 따라 천천히 변하기 때문에, 0.2 ~ 0.3 정도의 타임스텝에서 추출한 특징은 최종 이미지의 품질과 높은 상관관계를 가진다.

프로브는 (1) 특징 추출 단계에서 선택된 블록의 활성화를 전역 풀링 후 차원 축소하는 작은 비전 인코더(g_ϕ)와, (2) 타임스텝 임베딩을 결합한 뒤 스칼라 점수를 출력하는 MLP(p_ϕ)로 구성된다. 학습 목표는 두 가지 손실의 가중합이다. 첫 번째는 리스트와이즈 랭킹 손실(L_list)으로, 동일 배치 내 시드들의 외부 평가자 점수(y_i)를 기준으로 예측 점수(ŷ_iₜ)의 순위를 맞추도록 한다. 이는 절대값보다 상대적 순서에 초점을 맞추어 구조적 차이를 강조한다. 두 번째는 텍스트 정렬 손실(L_align)으로, 프로브 출력 uₜ와 고정된 텍스트 인코더(CLIP)의 프롬프트 임베딩 e_p 사이에 InfoNCE 대비 손실을 적용한다. 이를 통해 프로브가 프롬프트 의도를 반영한 품질 추정을 학습한다. 전체 손실은 L = L_list + λ·L_align 형태이며, λ는 두 손실 간 균형을 조절한다.

실험에서는 Stable Diffusion 2, Stable Diffusion 3.5 (Medium, Large) 및 FLUX.1‑dev 등 다양한 백본에 동일 프로브를 적용하였다. COCO 캡션 데이터에서 5개의 시드를 생성하고, 20 % 타임스텝(≈ S·0.2)에서 프로브 점수를 계산한 뒤 상위 1개만 계속 진행했다. 결과는 Spearman 상관계수가 0.7 ~ 0.99 수준으로 매우 안정적이며, 전체 샘플링 비용은 0.36배(≈ 64 % 절감)로 감소한다. 특히 ImageReward 점수는 SD2 기준 0.49→1.59, HPSv2.1은 26.95→29.03으로 크게 향상되었다. 이는 조기 종료가 저품질 시드를 걸러내고, 남은 시드에 더 많은 연산을 할당함으로써 품질-효율 트레이드오프를 최적화함을 의미한다.

ProbeSelect는 모델 자체나 스케줄을 변경하지 않으며, 단순히 중간 활성화를 읽어들이는 경량 모듈이므로 다른 확산·플로우‑매칭 프레임워크에도 손쉽게 적용 가능하다. 또한, 선택적 샘플링 외에도 품질 기반 적응형 스톱핑, 가이드 스케일 조정 등 다양한 파생 응용이 가능하다.