PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Zhiqi Li, Bo Zhu

arXiv:2609.37851v1cs.CVcs.LG

TL;DR

Few-step flow-map generators make long-range transport efficient, but their on-policy distillation is underexplored. FlowMap-OPD separates student-state collection from teacher–student distribution comparison and connects local supervision to the deployed map through flow–velocity consistency. In text-to-image experiments, it consolidates specialist teachers into one few-step student and surpasses multi-reward Flow-Map GRPO in task performance and convergence speed.

  • Problem

    On-policy distillation for few-step flow-map generators remains underexplored, despite their use of long-range transport to reduce sampling cost.

    Source: §1 · p. 2

  • Method

    FlowMap-OPD formulates comparisons over student-state marginals, separating rollout from supervision and using flow–velocity consistency to connect local supervision to the generated map.

    Source: §1 · p. 2, §3.1 · p. 4

  • Results

    In text-to-image experiments, the consolidated student matches or exceeds each specialist on its corresponding benchmark and surpasses multi-reward Flow-Map GRPO in task performance and convergence speed.

    Source: §1 · p. 2, §7 · p. 11

  • Takeaways & Limitations

    The experiments demonstrate consolidation of three specialist teachers into a single text-to-image student that retains few-step generation.

    Source: §7 · p. 11

    When a learned teacher map differs from the map induced by its velocity, reducing student consistency residual need not monotonically improve reproduction of the learned teacher map.

    Source: §4.5 · p. 7

Abstract

from arXiv · show

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.

1 Introduction

FlowMap-OPD separates how student states are collected from what teacher–student comparison is applied at those states. This expands few-step flow-map distillation beyond coupled local-transition comparisons and supports transfer across capacities and specialist teachers.

1 Introduction

  • Few-step flow-map generators use long-range transport to reduce sampling cost, but their on-policy distillation remains underexplored.Flow–velocity consistency motivates using different, compatible representations for rollout and teacher–student comparison.Source: §1 · p. 2
  • Rollout–kernel separation compares teacher and student over student-state marginals at selected comparison times, separating state acquisition from supervision.The framework characterizes conditions under which rollout replacement preserves the objective and expected update.Source: §1 · p. 2
  • The paper identifies instantaneous-velocity matching with independently tunable self-consistency as its most effective practical design.The design operates within a broader space that also includes flow-map and induced-velocity supervision.Source: §1 · p. 2
  • FlowMap-OPD consolidates three specialist teachers into one student with performance comparable to each specialist on its corresponding benchmark.The framework also improves generation quality in XL/2-to-B/2 distillation.Source: §1 · p. 2

2 Background

The background describes continuous-time generation through velocity-driven transport and its corresponding flow map. It contrasts local-transition OPD with FlowMap-OPD’s separate long-range rollout and teacher–student comparison.

2 Background

  • Continuous-time generative models transport noise toward data through a probability path governed by an instantaneous velocity field.The state evolves from t = 1 toward t = 0.Source: §2 · p. 2
  • A flow map transports a state along deterministic ODE dynamics and pushes its distribution from one time to another.The ODE flow map also represents the time marginals of the marginal-preserving SDE.Source: §2 · p. 2
  • FlowMap-OPD uses long-range map rollouts to collect states and separately compares teacher–student distributions using direct velocities, flow maps, or induced velocities.Its comparison uses a KL between Gaussian distributions with means constructed from the selected representation.Source: §2 · p. 3
  • Conventional diffusion and flow-matching sampling discretizes local dynamics into sequential steps, whereas few-step models directly parameterize long-range transport.The supplied background notes that diffusion and flow OPD use discretized stochastic dynamics to obtain tractable transition kernels.Source: §2 · pp. 2–3

3 Rollout–Kernel Separation for FlowMap Distillation

