pi-Flow: Generación Basada en Políticas en Pocos Pasos mediante Distilación por Imitación

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

October 16, 2025

Autores: Hansheng Chen, Kai Zhang, Hao Tan, Leonidas Guibas, Gordon Wetzstein, Sai Bi

cs.AI

Resumen

Los modelos generativos basados en difusión o flujo de pocos pasos suelen destilar un profesor que predice la velocidad en un estudiante que predice un atajo hacia datos sin ruido. Esta discrepancia de formato ha llevado a procedimientos de destilación complejos que a menudo sufren de un equilibrio entre calidad y diversidad. Para abordar esto, proponemos modelos de flujo basados en políticas (pi-Flow). pi-Flow modifica la capa de salida de un modelo de flujo estudiantil para predecir una política sin red en un paso de tiempo. La política luego produce velocidades de flujo dinámicas en subpasos futuros con un costo insignificante, permitiendo una integración rápida y precisa de EDO en estos subpasos sin evaluaciones adicionales de la red. Para igualar la trayectoria de EDO de la política con la del profesor, introducimos un enfoque novedoso de destilación por imitación, que iguala la velocidad de la política con la del profesor a lo largo de la trayectoria de la política utilizando una pérdida estándar de coincidencia de flujo ell_2. Al simplemente imitar el comportamiento del profesor, pi-Flow permite un entrenamiento estable y escalable y evita el equilibrio entre calidad y diversidad. En ImageNet 256^2, alcanza un FID de 1-NFE de 2.85, superando a MeanFlow de la misma arquitectura DiT. En FLUX.1-12B y Qwen-Image-20B a 4 NFEs, pi-Flow logra una diversidad sustancialmente mejor que los métodos de pocos pasos más avanzados, manteniendo la calidad a nivel de profesor.

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.