U-REPA: Aligning Diffusion U-Nets to ViTs

Yuchuan Tian, Hanting Chen, Mengyu Zheng, Yuchen Liang, Chao Xu, Yunhe Wang

Peking University· Huawei Noah's Ark Lab· University of Sydney

Representation Alignment U-Net Transformer Diffusion

Abstract

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:

different block functionalities necessitate revised alignment strategies;
spatial-dimension inconsistencies emerge from U-Net's spatial downsampling operations;
space gaps between U-Net and ViT hinder the effectiveness of tokenwise alignment.

To encounter these challenges, we propose 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.