기술 요약: FlowAWR

문제 정의

연속 공간상에서 생성형 플로우 모델(generative flow models)을 온라인 강화 학습(RL)으로 정렬하는 것은 근본적인 구조적 제약에 직면해 있습니다: 바로 궤적 가능도(trajectory likelihoods)의 계산 불가능성입니다. 이산적인 토큰 확률을 정확하게 계산할 수 있는 자기회귀(AR) 모델과 달리, 연속형 플flow 모델은 멱단계별 전이 밀도(step-wise transition densities)를 다룰 수 없습니다.

기존 접근 방식들은 두 가지 주요 패러다임을 통해 이를 우회하려고 시도하지만, 두 방식 모두 상당한 한계를 가집니다:

1. 밀도 근사 정책 경사법 (예: GRPO 기반 방법): 이 방법들은 연속적인 과정을 마르코프 결정 과정(MDP)으로 이산화하고, 트래커블한 가우시안 전이 커널을 구축하기 위해 확률 미분 방정식(SDE) 샘플러를 통해 확률성을 주입합니다. 이는 훈련-추론 불일치(stochastic training vs. deterministic inference)를 유발하며, 생성 품질을 유지하기 위해 추론 시 **Classifier-Free Guidance (CFG)**를 필수적으로 사용하게 되어 계산 오버헤드를 발생시킵니다.
2. 암묵적 최적화 (예: DiffusionNFT): 이러한 프레임워크는 명시적인 밀도 추정 없이 순방향 프로세스의 속도장(velocity fields)을 직접 최적화합니다. 그러나 이들은 휴리스틱한 고정 크기 보정에 의존하며, 그룹 내 샘플을 단순히 스칼라 정규화된 보상을 계산하는 데에만 활용합니다. 이는 최적화를 개별 샘플 수준에서의 이진적 "push-pull" 메커니즘으로 제한하며, 속도 업데이트의 크기를 동적으로 결정하기 위해 그룹 내 상대적 품질을 활용하지 못합니다.

방법론: Flow Advantage-Weighted Rectification (FlowAWR)

저자들은 FlowAWR을 제안하며, 이는 연속 생성 정책 최적화를 이론적으로 최적인 속도장을 향한 지도 회귀(supervised regression) 작업으로 재구성하는 온라인 RL 패러다임입니다. 이 방법은 KL 제약 조건이 있는 보상 최대화 문제의 정확한 폐쇄형 최적 정책으로부터 유도되었습니다.

이론적 유도

1. 최적 정책: KL 제약 보상 최대화 목적 함수로부터 시작하여, 저자들은 참조 정책 $\pi_{old}$의 지수 가중 버전인 폐쇄형 최적 터미널 분포 $\pi^*$를 유도합니다.
2. 최적 속도장: 이 최적 터미널 분포를 확률 경로를 따라 중간 시간 단계 $t$로 전파함으로써, 저자들은 그에 상응하는 최적 속도장 $v^*$를 분석적으로 도출합니다.
3. Advantage-Weighted Rectification: Tweedie의 공식과 주변 스코어(marginal scores) 및 사후 평균(posterior means) 사이의 관계를 사용하여, 최적 속도장은 참조 속도장 $\text{v}_{old}$를 **어드밴티지 가중 기대값(advantage-weighted expectation)**에 의해 교정(rectifying)하는 것과 동일함을 보여줍니다:
   $$v^*(x_t, c, t) = v_{old}(x_t, c, t) + \mathbb{E}_{x_1 \sim p_{old}} \left[ A(x_1, x_t) \cdot (u_t(x_t|x_1) - v_{old}(x_t, c, t)) \right]$$
   여기서 $A(x_1, x_t)$는 그룹 내 샘플의 상대적 품질을 정량화하는 중심화된 어드밴티지(Centered Advantage) 함수입니다. 이는 지수 보상을 분할 함수(partition function, 그룹 평균)에 대해 정규화하고 베이스라인 1을 빼서 구성됩니다.

실질적 알고리즘

이 프레임워크는 다음을 통해 이론을 구현합니다: Advantage-Weighted Rectification (AWR) 손실

