PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Jaeyo Shin, Jiwook Kim, Hyunjung Shim

arXiv:2603.14366v1cs.CVcs.LG

TL;DR

REPA can hinder pixel-space JiT training because high-dimensional denoising is aligned directly to a compressed semantic target, creating a shortcut that harms diversity and quality. PixelREPA transforms the target and constrains alignment with a Masked Transformer Adapter, improving convergence and final ImageNet performance, including over 2× faster convergence and FID reductions for JiT-B/16 and JiT-H/16.

  • Problem

    REPA degrades pixel-space JiT as training progresses because direct regression to a compressed semantic target can cause feature hacking and collapse diversity among images with similar semantic features.

  • Method

    PixelREPA replaces the MLP alignment projection with a shallow Transformer adapter and partial token masking that constrain alignment to the transformed target space.

  • Results

    FID decreases from 3.66 to 3.17 for PixelREPA-B/16, convergence is over 2× faster, and PixelREPA-H/16 reaches FID 1.81 on ImageNet 256×256.

  • Takeaways & Limitations

    PixelREPA improves both convergence speed and generation quality for high-resolution pixel-space JiT across the reported ImageNet backbones.

Abstract

from arXiv · show

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 \times 256$, while achieving $> 2\times$ 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.

1 Introduction

Applying REPA to pixel-space JiT unexpectedly degrades performance because denoising operates in a high-dimensional image space while the semantic target is compressed. PixelREPA addresses this mismatch with a transformed target and masked Transformer alignment, improving convergence and ImageNet quality.

  • Motivation: REPA unexpectedly degrades JiT as pixel-space training progresses, unlike its acceleration benefits in latent-space diffusion.JiT+REPA underperforms vanilla JiT on ImageNet 256 × 256 as training proceeds.
  • Failure mechanism: Pixel-space denoising has O(H × W) degrees of freedom, whereas the semantic encoder produces a compact target, creating an information asymmetry.At higher resolution, many visually distinct images map to similar semantic-feature regions.
  • PixelREPA: PixelREPA transforms intermediate JiT features with a shallow Transformer adapter before aligning them with frozen semantic-encoder features.The adapter preserves information needed for later JiT blocks to map representations back to pixels while injecting semantic structure.
  • PixelREPA: Partial token masking prevents trivial token-wise alignment by forcing contextual prediction under partial observation and constraining the pixel-side information.The adapter and masking together form the Masked Transformer Adapter, applied only to the training-time alignment branch.
  • Results: FID falls from 3.66 to 3.17 for PixelREPA-B/16, while PixelREPA achieves over 2× faster convergence and reaches FID 1.81 for H/16.On ImageNet 256×256, PixelREPA-H/16 reaches FID 1.81 versus 1.86 for vanilla JiT-H/16 and 1.82 for JiT-G/16.
  • Failure mechanism: Direct regression to the compressed semantic target causes feature hacking, reducing diversity among visually diverse images with similar semantic features.The degradation is especially evident on image subsets tightly clustered in semantic feature space.

2 Preliminaries

Diffusion models learn denoising objectives in either latent or pixel spaces, while related methods use flow matching or representation alignment to improve generation and training.

  • DDPM: DDPMs corrupt data through a forward noising process and train a reverse process to predict the added noise.The noise level follows a variance schedule β_t.
  • Flow-based Generative Models: Flow-based models deterministically transport samples from noise to clean data by learning a velocity field with velocity matching.Under linear schedules, the target velocity is v = x − ϵ.
  • Pixel-space Diffusion: Latent diffusion models denoise in compressed autoencoder latent spaces, reducing computation and memory for high-resolution generation.This efficiency comes with a reconstruction bottleneck that can remove fine textures and small structures.
  • Pixel-space Diffusion: JiT achieves strong pixel-space diffusion performance with a pure Transformer trained directly on raw images, without a latent tokenizer or auxiliary losses.JiT uses x-prediction and a velocity-based training objective.
  • Representation Alignment: REPA aligns intermediate diffusion features with representations from a frozen pretrained semantic encoder using a cosine-similarity objective.The alignment is computed patch-wise through a projection function.
  • Representation Alignment: Subsequent work extends representation alignment to joint VAE–diffusion tuning, early termination, attention alignment, video generation, 3D-aware generation, and unified model training.These applications build on REPA’s reported simplicity and effectiveness.

3 Motivation

REPA helps JiT at low resolution but degrades high-resolution pixel-space training as representation dimensionality increases. The failure is structured: compressed semantic targets harm diversity most for semantically clustered images.

  • Dimensionality and perceptual compression: Pixel-space diffusion retains O(H × W) degrees of freedom and fine-grained details that latent representations compress or attenuate.The mismatch grows with resolution because pixel space preserves textures and micro patterns weakly tied to semantics.
  • Dimensionality: REPA improves JiT at 32×32 resolution but degrades performance as training progresses at 256 × 256.This supports the finding that REPA failure emerges as representation-space dimensionality increases.
  • Information asymmetry: High-resolution pixel-space diffusion aligns a high-dimensional denoising space with a compact semantic representation, allowing many visually distinct images to share similar features.This ambiguity increases with resolution and makes direct alignment vulnerable to feature hacking.
  • Perceptual compression: On the Most Similar 100 subset, vanilla JiT achieves lower FID than JiT+REPA, whereas JiT+REPA performs better on the Least Similar 100 subset.The subsets are formed around class centroids in the external semantic feature space; the most similar images share global structure while differing mainly in fine details.
  • Feature hacking: REPA’s failure is structured: it harms generation quality specifically where semantic features are highly similar and ambiguous.The paper names this mechanism feature hacking, in which alignment drives pixel-diverse samples toward a narrow semantic region.