FlowMap-OPD separates student-state acquisition from teacher–student comparison by formulating the objective over state marginals and comparison times. It uses this separation to pair native flow-map rollouts with distinct local supervision kernels, subject to preserving the relevant marginals and time weighting.

3 Rollout–Kernel Separation for FlowMap Distillation

  • Rollout–kernel separation pairs a student rollout with a local comparison loss over its supplied state marginals.The comparison can match flow-map predictions, map-induced velocities, or instantaneous velocities with a consistency penalty.Source: §3 · p. 3, §3.2 · p. 4

3.1 From trajectories to state marginals

  • Trajectory-level OPD uses the same student transition both to generate the rollout and to define the transition compared with the teacher.This coupling motivates expressing the objective through state marginals and separating the two roles.Source: §3.1 · p. 3
  • The objective’s expectation depends on the student-generated state marginals and comparison-time distribution, rather than the full cross-time coupling of the rollout.The local loss is additive, so each teacher–student discrepancy can be evaluated from its query-state marginal.Source: §3.1 · pp. 3–4
  • Preserving state marginals at each query time preserves the objective for a fixed local comparison and time weighting, allowing rollout and optimization kernel to be chosen separately.Keeping the original comparison recovers the original objective; changing the comparison defines a new distillation objective.Source: §3.1 · p. 4

3.2 Optimizing with separated rollouts

  • A rollout can use few-step long-range flow-map updates or numerically integrate instantaneous velocity, provided it retains the required state marginals.The visited states become training queries for the separately specified optimization kernel.Source: §3.2 · p. 4
  • Changing the rollout preserves the objective and expected update when the required marginals and comparison-time sampling rules are preserved.The result assumes the local comparison and time weighting are fixed; integrability is required for equality of the expected local loss and update.Source: §3.2 · p. 4

4 Supervision Interfaces and Consistency

FlowMap-OPD links local velocity supervision to the deployed long-range flow map through consistency, while separating student-state acquisition from teacher–student comparisons. It supports several supervision interfaces paired with native few-step rollouts, with instantaneous-velocity supervision and independently tunable consistency performing strongest overall across the reported ImageNet reward settings.

4.1 Flow–velocity consistency

  • Flow–velocity consistency connects local velocity supervision to long-range transport, and exact consistency makes transport-error bounds depend only on velocity error.Approximate velocity matching and consistency can jointly control transport error through either the source or destination relation.Source: §4 · p. 4, §4.1 · p. 6

4.2 Native rollout families

  • Native few-step flow-map rollouts collect training states for fixed-endpoint consistency-model maps and two-time MeanFlow maps.The paper retains each model family’s native rollout for state collection rather than using multi-step velocity integration.Source: §4.2 · p. 6

4.3 Flow-map comparisons

  • Flow-map supervision compares stochastic student and teacher transition kernels with a conditional KL at states supplied by the selected rollout.The construction also describes a scaled zero-noise limit, not a KL between point masses.Source: §4.3 · pp. 6–7

4.4 Induced-velocity comparisons

  • Induced-velocity supervision applies the velocity-kernel KL to source- or destination-induced student predictions, thereby supervising the long-range map through consistency.The source kernel is evaluated at (x, t), while the destination kernel is evaluated at (ψθ t→r(x), r).Source: §4.4 · p. 7

4.5 Instantaneous velocity supervision with tunable consistency

  • Instantaneous-velocity supervision weights teacher-velocity error and consistency error separately, with λc ≥ 0 independently selectable and λc = 0 giving velocity-only supervision.The consistency term connects local velocity learning to the long-range flow map.Source: §4.5 · p. 7

4.6 Algorithm and implementation

  • Training repeatedly collects native student states, applies the chosen comparison to detached queries, and refreshes the states after updates.The rollout interval, velocity-kernel length, and anchor interval are chosen separately; ImageNet uses four steps and text-to-image uses five.Source: §4.6 · p. 8