• 그룹 기반 추정: 각 프롬프트에 대해 $G$개의 이미지를 생성합니다. 어드밴티지 $A$는 경직된 이진 가중치 대신 크기 인지형 보정(magnitude-aware corrections)을 사용하여 그룹의 보상 통계치를 통해 경험적으로 계산됩니다.
• 지도 회귀: 정책은 참조 필드를 어드밴티지 가중 잔차(advantage-weighted residual)로 교정한 확률적 타겟을 향해 매개변수화된 속도 $v_\theta$를 회귀하도록 훈련됩니다.
• 주요 특성:
  • SDE-Free: 훈련 중 SDE 샘플러를 사용할 필요가 없어 순방향 일관성(forward consistency)을 보장합니다.
    에는 CFG-Free: 가이던스 신호가 속도장에 암시적으로 증류(distilled)되어, 추론 시 CFG의 필요성을 제거합니다.
  • 적응형 스케일링: KL 제약 강도($\gamma$)는 현재 배치의 보상 표준 편차를 사용하여 적응적으로 스케일링됩니다.
  • 부드러운 온라인 진화: 어드밴티지 추정을 안정화하기 위해 참조 정책에 지수 이동 평균(EMA) 업데이트가 적용됩니다.

주요 기여

1. FlowAWR 프레임워크: SDE 시뮬레이션 없이 정확한 정책 업데이트를 수행하는 온라인 플로우 RL 프레임워크입니다. 이는 DiffusionNFT의 휴리스틱한 보정 방향을 검증하는 동시에, DiffusionNFT가 제안된 연속 포뮬레이션의 **이진 양자화된 특수 사례(binary quantized special case)**임을 입증합니다.
2. 이론적 규명: 저자들은 최적의 속도장을 중간 분포로 전파함으로써 최적 속도장을 규명하고, 이를 그룹 기반 근사 어드밴티지 함수을 통해 이론적 최적치를 구현합니다. 이는 플로우 기반 RL과 지도 학습 기반 미세 조정(SFT)을 연결하는 통합된 관점을 제공합니다.
3. 경험적 검증: SD3.5-Medium 모델에 대한 포괄적인 평가를 통해, FlowAWR이 베이스라인 대비 수렴 속도를 가속화하고 정렬 성능을 개선함을 입 ใน증합니다.

실험 결과

실험은 단일 보상 및 다중 보상 설정에서 SD3.5-Medium (2.5B 파라미터) 모델을 대상으로 수행되었습니다.

• 수렴 속도: FlowAWR은 DiffusionNFT에 비해 2배에서 5배 빠른 수렴을 달s합니다.
  • 예시 (PickScore): FlowAWR은 1.2k 스텝에서 24.12 점에 도달하는 반면, DiffusionNFT는 2.0k 스텝이 걸려 23.82에 도달하며, FlowGRPO는 CFG에 의존하면서도 23.50에 도달하기 위해 4k 이상의 스텝이 필요합니다.
  • 예시 (OCR): FlowAWR은 0.26k 스텝에서 0.97의 OCR 점수를 달성하며, 이는 DiffusionNFT가 0.52k 스텝에서 0.89를 달성한 것과 대조적입니다.
• 다중 보상 안정성: 미적(aesthetic) 규칙과 구조적 규칙을 결합한 다중 보상 최적화로 인해 발생하는 복잡한 분포 변화 하에서도, FlowAWR은 미적 영역 외 탈락(out-of-domain degradation) 없이 생성 품질을 유지하고 구조적 규칙을 준수합니다.
• CFG 제거: 이 방법은 CFG 없이도 고충실도 생성을 가능하게 하여 추론 비용을 크게 절감합니다.

의의 및 주장

본 논문은 FlowAWR이 플로우 기반 RL과 지도 학습 기반 미세 조정을 연결하는 통합적 분석 관점을 제공한다고 주장합니다. 제1원리로부터 최적 속도장을 유도함으로써, 이 방법은 기존의 휴리스틱 방법(DiffusionNFT 등)과 정책 최적화 사이의 관계를 명확히 하고, 이들을 연속적 프레임워크 내의 이산적 근사로 밝혀냅니다.

저자들은 FlowAWR이 Likelihood-free 접근 방식의 구조적 이점(순방향 일관성, 솔버 유연성)을 보존하면서도 기존 방법들의 한계를 극복한다고 강조합니다. 구체적으로, 경직된 이진 연산을 상대적 그룹 내 품질을 활용하는 **크기 인지형 교정(magnitude-aware rectifications)**으로 대체함으로써, 이론적 근거가 탄탄하고 수렴 속도 및 정렬 성능 면에서 경험적으로 우월하며 CFG가 필요 없는 안정적인 최적화를 실현했습니다.