RankE: 디스크리트 텍스트-투-이미지 생성을 위한 엔드-투-엔드 포스트 트레이닝 및 디코더 공진화

디스크리트 AR T2I에서 VQ decoder를 고정하면 Latent Covariate Shift가 발생한다. RankE는 정책과 디코더를 교대 업데이트하는 엔드-투-엔드 포스트 트레이닝으로 이 시프트를 흡수하고 픽셀 공간의 품질과 텍스트-이미지 정합성을 동시에 향상한다. LlamaGen-XL과 Janus-Pro에서 FID와 CLIP 모두 개선되며 GenEval에서도 강건한 성능이 확인된다.

단계적 배경: 디스크리트 AR T2I는 토큰화를 담당하는 VQ 토크나이저와 이를 예측하는 AR 정책으로 구성된다. 기존 포스트-트레이닝은 디코더를 고정한 채 정책만 업데이트하는 방식으로, RL이 진행될수록 토큰 분포가 디코더의 학습 분포와 멀어져 fidelity와 정합성 간의 무게추가 벌어진다. RankE는 두 모듈의 파라미터를 교대로 업데이트하는 이중 최적화로 이 간극을 오차 신호를 통해 메우고, 토큰 분포의 drift가 픽셀 공간의 재구성으로 전달되도록 한다. 정책은 토큰 레벨의 GRPO로 업데이트되고, 디코더는 픽셀 레벨의 Rank-GAN과 EMA 안정성 정규화로 업데이트된다. 이 교대 학습(GEM 관점)은 축적된 Latent Covariate Shift를 흡수하고 fidelity와 alignment를 동시에 끌어올리는 효과를 낸다. RankE의 핵심 효과는 두 가지 랭킹 신호의 보완적 속성과 매니폴드 앵커의 안정성으로, decoder가 evolving token distribution에 맞춰 재구성되도록 하는 점이다.

입력으로 프롬프트 y가 주어지고 z는 πθ(·|y)에서 샘플링된다. x = Dϕ(z)로 디코딩하고 보상 r(Dϕ(z), y)을 얻는다. RankE는 교대 업데이트를 수행한다. 입력: y, z; z_gt은 ground-truth 코드; x̂은 Dϕ(z)로 재구성한다. Stage 1: token-level ranking via GRPO. G개의 롤아웃에 대해 Ai = (ri − μr)/σr로 표준화된 advantage를 계산하고, ρi = πθ(zi|y)/π_old(zi|y)로 중요도 비를 정의한다. Lπ(θ) = − Ey [ (1/G) Σi min(ρi Ai, clip(ρi, 1±ε)) Ai ] − β DKL(πθ ∥ πref). Stage 2: pixel-level ranking. ẑ_i ∼ πEMA θ(·|y)로 샘플링하고 x̂_ema_i = Dϕ(ẑ_i)로 이미지 시퀀스를 얻어 Rank-GAN으로 가중치를 적용한 Ldec를 계산한다. Ldec = λr Lreward + λg LRank-GAN + λc Lconsist + λd Lrecon. 두 단계는 서로 다른 파라미터 공간에 대해 상승하는 최적화이며, GEM 수렴 보장을 따른다. Decoder의 안정화: Lrecon으로 ground-truth 코드에 대한 재구성을 유지하고, Lconsist로 정책 코드에 대한 EMA 타깃의 일관성을 보장한다. 이를 통해 Latent Covariate Shift를 흡수하고 정책-디코더 간 피드백 루프를 통해 픽셀 공간에서의 품질과 정합성을 동시에 향상시킨다.

주요 벤치마크에서 RankE는 frozen-decoder RL 대비 성능 향상을 보인다. LlamaGen-XL(775M)에서 CLIP 33.76, FID 15.21로 나타나고, Std. RL(CLIP)와 비교해 CLIP은 더 높고 FID는 더 낮다. GenEval 평균도 RankE가 0.425로 증가하고 RL baseline 대비 향상된다. Janus-Pro-1B에서 RankE의 CLIP은 33.86으로 증가하고 FID는 25.19로 감소하며 GenEval도 0.750으로 상승한다. HPSv2 기반 평가에서도 RankE의 GenEval가 상승하고 CLIP/GenEval의 균형이 개선된다. 4개의 이미지 인사이트에서 RankE는 색상/구성의 정확성 및 픽셀 품질을 유지하면서 Prompts의 속성(색상, 위치, 수, 관계)을 더 정확히 반영한다. 추적 지표로 KL(πθ ∥ preal)은 RankE에서 초기 수준을 유지하고 토큰 엔트로피는 real-image 수준에 가깝게 유지되며 Latent Covariate Shift를 흡수하는 메커니즘이 확인된다. 4개 이미지 비교에서 디코더 적응이 실질적인 픽셀 품질 개선과 정합성 향상에 기여한다.

단계: 전체 아키텍처는 AR policy πθ(z|y)와 VQ decoder Dϕ로 구성된다. 목표는 Ey,z∼πθ[r(Dϕ(z), y)]를 최대화하는 것인데, z 샘플링과 VQ 양자화로 인해 경사 경로가 차단된다. RankE는 두 단계의 교대 최적화를 통해 신호를 간접적으로 전달한다. 단계 1: token-level ranking via Group Relative Policy Optimization(GRPO). 각 프롬프트 y에 대해 G개의 롤아웃 z_i를 샘플링하고 r로 점수를 매긴 뒤 Ai=(ri−μr)/σr로 표준화된 advantage를 계산한다. 중요도 비 ρi=πθ(zi|y)/πold(zi|y)를 정의하고 Lπ(θ)=-Ey[(1/G)Σi min(ρiAi, clip(ρi,1±ε)) Ai]−βDKL(πθ∥πref)로 업데이트한다. 단계 2: pixel-level ranking. 정책을 고정하고 디코더를 업데이트하며, ẑi∼πEMA θ(·|y)로 샘플링하고 x̂i′=Dϕ(ẑi)로 이미지 시퀀스를 얻어 Rank-GAN으로 가중치를 적용한 LRank-GAN(ϕ)를 계산한다. 또한 Lreward(ϕ)로 differentiable 보상을 직접 역전파하는 경로를 제공하고, ground-truth 코드 zgt에 대한 재구성 손실 Lrecon(ϕ)과 EMA 기반 안정화 Lconsist(ϕ)를 추가한다. 이 두 손실은 원래의 ground-truth 매니폴드에 대한 안정성을 보장한다. RankE는 GEM 수렴 보장을 따르며, 각 스테이지가 서로를 개별적으로 개선하는 것을 전제로 한다.