4.5 Instantaneous velocity supervision with tunable consistency

  • Table 1 compares three separate XL/2-to-B/2 reward-transfer settings, with all students at epoch 300 using native four-step sampling; it reports MMD-teacher FID, classifier target log-probability, and DINO cosine similarity.These are different reward settings, not three scores of one student.Source: Table 1 · p. 8

5 Experiments

Across ImageNet and text-to-image experiments, FlowMap-OPD selects instantaneous-velocity supervision with tunable consistency and transfers multiple specialist capabilities to one student. In the two-task comparison, OPD outperforms mixed-reward Flow-Map GRPO in both task scores and training updates.

5 Experiments

  • The experiments use MeanFlow, compare XL/2 teachers with B/2 students across MMD, classifier, and DINO rewards, then apply the selected velocity-supervision family to text-to-image consolidation.The ImageNet comparison uses the same native rollout and optimization-step count across supervision families.Source: §5 · p. 9, §5.1 · p. 9

5.1 Design-space exploration on ImageNet

  • Instantaneous-velocity supervision with λc = 0.01 ranks first for MMD and classifier transfer and second for DINO, while DINO recall is 0.1160 versus 0.0225 for destination-induced supervision.It achieves MMD FID 18.23 and classifier log-probability −0.4720, improving on the base’s 22.68 and −0.5448.Source: §5.1 · p. 9

5.2 Specialist-teacher consolidation

  • In 300 steps, one student exceeds its corresponding specialist scores on all three benchmarks: GenEval 0.8580 vs 0.8454, OCR 0.8830 vs 0.8504, and PickScore 23.1502 vs 23.0772.The five-step student is evaluated at update 300.Source: §5.2 · p. 9, Table 2 · p. 10
  • At update 300, OPD reaches OCR 0.8629 and PickScore 23.1677 versus GRPO’s update-1,400 scores of 0.8158 and 22.9825, exceeding both by update 200.Figure 7 plots EMA task scores for OPD and equal-weight OCR/PickScore Flow-Map GRPO.Source: §5.2 · pp. 9–10, Figure 7 · p. 10

5.3 Consistency strength and teacher transfer

  • A small consistency penalty improves teacher transfer, but stronger regularization lowers the consistency residual at the expense of transfer quality; a smaller residual alone does not guarantee better transfer.The consistency sweep compares consistency and teacher-map error across λc.Source: §5.3 · p. 10, Figure 8 · p. 10, Table 3 · p. 10

6 Related Work

Prior work studies on-policy teacher feedback, flow-map learning, and the separation of sampling from learning objectives. This paper applies these ideas to few-step flow-map students, separating native long-range rollouts from compatible teacher-comparison representations.

6 Related Work

  • Policy distillation work varies sampling distributions and comparison objectives, while DiffusionOPD and Flow-OPD compare local conditional transitions [4] [13].Related approaches include GKD [1], MiniLLM [9], and DistiLLM [12], which address student-generated sequences, teacher divergences, or sample reuse.Source: §6 · p. 10
  • For few-step flow-map students, native sampling uses long-range transport while teacher supervision can compare local dynamics or a separately constructed flow-map transition.The paper characterizes when rollout changes preserve additive local OPD objectives and connects supervision to the deployed map through flow–velocity consistency.Source: §6 · p. 11
  • Flow-map work connects long-range transport to local dynamics, and post-training methods use induced instantaneous predictors or anchored stochastic transitions [2] [6] [10] [17] [18] [20].The paper uses these complementary interfaces in a common OPD formulation, with native flow-map rollouts supplying queries and velocity-space or flow-map kernels specifying comparisons.Source: §6 · p. 11
  • DiffusionNFT and DMSampler explore sampler–learner combinations for diffusion reinforcement learning; FlowMap-OPD extends that separation to teacher supervision on student-rollout states [21].Generation and supervision can use different compatible representations connected by flow–velocity consistency.Source: §6 · p. 11

7 Conclusion

