pi-Flow: Генерация на основе политик за несколько шагов через имитационное дистиллирование

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

October 16, 2025

Авторы: Hansheng Chen, Kai Zhang, Hao Tan, Leonidas Guibas, Gordon Wetzstein, Sai Bi

cs.AI

Аннотация

Модели генерации на основе диффузии или потоков с малым количеством шагов обычно преобразуют учителя, предсказывающего скорость, в ученика, который предсказывает кратчайший путь к очищенным данным. Это несоответствие форматов привело к сложным процедурам дистилляции, которые часто страдают от компромисса между качеством и разнообразием. Чтобы решить эту проблему, мы предлагаем модели потоков на основе политик (pi-Flow). pi-Flow модифицирует выходной слой модели потока ученика, чтобы предсказывать политику, не требующую сети, на одном временном шаге. Затем эта политика генерирует динамические скорости потока на последующих подшагах с минимальными накладными расходами, что позволяет быстро и точно интегрировать ОДУ на этих подшагах без дополнительных вычислений сети. Чтобы согласовать траекторию ОДУ политики с траекторией учителя, мы вводим новый подход к дистилляции через имитацию, который согласует скорость политики со скоростью учителя вдоль траектории политики с использованием стандартной функции потерь для согласования потоков ell_2. Просто имитируя поведение учителя, pi-Flow обеспечивает стабильное и масштабируемое обучение и избегает компромисса между качеством и разнообразием. На ImageNet 256^2 модель достигает FID 2.85 при 1-NFE, превосходя MeanFlow с той же архитектурой DiT. На FLUX.1-12B и Qwen-Image-20B при 4 NFEs pi-Flow демонстрирует значительно лучшее разнообразие по сравнению с современными методами с малым количеством шагов, сохраняя при этом качество на уровне учителя.

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.