Flow-OPD: Flow Matching 모델을 위한 온-폴리시 증류 기법

기존 Flow Matching 기반 이미지 생성 모델은 여러 작업을 동시에 학습할 때 성능이 상충하는 시소 효과와 보상 해킹 문제에 시달렸다. Flow-OPD는 거대 언어 모델의 온-폴리시 증류 기법을 이미지 생성에 최초로 도입하여, 텍스트 렌더링과 미적 품질을 동시에 극대화하는 새로운 정렬 패러다임을 제시한다.

기존의 Flow Matching 모델은 이미지 생성 시 노이즈에서 데이터로 가는 연속적인 속도장(velocity field)을 학습한다. 하지만 이를 특정 인간의 선호도나 OCR 정확도 같은 목표에 맞추기 위해 강화학습(RL)을 적용하면, 보상이 희소하고 여러 목표 간의 그래디언트 간섭이 발생해 한쪽 성능이 오르면 다른 쪽이 떨어지는 시소 효과가 나타난다. 이는 모델이 단순히 스칼라 값의 보상을 최대화하려다 발생하는 정보의 병목 현상 때문이다.

Flow-OPD는 이 문제를 해결하기 위해 모델이 현재 정책으로 직접 생성한 샘플(On-policy)에 대해 여러 전문가 모델(Teacher)이 밀집된 궤적 수준의 가이드를 제공하게 한다. 이는 단순히 잘했다/못했다는 점수만 주는 것이 아니라, 매 단계마다 어느 방향으로 속도 벡터를 수정해야 하는지 상세한 지도를 그려주는 것과 같다.

결과적으로 모델은 복잡한 다중 작업 환경에서도 각 전문가의 강점을 흡수하며, 단순히 교사를 따라가는 것을 넘어 여러 지식이 교차 오염(Cross-pollination)되면서 교사의 성능을 추월하는 현상까지 보여준다. 이는 파라미터 공간 내에서 서로 충돌하던 작업들이 밀집된 감독 신호를 통해 조화롭게 통합되었음을 의미한다.

Flow-OPD는 2단계 정렬 전략을 채택한다. 첫 번째 단계에서는 단일 보상 기반의 GRPO를 통해 특정 도메인(예: OCR, 미적 품질)에 특화된 전문가 교사 모델들을 개별적으로 육성한다. 이후 두 번째 단계에서 이 전문가들의 지식을 하나의 학생 모델로 통합하는 온-폴리시 증류를 수행한다.

온-폴리시 샘플링 과정에서는 결정론적인 ODE를 확률 미분 방정식(SDE)으로 변환하여 무작위성을 주입한다. [현재 상태 xt와 속도 vθ를 입력으로] → [SDE 기반의 Euler-Maruyama 이산화를 수행해] → [확률적인 전이 궤적을 얻고] → [이를 통해 학생 모델이 탐색할 수 있는 행동 공간을 확보한다].

핵심 보상 신호는 Dense KL Reward를 통해 계산된다. 학생과 교사의 전이 정책이 동일한 등방성 공분산을 공유한다는 점을 이용해, 복잡한 KL 발산을 두 모델의 속도 벡터 간 L2 거리로 단순화한다. [학생의 속도 vθ와 교사의 속도 vtarget을 입력으로] → [시간 적응형 가중치 w(t)를 곱한 L2 차이를 계산해] → [밀집된 보상 r_t를 얻고] → [이 값이 작아지도록 PPO 클리핑 메커니즘을 통해 정책을 업데이트한다].

마지막으로 Manifold Anchor Regularization(MAR)을 적용한다. [학생 모델의 출력과 고정된 미적 교사 모델의 속도장 vbase를 입력으로] → [전체 데이터에 대해 KL 패널티를 계산해] → [총 손실 함수 LTotal에 합산하고] → [모델이 특정 작업에 과적합되어 미적 품질이 붕괴되는 것을 방지한다].

Stable Diffusion 3.5 Medium을 기반으로 실험한 결과, Flow-OPD는 GenEval 점수를 63점에서 92점으로, OCR 정확도를 59%에서 94%로 대폭 향상시켰다. 이는 기본 GRPO 방식보다 약 10점 높은 수치이며, 특히 여러 작업을 동시에 수행할 때 발생하는 성능 저하 문제를 완벽히 해결했다.

Ablation Study에서는 콜드 스타트 전략의 중요성이 확인됐다. 모델 머징을 통한 초기화 방식(Ours-Merge)이 평균 0.9044의 점수를 기록하며 가장 우수한 성능을 보였으며, 이는 개별 전문가 모델들의 성능 한계를 뛰어넘는 'Teacher-surpassing' 효과를 입증했다.

또한 Manifold Anchor Regularization(MAR)의 도입으로 배경 붕괴나 시각적 중복 현상을 억제했다. MAR을 적용했을 때 ImageReward와 Aesthetic 점수가 각각 1.36, 6.23으로 나타나, 정규화가 없는 경우보다 시각적 무결성과 인간 선호도 정렬이 월등히 개선됨을 확인했다.

Flow-OPD는 Flow Matching의 연속적인 동역학을 강화학습의 마르코프 결정 과정(MDP)으로 재구성한다. 각 타임스텝 t에서의 상태 전이를 가우시안 정책으로 모델링하여, 연속 시간 영역에서의 KL 발산을 속도 벡터 간의 가중치 기반 L2 회귀 문제로 변환한 것이 수학적 핵심이다.

구현 측면에서는 Task-routing labeling 메커니즘을 사용해 텍스트 조건 c에 따라 적절한 전문가 교사를 동적으로 할당한다. 이를 통해 서로 다른 도메인의 그래디언트가 충돌하는 것을 방지하고, 학생 모델이 각 작업에 대해 독립적이고 밀집된 감독 신호를 받을 수 있도록 설계했다.

학습 효율성을 위해 LoRA(rank=32, alpha=64)를 사용했으며, 8개의 H800 GPU로 구성된 노드 4개에서 약 50시간 동안 학습을 진행했다. 샘플링 타임스텝 T=10, 평가 타임스텝 T=40 설정을 통해 추론 효율과 정렬 품질 사이의 균형을 맞추었다.

Flow-OPD는 교사 모델과 학생 모델 사이의 아키텍처적 동질성(Architectural homogeneity)을 요구한다. 이는 미세한 단계별 감독을 위해 두 모델의 속도장 구조가 일치해야 하기 때문이며, 서로 다른 구조를 가진 모델 간의 증류는 향후 과제로 남아 있다. 또한 교사 모델 자체가 가진 성능 한계가 학생 모델의 잠재적 상한선으로 작용할 수 있다.