This section summarizes FlowMap-OPD’s separation of student-state acquisition from distribution comparison, its use with different supervision targets and rollouts, and the reported ImageNet and text-to-image applications. It also limits the path-KL interpretation to conditional KL sums, noting that a Markov-chain interpretation requires the full state and query-distribution matching condition.

7 Conclusion

  • FlowMap-OPD separates student-state collection from distribution comparison, enabling native flow-map rollouts with flow-map, induced-velocity, or instantaneous-velocity supervision.The theory establishes when different rollouts preserve the training objective and connects local supervision to the deployed generator through flow–velocity consistency.Source: §7 · p. 11
  • The consolidated text-to-image student matches or exceeds each specialist on its corresponding benchmark while retaining few-step generation.Experiments also demonstrate transfer across model capacities on ImageNet.Source: §7 · p. 11

A.1 Sampled states, comparison times, and rollout distributions

  • The sampled objective is a local comparison averaged over queries drawn from rollout state marginals, retaining the state and comparison times needed by the loss.The formulation allows deterministic or history-dependent rollouts and does not require state densities.Source: §A.1 · pp. 13–14

A.2 Marginal formulation of the objective

  • Under the stated measurability and integrability assumptions, the objective depends on the query marginals rather than the full trajectory distribution.For the same local loss and query rules, integrating against coordinate state marginals gives the same objective as integrating over sampled paths.Source: §A.2 · p. 14

A.3 Discrete and continuous state marginals

  • Discrete marginal objectives converge to the continuous-time objective under the stated regularity and marginal-convergence conditions, while imperfect flow-map composition can make actual rollout marginals differ from direct-map marginals.ODE and compatible SDE dynamics have matching continuous marginals under the stated assumptions; finite-step discretizations introduce approximation error.Source: §A.3 · pp. 15–16

A.4 Rollout replacement and mismatch

  • Matching conditional state marginals preserves the expected local loss and gradient without requiring equal path distributions, but cross-time correlation can still affect estimator variance.Pairwise losses require matching pairwise query distributions, and batch-dependent normalization must be included in the comparison.Source: §A.4 · p. 17

A.5 Comparison chains and path KL

  • For conditional-KL comparisons, matching rollout marginals at the compared indices yields the comparison-chain path-KL identity without requiring equal rollout and comparison-chain path distributions.The result assumes the same initial distribution and finite expected unnormalized sum of conditional KLs.Source: §A.5 · p. 17
  • The theorem concerns the sum of conditional KLs, and data processing bounds mixed path-distribution KL by joint grid–path KL; this fixed-parameter value identity does not extend to full gradients.The gradient caveat applies when the query distribution is held fixed.Source: §A.5 · p. 18
  • An auxiliary comparison rule alone does not define a Markov chain; anchors, extra variables, or histories needed for transitions must be retained in the declared chain state.The query-distribution matching condition applies to that full state.Source: §A.5 · p. 18
  • Without full-state query-distribution matching, the conditional loss remains a valid local objective but lacks the path-KL interpretation; unequal Dirac transitions are treated as regression.Deterministic comparisons would otherwise have infinite KL.Source: §A.5 · p. 18

A.6 Gradients with fixed sampled states and times

The appendix specifies how fixed sampled states and times support expected-update matching, random-time estimation, and flow–velocity consistency, including bounds for transferring errors from velocity and consistency residuals to the generated map.

A.6 Gradients with fixed sampled states and times

  • Matching the complete query distribution matches expected updates for a shared update rule, while fixed-state differentiation requires integrable losses and derivative bounds.The appendix also distinguishes exact differentiation of the local loss from the stopped-JVP and stopped-target update conventions.Source: §A.6 · p. 18
  • A finite detached buffer approximates the population query distribution; differentiation is exact for its empirical objective, but population guarantees need population losses or a separate generalization argument.Within-rollout query correlation affects variance, while averages over independent sampled rollouts estimate the population objective.Source: §A.6 · p. 18

