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

Technical Summary: Poly-OPD

Problem Statement
Current open text-to-image models often exhibit complementary strengths: some excel in human preference-aligned aesthetics (e.g., FLUX.1-dev), while others strictly follow compositional instructions such as counting, spatial relations, and attribute binding (e.g., Z-Image). However, consolidating these capabilities into a single model is hindered by architectural heterogeneity. These specialists utilize different autoencoders, denoisers, and noise schedules, rendering their latent spaces mutually unintelligible. Existing distillation methods (e.g., score, velocity, or trajectory matching) require a shared coordinate system and are thus inapplicable. Furthermore, standard off-policy training on teacher-generated images suffers from a train-inference mismatch, while joint training on multiple capabilities often leads to destructive interference between conflicting gradient directions.

Methodology: Poly-OPD
The authors propose Poly-OPD (Poly-Objective On-Policy Distillation), a framework designed to distill heterogeneous, latent-incompatible teachers into a single, compact flow-matching student (specifically a 2.5B SD3.5-Medium model). The framework addresses three core challenges:

1. Heterogeneous On-Policy Distillation via a Pixel Bridge:
   To bridge incompatible latent spaces, Poly-OPD employs a "pixel bridge." The student generates a sample, which is decoded to pixel space (the only coordinate system shared by heterogeneous models). This image is then re-encoded by the active teacher's VAE encoder. The teacher refines this bridged latent from a noise level matched by magnitude rather than timestep index, producing a corrected target image. Supervision is applied not in the teacher's latent space, but in a frozen DINOv2 feature space (CLS embedding). This semantic feature space is invariant to model-specific pixel statistics, allowing the student to learn the semantic content of the teacher's correction without requiring compatible latents. This process is "on-policy" because the target is derived from the student's own generated sample, mitigating the train-inference mismatch.

2. Capability-Selectable Adapters with Gradient Compatibility:
   To prevent destructive interference between distinct capabilities (e.g., aesthetics vs. composition), Poly-OPD utilizes a modular adapter design. The authors introduce a gradient compatibility diagnostic to empirically determine which parameters should be shared and which should remain isolated. Measurements reveal that attention gradients are aligned across modes, whereas feed-forward network (FFN) gradients often conflict. Consequently, Poly-OPD shares a single attention LoRA across all teachers but isolates FFN adapters for each specific capability. This allows a single frozen backbone to host multiple capabilities, with inference-time switching requiring only an adapter swap.

3. Gap-Aware Adaptive Sampling:
   Recognizing that compositional subskills (e.g., counting, position) saturate at different rates, the framework employs a curriculum that allocates training budget based on the remaining teacher-student gap rather than raw prompt difficulty. A probe set is used to estimate the performance gap for each category. The sampling distribution is dynamically adjusted to focus training on categories where the student still lags significantly behind the teacher, annealing toward uniform sampling as gaps close.

Key Contributions

• Formulation: The paper formulates the problem of multi-capability consolidation from architecturally heterogeneous, latent-incompatible teachers, a setting where existing latent-space distillation fails.
• Framework: Poly-OPD introduces a pixel-bridge mechanism for on-policy supervision, a gradient-measurement-driven adapter sharing strategy, and a gap-aware sampling curriculum.
• Performance: The method successfully consolidates the strengths of FLUX.1-dev (12B) and Z-Image (6B) into a single 2.5B SD3.5-Medium student.

Experimental Results
Evaluated on standard benchmarks, the Poly-OPD student outperforms both larger teacher models in their respective domains:

• Composition (GenEval): The student achieves 73.3, surpassing Z-Image (69.4) and FLUX.1-dev (65.2). Significant gains are observed in structured categories like two-object composition (+14.93) and attribute binding (+8.25).
• Aesthetics (DrawBench): The student achieves an HPSv3 score of 11.35 and an ImageReward of 1.168, outperforming the SD3.5-Medium base and recovering a significant portion of the gap to the 12B FLUX.1-dev teacher.
• Ablation Studies: Removing the capability-selectable adapters (using a shared FFN) causes a drastic drop in performance (e.g., -11.0 GenEval points), confirming the necessity of isolating conflicting gradients. Similarly, removing the warm-start stage or the gap-aware sampling significantly degrades results, validating the proposed training pipeline.

Significance and Claims
The paper claims that Poly-OPD enables the consolidation of complementary capabilities into a single, switchable model without the memory and latency costs of deploying multiple large models. A key finding highlighted by the authors is that the distilled student surpasses its own teachers in specific metrics (e.g., GenEval and ImageReward). The authors attribute this "surpassing" capability uniquely to the on-policy nature of the training; unlike off-policy methods that are bounded by the distribution of teacher-generated samples, on-policy distillation allows the student to refine its own trajectory based on teacher corrections, potentially exceeding the quality of the raw teacher outputs. The framework demonstrates that heterogeneous distillation is feasible without shared latent spaces, provided supervision is anchored in a shared semantic feature space and training is managed through gradient-aware architecture and curriculum design.