4 PixelREPA: REPA for Pixel Space Diffusion Models

PixelREPA addresses JiT’s alignment failure by transforming the semantic target and constraining the alignment branch with a Masked Transformer Adapter. The adapter preserves the denoising pathway while using contextual aggregation and partial masking to prevent shortcut regression.

  • Target transformation: PixelREPA replaces direct JiT-to-semantic alignment with alignment in a transformed target space induced by a dedicated adapter.The adapter learns to extract semantic content from intermediate JiT features rather than forcing those features themselves to match the compressed target.
  • Masked Transformer Adapter: The Masked Transformer Adapter combines a shallow two-block self-attention transformer with partial masking of intermediate JiT tokens.The transformer provides contextual aggregation, while masking requires prediction of the full semantic target from partial observations.
  • Adapter architecture: Contextual aggregation replaces isolated point-wise mapping, producing a more structured alignment pathway and weakening the trivial regression route.The adapter uses neighboring-token information before matching the external semantic features.
  • Masking motivation: An unmasked adapter reaches FID 4.68 at 200 epochs, improving over JiT+REPA’s 5.14 but remaining worse than vanilla JiT’s 4.37.This intermediate result motivates masking as an additional constraint on the alignment branch.
  • Masking mechanism: Masking breaks per-token correspondence and narrows the alignment branch’s information gap while leaving the main denoising pathway unmasked.The effective adapter-input degrees of freedom decrease from O(N · d) to O((1 − r) · N · d).
  • Objective: PixelREPA combines its alignment loss with the JiT objective using a positive regularization weight λ.The alignment loss uses masked intermediate features passed through the adapter and compares them with frozen semantic features via cosine similarity.

5 Experiments

Experiments on ImageNet 256×256 show that PixelREPA improves JiT quality and convergence across model scales, while partial masking is essential for avoiding REPA’s degradation.

  • Overall comparison: FID 3.66 to 3.17 for B/16, a 13.4% improvement over JiT, with gains also reported for larger L/16 and H/16 models.PixelREPA-H/16 reaches FID 1.81 and surpasses JiT-G/16 despite having nearly twice as many parameters.
  • Masking ablation: Partial masking is essential: the unmasked adapter beats JiT+REPA but remains worse than the JiT baseline.The ablation attributes this benefit to discouraging shortcut learning and reducing overfitting to the external semantic feature.
  • Masking ratio: The best masking ratio is r = 0.2, whereas increasing it to 0.5 degrades performance by blocking supervision and gradient signals.The experiments therefore use r = 0.2 for all PixelREPA models.
  • Training convergence: More than 2× faster convergence is achieved by PixelREPA, which maintains an 8.5% improvement over vanilla JiT at 200 epochs.At 200 epochs, JiT+REPA instead shows a 17.6% FID degradation, reaching 5.14 versus 4.37 for the baseline.
  • Scalability: PixelREPA achieves lower FID as model size increases and consistently outperforms vanilla JiT at matched sizes and across training epochs.The scalability result indicates improved sample quality and diversity for larger models.

6 Conclusion

The paper identifies high-resolution REPA failure in JiT as feature hacking caused by alignment to a compressed semantic target. PixelREPA addresses this with transformed targets and a constrained Masked Transformer Adapter.

  • Failure mode: Standard REPA can cause feature hacking and degraded training when JiT aligns with a compressed semantic target at high resolution.The conclusion describes this as a failure mode of representation alignment for pixel-space diffusion.
  • Proposed solution: PixelREPA combines a shallow Transformer adapter with partial token masking to transform the alignment target and constrain the alignment pathway.The resulting Masked Transformer Adapter is reported to stabilize optimization and scale with model size.
  • Reported results: FID decreases from 3.66 to 3.17 for B/16, while H/16 reaches FID 1.81 across JiT backbones on ImageNet 256 × 256.These results summarize the reported high-resolution improvements.

A Implementation Details

The implementation follows JiT configurations while adding a fixed Masked Transformer Adapter, and JiT’s in-context concatenation begins only after a predefined block.

  • Optimization: PixelREPA uses Adam with learning rate 2 × 10^-4 and β1, β2 = (0.9, 0.95) under the original JiT configuration details.The same settings are used to evaluate the effectiveness of the Masked Transformer Adapter.
  • JiT block: The JiT architecture uses AdaLN-Zero modulation and concatenates condition embeddings with previous-block tokens only after a predefined in-context start block.This distinguishes JiT’s in-context concatenation from the cited DiT and SiT designs.

B Qualitative Results

The appendix provides uncurated qualitative samples from PixelREPA-H on ImageNet 256×256 across various classes under the same classifier-free guidance setting.

  • Qualitative samples: Uncurated qualitative results are shown for various ImageNet classes using PixelREPA-H at 256×256 resolution.The samples appear in Figures 9–12.
  • Evaluation setting: The qualitative figures share the same classifier-free guidance setting across the displayed samples.The passage does not state a qualitative comparison outcome.
Loading 2603.14366v1…