Test-Time Alignment of Text-to-Image Diffusion Models via Null-Text Embedding Optimisation

기술 요약: 텍스트-이미지 확산 모델의 Null-Text 테스트 시점 정렬 (Null-Text Test-Time Alignment)

문제 정의
사전 학습된 텍스트-이미지 확산 모델은 종종 정렬되지 않았거나 바람직하지 않은 콘텐츠가 포함된 대규모 미정제 웹 데이터셋으로 학습됩니다. 이러한 모델을 추론 단계에서 특정 인간의 가치나 목표 목적(예: 미적 품질, 특정 선호도)에 맞게 정렬하는 것은 신뢰할 수 있는 배포를 위해 매우 중요합니다. 기존의 정렬 방식은 크게 두 가지 범주로 나뉘며, 두 방식 모두 상당한 한계를 가집니다:

1. 미세 조정(Fine-tuning) 방법: 보상 함수를 최적화하기 위해 모델 파라미터를 업데이트합니다. 효과적이긴 하지만, 계산 비용이 많이 들고 보상 과최적화(reward over-optimization), 즉 모델이 프록시 보상에 과적합되어 다른 보상에 대한 일반화 능력과 출력 다양성을 상실하는 보상 해킹(reward hacking) 현상이 발생하기 쉽습니다.
2. 테스트 시점 정렬(Test-Time Alignment, TTA) 방법: 모델 가중치를 업데이트하지 않고 잠재 변수(latent)나 노이즈 변수를 최적화하거나, 보상 조건부 사후 분포에서 샘플링하거나, 이산 탐색 알고리즘을 사용합니다. 그러나 이러한 방법들은 잠재 및 노이즈 공간의 고도로 비구조적인 특성으로 인해 보상 과소 최적화 또는 과최적화 문제를 겪는 경우가 많습니다. 이러한 공간에서의 최적화는 모델이 시각적 품질, 의미적 일관성, 그리고 다른 지표에서의 성능을 저하시키면서 타겟 보상을 극대화하기 위해 비의미적인 노이즈 패턴을 악용하는 "보상 해킹"으로 이어질 수 있습니다.

방법론: Null-Text 테스트 시점 정렬 (Null-TTA)
저자들은 레이턴트나 노이즈 변수를 조작하는 대신, Classifier-Free Guidance (CFG) 내의 무조건부(null-text) 임베딩을 최적화함으로써 확산 모델을 정렬하는 훈련이 필요 없는(training-free) 프레임워크인 Null-TTA를 제안합니다.

