pi-Flow: Geração Baseada em Políticas em Poucos Passos via Distilação por Imitação

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

Resumo

Modelos generativos baseados em difusão ou fluxo com poucos passos geralmente destilam um professor que prevê a velocidade em um aluno que prevê um atalho em direção aos dados sem ruído. Essa incompatibilidade de formato levou a procedimentos de destilação complexos que frequentemente sofrem com uma troca entre qualidade e diversidade. Para resolver isso, propomos modelos de fluxo baseados em política (pi-Flow). O pi-Flow modifica a camada de saída de um modelo de fluxo aluno para prever uma política sem rede em um passo de tempo. A política então produz velocidades de fluxo dinâmicas em subpassos futuros com sobrecarga insignificante, permitindo uma integração rápida e precisa de EDO nesses subpassos sem avaliações extras de rede. Para alinhar a trajetória da EDO da política à do professor, introduzimos uma nova abordagem de destilação por imitação, que alinha a velocidade da política à do professor ao longo da trajetória da política usando uma perda padrão de correspondência de fluxo ell_2. Ao simplesmente imitar o comportamento do professor, o pi-Flow permite um treinamento estável e escalável e evita a troca entre qualidade e diversidade. No ImageNet 256^2, ele atinge um FID de 1-NFE de 2,85, superando o MeanFlow da mesma arquitetura DiT. No FLUX.1-12B e Qwen-Image-20B com 4 NFEs, o pi-Flow alcança uma diversidade substancialmente melhor do que os métodos state-of-the-art de poucos passos, mantendo a qualidade no nível do professor.

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.