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 and improves Inception Score (IS) from 275.1 to 284.6 on ImageNet 256×256, while achieving >2× faster convergence. Finally, PixelREPA-H/16 achieves FID=1.81 and IS=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 ] ( ↓ ) 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× faster convergence than the vanilla JiT. All evaluated models utilize JiT-B /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×32 and (b) ImageNet 256×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.
Model
params
FID ↓
IS ↑
Latent-space Diffusion
DiT-XL /2 [ 34 ]
675+49M
2.27
278.2
SiT-XL /2 [ 29 ]
675+49M
2.06
277.5
REPA [ 47 ] , SiT-XL /2
675+49M
1.42
305.7
LightningDiT-XL / 2 [ 46 ]
675+49M
1.35
295.3
DDT-XL /2 [ 43 ]
675+49M
1.26
310.6
Table 1: Quantitative comparison of diffusion models on ImageNet 256 × 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-B
PixelREPA-L
PixelREPA-H
Architecture
depth
12
24
32
hidden dim
768
1024
1280
heads
12
16
16
image size
256
patch size
(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.
Recent representation alignment (REPA) methods accelerate diffusion transformer training by aligning projections of the transformer's hidden states with representations from pretrained visual encoders. In this work, we explore a reverse and complementary direction to REPA: rather than projecting diffusion representations into the encoder's space, we inject encoder representations into the diffusion transformer, allowing them to actively participate in the denoising process. To this end, we introduce REPresentation Injection (REPI), a training framework based on a scaffold-to-internalization strategy, in which projected encoder representations initially serve as a temporary scaffold and are then progressively internalized by the diffusion transformer. REPI outperforms REPA across a wide range of backbones and is highly complementary to it: combining the two yields substantial gains over either alone. Notably, with only 160K training steps, REPI + REPA matches vanilla SiT trained for 7M steps, a speedup of over 43.5×. Code will be available at https://jeneveuxpas.github.io/REPI
Han Fu, Jiacheng Chen, Baoquan Zhao +4
Sun Yat-sen University · Video Rebirth · The Hong Kong Polytechnic University
Representation-alignment methods effectively accelerate the diffusion transformer training. REPA, for example, uses a pretrained vision encoder to supervise an intermediate layer through an MLP projection head which is discarded after training. It aims to import semantics from the vision encoder to the diffusion backbone. While this objective is achieved at early layers by back-propagating the alignment loss, there is no guarantee that later layers, supervised by the velocity prediction loss alone, can well access such semantics. We propose a simple solution, representation conditioning (RECON), to improve semantics in later layers. During training, we feed the pretrained vision encoder tokens into later layers; during inference, because we no longer have the clean image, we reuse the MLP head (as a proxy to the vision encoder) and project its output tokens to the later layers. Our method adds no parameters and leads to much faster convergence: RECON is 63× faster than SiT and 25× faster than REPA. Code is available at https://github.com/saeidrazavi/RECON
Recent advances in Diffusion Transformers (DiTs) demonstrate that aligning noisy latent states with well-trained semantic features-as pioneered by Representation Alignment (REPA)-can substantially accelerate training and improve generation fidelity. Subsequent analysis(e.g., iREPA) suggests that these gains arise primarily from transferring spatial structure contained in pre-trained vision representations. However, mostly existing alignment methods employ point-wise matching objectives or rely on implicit architectural tweaks, which fail to explicitly model the spatial relational geometry inherent in vision foundation models. We argue that such element-wise supervision is insufficient to capture the rich spatial topology of visual representations, and that effective alignment for generation should instead be formulated as an explicit structural constraint. To this end, we propose sREPA, a structural REPresentation Alignment framework to enforce consistency in the relational geometry of feature maps, rather than merely matching individual feature points. By encouraging the model to internalize holistic spatial layouts and structural correlations from pre-trained features, sREPA achieves faster and more stable convergence, along with improved sample quality, compared to state-of-the-art alignment strategies. Our code and models will be released.