PaperlayerView exact revision on arXiv

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.08063v5cs.CVcs.AI

TL;DR

Multi-task alignment in Flow Matching models is hindered by sparse scalar rewards and conflicting objectives. Flow-OPD uses dense on-policy distillation with manifold anchoring to consolidate specialized expertise, resolving cross-domain interference while preserving visual fidelity and producing teacher-surpassing results.

  • Problem

    Flow Matching models lack effective multi-task alignment methods for harmonizing heterogeneous objectives without sparse-reward seesaw effects and cross-domain interference.

  • Method

    Flow-OPD distills specialized teachers into one student through on-policy sampling, task routing, dense trajectory supervision, and task-agnostic manifold anchoring.

  • Results

    Flow-OPD consistently matches or surpasses specialized teachers across benchmarks, resolving cross-domain interference while achieving an optimal multi-task trade-off and teacher-surpassing cases.

  • Takeaways & Limitations

    Flow-OPD provides a scalable paradigm for generalist text-to-image models with consolidated composition and typography expertise and high visual fidelity.

  • Takeaways & Limitations

    Flow-OPD depends on teacher quality and requires architectural homogeneity between teacher and student for fine-grained step-wise supervision.

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 Introudction

Flow-OPD addresses reward sparsity and gradient interference in multi-task Flow Matching alignment by transferring On-Policy Distillation from LLMs to vision. Its two-stage framework separates expertise acquisition from model unification and improves performance and generalization over GRPO.

  • Motivation: Multi-task Flow Matching alignment is hindered by heterogeneous, conflicting feature representations and sparse scalar rewards that induce a zero-sum seesaw effect.These limitations make single-reward RL methods inadequate for harmonizing multiple objectives.
  • Motivation: On-Policy Distillation has enabled LLMs such as DeepSeek-V4 [9], Mimo v2 [20], and GLM-5 [19] to harmonize complex multi-domain capabilities by distilling specialized experts.This success motivates applying OPD to integrate diverse teacher strengths in Flow Matching models.
  • Method: Flow-OPD is the first proposed OPD paradigm for Flow Matching post-training, decoupling expertise acquisition from model unification to provide fine-grained supervision.The framework addresses failure modes identified in GRPO-based multi-task training, particularly reward sparsity and gradient interference.
  • Method: Manifold Anchor Regularization provides task-agnostic guidance to preserve global generative quality during Flow-OPD alignment.The method combines fine-grained supervision with a regularization mechanism intended to maintain generation quality across tasks.
  • Results: A 10-point improvement over the GRPO baseline is achieved across four mainstream benchmarks, while the unified student matches or surpasses specialized teachers in-domain and generalizes OOD.These results establish superior performance and generalization for Flow-OPD relative to GRPO.

2 Related Work

Prior work has applied reinforcement learning and preference optimization to align text-to-image models, while on-policy distillation addresses the mismatch between fixed supervision and evolving student trajectories. In language models, OPD has developed through frameworks and KL-based refinements targeting exposure bias and mode-seeking behavior.

  • RL for T2I Models: RL-based alignment for text-to-image generation evolved from diffusion policy optimization and preference alignment toward newer GRPO-style methods.DDPO, DPOK, and ImageReward/ReFL [25] optimize rewards for aesthetics, human preference, or text-image alignment, while Diffusion-DPO uses preference pairs.
  • On-Policy Distillation: On-Policy Distillation couples teacher supervision to the student’s evolving exploration space, unlike traditional offline distillation based on fixed datasets.The on-policy formulation dynamically adapts supervision to student trajectories.
  • On-Policy Distillation: In language models, OPD development includes GKD [34] for mitigating exposure bias and MiniLLM and DistiLLM [36] for refining mode-seeking behavior through Reverse and Skewed KL.These methods illustrate the rapid expansion of on-policy distillation techniques in the LLM domain.

3 Preliminaries

Flow Matching defines generation as an ODE transporting noise to data, while its discretized integration can be viewed as a Markovian denoising trajectory for reinforcement learning. On-policy distillation transfers teacher behavior on student-generated trajectories to reduce distribution shift and exposure bias.

  • Flow Matching Models: Flow Matching maps noise p0 to data pdata through the ODE dxt = vt(xt, t)dt, with the OT path xt = (1−t)x0+tx1 and constant velocity x1 − x0.The model vθ learns this constant velocity under the Optimal Transport formulation.
  • Flow Matching Models: Discretized ODE integration can be formulated as sequential Markovian denoising, connecting continuous generation dynamics to step-wise reinforcement-learning policy optimization.Each transition xt → xt+∆t is treated as a Markovian state step, defining a formal trajectory.
  • On-Policy Distillation: On-policy distillation minimizes student–teacher output divergence using teacher supervision on trajectories τ ∼ pθ(τ) generated by the student, mitigating distribution shift.For autoregressive models, the objective is formulated as reverse Kullback–Leibler divergence.
  • On-Policy Distillation: Single-reward GRPO severely compromises cross-task generalization, degrading capabilities on non-target metrics.This cross-task evaluation uses baseline setups that strictly follow the official Flow-GRPO implementation; on-policy alignment is described as suppressing exposure bias and improving robustness in interactive or iterative generation.