• 핵심 통찰: CFG에서 무조건부 임베딩($\phi$)은 모델의 생성 분포를 위한 기하학적 앵커 역할을 합니다. 구조화되지 않은 노이즈 공간과 달리, 텍스트 임베딩 공간은 구조화된 의미론적 성격을 가집니다. 이 공간 내에서 최적화를 수행하면 정렬이 의미적으로 일관된 매니폴드(manifold) 상에서 이루어짐을 보장하며, 모델이 비의미적인 노이즈 아티팩트를 악용하는 것을 방지합니다.
• 목적 함수: 이 방법은 모델의 생성 분포를 타겟 보상 쪽으로 유도하는 동시에 사전 학습된 모델로부터의 분포 드리프트(distributional drift)를 억제하는 정규화된 목적 함수를 공식화합니다. 이 목적 함수는 기대 보상 $R(x_0)$를 최대화하면서, 최적화된 임베딩 $\phi'$에 의해 유도된 분포와 원래의 $\phi$에 의해 유도된 분포 사이의 KL 발산을 최소화합니다:
  $$\max_{\phi'} \left( \lambda_1 \mathbb{E}_{p(x_0|\phi')}[R(x_0)] - \lambda_2 \text{KL}(p(x_{0:T}, \phi') || p(x_{0:T}, \phi)) \right)$$
  여기서 KL 항은 연속적인 디노이징 단계 간의 국소적 발산과 임베딩 분포 간의 발산으로 분해됩니다. 이는 예측된 노이즈($\tilde{\epsilon}$)의 차이와 임베딩의 차이($\phi' - \phi$)를 포함하는 폐쇄형(closed-form) 표현식으로 단순화됩니다.
• 최적화 전략:
  • 적응형 정규화(Adaptive Regularization): 정규화 가중치 $\lambda_2$는 디노이징 과정이 진행됨에 따라($t \to 0$) 적응적으로 어닐링(annealing)됩니다. 초기 단계에서는 강한 정규화를 통해 최적화를 안정화하고, 후기 단계에서는 더 약한 정규화를 통해 미세한 정렬을 허용합니다.
  • 탐욕적 탐색(Greedy Search): 각 디노이징 단계에서 이 방법은 경량화된 탐욕적 탐색을 사용합니다. 전이 커널(transition kernel)로부터 $K$개의 후보 레이턴트를 샘플링하고, 트위디(Tweedie) 공식을 통해 깨끗한 샘플을 추정하며, 보상 모델로 점수를 매긴 후, 다음 단계를 위해 가장 높은 점수를 가진 후보를 결정론적으로 선택합니다.
  • 비미분적 보상(Non-Differentiable Rewards): 보상 함수가 블랙박스(비미분)인 시나리오의 경우, Null-text 임베딩의 무작위 섭동을 사용하는 영차(zeroth-order) 그래디언트 추정법을 사용합니다.

주요 기여

1. Null-TTA 프레임워크: 텍스트 컨디셔닝의 구조화된 의미 공간에서 정렬을 수행하여, 구조화되지 않은 레이턴트/노이즈 최적화의 함정을 피하는 새로운 훈련 불필요(training-free) TTA 프레임워크를 제시합니다.
2. 원칙적인 목적 함수: 모델의 생성 분포 자체를 타겟 보상 쪽으로 재정향시키는 KL 정규화된 목적 함수를 통해, 파라미터 업데이트 없이 안정적이고 효율적인 정렬을 보장합니다. 이는 보상 해킹을 방지하고 의미적 일관성을 유지합니다.
3. 최첨단 성능: Null-TTA가 타겟 보상 성능을 높이면서도 다른 보상에 대한 교차 보상 일반화 및 지각적 품질을 효과적으로 유지함을 입증하는 광범위한 실험 결과를 보여줍니다.

실험 결과
저자들은 다양한 타겟 보상(HPSv2, PickScore, Aesthetic Score, ImageReward)과 비미분적 목적 함수(JPEG 압축률)에 대해 Stable Diffusion v1.5 및 SDXL을 사용하여 Null-TTA를 평가했습니다.

• 보상 최적화 vs. 일반화: Null-TTA는 베이스라인(DNO, DAS, DSearch, TDPO 포함)과 비교하여 파레토 프런트(Pareto front)를 지속적으로 바깥쪽으로 확장시킵니다. 이는 타겟 보상을 더 높게 달성하면서도 경쟁 모델들보다 홀드아웃(held-out) 보상의 저하를 훨씬 적게 일으켰음을 의미하며, 성공적인 정렬과 과최적화 방지를 나타냅니다.
• 질적 성능: 까다로운 카테고리(개수 세기, 구성성, 공간적 추론, 특이한 색상, 스타일 전이, 불가능한 장면)에서 Null-TTA는 프롬프트 제약을 더 충실히 만족하면서도 전역적 일관성을 유지하며, 잘못된 객체 수나 엉킨 관계와 같은 실패 모드를 피하는 이미지를 생성합니다.
• 사용자 스터디: 16명의 참가자가 참여한 블라인드 사용자 스터디 결과, Null-TTA가 베이스라인보다 통계적으로 유의미하게 선호되는 것으로 나타났으며, 이는 보상 이득이 지각적 품질로 이어진다는 것을 확인시켜 줍니다.
• 계산 효율성: Null-TTA는 더 낮은 GPU 메모리 사용량과 더 빠르거나 대등한 벽시계 시간(wall-clock time)으로 더 강력한 타겟 보상 성능을 달성합니다. 이러한 효율성은 전체 U-Net을 업데이트하는 대신 크로스 어텐션 레이어를 통해서만 역전파(backpropagate)하여 임베딩을 업데이트하기 때문입니다.
• 비미분적 보상: 이 방법은 영차 그래디언트를 사용하여 JPEG 파일 크기(블랙박스 목적 함수)를 성공적으로 최적화하였으며, 시각적 품질을 잘 보존하면서도 그래디언트 기반 베이스라인보다 경쟁력 있는 압축률을 달성했습니다.

의의 및 주장
본 논문은 **의미 공간 최적화(semantic-space optimization)**가 테스트 시점 정렬(TTA)을 위한 효과적이고 원칙적인 새로운 패러다임임을 입증합니다. 텍스트 임베딩 공간의 구조화된 특성을 활용함으로써, Null-TTA는 보상 해킹을 방지하고 모델의 사전 학습된 지식을 유지하는 내재적인 정규화를 제공합니다. 저자들은 이 접근 방식이 기존 TTA 방법들에 대한 단순하면서도 강력한 대안을 제공하며, 미세 조정의 계산 비용 없이 확산 모델의 신뢰할 수 있고 효율적이며 일반화 가능한 추론 시점 정렬을 가능하게 한다고 주장합니다.