Poly-OPD: Heterogeneous Multi-Teacher On-Policy Distillation for Capability-Selectable Flow Models
Abstract
Leading open text-to-image models often carry complementary strengths: one may lead on preference-aligned aesthetics while another follows compositional instructions more faithfully. However, differences in their autoencoders and noise schedules make it difficult to transfer these strengths across models. In this paper, we present Poly-OPD, a framework that can consolidate complementary strengths of heterogeneous teachers into a single compact flow-matching student. To bridge the incompatible latent spaces of different teachers, Poly-OPD performs on-policy distillation through a pixel bridge. Each student-generated image is re-encoded by a selected teacher's encoder and refined from a noise level matched by magnitude under the teacher's noise schedule. The resulting target is further matched to the student in frozen DINOv2 space, enabling supervision across incompatible latent spaces. To retain complementary capabilities without cross-teacher interference, Poly-OPD uses a gradient compatibility diagnostic to organize its adapters: attention LoRA modules are shared across teachers, whereas feed-forward adapters remain teacher-specific. During distillation, a gap-aware curriculum devotes more training to compositional categories where the student still falls short of the teacher. As each gap narrows, training shifts toward categories with larger remaining gaps. By distilling FLUX.1-dev and Z-Image into a 2.5B SD3.5-Medium student, Poly-OPD improves GenEval from 67.3 to 73.3, surpassing both larger teachers, and raises DrawBench HPSv3 from 9.34 to 11.35, consolidating both strengths within a switchable model.
Community
Leading open text-to-image models often carry complementary strengths: one may lead on preference-aligned aesthetics while another follows compositional instructions more faithfully. However, differences in their autoencoders and noise schedules make it difficult to transfer these strengths across models. In this paper, we present Poly-OPD, a framework that can consolidate complementary strengths of heterogeneous teachers into a single compact flow-matching student. To bridge the incompatible latent spaces of different teachers, Poly-OPD performs on-policy distillation through a pixel bridge. Each student-generated image is re-encoded by a selected teacher's encoder and refined from a noise level matched by magnitude under the teacher's noise schedule. The resulting target is further matched to the student in frozen DINOv2 space, enabling supervision across incompatible latent spaces. To retain complementary capabilities without cross-teacher interference, Poly-OPD uses a gradient compatibility diagnostic to organize its adapters: attention LoRA modules are shared across teachers, whereas feed-forward adapters remain teacher-specific. During distillation, a gap-aware curriculum devotes more training to compositional categories where the student still falls short of the teacher. As each gap narrows, training shifts toward categories with larger remaining gaps. By distilling FLUX.1-dev and Z-Image into a 2.5B SD3.5-Medium student, Poly-OPD improves GenEval from 67.3 to 73.3, surpassing both larger teachers, and raises DrawBench HPSv3 from 9.34 to 11.35, consolidating both strengths within a switchable model.
Really cool papers, congrats! I'm trying something similar and have a few questions!=)
- Actually my case isn't really "big teacher → small student" like in your papers — I'm adapting one model to a completely
different VAE. The teacher is the original model with its native VAE, and the student is the same transformer being
retrained to live in a new, much more compressed latent space. So there's no capacity gap in the transformer at all, only
the latent space changes. Your framework seems to cover this by construction (incompatible latents is the whole point), but
you never test this "VAE swap" setup — do you see any reason it wouldn't work, or anything you'd do differently for it? - My student is a 6-step distilled model, not 20-40 steps like yours. Will the noise-level matching still work on such a
coarse schedule, with gradients only flowing through the last 1-2 steps? - My student's VAE has way stronger compression than f8 — it's f128 with 512 channels, so a 1024px image is only 64 tokens.
I worry the teacher will paint in fine detail that the student's latent just can't hold. Did you ever hit this, and how
would I even tell — does the OPD loss just plateau? - Your depth ablation is about composition, where deep re-noising wins. What about the opposite case — I mostly want to fix
textures and fine detail without touching the structure. Is shallow refinement (r around 1-5) worth trying, and how did you
evaluate texture quality, since GenEval doesn't see it? - What value of λ_WS did you use in stage 2? Couldn't find it in the paper.
- Why stop at 800 steps? Does it start overfitting to DINOv2 features (some kind of style drift) if you go longer?
- Backprop goes through the student's decoder, which for me is a big one. Did you keep it in bf16, use checkpointing
through it? Any stability issues from that? - The noise you add before the teacher's refinement pass — new random noise every iteration, or fixed? You showed
stochastic rollouts hurt, so I wonder if this matters too. - What CFG did the teacher (FLUX.1-dev) run during refinement? High guidance gives oversaturated images — doesn't the
student inherit that? - Does it work with batch size 1? At higher resolutions I can't fit more. And is the final model EMA weights or raw?
- Any plans to release the code? Would save me a lot of guessing :)
Thanks!
Get this paper in your agent:
hf papers read 2608.04349 Don't have the latest CLI?
curl -LsSf https://hf.co/cli/install.sh | bash Models citing this paper 0
No model linking this paper
Datasets citing this paper 0
No dataset linking this paper
Spaces citing this paper 0
No Space linking this paper
Collections including this paper 0
No Collection including this paper