让扩散模型的U-Net与ViT对齐，加速训练并提升生成质量。
U-REPA: Aligning Diffusion U-Nets to ViTs
- 通过跳接结构选择中间层对齐，适配U-Net特性。
- 用MLP上采样特征，解决空间维度不一致问题。
- 引入流形损失增强样本间相似性，适合图像生成任务。
表示对齐（REPA）在扩散Transformer（DiT）训练中表现优异，但尚未在收敛更快的扩散U-Net架构上验证。由于块功能差异、空间下采样导致的维度不匹配以及U-Net与ViT之间的空间间隙，直接应用REPA面临挑战。为此，本文提出U-REPA：首先基于跳接观察，选定U-Net中间阶段为最佳对齐位置；其次，对经过MLP处理的U-Net特征进行上采样以对齐空间维度；最后，发现逐标记相似性对齐困难，引入流形损失正则化样本间的相对相似性。实验表明，U-REPA可实现优异生成质量，并显著加速收敛：在ImageNet 256×256上，使用CFG引导间隔，200轮或100万次迭代内达到FID<1.5，且所需总轮数仅为REPA的一半即可超越sd-vae-ft-ema。代码已开源。
原文摘要 · Abstract (English)
Representation Alignment (REPA) that aligns Diffusion Transformer (DiT) hidden-states with ViT visual encoders has proven highly effective in DiT training, demonstrating superior convergence properties, but it has not been validated on the canonical diffusion U-Net architecture that shows faster convergence compared to DiTs. However, adapting REPA to U-Net architectures presents unique challenges: (1) different block functionalities necessitate revised alignment strategies; (2) spatial-dimension inconsistencies emerge from U-Net's spatial downsampling operations; (3) space gaps between U-Net and ViT hinder the effectiveness of tokenwise alignment. To encounter these challenges, we propose \textbf{U-REPA}, a representation alignment paradigm that bridges U-Net hidden states and ViT features as follows: Firstly, we propose via observation that due to skip connection, the middle stage of U-Net is the best alignment option. Secondly, we propose upsampling of U-Net features after passing them through MLPs. Thirdly, we observe difficulty when performing tokenwise similarity alignment, and further introduces a manifold loss that regularizes the relative similarity between samples. Experiments indicate that the resulting U-REPA could achieve excellent generation quality and greatly accelerates the convergence speed. With CFG guidance interval, U-REPA could reach $FID<1.5$ in 200 epochs or 1M iterations on ImageNet 256 $\times$ 256, and needs only half the total epochs to perform better than REPA under sd-vae-ft-ema. Codes: https://github.com/YuchuanTian/U-REPA
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。