Flow-OPD: On-Policy Distillation for Flow Matching Models

Zhen Fang; Wenxuan Huang; Yu Zeng; Yiming Zhao; Shuang Chen; Kaituo Feng; Yunlong Lin; Lin Chen; Zehui Chen; Shaosheng Cao; Feng Zhao

Source-linked AI summary

Flow-OPD: On-Policy Distillation for Flow Matching Models

Zhen Fang, Wenxuan Huang, Yu Zeng, Yiming Zhao, Shuang Chen, Kaituo Feng, Yunlong Lin, Lin Chen, Zehui Chen, Shaosheng Cao, Feng Zhao

arXiv:2605.08063v52026년 5월 8일cs.CVcs.AI

TL;DR

Flow Matching 모델의 multi-task alignment는 희소한 scalar reward와 상충하는 objective로 인해 제약된다. Flow-OPD는 manifold anchoring을 결합한 dense on-policy distillation으로 전문화된 expertise를 통합하고, 도메인 간 interference를 해소하면서 visual fidelity를 보존하고 teacher-surpassing 성능을 달성한다.

문제
Flow Matching 모델에는 희소 reward로 인한 seesaw effect와 도메인 간 interference 없이 heterogeneous objective를 조화시키는 효과적인 multi-task alignment 방법이 부족하다.
방법
Flow-OPD는 on-policy sampling, task routing, dense trajectory supervision, task-agnostic manifold anchoring을 통해 전문화된 teacher들을 하나의 student로 distill한다.
결과
Flow-OPD는 benchmark 전반에서 전문화된 teacher와 일관되게 동등하거나 우수한 성능을 보이며, 도메인 간 interference를 해소하고 최적의 multi-task trade-off와 teacher-surpassing 사례를 달성한다.
시사점 및 한계
Flow-OPD는 composition 및 typography expertise를 통합하고 높은 visual fidelity를 유지하는 generalist text-to-image 모델을 위한 확장 가능한 paradigm을 제공한다.
시사점 및 한계
Flow-OPD는 teacher quality에 의존하며, 세밀한 step-wise supervision을 위해 teacher와 student 간 architectural homogeneity를 요구한다.

Abstract

from arXiv · show

Existing Flow Matching (FM) text-to-image models suffer from two critical bottlenecks under multi-task alignment: the reward sparsity induced by scalar-valued rewards, and the gradient interference arising from jointly optimizing heterogeneous objectives, which together give rise to a 'seesaw effect' of competing metrics and pervasive reward hacking. Inspired by the success of On-Policy Distillation (OPD) in the large language model community, we propose Flow-OPD, the first unified post-training framework that integrates on-policy distillation into Flow Matching models. Flow-OPD adopts a two-stage alignment strategy: it first cultivates domain-specialized teacher models via single-reward GRPO fine-tuning, allowing each expert to reach its performance ceiling in isolation; it then establishes a robust initial policy through a Flow-based Cold-Start scheme and seamlessly consolidates heterogeneous expertise into a single student via a three-step orchestration of on-policy sampling, task-routing labeling, and dense trajectory-level supervision. We further introduce Manifold Anchor Regularization (MAR), which leverages a task-agnostic teacher to provide full-data supervision that anchors generation to a high-quality manifold, effectively mitigating the aesthetic degradation commonly observed in purely RL-driven alignment. Built upon Stable Diffusion 3.5 Medium, Flow-OPD raises the GenEval score from 63 to 92 and the OCR accuracy from 59 to 94, yielding an overall improvement of roughly 10 points over vanilla GRPO, while preserving image fidelity and human-preference alignment and exhibiting an emergent 'teacher-surpassing' effect. These results establish Flow-OPD as a scalable alignment paradigm for building generalist text-to-image models. The codes and weights will be released in: https://github.com/CostaliyA/Flow-OPD .

1 서론

Flow-OPD는 LLM의 On-Policy Distillation을 vision으로 이전해 multi-task Flow Matching alignment의 reward sparsity와 gradient interference를 해결한다. 두 단계 framework는 전문성 획득과 모델 통합을 분리해 GRPO보다 성능과 일반화 성능을 향상한다.

