Representation Alignment for Just Image Transformers is not Easier than You Think
Abstract
Representation Alignment (REPA) has emerged as a simpleway to accelerate Diffusion Transformers training in latent space. At thesame time, pixel-space diffusion transformers such as Just image Trans-formers (JiT) have attracted growing attention because they remove adependency on a pretrained tokenizer, and then avoid the reconstruc-tion bottleneck of latent diffusion. This paper shows that the REPA canfail for JiT. REPA yields worse FID for JiT as training proceeds andcollapses diversity on image subsets that are tightly clustered in the rep-resentation space of pretrained semantic encoder on ImageNet. We tracethe failure to an information asymmetry: denoising occurs in the high di-mensional image space, while the semantic target is strongly compressed,making direct regression a shortcut objective. We propose PixelREPA,which transforms the alignment target and constrains alignment witha Masked Transformer Adapter that combines a shallow transformeradapter with partial token masking. PixelREPA improves both train-ing convergence and final quality. PixelREPA reduces FID from 3.66 to3.17 for JiT-B/16 and improves Inception Score (IS) from 275.1 to 284.6on ImageNet 256×256, while achieving > 2× faster convergence. Finally,PixelREPA-H/16 achieves FID= 1.81 and IS= 317.2.