pi-Flow : Génération en quelques étapes basée sur des politiques via la distillation par imitation

pi-Flow: Policy-Based Few-Step Generation via Imitation Distillation

October 16, 2025

papers.authors: Hansheng Chen, Kai Zhang, Hao Tan, Leonidas Guibas, Gordon Wetzstein, Sai Bi

cs.AI

papers.abstract

Les modèles génératifs basés sur la diffusion ou les flux en quelques étapes distillent généralement un enseignant prédictif de vitesse en un étudiant qui prédit un raccourci vers des données débruitées. Ce décalage de format a conduit à des procédures de distillation complexes qui souffrent souvent d'un compromis entre qualité et diversité. Pour résoudre ce problème, nous proposons des modèles de flux basés sur des politiques (pi-Flow). pi-Flow modifie la couche de sortie d'un modèle de flux étudiant pour prédire une politique sans réseau à un pas de temps. La politique produit ensuite des vitesses de flux dynamiques aux sous-étapes futures avec un surcoût négligeable, permettant une intégration rapide et précise des équations différentielles ordinaires (EDO) sur ces sous-étapes sans évaluations supplémentaires du réseau. Pour faire correspondre la trajectoire EDO de la politique à celle de l'enseignant, nous introduisons une nouvelle approche de distillation par imitation, qui aligne la vitesse de la politique sur celle de l'enseignant le long de la trajectoire de la politique en utilisant une perte de correspondance de flux standard ell_2. En imitant simplement le comportement de l'enseignant, pi-Flow permet un entraînement stable et évolutif et évite le compromis qualité-diversité. Sur ImageNet 256^2, il atteint un FID à 1-NFE de 2,85, surpassant MeanFlow de la même architecture DiT. Sur FLUX.1-12B et Qwen-Image-20B à 4 NFEs, pi-Flow obtient une diversité nettement meilleure que les méthodes de pointe en quelques étapes, tout en maintenant une qualité équivalente à celle de l'enseignant.

English

Few-step diffusion or flow-based generative models typically distill a velocity-predicting teacher into a student that predicts a shortcut towards denoised data. This format mismatch has led to complex distillation procedures that often suffer from a quality-diversity trade-off. To address this, we propose policy-based flow models (pi-Flow). pi-Flow modifies the output layer of a student flow model to predict a network-free policy at one timestep. The policy then produces dynamic flow velocities at future substeps with negligible overhead, enabling fast and accurate ODE integration on these substeps without extra network evaluations. To match the policy's ODE trajectory to the teacher's, we introduce a novel imitation distillation approach, which matches the policy's velocity to the teacher's along the policy's trajectory using a standard ell_2 flow matching loss. By simply mimicking the teacher's behavior, pi-Flow enables stable and scalable training and avoids the quality-diversity trade-off. On ImageNet 256^2, it attains a 1-NFE FID of 2.85, outperforming MeanFlow of the same DiT architecture. On FLUX.1-12B and Qwen-Image-20B at 4 NFEs, pi-Flow achieves substantially better diversity than state-of-the-art few-step methods, while maintaining teacher-level quality.