Source-linked AI summary

PE-OPSD: Internalizing Prompt Enhancement into Flow-matching Models via On-Policy Self-Distillation

Mingfeng Lin, Chengfei Cai, Lin Xu, Chengqian Ma, Yuxiang Wei, Liang Han

arXiv:2609.36638v1cs.LGcs.CV

TL;DR

Short, underspecified prompts leave a gap between user instructions and the detailed conditions text-to-image models benefit from, while inference-time prompt rewriting adds latency and computation. PE-OPSD uses an enhanced-prompt teacher to supervise a raw-prompt student on its own trajectories, then removes the teacher and PE at inference. Across three main backbones, it improves aggregate prompt fidelity over Base and exceeds evaluated alternatives while retaining base-model latency.

  • Problem

    Concise prompts leave visual choices unspecified, while generative models benefit from detailed textual conditions; existing prompt enhancement also remains an inference-time operation with added latency and computation.

    Source: §1 · p. 1
  • Method

    PE-OPSD uses enhanced prompts as privileged teacher conditions to supervise vector fields at states visited by a raw-prompt student, without requiring additional text–image pairs.

    Source: §1 · p. 2
  • Results

    PE-OPSD improves aggregate prompt fidelity by 13.35%-16.07% over Base across three main backbones and exceeds inference-time PE, SFT, and off-policy distillation while retaining base-model latency.

    Source: §5 · p. 10
  • Takeaways & Limitations

    After training, the PE and teacher are removed, allowing the student to generate directly from raw prompts at base-model inference latency.

    Source: §1 · p. 2, §5 · p. 10

    Performance depends on the selected PE’s quality and semantic faithfulness, so inaccurate or overly specific rewrites may introduce undesirable supervision.

    Source: H LIMITATIONS AND FUTURE WORK · p. 24

Abstract

from arXiv · show

Text-to-image users often provide concise and underspecified prompts, whereas generative models benefit from detailed textual conditions for reliable instruction following. Existing systems bridge this gap with Prompt Enhancers (PEs) that rewrite raw prompts at inference time, introducing additional latency and leaving prompt elaboration external to the generator. We instead view enhanced prompts as privileged training information and ask whether their benefits can be internalized. We propose Prompt-Enhanced On-Policy Self-Distillation (PE-OPSD) for text-to-image flow-matching models. During training, a raw-prompt student follows its own generation trajectory, while an enhanced-prompt teacher provides vector-field targets at the states visited by the student. This on-policy supervision distills the behavior induced by enhanced prompts into the raw-prompt student without requiring additional text--image pairs. At inference, both the PE and teacher are removed, and the student generates directly from raw prompts. Across multiple model families, PEs, and benchmarks, PE-OPSD achieves the strongest aggregate prompt fidelity among the evaluated baselines, yields positive aggregate visual appeal gains, and retains the base-model inference efficiency.

1 INTRODUCTION

Concise, underspecified prompts leave visual choices open even though generative models benefit from detailed conditions, creating a gap between user input and reliable instruction following. PE-OPSD treats enhanced prompts as training-time privileged information, transferring their generation behavior to a raw-prompt student that avoids inference-time rewriting.

1 INTRODUCTION

  • Concise prompts often leave pose, lighting, composition, and style unspecified, while detailed textual conditions can help models produce coherent outputs and follow constraints reliably.This mismatch makes reliable intent understanding difficult for short, underspecified prompts.Source: §1 · p. 1
  • Inference-time Prompt Enhancers rewrite raw prompts into longer descriptions, but leave prompt enhancement outside the generator and add latency and computation.Longer prompts also increase text-processing and conditioning overhead.Source: §1 · p. 1
  • The paper frames enhanced prompts as privileged information: they elaborate raw prompts while preserving explicit semantic constraints, but are used only to define the teacher condition.The student learns to reproduce the corresponding generation behavior from the raw prompt alone.Source: §1 · p. 2
  • PE-OPSD uses an enhanced-prompt teacher to supervise vector fields at states visited by a raw-prompt student, transferring enhanced-prompt behavior without additional text–image pairs.The student follows its own trajectories and is conditioned only on the raw prompt; after training, both PE and teacher are removed.Source: §1 · p. 2
  • PE-OPSD outperforms inference-time PE, SFT, and off-policy distillation in aggregate prompt fidelity on all three main backbones, with positive aggregate visual-appeal changes and base-model inference latency.Figure 1 also reports higher GenEval scores than Vanilla and PE while preserving inference efficiency.Source: §1 · p. 2, Figure 1 · p. 2

