기술 요약: Diffusion Transformer를 위한 품질 인식 변조 (Quality-Aware Modulation)

문제 정의

현대적인 텍스트-이미지 확산 모델, 특히 Stable Diffusion 3.5 (SD3.5)와 같은 Diffusion Transformer (DiT)는 적응형 레이어 정규화(Adaptive Layer Normalization, AdaLN)를 통해 디노이징 과정을 조절하기 위해 조건 신호(특히 타임스텝 임베딩과 풀링된 프롬프트 임베딩)에 의존합니다. 이러한 신호들은 현재의 노이즈 수준과 의미적 의도를 효과적으로 전달하지만, **품질 인식 정보(quality-aware information)**는 결여되어 있습니다. 결과적으로, 변조 메커니즘은 잠재 이미지(latent image)가 진화함에 따라 발생하는 오류 누적, 잘못된 프롬프트 해석, 또는 환각 현상(hallucinated features)과 같은 샘플 특유의 실패 모드에 직접적으로 반응할 수 없습니다. 베이스라인 조건화는 현재의 잠재 상태와 독립적($\partial M_{base} / \partial x_t = 0$)이므로, 중간 단계의 실제 품질을 기반으로 피델리티(fidelity) 문제를 수정하는 것이 불가능합니다.

방법론

저자들은 사전 학습된 백본을 수정하지 않고 디퓨전 과정에 품질 인식 신호를 주입하기 위해 설계된 경량 트랜스포머 기반 어댑터인 **품질 표현 모듈(Quality Representation-Module, QRM)**을 제안합니다.

아키텍처

QRM은 베이스라인 AdaLN 파라미터($M_{base}$)에 가산적인 변조 오프셋($M_{qrm}$)을 예측하는 트랜스포머 디코더로 구현됩니다.

• 입력: 모듈은 현재 잠재 이미지($x_t$), 풀링된 프롬프트 임베딩($p$), 타임스텝 임베딩($t$), 그리고 베이스라인 변조 파라미터($M_{base}$)를 받습니다.
• 프로세싱:
  1. 잠재 이미지는 학습된 $1\times1$ 컨볼루션과 적응형 풀링을 통해 잠재 토큰으로 투영됩니다.
  2. 베이스라인 변조 파라미터는 정규화된 후 쿼리 토큰($Z_0$)을 형성하기 위해 투영됩니다.
  3. 잠재 토큰, 프롬프트, 타임스텝 임베딩으로부터 키-값(Key-Value, KV) 메모리가 구축됩니다.
  4. 크로스 어텐션(cross-attention) 및 셀프 어텐션(self-attention) 레이어를 통해 쿼리 토큰은 KV 메모리에 어텐션을 수행하여 정제된 표현을 생성합니다.
  5. 출력은 네이티브 AdaLN 너비로 다시 투영되어 오프셋 벡터 $M_{qrm}$을 생성합니다.
• 적용: 트랜스포머 블록에 적용되는 최종 변조는 $M_t = M_{base,t} + M_{qrm,t}$입니다.

학습 전략 (ReFL)

QRM은 보상 피드백 학습(Reward Feedback Learning, ReFL) 전략을 사용하여 학습됩니다:

• 고정된 구성 요소: 디퓨전 백본(디노이저, 텍스트 인코더, VAE), 샘플러, 그리고 보상 모델은 고정된 상태로 유지됩니다. 오직 QRM 파라미터만 업데이트됩니다.
• 목적 함수: 모듈은 중간 디노이징 단계에서 디코딩된 이미지에 적용된 고정된 보상 모델(예: HPSv2.1 또는 CLIPScore)로부터 유도된 스칼라 보상 점수를 최대화하도록 학습됩니다.
• 손실 함수: 논문은 보상 최대화 손실 또는 마진 추구(margin-seeking) 손 loss를 사용하며, 이는 QRM이 변조된 이미지가 베이스라인보다 특정 마진만큼 더 우수하도록 유도합니다.
• 타임스텝 선택: 효율성을 높이고 구조적 형성에 집중하기 위해, QRM은 전체 디퓨전 스케줄이 아닌 정의된 하위 집합의 고노이즈 타임스텝(예: $t \in \{50, \dots, 26\}$) 동안에만 활성화됩니다.

추론

