PG-MAP: Joint MAP Optimization for Inference-Time Alignment of Diffusion and Flow-Matching Models

기술 요약: PG-MAP: 확산(Diffusion) 및 플로우 매칭(Flow-Matching) 모델의 추론 시 정렬을 위한 결합 MAP 최적화

문제 정의
사전 학습된 텍스트-이미지 모델(Diffusion 및 Flow-Matching)을 위한 현재의 추론 시 정렬 방법들은 일반적으로 단일 제어 축을 따라 작동합니다. 컨디셔닝 측면의 방법들(예: 어텐션 에디팅, 프롬프트 임베딩 학습)은 텍스트 임베딩 $c$를 수정하며, 잠재 변수 측면의 방법들(예: 보상 유도 잠재 변수 섭동)은 잠재 상태 $z_t$를 수정합니다. 기존 프레임드워크들은 컨디셔닝과 잠재 변수 사이의 결합 의존성을 모델링하지 못하고, 업데이트를 협응된 과정이 아닌 가산적인 과정으로 취급합니다. 또한, 이러한 방법들은 고정된 전략을 궤적의 시작점($z_T$)이나 오프라인 단계에서 설정함으로써, 고노이즈(전역 레이아웃)에서 저노이즈(지역적 세부 사항) 단계로 진화하는 과정에서 디노이저가 필요로 하는 동적인 정보 요구를 무시하는 정적인 특성을 가집니다. 이러한 한계는 서로 다른 생성 트랜스포트(Diffusion vs. Flow-Matching) 간의 정렬 기술 전이를 저해합니다.

방법론: PG-MAP
저자들은 PG-MAP(Preference-Guided Adaptive MAP)을 제안합니다. 이는 추론 시 정렬을 궤적 수준의 Gibbs-MAP / proximal energy 최적화로 공식화한 학습 불필요(training-free) 프레임드워크입니다. 핵심 혁신은 매 디노이징 단계마다 컨디셔닝 $c$와 잠재 상태 $z_t$를 공동으로 최적화하는 **순방향 일관성 결합(forward-consistency coupling)**입니다.

• 목적 함수: 이 프레임드워크는 세 가지 구성 요소를 포함하는 단계별 에너지 함수 $J_t(c, z_t)$를 정의합니다:

  1. 순방향 일관성 잔차(Forward-Consistency Residual): 현재의 잠재 변수 $z_t$와 모델이 예측한 이전 단계의 상태 사이의 편차를 처벌하여, 궤적이 디노이저의 순방향 커널과 일관성을 유지하도록 하는 항 $\ell_t(c, z_t)$입니다.
  2. 가우시안 앵커링 사전 확률(Gaussian Anchoring Priors): 최적화된 변수들을 원래의 값($c_0$ 및 DDIM 궤적 지점 $z^{ddim}_t$)에 고정하는 항 $R_c(c)$와 $R_z(z_t)$입니다. 결정적으로, 잠재 사전 확률 스케일 $\sigma_z(t)$는 확산 과정의 한계 노이즈 스케일($\sigma_z(t) = \gamma\sqrt{1-\bar{\alpha}_t}$)을 추적하도록 **스케줄 적응형(schedule-adaptive)**으로 설계되어, 궤적이 데이터 매니폴드에 접근함에 따라 조여지는 스케일 불변 신뢰 영역을 제공합니다.
  3. 선호도 보상 틸트(Preference Reward Tilt): 모델의 단계별 깨끗한 이미지 프리뷰를 기반으로 인간의 선호도(예: PickScore, HPS)를 향하도록 유도하는 항 $\lambda Q(\hat{x}_0, y)$입니다.

• 적응형 활성 집합(Adaptive Active Set): 프레임드워크는 단계별 활성 집합 $A_t \subseteq \{c, z_t\}$와 보상 게이트 $\lambda_t$를 사용합니다.

  • 확산 모델 (SD 1.5, SDXL): 프레임드워크는 $c$와 $z_t$를 공동으로 정제하며, 주로 고노이즈 구간에서 작동하여 궤적이 전개됨에 따라 컨디셔닝과 잠재 변수가 서로 정보를 주고받을 수 있도록 합니다.
  • 플로우 매칭 모델 (SD3.5): ODE의 결정론적 특성과 결합된 텍스트 인코더의 높은 용량으로 인해, 프레임드워크는 **잠재 변수 전용 변형(UG-FM)**으로 축소됩니다. 여기서 활성 집합은 데이터 측 윈도우(저노이즈)에 제한된 $\{z_t\}$로 붕괴되며, 이는 컨디셔닝 업데이트가 효과적이지 않고 노이즈 측 섭동이 ODE 역학에 의해 과도하게 증폭되기 때문입니다.

