Flow-matching models learn to improve prompts without extra steps
PE-OPSD: Internalizing Prompt Enhancement into Flow-matching Models via On-Policy Self-Distillation
Machine LearningComputer Vision and Pattern Recognition
Summary
Text-to-image AI often needs detailed prompts to create good pictures, but people usually give short prompts. To fix this, some methods rewrite prompts each time, making the process slower. The authors propose teaching the AI during training to understand and improve short prompts on its own, so it doesn't need extra rewriting later. This leads to better pictures from simple prompts without slowing down the generation.
What this means in practice
- •For ai developers: Build faster text-to-image generators that handle simple prompts with improved image quality without extra prompt rewriting.
- •For digital content creators: Use text-to-image tools that generate better visuals from brief prompts directly, streamlining creative workflows.
Authors
Mingfeng Lin, Chengfei Cai, Lin Xu, Chengqian Ma, Yuxiang Wei, Liang Han
Abstract
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.