cs.CVMar 15, 2026

Representation Alignment for Just Image Transformers is not Easier than You Think

Authors: Jaeyo Shin, Jiwook Kim, Hyunjung Shim

Organizations: KAIST AI, 291 Daehak-ro, Yuseong-gu, Daejeon 34141, Republic of Korea

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Representation Alignment (REPA) has emerged as a simple way to accelerate Diffusion Transformers training in latent space. At the same time, pixel-space diffusion transformers such as Just image Transformers (JiT) have attracted growing attention because they remove a dependency on a pretrained tokenizer, and then avoid the reconstruction bottleneck of latent diffusion. This paper shows that the REPA can fail for JiT. REPA yields worse FID for JiT as training proceeds and collapses diversity on image subsets that are tightly clustered in the representation space of pretrained semantic encoder on ImageNet. We trace the failure to an information asymmetry: denoising occurs in the high dimensional 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 with a Masked Transformer Adapter that combines a shallow transformer adapter with partial token masking. PixelREPA improves both training convergence and final quality. PixelREPA reduces FID from 3.66 to 3.17 for JiT-B/16/16 and improves Inception Score (IS) from 275.1 to 284.6 on ImageNet 256×256256 \times 256, while achieving >2×> 2\times faster convergence. Finally, PixelREPA-H/16/16 achieves FID=1.81=1.81 and IS=317.2=317.2. Our code is available at https://github.com/kaist-cvml/PixelREPA.

Figures & tables

Figure 1: REPA degrades JiT performance. As training progresses, JiT ++ REPA yields higher FID [ 14 ] ( ↓\downarrow ) than vanilla JiT, indicating that REPA hinders pixel space diffusion training. PixelREPA prevents overfitting to the external semantic feature target, which accelerates convergence in JiT training. Remarkably, PixelREPA achieves >2×>2\times faster convergence than the vanilla JiT. All evaluated models utilize JiT-B /16/16 .
Figure 2: Overall Framework of PixelREPA. PixelREPA masks a subset of tokens in an intermediate diffusion feature map. The full token sequence, with only a subset masked, are then transformed by a shallow Transformer adapter and aligned to features from a frozen pretrained semantic encoder. This transforms the alignment target and reduces overfitting to the external semantic representation.
Figure 3: REPA accelerates pixel diffusion training at low resolution, whereas it degrades training at high resolution. This figure illustrates FID scores across different resolutions comparing JiT and JiT ++ REPA. Results show (a) ImageNet 32×3232\times 32 and (b) ImageNet 256×256256\times 256 with varying training epochs.
Figure 4
Figure 6: Most Similar 100 and Least Similar 100 images in the external semantic feature space. (a) Most Similar 100 images to the class centroid in feature space of the external semantic encoder. (b) Least Similar 100 images to the centroid.
ModelparamsFID ↓\downarrowIS ↑\uparrow
Latent-space Diffusion
DiT-XL /2/2 [ 34 ]675+49M2.27278.2
SiT-XL /2/2 [ 29 ]675+49M2.06277.5
REPA [ 47 ] , SiT-XL /2/2675+49M1.42305.7
LightningDiT-XL // 2 [ 46 ]675+49M1.35295.3
DDT-XL /2/2 [ 43 ]675+49M1.26310.6
Table 1: Quantitative comparison of diffusion models on ImageNet 256 ×\times 256. FID and IS are evaluated with 50K samples.
Table 7
Figure 7: Scalability. As model size grows, PixelREPA shows better performance.
Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

 PixelREPA-BPixelREPA-LPixelREPA-H
Architecture
depth122432
hidden dim76810241280
heads121616
image size256
patch size(image size)/16\text{(image size)}/16
Appendix
Table 4: Model configuration details.
Figure 8: Implementation of the JiT block. The in-context concatenation is operated only after the predefined in-context start block.
class n01443537
class n01735189
class n02317335
class n03207743
More in this topic

Diffusion Models

Diffusion models learn to reverse gradual data corruption, enabling flexible generation, editing, forecasting, and inference.

See all 1,169 papers

Explore similar work

CardsList

  1. Scaffold Then Internalize: Representation Injection for Diffusion Transformers

    Sep 28, 2026Han Fu, Jiacheng Chen, Baoquan Zhao +4Diffusion TransformerDiffusion Model Acceleration

    Sun Yat-sen University · Video Rebirth · The Hong Kong Polytechnic University

  2. Representation Conditioning for Generation

    Oct 14, 2026Saeed Razavi, Xingjian Leng, Liang ZhengDiffusion-Based Image GenerationDiffusion Transformer

    Australian National University

  3. Beyond Point-Wise Matching: Structural Representation Alignment for Accelerating Diffusion Transformers

    May 16, 2026Shaodong Xu, Zhendong Wang, Litong Gong +4Diffusion Model AlignmentDiffusion Transformer

    University of Science and Technology of China