스텝 인식 이점으로 흐름 매칭 모델을 빠르고 정확하게 정렬하기

초록

본 논문은 흐름 매칭 기반 텍스트‑투‑이미지 생성 모델을 몇 단계만 사용해도 인간 선호와 높은 정렬성을 유지하도록 설계한 TAFS‑GRPO 프레임워크를 제안한다. 온도‑annealing 방식으로 한 단계 샘플에 적응형 노이즈를 반복 주입해 스텝마다 의미 있는 이미지와 풍부한 보상을 얻고, GRPO 기반 정책 최적화에 스텝‑인식 어드밴티지를 결합해 보상 희소성을 해소한다. 실험 결과 4‑step·8‑step 설정에서 기존 Reward‑Instruct 대비 품질·정렬 모두 크게 향상됨을 확인하였다.

상세 분석

TAFS‑GRPO는 흐름 매칭(Flow Matching) 모델의 근본적인 한계인 “다수의 디노이징 단계 필요”와 “보상 신호의 희소성”을 동시에 해결한다는 점에서 혁신적이다. 기존의 RL‑기반 정렬 방법은 전체 ODE/ODE‑SDE 경로를 모두 샘플링해야 보상을 계산할 수 있었으며, 이는 20~40 단계의 고비용을 초래한다. 저자들은 이를 “온도‑annealed few‑step sampling”이라는 새로운 샘플링 메커니즘으로 대체한다. 구체적으로, 초기 가우시안 노이즈 x_T에서 한 단계 흐름 모델 v_θ를 이용해 x₁⁰을 얻고, 이후 τ = T/N 로 정의된 작은 시간 간격마다 적응형 노이즈 ε_{T‑τ·k}를 삽입한다. 각 단계에서 x_{k+1}⁰ = x_k⁰ + ε_{T‑kτ} + (T‑kτ)·v_θ(x_{k} + ε_{T‑kτ}, T‑kτ) 로 업데이트되며, 이는 기존 디노이징 단계와 달리 “한 번의 샘플링” 안에서 여러 의미 있는 중간 이미지가 생성된다. 이 과정은 확률적 SDE 형태를 유지하면서도 각 중간 이미지가 충분히 의미론적 정보를 담게 하여, 보상 함수를 단계별로 적용할 수 있게 만든다.

스텝‑인식 어드밴티지 통합은 GRPO의 핵심인 그룹 기반 정책 업데이트에 직접 연결된다. 기존 GRPO는 동일한 보상 r을 모든 타임스텝에 동일하게 할당해 보상의 밀도가 낮았다. TAFS‑GRPO는 각 중간 이미지에 대해 별도의 보상 R(x_k⁰)·를 계산하고, 이를 정규화한 어드밴티지 ĤA_{k} = (r_k ‑ mean(r))/std(r) 로 변환한다. 이후 정책 비율 r_{i,k}=π_θ(o_{i,k}|·)/π_{θold}(o_{i,k}|·)와 KL 정규화 항을 포함한 손실 J_{GRPO}에 삽입해 스텝별로 차등된 정책 그라디언트를 얻는다. 이때 보상 함수는 미분 가능할 필요가 없으며, 인간 선호를 반영한 라벨링 혹은 LLM‑생성 점수 등 임의의 블랙박스 평가를 그대로 사용할 수 있다.

실험 설계는 4‑step·8‑step 두 가지 설정에서 기존 최첨단 Reward‑Instruct와 동일한 학습 환경을 유지하면서 비교하였다. 평가 지표는 FID, IS, CLIP‑Score 및 인간 평가(선호도 비율)이다. 결과는 TAFS‑GRPO가 FID를 평균 12% 개선하고, CLIP‑Score와 인간 선호도에서 각각 0.07·0.09 포인트 상승했으며, 특히 4‑step 상황에서 품질 저하 없이 8‑step 대비 2배 가량 빠른 추론 속도를 달성했다는 점을 강조한다. 또한, 온도‑annealing 파라미터 τ와 노이즈 스케줄링을 다양하게 변형해도 성능이 크게 변하지 않아 방법의 견고함을 입증한다.

이 논문의 주요 기여는 다음과 같다. 첫째, 흐름 매칭 모델에 온도‑annealed 샘플링을 도입해 단계별 의미 보존과 stochasticity를 동시에 확보하였다. 둘째, 스텝‑인식 어드밴티지를 GRPO와 결합해 보상 희소성을 근본적으로 해소하고, 미분 불가능한 보상 함수도 활용 가능하도록 했다. 셋째, 실험을 통해 몇 단계만으로도 고품질·고정렬 이미지 생성이 가능함을 입증했으며, 코드와 모델을 공개해 재현성을 보장한다. 향후 연구는 더 복잡한 멀티모달 프롬프트, 비디오 생성, 그리고 실제 서비스 환경에서의 latency‑aware 배포 등에 적용할 여지를 남긴다.