2 RELATED WORK

Prior prompt-enhancement work improves text–image generation through external prompt rewriting, while on-policy distillation and privileged-information methods provide related training paradigms. PE-OPSD instead uses enhanced prompts as training-time privileged information to transfer their effect into a raw-prompt flow-matching generator.

Prompt enhancement for text-to-image generation

  • Promptist, BeautifulPrompt, and PromptEnhancer improve prompts externally, whereas PE-OPSD uses enhanced prompts only during training and distills their effect into the generator.The resulting generator can operate directly from raw prompts without inference-time prompt enhancement.Source: §2 · p. 3
  • Prompt rewriting can improve human preference, text–image relevance, attribute binding, and compositional relationships, but keeps enhancement in an external module.PE-OPSD shifts that role to training-time supervision rather than modifying the inference input.Source: §2 · p. 3

On-policy distillation

  • On-policy distillation supervises a student on its own policy’s samples, while PE-OPSD applies enhanced prompts as privileged teacher information without requiring inference-time enhancement.Unlike D-OPSD, PE-OPSD does not require paired image–text data or a generator that accepts image-conditioned inputs.Source: §2 · p. 3

3 METHODOLOGY

Prompt enhancement supplies richer conditioning that can improve compositional scores, while PE-OPSD treats enhanced prompts as training-only privileged information. It trains a raw-prompt student on its own rollout states using enhanced-prompt teacher targets, then retains the student for inference.

3.1 PRELIMINARIES

  • On-policy distillation evaluates the teacher on samples from the current student distribution, while on-policy self-distillation derives asymmetric teacher and student views from the same base model.The teacher may be a stop-gradient base model or an EMA copy of the student.Source: §3.1 · pp. 3–4

3.2 FROM PROMPT ENHANCEMENT TO PRIVILEGED INFORMATION

  • With the same frozen generator, enhanced prompts improve reported compositional scores while substantially increasing prompt length.The figure compares raw and enhanced prompts, with enhanced prompts providing denser, more detailed textual information.Source: §3.2 · p. 4, Figure 3 · p. 4
  • PE-OPSD treats enhanced prompts as textual privileged information: the teacher receives the elaborated prompt while the student sees only the raw prompt.The approach distills generation dynamics rather than reconstructing the enhanced prompt, and requires neither paired target images nor changes to the conditioning interface.Source: §3.2 · p. 4

3.3 PE-OPSD: PROMPT-ENHANCED ON-POLICY SELF-DISTILLATION

  • PE-OPSD trains the student on raw-prompt rollouts and matches enhanced-prompt teacher predictions at the states visited by that student.The student and teacher share an architecture and start from the same pretrained parameters.Source: §3.3 · pp. 4–5
  • For deterministic flow trajectories, PE-OPSD directly regresses teacher vector fields and unifies velocity, transition, and clean-latent matching as weighted velocity mismatch.The variants share a pointwise velocity-space optimum but differ in timestep weighting.Source: §3.3 · p. 5
  • After training, only the student remains, generating directly from raw prompts without invoking the prompt enhancer or teacher.Prompt pairs are constructed offline, and the teacher is updated by EMA during training.Source: §3.3 · pp. 5–6

4 EXPERIMENTS

The experiments compare PE-OPSD with base, inference-time enhancement, SFT, and off-policy distillation across multiple models, benchmarks, and training choices. PE-OPSD leads in aggregate prompt fidelity while retaining base-model inference latency, with further tests examining generalization and ablations.

4.1 EXPERIMENTAL SETUPS

  • The setup compares SD3.5-M, Z-Image, and Z-Image-Turbo against Base, Base+PE, SFT, and off-policy distillation, using shared training initialization, data, optimization, and steps.Evaluation covers GenEval and GenEval2 metrics for prompt fidelity and visual appeal, plus DPG-Bench, T2I-CompBench++, and EvalMuse for out-of-domain evaluation.Source: §4.1 · pp. 6–7