4 Motivation

Online GRPO can surpass offline SFT by exploring self-generated trajectories, but single-reward optimization damages orthogonal capabilities through gradient interference. Controlled experiments show that mixing scalar rewards is fundamentally unstable and unscalable because conflicting objectives are compressed into a zero-sum advantage.

  • Online Exploration: GRPO overcomes offline SFT’s static-data ceiling by sampling groups from its current policy, evaluating self-generated states with relative advantages, and discovering novel high-reward trajectories.Standard flow matching relies on offline reconstruction and cannot directly optimize non-differentiable preferences, whereas online exploration supplies policy-driven training signals.
  • Gradient Interference: Single-reward GRPO severely degrades orthogonal capabilities because sparse scalar rewards induce unconstrained gradient interference within the shared parameter space.Conflicting task gradients can have negative inner products, allowing optimization for one target to exploit unmonitored degrees of freedom and dismantle pretrained synergies.
  • Mixing Multiple Rewards: 5% GenEval degradation follows adding OCR after +GenEval, showing that stacking scalar rewards triggers catastrophic forgetting instead of stable multi-reward optimization.The experiment progressively stacks GenEval, OCR, PickScore, and DeQA on SD-3.5-M; the table reports capability degradation under multi-reward optimization.
  • Mixing Multiple Rewards: Scalar reward mixing creates a zero-sum game in which optimizing one objective can overwrite representations needed by another, such as geometric precision under aesthetic stylization.The passage identifies conflicting gradients as the mechanism making scalar reward mixing fundamentally unscalable.

5 Method: Flow-OPD

Flow-OPD aligns Flow Matching models through dense supervision on self-generated trajectories, combining expert teachers, cold-start initialization, task-specific routing, and Manifold Anchor Regularization. Its continuous formulation analytically replaces high-variance policy-gradient estimation with an equivalent vector-field regression objective while preserving the visual manifold.

  • Flow-OPD Overview: The framework trains domain-expert teachers, initializes the student through cold-start SFT or model merging, and distills expertise through dynamically routed online trajectories.SFT transfers specialized teachers’ knowledge distributions, whereas model merging superposes divergent teachers’ anisotropic priors into one parameter state.
  • Task-Specific Teacher Labeling: Hard task routing activates one domain expert per condition, providing a localized reference velocity field as the definitive target for each on-policy trajectory.The student samples its own distribution through stochastic SDE-based rollouts, while routing prevents inter-domain gradient interference.
  • Bridging OPD and Flow Matching: Flow-OPD bypasses high-variance policy-gradient estimation by directly minimizing a closed-form vector-field MSE that is mathematically equivalent to LLM-style Policy-Gradient OPD.The continuous flow formulation avoids log-probability computation and PPO surrogate bounds, reducing gradient variance to zero while uniting on-policy exploration with regression efficiency.
  • Manifold Anchor Regularization: Manifold Anchor Regularization uses a frozen aesthetic teacher’s vector field to constrain functional alignment to a high-quality visual manifold and avert aesthetic degradation.The regularizer translates the Reverse KL penalty into a time-weighted vector-field distance, preserving visual quality during multi-teacher optimization.

6 Experiments

Flow-OPD is evaluated across four alignment tasks against monolithic- and hybrid-reward GRPO baselines, consistently matching or surpassing specialized teachers while improving multi-task trade-offs and generalization. Ablations support cold-start initialization and MAR as mechanisms for robust alignment, capability preservation, and image-quality control.

  • Experimental setup: Evaluation follows Flow-GRPO splits across GenEval, OCR, PickScore, and DeQA, using official expert checkpoints except for a DeQA teacher trained with a 4:6 DeQA-to-PickScore reward blend.Baselines include single-reward GRPO and GRPO-Mix with GenEval:OCR:PickScore weights of 3:1:1.
  • Main results: Flow-OPD consistently matches or surpasses specialized teachers across GenEval, OCR, PickScore, and DeQA, resolving cross-domain interference and sparse-reward optimization bottlenecks.Table 2 compares compositional generation, visual text rendering, and image quality; the PickScore teacher’s GenEval score drops to 0.51, illustrating specialization interference.
  • Qualitative results: Flow-OPD achieves an optimal multi-task trade-off, combining prompt fidelity and visual aesthetics while succeeding on edge cases where all individual teachers fail.These cases motivate the reported Teacher-Surpassing phenomenon, hypothesized to arise from knowledge cross-pollination within the latent flow manifold.
  • Cold-start ablation: Cold-start initialization rapidly establishes a robust training foundation; SFT offers scalable heterogeneous-teacher extraction, while model merging best exploits homogeneous teachers without additional training.These alternatives are compared in the cold-start ablation.
  • Out-of-domain generalization: On T2I-CompBench, Flow-OPD achieves state-of-the-art out-of-domain compositional generalization, whereas multi-task GRPO exhibits catastrophic forgetting in capabilities such as shape rendering and 3D spatial reasoning.The comparison uses the identical cold-start baseline for both methods.
  • Manifold Anchor Regularization: MAR prevents reward-induced background mode collapse and semantic redundancy by anchoring optimization to a high-fidelity manifold, improving structural diversity, semantic adherence, image quality, and human-preference alignment.The constraint is task-agnostic; teachers such as DeQA provide diverse samples but can struggle with instruction following.

