확산 모델 강화학습을 위한 라우팅 정책 최적화

핵심 요약

• 기존 강화학습 방식은 역방향 디노이징 궤적을 섭동하여 비용이 많이 드는 롤아웃 기반 최적화를 수행하는데, 생성 모델 내부에서 확률적 탐색을 도입하면서 원래의 샘플러를 보존할 수 있는지 묻는 새로운 접근을 제시했다.
• Routing Policy Optimization (RPO)는 어텐션 라우팅을 확률적 탐색 행동으로 변환하여 이를 답한다.
• 선택된 몇 개의 디노이징 단계에서 각 쿼리 블록은 자신의 블록을 유지하면서 생성 모델의 자연스러운 어텐션 질량에 비례하여 추가 키 블록을 샘플링하거나 밀집 어텐션을 유지한다.
• 샘플링된 경로는 해당 전이의 속도를 섭동하며, 밀집 속도 주변의 가우스 우도가 점수 함수 목표로 섭동에 점수를 매긴다.
• RPO는 제어되는 이중 롤아웃 추정기를 사용하는데, 두 롤아웃은 동일한 프롬프트, 초기 잠재, 라우팅된 디노이징 단계, 밀집 앵커 레이어를 공유하며 샘플링된 어텐션 경로만 다르다.
• 두 롤아웃의 보상 차이가 라우팅된 전이에 대한 어드밴티지 신호를 제공한다.
• 우도가 밀집 속도 자체를 미분하기 때문에 학습된 업데이트는 추론 시 라우팅을 제거한 후에도 효과적으로 유지된다.
• 추론 단계에서 RPO는 밀집 어텐션을 복원하고 원래의 결정적 샘플러를 유지한다.
• 텍스트-이미지 및 텍스트-비디오 생성 실험에서 RPO는 생성 품질을 개선하면서 Flow-GRPO 스타일 방법 대비 약 4배 높은 훈련 효율을 달성했다.