让流形生成模型一步完成图像转换,实现实时推理。
TReFT: Taming Rectified Flow Models For One-Step Image Translation
- 直接使用预训练模型预测的速度作为输出,解决对抗训练收敛难题。
- 在多个数据集上达到当前最优性能,且推理速度可达实时。
- 适合需要快速图像转换的应用场景,如实时风格迁移、图像编辑。
矩形流(RF)模型通过最优传输理论实现了高质量的图像与视频生成。然而,在图像到图像翻译任务中,仍依赖于耗时的多步去噪过程,限制了实时应用。尽管最近的循环对抗训练范式(CycleGAN-Turbo)能在预训练扩散模型中实现一步转换,但直接应用于RF模型会导致严重收敛问题。本文分析此挑战并提出TReFT:一种用于一步图像翻译的矩形流模型驯化方法。不同于以往工作,TReFT直接以预训练DiT或UNet预测的速度作为输出,设计简单却有效,解决了对抗训练下的收敛问题,并支持一步推理。该设计基于一个新观察:在去噪过程末期,预训练RF模型预测的速度趋于从原点指向最终干净图像的向量,这一性质通过理论分析加以验证。将TReFT应用于SD3.5和FLUX等大型预训练RF模型时,引入内存高效的潜在循环一致性与身份损失,以及轻量级结构简化以加速推理。经TReFT微调的预训练模型在多个图像翻译数据集上表现媲美最先进方法,同时实现实时推理。
原文摘要 · Abstract (English)
Rectified Flow (RF) models have advanced high-quality image and video synthesis via optimal transport theory. However, when applied to image-to-image translation, they still depend on costly multi-step denoising, hindering real-time applications. Although the recent adversarial training paradigm, CycleGAN-Turbo, works in pretrained diffusion models for one-step image translation, we find that directly applying it to RF models leads to severe convergence issues. In this paper, we analyze these challenges and propose TReFT, a novel method to Tame Rectified Flow models for one-step image Translation. Unlike previous works, TReFT directly uses the velocity predicted by pretrained DiT or UNet as output-a simple yet effective design that tackles the convergence issues under adversarial training with one-step inference. This design is mainly motivated by a novel observation that, near the end of the denoising process, the velocity predicted by pretrained RF models converges to the vector from origin to the final clean image, a property we further justify through theoretical analysis. When applying TReFT to large pretrained RF models such as SD3.5 and FLUX, we introduce memory-efficient latent cycle-consistency and identity losses during training, as well as lightweight architectural simplifications for faster inference. Pretrained RF models finetuned with TReFT achieve performance comparable to sota methods across multiple image translation datasets while enabling real-time inference.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。