4.2 MAIN EXPERIMENTS

  • PE-OPSD improves aggregate prompt fidelity over Base by 13.35% on SD3.5-M, 16.07% on Z-Image, and 15.22% on Z-Image-Turbo, with positive aggregate visual-appeal changes.It achieves the highest GE, GE2GM, and IPF across all three backbones, outperforming SFT, off-policy distillation, and Base+PE.Source: Table 2 · p. 7, §4.2 · p. 7
  • On Z-Image-Turbo, PE-OPSD remains effective after post-training, improving both fidelity and visual appeal while SFT yields smaller fidelity gains and a 1.40% visual-appeal decline.This supports applicability to models that have already undergone post-training.Source: §4.2 · p. 7
  • On Z-Image, PE-OPSD retains Base latency and is 1.96× and 4.51× faster than PromptEnhancer-7B and PromptEnhancer-32B; human raters prefer it over Base in 66% of fidelity and 58% of appeal comparisons.The latency comparison holds under both evaluated PEs; the human preference rates are the highest among evaluated methods.Source: Figure 5 · p. 7, Table 3 · p. 8, §4.2 · p. 8

4.3 GENERALIZATION

  • Across larger models, alternative PEs, and out-of-domain benchmarks, PE-OPSD improves aggregate fidelity; it largely preserves base performance on detailed DPG-Bench prompts.It has the highest IPF and IVA with each of three tested PEs; on T2I-CompBench++ and EvalMuse, all three backbones improve Overall, with consistent spatial-relation gains on T2I-CompBench++.Source: §4.3 · pp. 8–9, Table 4 · p. 8, Table 5 · p. 8

4.4 ABLATION STUDIES

  • In Z-Image-Turbo, four rollout steps achieve the highest GE2 scores at about half the iteration time of eight steps, with nearly identical GE across tested settings.For training data, GenEval-only and GenEval2-only training specialize toward their respective benchmarks, whereas MixDataset gives the best balance across both.Source: Table 8 · p. 9, §4.4 · p. 9, Figure 7 · p. 9
  • Among timestep-weighting losses, µ-loss gives the highest GE2 scores while staying close to v-loss on GE; omitting CFG roughly halves training time and improves GE but lowers GE2.Detaching the unconditional branch partially reduces CFG's computational overhead; CFG is omitted by default for training efficiency.Source: Table 6 · p. 9, §4.4 · p. 9, Table 7 · p. 9

5 CONCLUSION

The conclusion presents PE-OPSD as a way to transfer enhanced-prompt behavior to a generator conditioned on raw prompts. The appendix explains its on-policy trajectory objective and deterministic matching surrogate, and describes how prompt pairs and benchmark data are prepared.

5 CONCLUSION

  • Across three main backbones, PE-OPSD improves aggregate prompt fidelity by 13.35%–16.07% over Base and exceeds inference-time PE, SFT, and off-policy distillation while retaining base-model latency.Alternative PEs, larger models, and out-of-domain benchmarks further support its applicability.Source: §5 · p. 10

A FROM TRAJECTORY KL TO DETERMINISTIC MATCHING

  • The appendix interprets the deterministic training objective through trajectory-level distillation, decomposing it into teacher matching at states induced by the raw-prompt student.The teacher is evaluated on student-visited states rather than states from its own enhanced-prompt trajectory.Source: A FROM TRAJECTORY KL TO DETERMINISTIC… · p. 13
  • The deterministic objective is not obtained by directly taking a zero-variance KL limit; it retains the shared-covariance Gaussian objective's pointwise optimizer through deterministic L2 matching.Rollout states and teacher predictions are detached during each optimization step, avoiding transition noise during training.Source: A FROM TRAJECTORY KL TO DETERMINISTIC… · p. 13

B.1 TRAINING DATA

  • GenEval and GenEval2 training prompts are combined into MixDataset, while evaluation prompts are never used for training.Prompt pairs are constructed offline by applying each PE to raw training prompts, with GPT-5.6 Sol outputs checked for consistency and rejected prompts regenerated.Source: §B.1 · p. 14

B.2 TRAINING CONFIGURATIONS.

The experiments use shared training settings within each backbone, with reduced-step student rollouts and backbone-specific evaluation configurations. The baselines differ in when enhancement is applied or whether enhanced-prompt behavior is transferred through images or distillation, and evaluations span in-domain and out-of-domain benchmarks.

B.2 TRAINING CONFIGURATIONS.

  • Training rollouts use 4, 10, or 14 denoising steps by backbone, while evaluation uses backbone-specific inference-step and CFG settings.All methods within a backbone share initialization, training data, and hyperparameters; ablations vary specified settings against the default configuration.Source: §B.2 · p. 14, Table 9 · p. 15

B.3 MODELS

  • Table 10 lists the models and checkpoints used in the experiments.The supplied passage identifies the table's scope but does not specify checkpoint names.Source: §B.3 · p. 15, Table 10 · p. 15

