基于表征自编码器的 MeanFlow Transformer
MeanFlow Transformers with Representation Autoencoders
将 MeanFlow 的潜空间训练迁移到 DINO 等预训练视觉编码器的表征自编码器潜空间,并给出稳定训练的两阶段蒸馏方案,为视觉生成模型提供更高效、无需引导的少步采样路径。
原文依据 we develop an efficient training and sampling scheme for MF in the latent space of a Representation Autoencoder (RAE)
摘要译文
由模型翻译 · 术语统一 · 请对照原文核对关键结论MeanFlow(MF)是一种受扩散启发的生成模型,通过学习从噪声到数据的长跳跃来实现高效的少步生成。实践中,它常被用作潜空间 MF,利用预训练的 Stable Diffusion 变分自编码器(SD-VAE)进行高维数据建模。然而,MF 的训练仍然计算量大且常常不稳定。推理时,SD-VAE 解码器主导了生成成本,且 MF 依赖复杂的引导超参数进行类别条件生成。在本工作中,我们为 MF 在表征自编码器(RAE)的潜空间中开发了一种高效的训练与采样方案,其中预训练的视觉编码器(如 DINO)提供语义丰富的潜变量,并搭配轻量级解码器。我们观察到,在 RAE 潜空间中进行朴素的 MF 训练会遭遇严重的梯度爆炸。为了稳定并加速训练,我们采用一致性中训练进行轨迹感知初始化,并使用两阶段方案:先从预训练的流匹配教师模型蒸馏以加速收敛并降低方差,随后可选地使用单点速度估计器进行自举阶段,以进一步减小与理想平均流的偏差。该设计消除了对引导的需求,简化了训练配置,并减少了训练和采样中的计算量。实验上,我们的方法在 ImageNet 256 上实现了 1 步 FID 2.03,优于原始 MF 的 3.43,同时将采样 GFLOPS 降低 38%,总训练成本降低 83%。我们进一步将方法扩展到 ImageNet 512,以所有基线中最低的 GFLOPS 实现了具有竞争力的单步 FID 3.23。
英文原文摘要 1702 字符
MeanFlow (MF) is a diffusion-motivated generative model that enables efficient few-step generation by learning long jumps directly from noise to data. In practice, it is often used as a latent MF by leveraging the pre-trained Stable Diffusion variational autoencoder (SD-VAE) for high-dimensional data modeling. However, MF training remains computationally demanding and is often unstable. During inference, the SD-VAE decoder dominates the generation cost, and MF depends on complex guidance hyperparameters for class-conditional generation. In this work, we develop an efficient training and sampling scheme for MF in the latent space of a Representation Autoencoder (RAE), where a pre-trained vision encoder (e.g., DINO) provides semantically rich latents paired with a lightweight decoder. We observe that naive MF training in the RAE latent space suffers from severe gradient explosion. To stabilize and accelerate training, we adopt Consistency Mid-Training for trajectory-aware initialization and use a two-stage scheme: distillation from a pre-trained flow matching teacher to speed convergence and reduce variance, followed by an optional bootstrapping stage with a one-point velocity estimator to further reduce deviation from the oracle mean flow. This design removes the need for guidance, simplifies training configurations, and reduces computation in both training and sampling. Empirically, our method achieves a 1-step FID of 2.03, outperforming vanilla MF's 3.43, while reducing sampling GFLOPS by 38% and total training cost by 83% on ImageNet 256. We further scale our approach to ImageNet 512, achieving a competitive one-step FID of 3.23 with the lowest GFLOPS among all baselines.