• 최적화: 정제된 쌍 $(c^\star_t, z^\star_t)$는 에너지 함수에서 유도된 그래디언트를 사용하여 적은 수의 공동 상승 단계(통상적으로 $K=2$)를 통해 얻어집니다. 이 과정은 CFG(Classifier-Free Guidance)와 같은 기존 방법들과 결합 가능합니다. CFG가 벡터 필드를 수정한다면, PG-MAP은 고정된 벡터 필드 하에서 쿼리 포인트를 이동시킵니다.

주요 기여

1. 결합 (c, $z_t$) MAP 프레임워크: 디노이저의 순방향 커널을 통해 컨디셔닝과 잠재 변수를 결합하여, 구성 실패(c-side)와 질감 실패(z-side)를 동시에 해결하는 최초의 추론 시 프레임드워크입니다.
2. 통합된 목적 함수: 목적 함수 $J_t$는 기존의 단일 축 방법들(MAP-c, Reward-z)을 특수 사례로서 분석적으로 회복하며, CFG는 특수 사례가 아닌 결합 가능한 외부 메커니즘으로서 유지됩니다.
3. 스케줄 적응형 최적화: 노이즈 스케줄에 적응하여 고노이즈에서는 제약을 완화하고 데이터 종단부 근처에서는 제약을 조이는 시간 의존적 신뢰 영역과 활성 집합을 도입합니다.
4. 트랜스포트 의존적 분석: 최적의 활성 집합이 생성 트랜스포트에 따라 달라짐을 입증합니다. 확산 모델은 고노이즈에서 공동 정제의 이점을 얻는 반면, 플ow-Matching은 데이터 엔드에서 잠재 변수 전용 정제의 이점을 얻습니다. 이는 국소 섭동 분석과 경험적 검증을 통해 뒷받-되었습니다.

결과

• 확산 백본 (SD 1.5, SDXL): PG-MAP은 정렬 지표를 일관되게 개선합니다. SDXL에서, 정적 베이스라인 대비 56.4%의 PickScore 승률과 56.2%의 Aesthetic 승률을 달연했습니다. PG-MAP을 튜닝된 CFG와 결합하면 가장 강력한 성능을 보이며, 64.6% HPS 및 56.5% Aesthetic에 도달합니다.
• 플로우 매칭 백본 (SD3.5-medium): 잠재 변수 전용 변형(UG-FM)은 정적 베이스라인 대비 91.9% PickScore 및 75.7% HPS 승률을 달성합니다. 통제된 실험을 통해 노이즈 관련 아티팩트를 배제함으로써, 이러한 이득이 구조화된 저주파 잠재 섭동으로부터 기인함을 확인했습니다.
• 인간 평가: SDXL에 대해, PG-MAP은 쌍체 비교 인간 평가에서 정적 베이스라인, 튜닝된 CFG, 그리고 연산량이 동일한 Universal Guidance보다 선호되었습니다 (60–67% 승률).
• 오라클 라우팅(Oracle Routing): 오라클 분석 결과, 서로 다른 프롬프트 유형(예: 속성 결합 vs. 분위기 묘사 장면)은 서로 다른 최적화 변형의 이점을 얻습니다. 프롬프트별 선택기(selector)를 사용하면 지표를 추가로 5–14 퍼센트 포인트 향상시킬 수 있으며, 이는 향후 라우팅 메커니즘의 상당한 발전 가능성을 시사합니다.

의의 및 주장
본 논문은 PG-MAP이 정적인 단일 축 가이드 휴리스틱에서 동적이고 궤적을 인식하는 최적화로의 패러다임 전환을 나타낸다고 주장합니다. 순방향 일관성 결합을 갖춘 근사 에너지 최적화로 정렬을 공식화함으로써, 이 프레임워크는 확산 및 플로우 매칭 트랜스포트 모두와 호환되는 모달리티 간의 협응된 업데이트를 가능하게 합니다. 저자들은 "핵심적인(load-bearing)" 요소가 단순히 더 큰 스텝 사이즈나 더 강한 보상 신호가 아니라, 결합된(joint) 결합과 비정상적(non-stationary) 스케줄링임을 강조합니다. 이 연구는 추론 시 정렬을 위한 유연하고 학습이 필요 없는 접근법을 제공하며, 컨디셔닝과 잠재 변수가 생성 과정 중에 어떻게 상호작용하는지에 대한 통합된 이론적 및 경험적 토대를 제공합니다.