B.4 BASELINES

  • The baselines include raw-prompt Base, inference-time Base+PE, SFT on raw-prompt/image pairs, and off-policy distillation using teacher-generated states.SFT transfers enhanced-prompt behavior through offline-generated pseudo-target images; off-policy distillation differs from PE-OPSD in its rollout distribution.Source: §B.4 · p. 15

B.5 BENCHMARKS

  • In-domain evaluation covers GenEval and GenEval2, while out-of-domain evaluation covers DPG-Bench, T2I-CompBench++, and EvalMuse.The reported measures include benchmark-specific correctness or alignment scores, plus CLIP score, normalized PickScore, and aesthetics on the in-domain benchmarks.Source: §B.5 · p. 16

C HUMAN EVALUATION PROTOCOL

Human evaluation compares Base with four alternatives on all 800 GenEval2 test prompts using Z-Image-Turbo and GPT-5.6 Sol as the PE. Three professional assessors judge prompt fidelity and visual appeal in blinded, randomized pairwise comparisons.

Evaluation data and comparisons

  • The evaluation compares Base against Base+PE, SFT, off-policy distillation, and PE-OPSD on all 800 GenEval2 test prompts.Z-Image-Turbo is the generator and GPT-5.6 Sol is the PE.Source: C HUMAN EVALUATION PROTOCOL · p. 16

Evaluation interface and aggregation

  • Assessors see randomized left-right image pairs with method identities hidden and the raw prompt displayed, then judge prompt fidelity and visual appeal.Each criterion is judged independently by all three assessors, with each pair's preference determined by majority vote.Source: C HUMAN EVALUATION PROTOCOL · p. 16

Assessors and quality control

  • Three compensated professional assessors receive criterion-specific instructions and complete a qualification test before participating.They are informed about the task and potential exposure to generated content.Source: C HUMAN EVALUATION PROTOCOL · p. 16

D.1 TRAINING EFFICIENCY

PE-OPSD keeps base-model inference latency by generating directly from raw prompts, while training-cost comparisons show it avoids SFT's offline image-generation and storage requirements. EMA experiments find that teacher adaptation and temporal stability matter for later-stage supervision.

D.1 TRAINING EFFICIENCY

  • PE-OPSD reduces training time from 7.0 to 3.1 hours on SD3.5-M and from 28.4 to 13.6 hours on Z-Image, while SFT is slightly faster on Z-Image-Turbo.SFT's pseudo-target generation dominates total cost on SD3.5-M and Z-Image; PE-OPSD and off-policy distillation have identical training costs and avoid SFT's additional image-generation and storage requirements.Source: §D.1 · p. 16, Table 11 · p. 17

D.2 INFERENCE LATENCY

  • 1.52×–12.21× speedups over PromptEnhancer-7B and 3.31×–50.67× over PromptEnhancer-32B accompany PE-OPSD's retention of Base latency.The largest gains occur on Z-Image-Turbo and FLUX.2-klein; Table 12 compares latency and speedup across settings.Source: Table 12 · p. 17, §D.2 · p. 17

D.3 EFFECT OF EMA TEACHER

  • The frozen teacher reaches a lower performance ceiling than the default setting, while γ = 0.99 performs slightly below γ = 0.999.The frozen teacher is similar to the default early in training; images with γ = 0.99 remain visually coherent.Source: §D.3 · p. 17, Figure 8 · p. 18, Figure 9 · p. 18
  • With γ = 0.9, scores fall below Base and images show pronounced artifacts despite rapidly decreasing training loss.The authors attribute this to an overly responsive teacher becoming too tightly coupled to the student and providing insufficiently stable targets; they use γ = 0.999.Source: §D.3 · p. 17, Figure 8 · p. 18, Figure 9 · p. 18

D.4 DETAILED RESULTS OF SCALABILITY

On larger models, PE-OPSD achieves the strongest aggregate prompt-fidelity improvement across three backbones and improves aggregate visual appeal over Base, with a noted exception among the FLUX variants.

D.4 DETAILED RESULTS OF SCALABILITY

  • PE-OPSD has the highest aggregate prompt-fidelity improvement across all three larger backbones and leads on GE and GE2GM, while Base+PE is marginally better in visual appeal on the two FLUX variants.On QwenImage-2512, PE-OPSD has the best aggregate results for both objectives; Table 13 reports the larger-model results.Source: §D.4 · p. 18, Table 13 · p. 18

D.5 MORE ANALYSIS IN TABLE 4

