pi-Flow: Beleidsgebaseerde generatie in enkele stappen via imitatiedistillatie

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

October 16, 2025

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

cs.AI

Samenvatting

Weinig-staps diffusie- of stroomgebaseerde generatieve modellen distilleren doorgaans een snelheid-voorspellend leraarmodel naar een studentmodel dat een kortere weg naar ontruiste data voorspelt. Dit formaatverschil heeft geleid tot complexe distillatieprocedures die vaak lijden onder een kwaliteit-diversiteit afweging. Om dit aan te pakken, stellen we policy-gebaseerde stroommodellen (pi-Flow) voor. pi-Flow past de uitvoerlaag van een student-stroommodel aan om een netwerkvrij policy te voorspellen op één tijdstap. Het policy produceert vervolgens dynamische stroomsnelheden op toekomstige substappen met verwaarloosbare overhead, waardoor snelle en nauwkeurige ODE-integratie op deze substappen mogelijk is zonder extra netwerkevaluaties. Om het ODE-traject van het policy af te stemmen op dat van de leraar, introduceren we een nieuwe imitatiedistillatiebenadering, die de snelheid van het policy afstemt op die van de leraar langs het traject van het policy met behulp van een standaard ell_2 stroommatcheringsverlies. Door simpelweg het gedrag van de leraar na te bootsen, maakt pi-Flow stabiele en schaalbare training mogelijk en vermijdt het de kwaliteit-diversiteit afweging. Op ImageNet 256^2 behaalt het een 1-NFE FID van 2,85, wat beter is dan MeanFlow van dezelfde DiT-architectuur. Op FLUX.1-12B en Qwen-Image-20B bij 4 NFEs bereikt pi-Flow aanzienlijk betere diversiteit dan state-of-the-art weinig-staps methoden, terwijl het kwaliteit op leraarniveau behoudt.

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.