FlowMap-OPD improves on-policy training for fast image generation models

FlowMap-OPD: Rollout--Kernel Separation for On-Policy Distillation of Few-Step Flow-Map Generators

Computer Vision and Pattern RecognitionMachine Learning

Summary

Generating images or similar data often involves repeatedly applying a learned process many times, which can be slow. Some new methods create models that take only a few steps to produce high-quality results, but training these models alongside the original ones is hard. The authors developed FlowMap-OPD, a way to teach the faster models by separating how data states are collected from how differences are measured, making training more effective. They tested their approach on image datasets and found it works better and faster than previous methods, especially for combining multiple specialized tasks.

What this means in practice

  • For image generation engineers: Train efficient few-step models that produce high-quality images faster by improving how these models learn from longer, slower ones.
  • For machine learning platform teams: Integrate multi-task few-step flow-map generators that consolidate different image generation styles or rewards more effectively.

Authors

Zhiqi Li, Bo Zhu

Abstract

Few-step flow-map generators, including MeanFlow and consistency models, enable efficient sampling through long-range transport, yet their on-policy distillation remains underexplored. We introduce FlowMap-OPD, an on-policy distillation framework that separates student-state acquisition from teacher--student distribution comparison. A formulation based on state marginals establishes this separation, while flow--velocity consistency connects local supervision to the deployed long-range map. Within this framework, we develop flow-map, induced-velocity, and instantaneous-velocity distribution supervision, each paired with a separately specified native flow-map rollout. Cross-capacity ImageNet experiments across three teacher rewards identify instantaneous-velocity distribution supervision with independently tunable student consistency as the most effective choice. In text-to-image experiments, FlowMap-OPD demonstrates strong multi-specialist consolidation capabilities and surpasses multi-reward Flow-Map GRPO in task performance and convergence speed.