A.7 Random-time training and unbiased estimation

  • Uniformly selecting among five sorted source indices yields a mixture of a point mass at 1 and a uniform interior-time component, so uniform continuous-time weighting requires adjusting the query weights.A minimum-spacing rejection rule changes the time distribution; averaging all starts and sampling one uniformly have the same expected loss but different cost and variance.Source: §A.7 · p. 19

B Flow–Velocity Consistency and Error Bounds

  • Flow–velocity consistency connects the learned map and velocity through source and destination derivative identities, with regularity assumptions required for the flow identification results.The appendix uses consistency residuals and a consistent teacher map in its error-bound analysis.Source: §B.1 · pp. 19–20, B Flow–Velocity Consistency and Error… · p. 19

B.1 Flow–Velocity Consistency

  • Zero weighted population losses identify the teacher velocity and map only when residuals are continuous and query distributions have full support with positive weights.Vanishing residuals on a fixed finite grid alone does not establish the differential identities between grid nodes.Source: §B.1 · p. 20

B.2 Source-consistency error bound

  • Source-consistency residual and velocity-error bounds transfer to a bound on map error over the interval.The appendix also gives an alternative source-consistency bound when the source residual is uniformly bounded.Source: §B.2 · p. 21

B.3 Destination-consistency error bound

  • Destination consistency and velocity-error bounds imply a map-error bound of h e^(L_T h)(ε_c + ε_v), under the stated teacher Lipschitz condition.With uniform bounds along arrival curves from native source states, the appendix extends the per-interval result across a valid grid.Source: §B.3 · p. 21

B.4 Error propagation through native rollouts

The appendices bound deployment error from native rollouts and population training losses, then derive compatible local kernels and supervision variants for flow-map and fixed-endpoint consistency models.

B.4 Error propagation through native rollouts

  • Under the stated hypotheses, native-rollout deployment error is bounded gridwise and in W2 for finite-second-moment output distributions.The bound also applies to a deployment grid sampled independently when the stated integrability conditions hold.Source: §B.4 · p. 22

B.5 Population training losses and coverage

  • Covered population losses bound expected output error, while the transfer constants reflect off-trajectory coverage and mismatch with the frozen acquisition distribution.A finite buffer or randomized times alone do not establish bounded coverage ratios; omitted endpoint intervals need an added error term or separate boundary guarantee.Source: §B.5 · pp. 23–24

C.1 Schedule-compatible local velocity kernels

  • With shared covariance, schedule-compatible Gaussian kernels compare teacher and student local velocities through their mean difference.The displayed Gaussian KL is exact for the stated kernels; marginal preservation additionally requires compatible dynamics and exact integration.Source: §C.1 · p. 24

C.2 Anchored flow-map kernels

  • Anchored flow-map kernels yield a map-error bound when the stated Lipschitz and positivity condition holds; the zero-noise deterministic regression uses a scaled limit.At zero noise, unequal point masses have infinite KL.Source: §C.2 · p. 25

C.3 Induced source and destination supervision

  • Induced source or destination supervision can establish teacher-flow consistency under the stated conditions, while separate velocity and consistency penalties control the induced-source residual.Matching only their sum can permit finite-error cancellation.Source: §C.3 · p. 25

C.4 Fixed-endpoint consistency models

  • For a fixed-endpoint consistency model, teacher velocity serves as the closure field; a vanishing source residual with the identity endpoint yields the teacher flow map.The appendix also describes endpoint prediction with re-noising and an endpoint-anchor comparison.Source: §C.4 · pp. 25–26

D Gradient Implementation

The implementation specifies which student and teacher gradients are retained for JVP-based objectives, and notes an additional error when a finite-difference fallback is needed.

