UniGRPO: 추론 기반 시각적 생성을 위한 통합 정책 최적화

기존 이미지 생성 모델은 복잡한 프롬프트를 정확히 반영하지 못하는 한계가 있었다. 이 논문은 모델이 이미지를 그리기 전 스스로 '생각'하는 과정을 거치도록 강화학습으로 훈련시켜, 텍스트와 이미지 사이의 정렬도를 획기적으로 높이고 복잡한 지시사항도 정확하게 시각화하는 방법을 제시한다.

기존의 텍스트-이미지 생성 모델은 프롬프트를 입력받으면 즉시 픽셀을 생성하기 시작한다. 이는 사람이 복잡한 그림을 그릴 때 미리 구도를 잡거나 세부 요소를 계획하는 과정이 빠진 것과 같다. UniGRPO는 모델이 이미지를 생성하기 전에 '사고의 사슬(Chain-of-Thought)' 토큰을 먼저 생성하도록 유도하여, 프롬프트의 논리적 구조를 먼저 파악하게 한다.

이 과정은 강화학습 알고리즘인 GRPO를 통해 최적화된다. 모델이 여러 개의 추론 경로와 그에 따른 이미지 샘플들을 생성하면, 최종 이미지의 품질과 텍스트 정렬도를 평가하여 그룹 내에서 상대적으로 우수한 샘플에 더 높은 보상을 준다. 이 보상 신호는 이미지 생성 능력뿐만 아니라, 그 이미지를 만들기 위해 수행했던 앞선 '생각' 과정까지 소급하여 강화한다.

결과적으로 모델은 더 정확하고 풍부한 추론 텍스트를 작성하게 되며, 이 텍스트는 다시 이미지 생성 과정을 정교하게 가이드하는 선순환 구조를 만든다. 특히 이미지 생성 방식인 Flow Matching을 확률적 과정으로 변환하여 강화학습이 가능하도록 설계함으로써, 텍스트와 이미지라는 서로 다른 양식을 하나의 논리적 루프 안에서 성공적으로 결합했다.

전체 생성 과정을 단일 MDP(S, A, P, R)로 정의한다. 텍스트 단계에서는 어휘 사전에서 토큰을 선택하는 이산적 액션을, 이미지 단계에서는 노이즈를 제거하여 다음 단계의 잠재 변수를 결정하는 연속적 액션을 수행한다. 보상은 이미지가 완전히 생성된 마지막 단계에서만 주어지는 희소 터미널 보상을 사용한다.

텍스트 최적화에는 GRPO를 적용한다. 프롬프트 c에 대해 G개의 추론 체인 {y_i}를 샘플링하고, 각 샘플의 보상 Ri에서 그룹 평균을 뺀 뒤 표준편차로 나눈 값인 Advantage A_i를 계산한다. [Ri - mean({Rj}) / std({Rj}) → Advantage A_i 산출 → 정책 업데이트] 순으로 연산이 이루어지며, 이를 통해 별도의 가치 모델 없이도 효율적인 학습이 가능하다.

이미지 최적화는 FlowGRPO를 기반으로 한다. 결정론적 ODE를 확률적 SDE로 변환하여 탐색을 가능하게 하며, 학습 시에는 Classifier-Free Guidance(CFG)를 제거하여 선형적인 롤아웃을 유지한다. 또한 보상 해킹을 방지하기 위해 속도장 v_θ와 참조 모델 v_ref 사이의 차이를 제곱하여 합산하는 MSE 패널티를 적용한다. [||v_θ - v_ref||² 계산 → 패널티 부여 → 원래의 생성 능력 보존] 과정을 통해 안정적인 최적화를 달성한다.

TA(Text Alignment) 벤치마크 결과, UniGRPO는 0.8381점을 기록하여 기본 모델인 Bagel(0.6810) 및 SFT(0.7769) 대비 큰 폭의 성능 향상을 보였다. 이는 추론 과정이 이미지와 텍스트 사이의 정렬도를 높이는 데 결정적인 역할을 함을 시사한다.

GenEval 벤치마크에서는 0.90점을 달성하며 기존의 ReFL, FPO, FlowGRPO 등 다양한 RL 기반 방법론들을 능가하는 SOTA 성능을 입증했다. 특히 객체 수 세기(Counting)와 위치 관계(Position) 등 논리적 판단이 필요한 항목에서 추론 기반 생성의 강점이 두드러졌다.

Ablation Study를 통해 학습 시 CFG를 제거하더라도 최종 성능에 부정적인 영향이 없음을 확인했으며, 오히려 연산 효율성을 확보하여 다회차 생성으로의 확장 가능성을 열었다. 또한 속도 기반 MSE 패널티가 기존 Latent KL 패널티보다 보상 해킹을 억제하고 이미지 질감을 사실적으로 유지하는 데 효과적임이 증명됐다.

UniGRPO는 Bagel 아키텍처를 기반으로 하며, 텍스트 생성을 위한 Autoregressive 모듈과 이미지 생성을 위한 Flow Matching 모듈이 통합된 구조를 가진다. 학습 시에는 FlowGRPO-Fast 변형을 사용하여 특정 시간 윈도우 내에서만 SDE 샘플링과 그래디언트 트래킹을 수행함으로써 연산 비용을 최적화했다.

보상 해킹 방지를 위해 도입된 Ratio Normalization(RatioNorm)은 디퓨전 모델 학습 시 중요도 샘플링 비율이 한쪽으로 치우치는 문제를 해결한다. 로그 중요도 비율을 표준화하여 분포의 중심을 0으로 맞춤으로써 클리핑(Clipping) 메커니즘이 정상적으로 작동하도록 돕는다.

학습 과정에서 CFG를 제거한 것은 향후 다회차 인터리브드(Interleaved) 생성으로 확장하기 위한 핵심 설계다. CFG는 조건의 수에 따라 연산량이 기하급수적으로 늘어나지만, UniGRPO는 선형적인 롤아웃을 유지하여 복잡한 멀티턴 상호작용 환경에서도 안정적인 학습과 추론이 가능하다.

현재 UniGRPO는 최종 이미지 결과물에만 의존하는 희소 보상을 사용하므로, 중간 추론 단계의 논리적 오류를 직접 교정하는 신용 할당(Credit Assignment) 능력이 부족할 수 있다. 또한 단일 라운드 생성에 대해서만 검증되었으며, 실제 다회차 대화형 생성에서의 성능은 추가 연구가 필요하다.