동기: Multi-task Flow Matching alignment는 이질적이고 상충하는 feature representation과 zero-sum seesaw effect를 유발하는 희소한 scalar reward로 인해 제약된다.이러한 한계로 인해 single-reward RL method만으로는 여러 objective를 조화시키기 어렵다.
동기: On-Policy Distillation은 DeepSeek-V4 [9], Mimo v2 [20], GLM-5 [19]와 같은 LLM이 specialized expert를 distilling해 복잡한 multi-domain capability를 조화시키도록 했다.이러한 성공은 Flow Matching model에서 다양한 teacher의 강점을 통합하기 위해 OPD를 적용하는 동기를 제공한다.
방법: Flow-OPD는 Flow Matching post-training을 위한 최초의 제안된 OPD paradigm으로, 전문성 획득과 모델 통합을 분리해 fine-grained supervision을 제공한다.이 framework는 GRPO 기반 multi-task training에서 확인된 failure mode, 특히 reward sparsity와 gradient interference를 해결한다.
방법: Manifold Anchor Regularization은 Flow-OPD alignment 동안 전반적인 generative quality를 보존하기 위해 task-agnostic guidance를 제공한다.이 방법은 task 전반에서 generation quality를 유지하도록 설계된 regularization mechanism과 fine-grained supervision을 결합한다.
결과: 네 가지 mainstream benchmark에서 GRPO baseline 대비 10-point improvement를 달성했으며, unified student는 in-domain에서 specialized teacher와 동등하거나 이를 능가하고 OOD에서도 일반화한다.이 결과는 GRPO 대비 Flow-OPD의 우수한 성능과 일반화 성능을 확립한다.

2 관련 연구

기존 연구는 text-to-image 모델 정렬에 reinforcement learning과 preference optimization을 적용해 왔으며, on-policy distillation은 고정된 supervision과 변화하는 student trajectory 간의 불일치를 다룬다. 언어 모델에서는 exposure bias와 mode-seeking behavior를 겨냥한 framework 및 KL 기반 refinement를 통해 OPD가 발전해 왔다.

T2I 모델을 위한 RL: text-to-image generation을 위한 RL 기반 정렬은 diffusion policy optimization과 preference alignment에서 출발해 최신 GRPO 계열 방법으로 발전했다.DDPO, DPOK, ImageReward/ReFL [25]은 aesthetics, human preference 또는 text-image alignment를 위한 reward를 최적화하며, Diffusion-DPO 는 preference pair를 사용한다.
On-Policy Distillation: On-Policy Distillation은 traditional offline distillation과 달리 teacher supervision을 student의 변화하는 exploration space에 결부한다.on-policy formulation은 student trajectory에 맞춰 supervision을 동적으로 조정한다.
On-Policy Distillation: 언어 모델에서 OPD의 발전에는 exposure bias 완화를 위한 GKD [34]와 Reverse 및 Skewed KL을 통해 mode-seeking behavior를 refinement하는 MiniLLM 및 DistiLLM [36]이 포함된다.이 방법들은 LLM 분야에서 on-policy distillation 기법이 빠르게 확장되고 있음을 보여준다.

3 사전 지식

Flow Matching은 noise를 data로 운반하는 ODE로 생성을 정의하며, 이산화된 적분은 reinforcement learning을 위한 Markovian denoising trajectory로 볼 수 있다. On-policy distillation은 student가 생성한 trajectory에서 teacher의 행동을 전달해 distribution shift와 exposure bias를 줄인다.

Flow Matching Models: Flow Matching은 ODE dxt = vt(xt, t)dt를 통해 noise p0를 data pdata로 매핑하며, OT path xt = (1−t)x0+tx1와 constant velocity x1 − x0를 사용한다.Optimal Transport formulation에서 model vθ는 이 constant velocity를 학습한다.
Flow Matching Models: 이산화된 ODE 적분은 sequential Markovian denoising으로 정식화할 수 있으며, continuous generation dynamics를 step-wise reinforcement-learning policy optimization과 연결한다.각 transition xt → xt+∆t는 Markovian state step으로 취급되어 formal trajectory를 정의한다.
On-Policy Distillation: On-policy distillation은 student가 생성한 trajectory τ ∼ pθ(τ)에서 teacher supervision을 사용해 student–teacher output divergence를 최소화하고 distribution shift를 완화한다.Autoregressive models에서는 objective를 reverse Kullback–Leibler divergence로 정식화한다.
On-Policy Distillation: Single-reward GRPO는 cross-task generalization을 심각하게 저해하여 non-target metrics에서 capability를 저하시킨다.이 cross-task evaluation은 official Flow-GRPO implementation을 엄격히 따르는 baseline setups를 사용한다. On-policy alignment는 exposure bias를 억제하고 interactive 또는 iterative generation에서 robustness를 향상하는 것으로 설명된다.

