pi-Flow: Richtlinienbasierte Generierung in wenigen Schritten durch Imitationsdistillation

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

Wenige-Schritt-Diffusions- oder flussbasierte generative Modelle destillieren typischerweise einen geschwindigkeitsvorhersagenden Lehrer in einen Schüler, der einen direkten Weg zu entrauschten Daten vorhersagt. Diese Formatinkongruenz hat zu komplexen Destillationsverfahren geführt, die oft unter einem Qualitäts-Diversitäts-Kompromiss leiden. Um dies zu beheben, schlagen wir policy-basierte Flussmodelle (pi-Flow) vor. Pi-Flow modifiziert die Ausgabeschicht eines Schüler-Flussmodells, um eine netzwerkfreie Policy in einem Zeitschritt vorherzusagen. Die Policy erzeugt dann dynamische Flussgeschwindigkeiten in zukünftigen Teilschritten mit vernachlässigbarem Overhead, was eine schnelle und präzise ODE-Integration auf diesen Teilschritten ohne zusätzliche Netzwerkauswertungen ermöglicht. Um die ODE-Trajektorie der Policy an die des Lehrers anzupassen, führen wir einen neuartigen Imitations-Destillationsansatz ein, der die Geschwindigkeit der Policy entlang ihrer Trajektorie an die des Lehrers anpasst, indem ein standardmäßiger ell_2-Fluss-Anpassungsverlust verwendet wird. Durch einfaches Nachahmen des Verhaltens des Lehrers ermöglicht pi-Flow ein stabiles und skalierbares Training und vermeidet den Qualitäts-Diversitäts-Kompromiss. Auf ImageNet 256^2 erreicht es einen 1-NFE-FID von 2,85 und übertrifft damit MeanFlow derselben DiT-Architektur. Auf FLUX.1-12B und Qwen-Image-20B bei 4 NFEs erzielt pi-Flow eine deutlich bessere Diversität als state-of-the-art Wenige-Schritt-Methoden, bei gleichbleibender Lehrer-Qualität.

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.