MAR-GRPO: Stabilized GRPO for AR-diffusion Hybrid Image Generation

1. 문제 정의 (Problem)

최근 텍스트 조건부 이미지 생성 분야에서 마스크 자동회귀 (Masked Autoregressive, MAR) 모델은 이산적 토큰의 양자화 한계를 극복하고 연속적 잠재 공간 (continuous latent space) 을 활용하여 뛰어난 성능을 보여주고 있습니다. MAR 모델은 AR 트랜스포머가 잠재 특징을 예측하고, 경량화된 확산 (Diffusion) 헤드가 이를 고해상도 이미지로 디코딩하는 하이브리드 구조를 가집니다.

그러나 이러한 MAR 모델에 강화 학습 (Reinforcement Learning, RL), 특히 GRPO (Group Relative Policy Optimization) 를 적용하는 것은 다음과 같은 근본적인 어려움에 직면해 있습니다:

• 불안정한 최적화 동역학: AR 모델과 확산 헤드가 결합된 복잡한 추론 과정에서, 확산 헤드의 확률적 (stochastic) 데노이징 과정이 AR 모델로 전달되는 기울기 (gradient) 에 큰 노이즈를 유발합니다.
• 기울기 분산 증폭: AR 잠재 공간과 이미지 공간 간의 매핑이 확산 헤드의 업데이트에 따라 계속 변하므로 (Non-stationary), 보상 신호가 불안정해지고 조기 수렴 (early saturation) 이나 성능 붕괴가 발생합니다.
• 기존 방법의 한계: 단순히 엔드 - 투 - 엔드 (End-to-End) 방식으로 GRPO 를 적용하면 구조적 불안정성과 세부 사항의 손실, 다양성 붕괴 (diversity collapse) 가 관찰됩니다.

2. 방법론 (Methodology)

저자들은 MAR 모델의 훈련 안정성을 확보하고 성능을 극대화하기 위해 MAR-GRPO 프레임워크를 제안합니다. 이 프레임워크는 확산 헤드의 확률적 특성을 보완하기 위한 세 가지 핵심 전략을 포함합니다.

가. 고정된 확산 헤드를 통한 안정화 (Stabilized Training Dynamics)

• 분석: 엔드 - 투 - 엔드 학습 시 확산 헤드의 파라미터 업데이트가 AR 모델의 최적화 목표를 지속적으로 변경시켜 불안정을 초래함을 발견했습니다.
• 해결: 확산 헤드를 고정 (Freeze) 하고 AR 트랜스포머만 최적화하는 방식을 채택했습니다. 이는 AR 모델에게 일관된 최적화 목표를 제공하여 훈련 동역학을 크게 안정화시킵니다.

나. 다중 궤적 기대값 (Multi-Trajectory Expectation, MTE)

• 문제: 확산 헤드를 고정하더라도, 단일 확산 궤적 (trajectory) 을 사용할 경우 확산 과정의 무작위성 (noise realization) 으로 인해 정책 로그 확률 추정의 분산이 여전히 높습니다.
• 해결: 동일한 AR 잠재 벡터에 대해 여러 개의 확산 궤적 (Multi-seed Diffusion Sampling) 을 생성하고, 이들에 대한 기대값 (Expectation) 을 통해 최적화 방향을 추정합니다.
  • 이는 확산으로 인한 기울기 노이즈를 평균화하여 줄여주며, 더 안정적이고 신뢰할 수 있는 정책 업데이트를 가능하게 합니다.

다. 불확실성 기반 토큰 선택 및 일관성 필터링

• 효율성 최적화: 모든 토큰에 대해 다중 궤적 계산을 수행하면 계산 비용이 증가하고 과도한 평활화 (over-smoothing) 로 인해 중요한 기울기 정보가 약화될 수 있습니다.
  • Top-k% 불확실성 선택: 여러 확산 궤적 간의 분산을 기반으로 토큰별 불확실성 (Uncertainty) 지도를 생성합니다. 구조가 모호하거나 미세한 세부 사항이 필요한 불확실성이 높은 상위 k% 토큰에만 다중 궤적 계산을 적용하고, 나머지는 단일 궤적 계산을 유지하여 효율성을 높입니다.
• 일관성 인식 토큰 선택 (Consistency-Aware Token Selection): AR 모델이 생성한 토큰 중 최종 이미지 생성에 부정적인 영향을 미치는 노이즈가 섞인 토큰을 필터링합니다.
  • 각 AR 단계에서 예측된 잠재 벡터와 최종 생성된 잠재 벡터 간의 코사인 유사도 변화를 측정하여, 일관되게 개선되는 토큰만 최적화에 포함시킵니다.

3. 주요 기여 (Key Contributions)

1. MAR 를 위한 최초의 GRPO 프레임워크: AR-확산 하이브리드 아키텍처에 강화 학습을 적용하는 첫 번째 체계적인 방법론을 제시했습니다.
2. 확산 헤드의 역할 규명: 하이브리드 구조에서 확산 헤드가 최적화 불안정성의 주요 원인임을 분석하고, 이를 해결하기 위한 다중 궤적 기대값 (MTE) 기법을 도입하여 확산으로 인한 기울기 분산을 효과적으로 줄였습니다.
3. 성능 및 안정성 동시 달성: 기존 GRPO 베이스라인 및 사전 RL 모델 대비 인간 선호도 (Human Preference), 프롬프트 충실도, 공간 구조 이해도 등 다양한 벤치마크에서 일관된 성능 향상을 입증했습니다.

4. 실험 결과 (Results)

• 훈련 안정성: 기존 GRPO 는 훈련 중 HPS (Human Preference Score) 가 급격히 떨어지는 '붕괴 (Collapse)' 현상을 보인 반면, MAR-GRPO 는 훈련 내내 안정적인 성능 향상을 보였습니다.
• 정량적 지표:
  • HPS, ImageReward, Aesthetic Score: 기존 GRPO 및 베이스라인 모델 (NOVA, Harmon 등) 대비 유의미한 점수 상승을 기록했습니다.
  • T2I-CompBench (공간 및 구성적 추론): 객체의 위치, 수, 배치 등 공간적 관계 이해 능력이 크게 향상되었습니다.
• 정성적 평가:
  • 구조적 일관성: 객체 간의 배치와 형태가 왜곡되지 않고 정확하게 생성됩니다.
  • 세부 사항: 텍스트 프롬프트에 대한 세부 사항 (예: 안경 쓴 캐릭터, 특정 개수의 객체) 이 더 정교하게 표현됩니다.
  • 다양성: 기존 방법에서 발생하던 다양성 붕괴가 개선되었습니다.

5. 의의 및 결론 (Significance)

이 연구는 AR-확산 하이브리드 모델에 강화 학습을 적용할 때 발생하는 확산 헤드의 확률적 노이즈 문제를 해결하는 중요한 통찰을 제공합니다. 확산 헤드를 고정하고, 불확실성이 높은 영역에 대해만 다중 궤적 계산을 적용하는 경량화된 노이즈 보정 전략은 추가적인 계산 오버헤드를 거의 증가시키지 않으면서 훈련 안정성과 최종 생성 품질을 동시에 극대화합니다. 이는 텍스트 - 이미지 생성뿐만 아니라 비디오 생성 등 다른 멀티모달 AR-확산 모델의 RL 최적화 연구에도 중요한 기반을 마련합니다.