4 동기

Online GRPO는 자체 생성 trajectory를 탐색해 offline SFT를 능가할 수 있지만, 단일 보상 최적화는 gradient interference를 통해 직교 능력을 손상시킨다. 통제 실험은 충돌하는 목적이 zero-sum advantage로 압축되기 때문에 scalar reward를 혼합하는 방식이 본질적으로 불안정하고 확장 불가능함을 보여준다.

Online Exploration: GRPO는 현재 policy에서 그룹을 샘플링하고, 자체 생성 state를 relative advantage로 평가하며, 새로운 high-reward trajectory를 발견함으로써 offline SFT의 static-data ceiling을 극복한다.Standard flow matching은 offline reconstruction에 의존해 미분 불가능한 선호를 직접 최적화할 수 없는 반면, online exploration은 policy-driven training signal을 제공한다.
Gradient Interference: Sparse scalar reward는 공유 parameter space 내에서 제약되지 않은 gradient interference를 유발하므로, 단일 보상 GRPO는 직교 능력을 심각하게 저하시킨다.충돌하는 task gradient는 음의 inner product를 가질 수 있으며, 그 결과 한 target을 위한 최적화가 모니터링되지 않은 자유도를 악용하고 pretrained synergy를 해체할 수 있다.
Mixing Multiple Rewards: +GenEval 이후 OCR을 추가하면 5% GenEval degradation이 발생하며, 이는 scalar reward를 누적할 때 안정적인 multi-reward optimization이 아니라 catastrophic forgetting이 일어남을 보여준다.이 실험은 SD-3.5-M에서 GenEval, OCR, PickScore, DeQA를 점진적으로 누적하며, 표는 multi-reward optimization에 따른 능력 저하를 보고한다.
Mixing Multiple Rewards: Scalar reward mixing은 한 objective를 최적화하는 과정에서 aesthetic stylization에 필요한 geometric precision과 같이 다른 objective에 필요한 representation을 덮어쓸 수 있는 zero-sum game을 만든다.이 대목은 conflicting gradient가 scalar reward mixing을 본질적으로 확장 불가능하게 만드는 메커니즘임을 지적한다.

5 방법: Flow-OPD

Flow-OPD는 자체 생성 궤적에 대한 dense supervision을 통해 Flow Matching 모델을 정렬하며, expert teacher, cold-start 초기화, task-specific routing, Manifold Anchor Regularization을 결합한다. 연속 formulation은 시각적 manifold를 보존하면서 고분산 policy-gradient 추정을 동등한 vector-field regression objective로 분석적으로 대체한다.

Flow-OPD 개요: 이 framework는 domain-expert teacher를 학습하고, cold-start SFT 또는 model merging을 통해 student를 초기화한 뒤, 동적으로 라우팅된 online trajectory를 통해 전문 지식을 distillation한다.SFT는 specialized teacher의 knowledge distribution을 전달하는 반면, model merging은 서로 다른 teacher의 anisotropic prior를 하나의 parameter state에 중첩한다.
Task-Specific Teacher Labeling: Hard task routing은 각 condition에서 하나의 domain expert를 활성화하여, 각 on-policy trajectory의 definitive target으로 국소화된 reference velocity field를 제공한다.student는 stochastic SDE 기반 rollout을 통해 자체 distribution에서 sampling하며, routing은 domain 간 gradient interference를 방지한다.
OPD와 Flow Matching의 연결: Flow-OPD는 LLM-style Policy-Gradient OPD와 수학적으로 동등한 closed-form vector-field MSE를 직접 최소화하여 고분산 policy-gradient 추정을 우회한다.연속 flow formulation은 log-probability 계산과 PPO surrogate bound를 피하고, gradient variance를 zero로 낮추면서 on-policy exploration과 regression efficiency를 통합한다.
Manifold Anchor Regularization: Manifold Anchor Regularization은 frozen aesthetic teacher의 vector field를 사용해 functional alignment를 고품질 visual manifold로 제한하고 aesthetic degradation을 방지한다.이 regularizer는 Reverse KL penalty를 time-weighted vector-field distance로 변환하여, multi-teacher optimization 중 visual quality를 보존한다.

6 실험