추론 시, 학습된 QRM 모듈은 플러그인 어댑터로서 작동합니다. 이는 현재 잠재 상태를 기반으로 변조 오프셋을 예측하고 이를 베이스라인 파라미터에 더합니다. 이를 통해 모델은 백본을 재학습하거나 샘플링 스케줄을 변경하지 않고도 디노이징 궤적을 더 높은 품질의 상태로 유도할 수 있습니다.

주요 기여

1. QRM 아키텍처: 기존 모델의 입력(잠재 상태, 프롬프트, 타임스텝)으로부터 품질 인식 표현을 학습하여 변조 파라미터를 정제하는 경량 트랜스포머 모듈을 도입했습니다.
2. 비침습적 통합: 기존의 디퓨전 트랜스포머 백본, VAE, 또는 샘플링 스케줄을 수정하지 않고 잠재 이미지 품질 신호로 변조를 증강하는 메커니즘을 제시했습니다.
3. 보상 가이드 변조: 예측된 변조 업데이트($M_{qrm}$)가 각 타임스텝에서 디노이징 과정을 효과적으로 정제하여 프롬프트 준수 능력과 이미지 피델리티를 개선할 수 있음을 입증했습니다.
4. 포괄적 평가: QRM 아키텍처 크기, 활성 타임스텝 범위, 학습 보상 지표가 이미지 품질에 미치는 영향을 분석하는 광범위한 절제 연구(ablation studies)를 수행했습니다.

실험 결과

저자들은 Parti-Prompts 벤치마크에서 **Stable Diffusion 3.5 (SD3.5)**를 사용하여 QRM을 평가했습니다.

• 정량적 성능:
  • SD3.5 베이스라인과 비교했을 때, HPSv2.1을 보상 신호로 사용하여 학습된 QRM이 평가된 모든 방법(프롬프트 정제, 파인튜닝, ReNeg 방식 포함) 중 가장 높은 점수를 달성했습니다.
  • 개선 사항: Aesthetic 점수에서 +0.20, HPSv2.1 점수에서 +0.82를 기록했습니다.
  • CLIPScore로 학습했을 경우, QML은 인간 선호도 정렬(HPSv2.1)을 개선하면서도 대등한 지각적 품질을 유지했습니다.
• 정성적 분석:
  • 디노이징 궤적 시각화를 통해, 초기 고노이즈 타임스텝에 QRM을 적용하면 베이스라인보다 의미적 구조가 훨씬 더 일찍 나타남을 확인했습니다.
  • 중간 또는 후기 단계에 QRM을 적용했을 때는 개선 효과가 작았으며, 이는 전역적 구조가 디퓨전 과정의 초기에 확립된다는 것을 뒷받침합니다.
• 절제 연구 결과:
  • 용량(Capacity): 더 큰 QRM 모델은 특정 학습 보상(예: HPSv2.1)에 더 강력하게 최적화되었으나, 다른 지표(예: CLIPScore)에서는 약간의 트레이드오프를 보였습니다.
  • 타임스텝 범위: 초기 타임스텝($50 \to 26$)에서의 변조가 후기 단계에서의 변조보다 정렬(alignment) 측면에서 일관되게 더 효과적이었습니다.
  • 보상 지표: 보상 모델(CLIPScore vs. HPSv2.1)의 선택은 QRM이 학습하는 지각적 편향을 직접적으로 결정했습니다.

의의 및 주장

본 논문은 **변조 레벨의 개입(modulation-level interventions)**이 거대 디퓨전 모델을 인간 선호 신호에 정렬하기 위한 확장 가능하고 효과적인 경로를 제공한다고 주장합니다. 전체 백본의 값비싼 파인튜닝(예: LoRA 또는 전체 파인튜닝)이 필요한 방법들과 달리, QRM은 진화하는 잠재 표현에 조건화된 가산적 변조 업데이트를 통해 일관된 품질 향상을 달 Achieves 합니다.

저자들은 QRM을 다양한 보상 신호에 의해 제어될 수 있고 다양한 디퓨전 트랜스포머 아키텍처에 적용될 수 있는 유연한 품질 인식 어댑터로 정의합니다. 또한 QRM이 학습 보상 모델의 편향을 상속한다는 점을 인정하며, 향-후 연구에서 견고성과 일반화 능력을 더욱 향상시키기 위해 멀티 보상 학습 및 적응형 타임스텝 선택을 탐구할 수 있음을 시사합니다. 이 연구는 사전 학습된 가중치를 변경하지 않고도 최첨단 모델의 이미지 피델리티와 프롬프트 정렬에서 상당한 이득을 얻을 수 있음을 보여줍니다.