Poly-OPD: Heterogeneous Multi-Teacher On-Policy Distillation for Capability-Selectable Flow Models

Thanks — I'll answer what we actually measured and be explicit about where we have no data, since several of your questions are about regimes we never ran.

0. VAE swap

We haven't tested it, so I won't guess at whether it works. What I can tell you is what the framework assumes, which may be the useful part: the teacher and student never exchange latents. The teacher only ever receives the student's decoded pixels, re-encodes them with its own encoder, refines, and decodes back to pixels; the comparison then happens in frozen DINOv2 feature space. Nothing in the pipeline reads the other model's latent coordinates.

Related evidence you may find relevant: in our perceptual-space ablation, replacing the DINOv2 feature comparison with a direct MSE in student latent coordinates collapses outright. Whatever else varies, the loss has to live in a shared, stable feature space rather than in either model's coordinates.

1. 6-step student

Not a regime we ran — our student uses a 20-step Euler sampler throughout and we never distilled into a few-step model, so I can't tell you how the noise-level matching behaves on a 6-step grid.

Our refinement-depth result is in (3), but note it separates by metric type rather than being uniformly "deeper is better", so I'd be careful extrapolating it to a coarse schedule.

2. Strongly compressed latent

We never hit this — both our VAEs are f8/16-channel, so we have no experience with a teacher painting detail the student's latent cannot hold, and no diagnostic to offer for spotting it.

Two facts that bear on it. Our loss is DINOv2 CLS cosine computed on images resized to 224px, so it is a semantic objective and largely blind to fine texture by construction; in our runs unreachable high-frequency detail never became visible in the loss. And, as above, the direct latent-space MSE variant collapsed, so we have no data on a loss that would be sensitive to latent capacity in the way you're worried about.

3. Shallow refinement and texture

This is the one place our ablation speaks fairly directly to your question, and it points the way you were hoping.

Varying the teacher refinement range with everything else fixed, preference metrics are largely insensitive: HPSv3 varies within 0.24 across the whole range, UR-Avg within 0.05, and UR-Avg is in fact marginally higher for the low-noise range [0,10] than for our main setting [15,20]. What moves sharply is composition — GenEval goes from 67.2 at [0,10] to 73.3 at [15,20]. Our reading is that compositional prompts need the teacher to intervene at high noise, where layout, counts and spatial relations can still be reorganized, whereas appearance quality is already served by a low-noise local polish. So if structure is explicitly what you don't want to touch, shallow refinement is not obviously the wrong regime in our data.

Caveat on what we measured: appearance was evaluated with ImageReward, HPSv3, PickScore and UnifiedReward, not with a texture-specific metric, and our objective is a 224px semantic loss. So "appearance metrics are insensitive to depth" is not the same claim as "texture is preserved at shallow depth", and we don't have the second one.

4. λ_WS

0.01.

5. Why 800 steps

For this paper it's a 500-step warm start followed by 800 OPD steps. The 800 is checkpoint selection on a plateau rather than a principled stopping rule.

On whether it's DINOv2 overfitting or style drift specifically: we did not run that experiment, so I have no evidence either way. Our ConvNeXt and SigLIP variants were run as perceptual-space ablations at fixed length — both stay competitive on preference metrics while costing about 7 GenEval points against DINOv2 — not as long-horizon drift tests, which is what your question would actually need.

6. Backprop through the decoder

The student VAE is kept in fp32 for the whole run even though training is otherwise fp16, and we force autocast off around every VAE call. We never attempted a bf16 decode backward, so I have no stability data on it.

The decode is gradient-checkpointed one image at a time — a loop over the batch with per-sample checkpointing, then concatenate — rather than checkpointing the batched decode. The teacher-side decode runs under no_grad.

One detail not in the paper: we clamp to [-1, 1] inside the checkpointed decode, so saturated pixels receive exactly zero gradient.

7. The teacher's noise

Fresh every iteration — new noise at each step, independent of the student's rollout noise, and the student's rollout also starts from fresh initial noise each step. We never tried fixing it, so I can't tell you what it costs.

Worth separating from the stochastic-rollout result you're referring to: that concerns noise injected into the student's own sampling trajectory during rollout, which is a different mechanism from the teacher's SDEdit noise. It doesn't carry over.

8. Teacher CFG

For FLUX.1-dev we pass guidance 3.5, but this is FLUX-dev's distilled guidance embedding — a single forward pass, no negative branch — so the usual classifier-free double-pass mechanism isn't in play. Our Z-Image teacher does use true CFG at 5.0.

On inheritance: against the SD3.5-Medium baseline, ImageReward rises from 0.922 to 1.168 and HPSv3 from 9.341 to 11.354, with PickScore essentially unchanged (0.869 vs. 0.866). So whatever the teachers' guidance imparts did not show up as a regression on preference metrics. We didn't run a teacher-guidance sweep, so that's an observation from the main runs rather than a controlled result.

9. Batch size 1, and EMA

Nothing in our loss uses batch statistics — the feature-space comparison is per-sample and then averaged — so there's no structural obstacle. We ran an effective batch of 28 (4 samples per GPU across 7 GPUs, no gradient accumulation) and never tried batch 1. One thing about our loop that would matter there: we draw a single refinement depth per step for the whole batch, and with no accumulation there is nothing averaging over that draw.

The reported model is EMA weights — checkpoints are saved with the EMA parameters swapped in, decay 0.999, updated every 8 optimizer steps.

10. Code release

Yes — the code is going through internal review at our company right now, and we'll open-source it as soon as that clears.

Hope this helps, and good luck with the VAE swap — I'd be curious to hear how the warm-start stage behaves in that setting.