Flow-OPD는 네 가지 alignment task에서 monolithic- 및 hybrid-reward GRPO baseline과 비교 평가되며, specialized teacher와 일관되게 동등하거나 더 나은 성능을 내는 동시에 multi-task trade-off와 일반화를 개선한다. Ablation 결과는 cold-start initialization과 MAR가 견고한 alignment, capability 보존, image-quality control을 가능하게 하는 메커니즘임을 뒷받침한다.

실험 설정: 평가는 GenEval, OCR, PickScore, DeQA에 걸친 Flow-GRPO split을 따르며, 4:6 DeQA-to-PickScore reward blend로 학습한 DeQA teacher를 제외하면 공식 expert checkpoint를 사용한다.Baseline에는 single-reward GRPO와 GenEval:OCR:PickScore 가중치가 3:1:1인 GRPO-Mix가 포함된다.
주요 결과: Flow-OPD는 GenEval, OCR, PickScore, DeQA 전반에서 specialized teacher와 일관되게 동등하거나 더 나은 성능을 내며, cross-domain interference와 sparse-reward optimization bottleneck을 해소한다.Table 2는 compositional generation, visual text rendering, image quality를 비교한다. PickScore teacher의 GenEval score가 0.51로 하락하는 현상은 specialization interference를 보여준다.
정성적 결과: Flow-OPD는 prompt fidelity와 visual aesthetics를 결합하는 최적의 multi-task trade-off를 달성하며, 개별 teacher가 모두 실패하는 edge case에서도 성공한다.이러한 사례는 보고된 Teacher-Surpassing phenomenon을 뒷받침하며, 이는 latent flow manifold 내부의 knowledge cross-pollination에서 비롯된 것으로 추정된다.
Cold-start ablation: Cold-start initialization은 견고한 training foundation을 빠르게 구축한다. SFT는 heterogeneous-teacher extraction을 확장 가능하게 제공하는 반면, model merging은 추가 학습 없이 homogeneous teacher를 가장 효과적으로 활용한다.이 대안들은 cold-start ablation에서 비교된다.
Out-of-domain generalization: T2I-CompBench에서 Flow-OPD는 state-of-the-art out-of-domain compositional generalization을 달성하는 반면, multi-task GRPO는 shape rendering과 3D spatial reasoning 같은 capability에서 catastrophic forgetting을 보인다.비교에는 두 방법에 동일한 cold-start baseline을 사용한다.
Manifold Anchor Regularization: MAR는 optimization을 high-fidelity manifold에 고정해 reward로 유발되는 background mode collapse와 semantic redundancy를 방지하고, structural diversity, semantic adherence, image quality, human-preference alignment를 개선한다.이 constraint는 task-agnostic하다. DeQA와 같은 teacher는 diverse sample을 제공하지만 instruction following에 어려움을 겪을 수 있다.

7 결론 · A More Details

Flow-OPD는 on-policy distillation을 Flow Matching 모델에 통합하고, dense trajectory-level supervision과 Manifold Anchor Regularization을 사용해 reward sparsity, gradient interference, aesthetic degradation을 해결하는 동시에 composition 및 typography 전문성을 통합한다. 추가 세부사항은 hybrid-training 설정과 선택된 hyperparameter를 제시한다.

7 결론: Flow-OPD는 on-policy distillation을 Flow Matching 모델에 통합하고, scalar reward를 dense trajectory-level supervision으로 대체하며, emergent “teacher-surpassing” effect와 함께 composition 및 typography 전문성을 통합한다.이 framework는 reward sparsity, gradient interference, 그리고 경쟁 metric 간 seesaw effect를 해결하는 것으로 제시된다.
A More Details: GRPO-mix는 GenEval, OCR, PickScore에 대해 3:1:1 epoch 비율을 사용하며, 각 data partition은 대응하는 reward model에서만 reward를 받는다.Training은 Flow-GRPO의 data 및 reward configuration을 따랐다.
A More Details: Training은 각각 8개의 H800 GPU를 탑재한 4개의 distributed node에서 약 50시간 동안 진행됐다.이는 hybrid setup에 대해 보고된 distributed training configuration을 설명한다.
A More Details: 구현에는 sampling timestep T = 10, evaluation timestep T = 40, group size G = 24, noise level a = 0.7, 그리고 512-resolution image가 사용된다.이 설정들은 고정된 GRPO hyperparameter 중 일부로 제시된다.
A More Details: Manifold Anchor Regularization은 KL ratio β = 0.02를 사용하며, LoRA는 α = 64 및 r = 32를 사용한다.해당 passage는 이러한 MAR 및 parameter-efficient fine-tuning 설정을 명시한다.
A More Details: Qwenvl Score의 경우, 이 방법은 Flow-GRPO prompt [16]를 변형해 사용하며 Qwen3-30B-A3B-Instruct-2507을 사용한다.변형된 prompt는 Fig. 6에 제시된다.

