用视觉编码器重构潜在空间,让生成模型更快更稳。
MeanFlow Transformers with Representation Autoencoders
- 用预训练视觉编码器构建语义丰富的潜在表示,替代传统VAE。
- 1步生成即达FID 2.03,采样效率提升38%,训练成本降83%。
- 无需引导参数,适合追求高效高质生成的开发者使用。
MeanFlow(MF)是一种受扩散启发的生成模型,通过直接学习从噪声到数据的长距离跳变实现高效少步生成。实际中常结合Stable Diffusion变分自编码器(SD-VAE)在潜在空间使用,但其训练仍计算昂贵且不稳定。推理时SD-VAE解码器占主导开销,且类条件生成依赖复杂引导超参数。本文在表示自编码器(RAE)的潜在空间中开发了高效的MF训练与采样方案:使用预训练视觉编码器(如DINO)提供富含语义的潜在表示,并搭配轻量解码器。我们发现直接在RAE潜在空间训练MF会导致严重梯度爆炸。为稳定加速训练,采用轨迹感知的一致性中期初始化,并设计两阶段方案:先通过预训练流匹配教师模型进行蒸馏以加速收敛并降低方差,再可选地引入单点速度估计器进行自举,进一步减少与理想均值流的偏差。该方法消除对引导的需求,简化配置,显著降低训练与采样计算量。实验表明,在ImageNet 256上,本方法实现1步生成FID 2.03,优于原始MF的3.43,采样GFLOPS降低38%,总训练成本减少83%;在ImageNet 512上达到竞争性1步FID 3.23,且所有基线中采样开销最低。
原文摘要 · Abstract (English)
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 1-step FID of 3.23 with the lowest GFLOPS among all baselines. Code is available at https://github.com/sony/mf-rae.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。