기술 요약: STAR (SpatioTemporal Adaptive Reward Allocation)

문제 정의

Flow-GRPO와 같은 기존의 텍스트-이미지(T2I) 생성용 강화 학습(RL) 사후 학습 방법들은 일반적으로 최종 이미지 수준의 보상을 단일 스칼라 어드밴티지(advantage)로 변환합니다. 이 스칼라는 전체 생성 궤적(모든 디노이징 단계 및 모든 공간 잠재 토큰)에 걸쳐 균일한 강도로 적용됩니다.

본 논문은 이러한 접근 방식에서 근본적인 **입도 불일치(granularity mismatch)**를 식별합니다:

1. 시간적 구조(Temporal Structure): 서로 다른 디노이징 단계는 서로 다른 생성 단계(예: 전역 레이아웃 vs. 세부 사항 정제)를 담당합니다.
2. 공간적 구조(Spatial Structure): 텍스트 정렬(예: 특정 객체, 속성 또는 공간적 관계)을 결정하는 콘텐츠는 종종 이미지의 특정 영역에만 나타나는 반면, 배경과 같은 다른 영역은 프롬프트의 의미론적 요구 사항에 기여하는 바가 적습니다.

스칼라 보상이 균일하게 적용될 때, 정책 업데이트는 프롬프트를 충족하는 데 중요한 영역과 그렇지 않은 영역을 구분하는 데 어려움을 겪습니다. 이는 모델이 실제 보상을 유도하는 특정 생성 구성 요소에 학습을 집중하는 것을 어렵게 만들어, 구성적 일반화(compositional generalization), 객체 카운팅, 텍스트 렌더링 성능을 저하시키는 원인이 됩니다.

방법론: STAR

저자들은 확산(diffusion) 및 플로우(flow) 모델의 구조화된 디노이징 궤적에 걸쳐 이미지 수준의 보상을 재할당하도록 설계된 프레임워크인 **STAR(SpatioTemporal Adaptive Reward Allocation)**를 제안합니다. STAR는 외부 보상 소스를 수정하지 않고, 대신 생성 모델의 내부 메커니즘을 활용하여 작동합니다.

1. 어텐션 가이드 할당 (Attention-Guided Allocation)

STAR는 디노이징 단계($t$)와 롤아웃($i$)에 따라 변화하는 동적 공간 할당 맵을 구축합니다. 과정은 다음과 같습니다:

• 텍스트 단위 추출: 프롬프트 $c$를 핵심 텍스트 단위 $K(c)$(예: 주체, 스타일, 속성)의 집합으로 분해합니다.
• 어텐션 매핑: 모델은 텍스트-이미지 어텐션 레이어(특히 Stable Diffusion 3.5의 결합 어텐션 백본)를 사용하여 라우팅 신호를 추출합니다. 각 잠재 공간 위치 $u$와 텍스트 토큰 $j$에 대해 어텐션 가중치 $\alpha(u, j)$를 계산합니다.
• 히트맵 구축: 각 텍스트 단위에 대해, 대응하는 텍스트 토큰들에 대한 어텐션 가중치를 집계하여 공간 맵을 생성합니다. 이 맵들은 프롬프트의 주요 구성 요소와 관련하여 가장 활발한 잠재 영역을 강조하는 궤적별 히트맵 $H_t^i$로 통합됩니다.
• 정규화: 히트맵은 정규화 및 스케일링되어 계수 맵 $C_t^i(u)$를 생성합니다. 이 맵은 프롬프트 관련 영역에는 높은 가중치를, 배경 영역에는 낮지만 0이 아닌 가중치를 할당합니다.

2. 공간적으로 해상된 정책 목적 함수 (Spatially Resolved Policy Objective)

STAR는 스칼라 그룹 상대적 어드밴티지 $A_i$를 공간적으로 해상된 훈련 신호 $A_t^i(u) = A_i \cdot C_t^i(u)$로 변환합니다.

• 정책 목적 함수는 스칼라 궤적 확률이 아닌 공간적 로그 확률 기여도 상에서 작동하도록 수정됩니다.
• 각 잠재 위치에 대해 정규화된 공간 비율 $\rho$가 계산됩니다.
• 최종 목적 함수(PPO 스타일의 클리핑 기반)는 모든 타임스텝과 공간 위치에 대해 공간 어드밴티지 맵 $A_t^i(u)$에 의해 가중치가 부여된 클리핑된 비율들을 합산합니다. 이를 통해 정책 업데이트가 핵심 프롬프트 구성 요소에 해당하는 영역에서 더 강력하게 이루어지도록 보장합니다.

주요 기여

1. STAR 프레임워크: 이미지 수준의 보상을 시간 및 잠재 공간에 분포된 훈련 신호로 변환하여, 주요 객체, 속성 및 관계에 대응하는 생성 영역에 업데이트를 집중할 수 있게 하는 T2I 확산/플로우 RL 사후 학습 방법입니다.
2. 통합 인터페이스: 내부 텍스트-이미지 상호작용을 사용하여 타임스텝별 공간 할당 맵을 구축함으로써, 새로운 보상 모델 없이도 기존 RL 목적 함수를 공간적으로 해상된 정책 목적 함수로 확장합니다.
3. 실증적 검증: Stable Diffusion 3.5 Medium을 베이스 모델로 사용하여 세 가지 별도의 작업(GenEval, OCR, PickScore)에서 STAR가 구성적 의미론적 정렬, 텍스트 렌더링 및 선호도 최적화를 개선함을 입증했습니다.

실험 결과

저자들은 Stable Diffusion 3.5 Medium (SD3.5-M)을 대상으로 Flow-GRPO 베이스라인 및 다른 최첨단 모델들과 비교하여 STAR를 평가했습니다.

• GenEval (구성적 생성): STAR는 0.9759의 전체 점수를 달성하여, Flow-GRPO(0.95)와 기본 SD3.5-M(0.63)을 능가했습니다. 특히 카운팅(0.98), 색상(0.96), 속성 결합(0.94)에서 눈에 띄는 개선이 있었습니다.
• OCR 텍스트 렌더링: STAR는 0.9757의 정확도를 달성하여, Flow-GRPO 베이스라인(0.92)을 넘어섰습니다.
• PickScore (선호도 최적화): STAR는 23.60의 점수를 기록하여, Flow-GRPO의 23.31과 비교되었습니다.
• 효율성: 이 방법은 계산 오버헤드를 최소화하며, 포워드 패스(forward pass)의 어텐션 맵을 재사용함으로써 훈련 시간을 2% 미만으로 추가하고 피크 메모리 사용량을 단 **1.56%**만 증가시킵니다.

의의 및 주장

본 논문은 STAR가 확산 생성의 시공간적 특성을 인정함으로써 T2I RL의 "신용 할당(credit assignment)" 문제를 해결한다고 주장합니다. STAR는 동일한 스칼라 보상 신호를 텍스트 정렬을 결정하는 특정 잠재 영역으로 라우팅함으로써, 더욱 선택적이고 효과적인 정책 업데이트를 가능하게 합니다.

저자들은 STAR가 외부 보상 소스를 변경하지 않고 혹은 추가적인 보상 모델을 요구하지 않고도 이러한 이득을 달성한다는 점을 강조합니다. 이는 이미지 수준의 피드백을 구조화된 궤적으로 할당하기 위한 "단순하고, 가볍고, 효과적인 인터페이스"를 제공합니다. 결과는 시공간적 적응형 할당이 생성 모델의 미세한 의미 이해, 로컬 콘텐츠 제어 및 선호도 정렬을 최적화하는 데 중요한 요소임을 시사합니다.