FlowCTS: On-policy Continuous Trajectory Supervision of Flow Models

2026-07-27Machine Learning

Machine LearningComputer Vision and Pattern Recognition
AI summary

The authors introduce Flow Continuous Trajectory Supervision (FlowCTS) as a new way to improve on-policy distillation (OPD), especially for flow models, which was less studied before. Their method works by aligning the student and reference model trajectories starting from the same state, using a mathematical way to match velocity fields over time. They find that FlowCTS performs better and faster than traditional KL-based OPD and also beats some reinforcement learning baselines on tasks like GenEval and OCR. Additionally, they identify a timing mismatch problem in previous methods and show that FlowCTS works well even outside the usual on-policy setting.

on-policy distillationflow modelstrajectory supervisionvelocity matchingKL divergencereinforcement learningGeneEvaloptical character recognition (OCR)stochastic differential equations (SDE)supervised fine-tuning
Authors
Kaiyang Ye, Yuan Ge, Junxiang Zhang, Bei Li, Ziming Zhu, Haishu Zhao, Xiaoqian Liu, Chenglong Wang, Jingbo Zhu, Zhengtao Yu, Tong Xiao
Abstract
While on-policy distillation (OPD) effectively addresses sparse rewards and exposure bias in large language model post-training, its extension to flow models remains underexplored. To this end, we propose Flow Continuous Trajectory Supervision (FlowCTS), which matches subsequent student and reference trajectories initialized from the same student-visited state. Using the integral relation between trajectories and velocity fields, we derive a temporally weighted velocity-matching upper bound and discretize it into practical objectives parameterized by the number of supervision steps. Under a multi-reference setup, single-state FlowCTS-OPD outperforms vanilla KL-based OPD with faster convergence. FlowCTS-OPD improves GenEval from 0.90 to 0.93, OCR from 0.90 to 0.92, and PickScore from 22.75 to 23.06, while outperforming a mixed-reward RL baseline across all target metrics. Further analysis reveals a clear temporal supervision mismatch in vanilla KL-based OPD arising from its auxiliary SDE transition kernels. Beyond on-policy setting,FlowCTS also consistently outperforms vanilla SFT , particularly on OCR, while increasing supervision steps exhibit a trade-off between richer trajectory information and greater optimization difficulty.