Source-linked AI summary
Any-OPD: Heterogeneous On-Policy Distillation for Flow-Matching Models via Representation-Space Bridging
Siming Fu, Zheming Fu, Ruizhe He, Hualiang Wang, Jie Huang, Xiaoxiao Ma, Mingchen Zhong, Weihu Huang, Xiaoxuan He, Haojun Xu
TL;DR
Existing on-policy distillation assumes shared latent spaces and timestep grids, leaving heterogeneous flow-matching models unsupported. Any-OPD bridges arbitrary teacher–student pairs through frozen vision representations, noise-level alignment, and anchoring; a 2.5B student then surpasses its 12B teacher on several preference metrics.
Problem
Heterogeneous flow-matching models lack shared representations and comparable timestep grids, so existing on-policy distillation methods cannot directly supervise them.
Method
Any-OPD couples teacher and student through a frozen model-agnostic representation, aligns trajectories by noise level, and anchors the student on the teacher’s manifold before on-policy correction.
Results
The distilled 2.5B student overtakes its 12B teacher on Aesthetic Score, ImageReward, and PickScore, while improving all six DrawBench metrics over baseline.
Takeaways & Limitations
Any-OPD demonstrates on-policy distillation across model families, with a smaller student matching or surpassing its larger teacher on several preference metrics.
Takeaways & Limitations
Anchoring fixes the student’s distribution relative to the teacher but does not correct how the student behaves on states along its own imperfect inference trajectory.
Abstract
from arXiv · showhide
On-policy distillation, in which a teacher corrects samples that the student itself generates, presupposes that the two models speak the same language: identical VAE latents, matching architectures, and a common timestep grid. We ask what happens when none of this holds, as when the strongest teacher available and the student one wishes to deploy come from different model families, and find that the standard recipes have no answer: teacher latents cannot serve as targets in a foreign coordinate system, per-pixel losses against a teacher that stochastically re-draws local detail degenerate into blur or divergence, and timestep indices lose their meaning across mismatched schedules. We present Any-OPD, to our knowledge the first framework for on-policy distillation between arbitrary pairs of latent flow-matching generators. Any-OPD treats the teacher purely as a black-box sampler and connects the two models at exactly one point: a frozen, model-agnostic vision representation in which their independently decoded outputs are compared, sidestepping every assumption about latents, features, or architecture. Trajectory correspondence is recovered by matching continuous noise levels instead of step indices, and a brief anchoring phase, in which teacher samples are re-encoded through the student's own VAE, ensures the on-policy gradient measures sample quality rather than domain mismatch. Distilling the 12B FLUX.1-dev into the 2.5B SD3.5-Medium, Any-OPD lifts the student's PickScore from 0.846 to 0.884 and HPSv3 from 9.12 to 10.97, rivaling the teacher at a fifth of its size, where direct latent regression fails to train at all.
1 Introduction
Any-OPD enables on-policy distillation between heterogeneous latent flow-matching generators that share neither representation nor timestep grid. Applied from 12B FLUX.1-dev to 2.5B SD3.5-Medium, it raises PickScore from 0.846 to 0.884 and HPSv3 from 9.12 to 10.97.
- Motivation: Large flow-matching generators such as FLUX.1-dev are difficult to deploy, motivating distillation into smaller students trained on their own trajectories.FLUX.1-dev carries 12B parameters in its denoising transformer; on-policy supervision avoids exposure bias from purely offline imitation.
- Method: An offline anchoring stage re-encodes teacher outputs with the student’s VAE before on-policy training, closing the distribution-level gap.This leaves the on-policy objective focused on per-sample correction rather than mismatch between the teacher’s and student’s latent domains.
- Problem: Heterogeneous distillation fails because teacher and student lack both a shared representation and a comparable timestep grid.Teacher latents are invalid targets in the student’s coordinates, while solver timestep indices are not comparable across different noise schedules.
- Method: Any-OPD couples arbitrary latent flow-matching teachers and students through a frozen, model-agnostic representation and aligns trajectories by continuous noise levels.Each model decodes with its own VAE, and the resulting images are compared using frozen DINOv2 CLS embeddings rather than model internals or solver indices.
- Results: 0.884 PickScore and 10.97 HPSv3 are achieved by the 2.5B SD3.5-Medium student after distillation from the 12B FLUX.1-dev teacher.The reported improvements are from 0.846 to 0.884 on PickScore and from 9.12 to 10.97 on HPSv3; naive latent regression across the VAE boundary collapses.
2 Related Work
Diffusion and flow matching provide the foundation for modern text-to-image generation, while distillation methods reduce sampling cost by transferring behavior from many-step teachers to few-step students. Unlike predominantly offline approaches, on-policy distillation trains on student-generated trajectories to reduce exposure bias through teacher feedback.
- Diffusion and Flow Matching: Flow matching learns continuous transport from noise to data and supports efficient deterministic sampling for high-resolution text-to-image systems.The formulation underlies systems including SD3, FLUX, and Z-Image.
- Diffusion Distillation: Diffusion distillation compresses sampling cost by training few-step students to match many-step teachers or align with teacher output preferences.Progressive and consistency-based methods match teacher behavior, whereas distribution matching methods align output distributions; these approaches are typically offline.
- Diffusion Distillation: On-policy distillation reduces exposure bias by letting students learn from teacher feedback on trajectories they generate themselves.The approach originated in language-model distillation and has recently been extended to iterative diffusion and flow image generation.
3 Method
Any-OPD distills heterogeneous latent flow-matching generators by anchoring teacher samples in the student’s VAE coordinates, then refining on-policy student outputs through a black-box teacher projection. It aligns supervision by continuous noise levels rather than solver indices and compares independently decoded images in an external vision representation.
- Problem setup: Heterogeneous distillation fails because teacher latents, architectural features, pixels, and timestep indices are not directly compatible across model families.Teacher refinement can stochastically resynthesize local detail, while differing shift functions make index-aligned states represent incomparable noise regimes.
- Noise-level routing: Noise-level matching routes each correction through the tightest student solver segment whose noise range contains the teacher projection’s range, avoiding incomparable index alignment.The correction must reach every student step operating across the affected noise scales so the student can express the edit.
- Anchoring: Anchoring re-encodes teacher-generated images with the student’s encoder and trains with the standard flow-matching objective to close the distribution-level gap offline.This transfers the teacher manifold into the student’s coordinate system without defining a loss across the two latent spaces.
- Anchoring limitation: Anchoring fixes the student’s distribution but not its behavior on imperfect inference trajectories, so the residual per-sample correction requires on-policy training.The student observes states interpolated from correct endpoints during anchoring, rather than the states its own rollout reaches.
- On-policy projection: The on-policy stage projects each deployed student sample through teacher noise-and-denoise refinement, treating the teacher as a sampler-only stochastic projection onto its image manifold.The student receives zero gradient when its decoded image and the projected teacher image are equivalent in the frozen representation.
4 Experiments · 4.1 Experimental Setup
The experiments distill FLUX.1-dev into SD3.5-Medium despite simultaneous differences in VAE, architecture, and noise schedule. They use a two-phase training protocol with frozen student components except LoRA adapters, followed by standardized evaluation across image-quality and compositional benchmarks.
- 4.1 Experimental Setup: The student is SD3.5-Medium (2.5B), while the teacher is FLUX.1-dev (12B), with different VAEs, architectures, and noise schedules.The student uses an SD3 VAE, static shift 3.0, and guidance 4.5; the teacher uses a FLUX VAE, dynamic shifting, and guidance 3.5.
- 4.1 Experimental Setup: Training updates only rank-32 LoRA adapters with α = 64 in the student transformer, while all other components remain frozen.This isolates adaptation to the student transformer during distillation.
- 4.1 Experimental Setup: Anchoring runs for 400 steps across M = 10 noise levels using teacher targets generated at 512 × 512 with 50-step Euler sampling.Anchoring uses prompts from the Pick-a-Pic training set, while Pick-a-Pic test prompts are held out for validation.
- 4.1 Experimental Setup: On-policy training runs for 800 steps at 512 × 512 with NS = NT = 20.Both anchoring and on-policy phases use learning rate 10−4 and per-GPU batch size 4.
- 4.1 Experimental Setup: Teacher refinement budget affects transferable quality: blur in NT = 10 targets propagates into the student, whereas the comparison uses NT = 10 versus 20.The figure attributes this effect to the teacher’s denoising budget bounding transferable quality.
- 4.1 Experimental Setup: Evaluation uses 1024 × 1024 images under standard inference settings: 50-step Euler for the teacher and 40-step Euler for the student and Any-OPD.DrawBench reports Aesthetic Score, ImageReward, PickScore, HPSv3, UnifiedReward, and UnifiedReward2 over five images per prompt; GenEval and DPG-Bench are also included.
4.2 Main Results
Any-OPD enables the 2.5B student to outperform its 12B teacher on several quality metrics while improving all six DrawBench metrics, without changing architecture or sampling budget.
- Main Results: The 2.5B distilled student overtakes its 12B teacher on Aesthetic Score, ImageReward, and PickScore.It also closes most of the gap on HPSv3 and UnifiedReward scores.
- Main Results: The student improves all six DrawBench metrics over the baseline at unchanged architecture and sampling budget.This demonstrates gains without increasing the deployment model’s architecture or sampling cost.
- Main Results: The student can exceed its teacher because supervision targets the teacher’s projection of a student sample.The target combines the student’s composition with the teacher’s refinement.
4.3 Ablation Studies
Ablations show that stable heterogeneous OPD depends on representation-space supervision, converged anchoring, and sufficiently strong teacher refinement. Deterministic rollouts and adequate denoising budgets further preserve transferable quality, whereas latent regression, weak refinement, and rollout noise degrade training or targets.
- Objective: Only DINOv2 representation-space objectives train stably; latent MSE collapses early, while LPIPS improves initially before degrading.The comparison includes DINOv2 CLS, multi-layer CLS, and CLS+Patch variants.
- Anchoring: Anchor-400 dominates Anchor-100 and no anchoring throughout training, showing OPD should begin only after anchoring converges.Anchor-100 never catches Anchor-400 under the same OPD budget, so the gap reflects initialization quality rather than total compute.
- Refinement strength: Refinement strength r ∈ [0, 10) underperforms every other setting, while [10, 15) leads PickScore and ImageReward and [15, 20] wins DPG-Bench and HPSv3.The adopted range is [15, 20] as the best overall balance.
- Refinement steps: NT = 10 leaves visible blur in refined targets, whereas NT = 20 produces sharp targets and a sharper student.The teacher’s denoising budget bounds the quality that can transfer to the student.
4.4 Generalization: Swapping the Teacher
Any-OPD allows the teacher to be swapped without changing the student, loss, hyperparameters, or evaluation procedure. Replacing FLUX.1-dev with Z-Image improves every reported metric, including compositional benchmarks, while inheriting the teacher’s strengths.
- Teacher replacement: Z-Image replaces FLUX.1-dev as teacher while the student, loss, hyperparameters, and evaluation remain unchanged.Z-Image is a 6B model with its own VAE, shift 6, and guidance 5.0.
- Results: +0.80 GenEval and +0.95 DPG-Bench are among the improvements achieved by the student with Z-Image as teacher.The student improves on every reported metric.
- Teacher-profile inheritance: 74.90 GenEval for Z-Image corresponds to Any-OPD-Z inheriting the teacher’s compositional strength alongside perceptual gains.The gains follow Z-Image’s profile, which is itself strong on GenEval.
5 Conclusion
Any-OPD extends on-policy distillation to heterogeneous model pairs by replacing shared-representation and shared-timestep assumptions with model-agnostic mechanisms. Future work includes combining complementary teachers and extending the framework beyond image modalities.
- Any-OPD extends on-policy distillation beyond the homogeneous setting by addressing the shared-representation and shared-timestep assumptions broken by heterogeneous pairs.
- The formulation avoids model internals: the teacher acts as a sampler through noise-and-denoise projection, while both models meet in a frozen external representation.
- Limitations and Future Work: Future work includes combining complementary teachers and extending Any-OPD to other latent generative modalities, since the formulation is not specific to images beyond its feature extractor.