Across three PEs, PE-OPSD consistently delivers stronger aggregate results than off-policy distillation, and its training remains stable across PEs and models. Its advantage over inference-time PE also varies with the teacher condition.

D.5 MORE ANALYSIS IN TABLE 4

  • Across three PEs, PE-OPSD improves IPF over off-policy distillation by 1.80–3.78 percentage points and IVA by 0.18–0.33 points.The two distillation methods share the teacher condition, objective, and optimization budget.Source: §D.5 · p. 19
  • PE-OPSD exceeds inference-time PE on both aggregate indices, with its largest fidelity margin over Base+PE for PromptEnhancer-7B.PromptEnhancer-7B has the weakest direct inference-time improvement; PromptEnhancer-32B gives the strongest fidelity supervision, while GPT-5.6 Sol yields the largest visual-appeal improvement.Source: §D.5 · p. 19

D.6 DIFFERENT PES

  • Across three PEs and three models, PE-OPSD loss decreases and stabilizes within 1,000 training steps.Figure 10 plots raw losses as faint lines and moving averages as bold lines.Source: §D.6 · p. 19, Figure 10 · p. 19
  • Figures 11–13 provide examples comparing raw prompts with enhanced versions from GPT-5.6 Sol, PromptEnhancer-7B, and PromptEnhancer-32B.The examples show the different PE outputs used in the experiments.Source: §D.6 · p. 19, Figure 11 · p. 19, Figure 12 · p. 20

E ADDITIONAL BENCHMARK RESULTS

The appendix details benchmark breakdowns, prompt-consistency filtering, enhanced-prompt instructions, and additional qualitative examples. The supplied passages describe these materials but do not report their detailed benchmark values.

E.1 GENEVAL DETAILS

  • GenEval breakdowns report Single Object, Two Object, Counting, Colors, Position, Attribute Binding, and Overall scores for GPT-5.6 Sol and other PEs on Z-Image.Tables 14 and 15 provide the detailed results.Source: §E.1 · p. 20, Table 14 · p. 21, Table 15 · p. 21

E.2 GENEVAL2 DETAILS

  • GenEval2 breakdowns cover Object, Attribute, Count, Position, and Verb, plus overall soft-tifa scores in GenEval2AM and GenEval2GM.Tables 16 and 17 include GPT-5.6 Sol and other PEs on Z-Image.Source: §E.2 · p. 20, Table 16 · p. 22, Table 17 · p. 22

E.3 EVALMUSE DETAILS

  • EvalMuse reports fine-grained scores for Attribute, Location, Color, Object, Material, Animal/Human, Food, Shape, Activity, Spatial, and Counting, alongside an Overall score.Table 18 marks the best Overall result in bold.Source: §E.3 · p. 23, Table 18 · p. 23

F.1 SYSTEM PROMPT FOR CONSISTENCY FILTERING

  • GPT-generated enhancements are screened by a consistency judge that checks whether critical prompt facts are explicitly preserved and returns a structured binary decision.The rules allow non-conflicting added composition, lighting, style, and background details, while warning against changed counts, identities, attributes, actions, or relations.Source: §F.1 · p. 23

F.2 SYSTEM PROMPTS FOR GENERATING ENHANCED PROMPTS.

  • The GPT-5.6 Sol enhancer is instructed to produce a hierarchical expanded caption while preserving specified objects, attributes, relations, rendered text, and intellectual-property subjects.Its prompt organizes descriptions from an overview through spatial and part-level detail to a concluding stylistic identification.Source: §F.2 · p. 24

G ADDITIONAL QUALITATIVE EXAMPLES

H LIMITATIONS AND FUTURE WORK

PE-OPSD depends on the quality and semantic faithfulness of the selected prompt enhancer, and inaccurate or overly specific rewrites may produce undesirable supervision. The paper also identifies training-compute reduction and extension beyond flow-matching text-to-image models as future directions.

H LIMITATIONS AND FUTURE WORK

  • The supervision signal depends on PE quality and semantic faithfulness; inaccurate or overly specific rewrites may introduce undesirable targets.Suggested future directions include confidence-aware filtering, agreement across PEs, and adaptive emphasis on reliable prompt elements.Source: H LIMITATIONS AND FUTURE WORK · p. 24
  • Reducing the shifted offline prompt-construction and post-training cost through selective timestep supervision or fewer teacher evaluations remains future work.The paper also proposes extending privileged-condition distillation beyond flow-matching text-to-image models.Source: H LIMITATIONS AND FUTURE WORK · p. 24
Loading 2609.36638v1…