7 Conclusion · A More Details

Flow-OPD integrates on-policy distillation into Flow Matching models, using dense trajectory-level supervision and Manifold Anchor Regularization to address reward sparsity, gradient interference, and aesthetic degradation while consolidating composition and typography expertise. Additional details specify the hybrid-training setup and selected hyperparameters.

  • 7 Conclusion: Flow-OPD integrates on-policy distillation into Flow Matching models, replaces scalar rewards with dense trajectory-level supervision, and consolidates composition and typography expertise with an emergent “teacher-surpassing” effect.The framework is presented as addressing reward sparsity, gradient interference, and the seesaw effect of competing metrics.
  • A More Details: GRPO-mix uses a 3:1:1 epoch ratio for GenEval, OCR, and PickScore, with each data partition receiving rewards only from its corresponding reward model.Training followed Flow-GRPO’s data and reward configurations.
  • A More Details: Training ran for about 50 hours on four distributed nodes equipped with eight H800 GPUs each.This describes the reported distributed training configuration for the hybrid setup.
  • A More Details: The implementation uses sampling timestep T = 10, evaluation timestep T = 40, group size G = 24, noise level a = 0.7, and 512-resolution images.These settings are listed among the fixed GRPO hyperparameters.
  • A More Details: Manifold Anchor Regularization uses KL ratio β = 0.02, while LoRA uses α = 64 and r = 32.The passage specifies these MAR and parameter-efficient fine-tuning settings.
  • A More Details: For Qwenvl Score, the method adapts the Flow-GRPO prompt [16] and uses Qwen3-30B-A3B-Instruct-2507.The adapted prompt is shown in Fig. 6.

B More Results · B.1 Qualitative results · B.2 Comparison with DiffusionNFT

Flow-OPD’s qualitative results show precise content generation, superior image quality, coherent layouts, and stronger alignment with human preferences. Compared with DiffusionNFT [49], it avoids limitations including incompatibility with Classifier-Free Guidance (CFG).

  • B.1 Qualitative results: Flow-OPD delivers precise content generation alongside superior image quality and coherent structural layouts.
  • B.1 Qualitative results: Figures 7, 8, and 9 provide additional qualitative results for Flow-OPD.
  • B.1 Qualitative results: Flow-OPD demonstrates stronger alignment with human preferences.
  • B.1 Qualitative results: Flow-OPD bridges functional accuracy with aesthetic excellence.
  • B.2 Comparison with DiffusionNFT: DiffusionNFT [49] introduces online reinforcement learning that integrates reward feedback into the forward diffusion process.
  • B.2 Comparison with DiffusionNFT: DiffusionNFT [49] enables policy optimization during the noise-injection phase and achieves competitive benchmark scores.
  • B.2 Comparison with DiffusionNFT: DiffusionNFT [49] is fundamentally incompatible with Classifier-Free Guidance (CFG), bottlenecking its performance upper bound.
  • B.2 Comparison with DiffusionNFT: DiffusionNFT [49] also suffers from pronounced reward hac.

B.3 Failure Cases and Limitations

Flow-OPD remains constrained by the performance ceiling of its teacher models and by architectural homogeneity requirements between teachers and students. Specialized-teacher errors can propagate through dense supervision, introducing noise that limits student improvement.

  • Limitations: Flow-OPD’s performance ceiling is constrained by the capabilities of its teacher models.
  • Failure Cases: Specialized-teacher failures propagate through dense supervisory signals, introducing noise into distillation and hindering the student from surpassing the teacher ensemble.This failure occurs when teachers generate semantically incorrect images.
  • Limitations: Teacher and student models must share architectural homogeneity for the distillation process.

Broader Impact

Flow-OPD offers societal benefits through more reliable, coherent, and OCR-capable AI-generated content, while its stronger image-generation realism introduces risks of sophisticated disinformation and deceptive visual content. The authors therefore emphasize careful consideration and mitigation through measures including digital watermarki.

  • Broader Impact: Flow-OPD combines positive societal contributions with potential risks that require careful consideration.The framework is presented as robust for multi-task alignment in generative models.
  • Positive Societal Impacts: Improved layout coherence and OCR accuracy could benefit graphic design, educational content, and assistive technologies for visually impaired people.The Multi-Teacher paradigm is also described as promoting a more balanced optimization objective and mitigating winner-takes-all bias.
  • Negative Societal Impacts: Higher proficiency in generating realistic, instruction-following images could enable sophisticated disinformation, fake documents, and misleading social-media assets.Although Flow-OPD inherits its foundation model’s safety filters, the authors advocate integrating digital watermarki to mitigate misuse.
Loading 2605.08063v5…