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는 학생 흐름 모델의 출력 계층을 수정하여 한 타임스텝에서 네트워크 없는 정책을 예측하도록 합니다. 이 정책은 미래의 하위 단계에서 동적 흐름 속도를 생성하며, 이는 추가 네트워크 평가 없이도 빠르고 정확한 ODE 통합을 가능하게 합니다. 정책의 ODE 궤적을 교사의 궤적과 일치시키기 위해, 우리는 새로운 모방 증류 접근법을 도입했습니다. 이 접근법은 정책의 궤적을 따라 정책의 속도를 교사의 속도와 일치시키기 위해 표준 ell_2 흐름 매칭 손실을 사용합니다. 단순히 교사의 행동을 모방함으로써, pi-Flow는 안정적이고 확장 가능한 훈련을 가능하게 하며 품질과 다양성 간의 트레이드오프를 피합니다. ImageNet 256^2에서, pi-Flow는 1-NFE FID 2.85를 달성하여 동일한 DiT 아키텍처의 MeanFlow를 능가합니다. 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.