B 추가 결과 · B.1 정성적 결과 · B.2 DiffusionNFT와의 비교

Flow-OPD의 정성적 결과는 정밀한 콘텐츠 생성, 우수한 이미지 품질, 일관된 레이아웃, 그리고 인간 선호와의 더 강한 정렬을 보여준다. DiffusionNFT [49]와 비교하면 Classifier-Free Guidance (CFG)와의 비호환성을 비롯한 여러 한계를 피한다.

B.1 정성적 결과: Flow-OPD는 우수한 이미지 품질과 일관된 구조적 레이아웃을 바탕으로 정밀한 콘텐츠 생성을 달성한다.
B.1 정성적 결과: Figures 7, 8, and 9는 Flow-OPD의 추가 정성적 결과를 제시한다.
B.1 정성적 결과: Flow-OPD는 인간 선호와의 더 강한 정렬을 보여준다.
B.1 정성적 결과: Flow-OPD는 기능적 정확성과 미적 탁월성을 결합한다.
B.2 DiffusionNFT와의 비교: DiffusionNFT [49]는 보상 피드백을 forward diffusion process에 통합하는 online reinforcement learning을 도입한다.
B.2 DiffusionNFT와의 비교: DiffusionNFT [49]는 noise-injection phase에서 policy optimization을 가능하게 하며 경쟁력 있는 benchmark 점수를 달성한다.
B.2 DiffusionNFT와의 비교: DiffusionNFT [49]는 Classifier-Free Guidance (CFG)와 근본적으로 비호환되므로 성능 상한이 제한된다.
B.2 DiffusionNFT와의 비교: DiffusionNFT [49]는 또한 뚜렷한 reward hac 문제를 겪는다.

B.3 실패 사례 및 한계

Flow-OPD는 teacher model의 성능 상한과 teacher와 student 간 아키텍처 동질성 요구로 제약된다. 전문 teacher의 오류는 dense supervision을 통해 전파되어 noise를 유입하고 student의 개선을 제한할 수 있다.

한계: Flow-OPD의 성능 상한은 teacher model의 역량에 의해 제한된다.
실패 사례: 전문 teacher의 실패는 조밀한 감독 신호를 통해 전파되어 distillation에 noise를 유입하고 student가 teacher ensemble을 능가하는 것을 방해한다.이 실패는 teacher가 의미적으로 잘못된 이미지를 생성할 때 발생한다.
한계: distillation 과정에서는 teacher와 student model이 아키텍처 동질성을 공유해야 한다.

더 넓은 영향

Flow-OPD는 더 신뢰할 수 있고 일관되며 OCR이 가능한 AI 생성 콘텐츠를 통해 사회적 편익을 제공하지만, 향상된 이미지 생성 사실성은 정교한 허위정보와 기만적 시각 콘텐츠의 위험을 높인다. 이에 저자들은 digital watermarki를 포함한 조치를 통해 신중한 검토와 완화를 강조한다.

더 넓은 영향: Flow-OPD는 긍정적 사회 기여와 신중한 검토가 필요한 잠재적 위험을 함께 지닌다.이 프레임워크는 생성 모델의 multi-task alignment에 강건한 방법으로 제시된다.
긍정적 사회적 영향: 향상된 레이아웃 일관성과 OCR 정확도는 그래픽 디자인, 교육 콘텐츠, 시각장애인을 위한 보조 기술에 기여할 수 있다.또한 Multi-Teacher 패러다임은 보다 균형 잡힌 최적화 목적을 촉진하고 winner-takes-all 편향을 완화하는 것으로 설명된다.
부정적 사회적 영향: 사실적이고 지시를 따르는 이미지를 생성하는 능력이 향상되면 정교한 허위정보, 위조 문서, 오해를 유발하는 소셜 미디어 콘텐츠를 만드는 데 악용될 수 있다.Flow-OPD는 기반 모델의 safety filter를 계승하지만, 저자들은 오용을 완화하기 위해 digital watermarki를 통합할 것을 주장한다.