D Gradient Implementation

  • The primary update detaches the entire source JVP output, including its direction dependence, while average-velocity and diagonal-velocity predictions remain differentiable.Teacher parameters remain frozen; an exact destination-residual gradient retains the teacher-input derivative.Source: D Gradient Implementation · p. 26
  • For text-to-image coordinates, the source JVP direction is (vθ, 1, 1); induced-destination supervision stops the teacher target, including its input derivative.The mathematical objective permits differentiation through the teacher input.Source: D Gradient Implementation · p. 26
  • A one-sided finite-difference fallback adds discretization error beyond the theoretical residual when a required operator lacks JVP support.The fallback uses a fixed direction and a valid signed step η.Source: D Gradient Implementation · p. 26

E Additional ImageNet Results

The ImageNet appendix evaluates cross-capacity reward transfer across supervision configurations and sampling grids, and reports a supplementary consistency-model comparison.

E Additional ImageNet Results

  • The reward-transfer study compares XL/2 teachers adapted to MMD, classifier, and DINO rewards with B/2 students initialized from the same pretrained model.All 36 configurations are compared at epoch 300.Source: E Additional ImageNet Results · p. 26
  • Each reward setting compares twelve B/2 supervision configurations, including velocity, stochastic-map, deterministic-map, and induced-velocity recipes.The configurations use the stated shared training schedule and implementation choices.Source: E Additional ImageNet Results · pp. 26–27
  • Evaluation uses 5,000 generated images across 100 ImageNet classes, with four- and five-step results plus fixed-checkpoint evaluations at 4, 5, 8, and 16 steps.Figure 9 marks recorded evaluations, frozen teachers, and the common epoch-300 comparison; its metric axes use logarithmic or symmetric-log scales.Source: E Additional ImageNet Results · p. 27, Figure 9 · p. 27, Table 5 · p. 28
  • Source-induced supervision achieves the lowest FID among compared consistency-model recipes under both four- and five-step generation.This supplementary comparison distills an MMD-adapted XL/2 teacher into a B/2 student at epoch 150.Source: Table 7 · p. 28, E Additional ImageNet Results · p. 28

F Text-to-Image Consistency Analysis

The text-to-image analysis examines how consistency strength relates to teacher matching and task performance, alongside paired qualitative comparisons across base and specialist models.

F Text-to-Image Consistency Analysis

  • Stronger λc reduces consistency residuals, but large weights increase teacher mismatch and reduce task performance, trading teacher fit against consistency.Table 8 reports consistency and teacher-matching errors after 200 training steps, averaged over the preceding 20 steps at the last supervised interval.Source: F Text-to-Image Consistency Analysis · p. 29, Table 8 · p. 29

G Paired Qualitative Comparisons

  • Paired qualitative comparisons show the base model, three specialist teachers, and the student under shared prompts, initial noise, sampling grids, and guidance.The illustrative comparisons use the λc = 0 student.Source: G Paired Qualitative Comparisons · p. 30

G.1 Cross-task Specialist–Student Comparisons

The panels present paired examples of specialist capabilities and their consolidation in one student, comparing outputs across OCR, visual-preference, and compositional-generation tasks under shared prompts and sampling conditions.

G.1 Cross-task Specialist–Student Comparisons

  • Across the three tasks, panels combine examples to show specialist capabilities alongside their consolidation in one student.Cross-task comparisons include text rendering, visual preferences, and object relations.Source: §G.1 · p. 30, Figure 10 · p. 30, Figure 11 · p. 31

G.2 OCR Capability Comparisons

  • OCR examples compare the base, three specialist teachers, and the same student for shared prompts, initial noise, and few-step sampling.Paired examples include a movie poster, a seed packet, a classroom scene, a reception area, and a campsite.Source: §G.1 · pp. 30–32, Figure 13 · p. 33, §G.2 · p. 33

G.3 Visual Preference Comparisons

G.4 Compositional Generation Comparisons

Related papers

Loading 2609.37851v1…