pi-Flow: Generazione Basata su Politiche in Pochi Passi tramite Distillazione per Imitazione

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

October 16, 2025

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

cs.AI

Abstract

I modelli generativi basati su diffusione o flusso con pochi passi tipicamente distillano un insegnante che predice la velocità in uno studente che predice una scorciatoia verso dati denoisati. Questo disallineamento di formato ha portato a procedure di distillazione complesse che spesso soffrono di un compromesso qualità-diversità. Per affrontare questo problema, proponiamo modelli di flusso basati su policy (pi-Flow). pi-Flow modifica lo strato di output di un modello di flusso studente per predire una policy senza rete in un singolo passo temporale. La policy produce poi velocità di flusso dinamiche nei sottopassi futuri con un overhead trascurabile, consentendo un'integrazione ODE rapida e accurata su questi sottopassi senza valutazioni aggiuntive della rete. Per allineare la traiettoria ODE della policy a quella dell'insegnante, introduciamo un nuovo approccio di distillazione per imitazione, che allinea la velocità della policy a quella dell'insegnante lungo la traiettoria della policy utilizzando una perdita standard di flusso ell_2. Semplicemente imitando il comportamento dell'insegnante, pi-Flow consente un addestramento stabile e scalabile ed evita il compromesso qualità-diversità. Su ImageNet 256^2, raggiunge un FID a 1-NFE di 2.85, superando MeanFlow della stessa architettura DiT. Su FLUX.1-12B e Qwen-Image-20B a 4 NFE, pi-Flow ottiene una diversità sostanzialmente migliore rispetto ai metodi all'avanguardia con pochi passi, mantenendo una qualità pari a quella dell'insegnante.

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.