PixelREPA: Pixel-Space Diffusion Alignment

• The paper introduces PixelREPA, which modifies the alignment branch with a Masked Transformer Adapter to prevent shortcut learning in pixel-space diffusion.
• It improves ImageNet 256×256 performance by reducing FID and enhancing IS metrics, effectively addressing the failure modes of vanilla REPA.
• PixelREPA employs a partial token masking strategy that preserves rich image details and ensures robust, faster convergence during training.

PixelREPA is a representation-alignment method for pixel-space diffusion transformers, developed to address the mismatch between standard REPA and direct pixel-space denoising in Just image Transformers (JiT). Its central claim is that naive alignment of intermediate diffusion features to a frozen semantic encoder can be beneficial in latent diffusion yet become counterproductive in pixel space, where denoising must retain substantially richer image detail. PixelREPA modifies the alignment branch rather than the denoising backbone, replacing direct regression with a constrained contextual prediction mechanism based on a Masked Transformer Adapter (MTA). In the reported ImageNet $256\times256$ setting, the method improves both convergence and final sample quality relative to JiT, with JiT-B/16 moving from FID $3.66$ to $3.17$ and IS from $275.1$ to $284.6$, while PixelREPA-H/16 reaches FID $1.81$ and IS $317.2$ (Shin et al., 15 Mar 2026).

1. Conceptual basis and problem setting

PixelREPA was introduced in response to a specific failure mode of Representation Alignment (REPA) in pixel-space diffusion. REPA had emerged as a simple means of accelerating Diffusion Transformers in latent space, but the same formulation does not transfer cleanly to JiT, which denoises directly in image space rather than in a tokenizer-compressed latent. The underlying issue is described as an information asymmetry: pixel-space denoising must model an image with $O(H \times W)$ degrees of freedom, whereas the semantic target supplied by a frozen encoder such as DINOv2 is strongly compressed (Shin et al., 15 Mar 2026).

This asymmetry makes direct alignment ill-posed. Many visually distinct images can map to similar semantic representations, so forcing intermediate diffusion features to match that compressed target can encourage the model to satisfy the alignment loss through shortcut behavior rather than through improved image synthesis. The paper characterizes this behavior as feature hacking: the denoiser over-optimizes the external semantic feature space and sacrifices diversity. This framing is central to PixelREPA, because the method is not presented as a generic alignment improvement but as a redesign of alignment specifically for high-dimensional pixel generation (Shin et al., 15 Mar 2026).

A further implication is resolution dependence. The reported evidence shows that REPA can help JiT on ImageNet $32\times32$, yet on ImageNet $256\times256$ it eventually degrades JiT as training proceeds. This suggests that the severity of the mismatch grows with pixel-space dimensionality. The method is therefore best understood not as a universal replacement for REPA in all diffusion settings, but as a response to the geometry of pixel-space diffusion training (Shin et al., 15 Mar 2026).

2. Failure of vanilla REPA in pixel-space JiT

Standard REPA aligns an intermediate diffusion feature to a frozen semantic encoder feature, typically from DINOv2. In latent diffusion this is relatively natural because both the diffusion state and the semantic target are already compressed. In JiT, however, the denoiser operates in pixel space, where fine-grained detail remains present throughout training. Direct alignment to a compact semantic target then becomes a shortcut objective rather than a faithful regularizer (Shin et al., 15 Mar 2026).

The empirical failures reported for vanilla REPA are specific and structured. On ImageNet $3.66$0, REPA may improve JiT early in training but later yields worse FID. The degradation is especially pronounced on subsets of images that are visually diverse but tightly clustered in semantic-feature space. The paper formalizes this analysis through “Most Similar 100” and “Least Similar 100” subsets: vanilla JiT performs well on the semantically clustered subset, JiT + REPA performs worse there, and PixelREPA attains the best FID on both subsets. This pattern is used to support the feature-hacking interpretation rather than a simpler claim that alignment is uniformly harmful in pixel space (Shin et al., 15 Mar 2026).

The ablation at 200 epochs makes the comparison explicit. In that table, JiT reaches FID $3.66$1, JiT + REPA reaches $3.66$2, PixelREPA without masking reaches $3.66$3, and full PixelREPA reaches $3.66$4. These numbers indicate that a plain REPA-style objective can be worse than the baseline in this setting, and that neither a Transformer adapter alone nor masking alone is sufficient to recover the full gain (Shin et al., 15 Mar 2026).

3. Method: Masked Transformer Adapter and alignment objective

PixelREPA addresses the mismatch by transforming the alignment target pathway instead of directly regressing JiT features to the compressed semantic space. The method replaces the simple MLP projection head used in REPA with a Masked Transformer Adapter. The MTA has two components: a shallow Transformer adapter and a partial token masking strategy. The alignment branch thus becomes a contextual prediction problem from incomplete intermediate features rather than a token-wise shortcut to semantic embeddings (Shin et al., 15 Mar 2026).

In procedural terms, the method takes an intermediate hidden representation from JiT, randomly masks a fraction of the tokens, feeds the masked tokens into a 2-block Transformer adapter, and aligns the adapter output with frozen semantic encoder features. The semantic encoder is DINOv2. The adapter uses self-attention for contextual aggregation. The mask ratio is fixed to $3.66$5 for all model sizes in the main experiments. Masking is applied only to the alignment path, not to the denoising pathway, and the MTA is used only during training, so the paper states that it adds no inference-time cost (Shin et al., 15 Mar 2026).

The method’s intuition has two parts. First, partial masking is intended to prevent shortcut learning: if all tokens are visible, the adapter can still learn a near-trivial token-wise mapping to the semantic target. Second, masking acts as an information bottleneck, reducing the adapter-side information content and making the two spaces more compatible, in a way that is described as similar in spirit to the compression already present in latent diffusion (Shin et al., 15 Mar 2026).

The paper gives the following objective for PixelREPA: $3.66$6 where $3.66$7 is a patch-wise mask, $3.66$8 is the Transformer adapter, and $3.66$9 denotes masked intermediate features. The full training objective is reported as

$3.17$0

with $3.17$1 as the alignment weight (Shin et al., 15 Mar 2026).

4. Architecture, training configuration, and placement within JiT

PixelREPA is architecturally conservative with respect to the main denoising model. The MTA is attached to an intermediate JiT representation at a layer directly before JiT’s in-context start block. This detail matters because it localizes the auxiliary alignment to a specific point in the network rather than distributing it across the entire denoising stack (Shin et al., 15 Mar 2026).

The training configuration reported for the main ImageNet experiments is specific. The optimizer is Adam with learning rate $3.17$2 and $3.17$3. The batch size is $3.17$4, the regularization weight is $3.17$5, and the mask ratio is $3.17$6. Training is performed on 8 NVIDIA H200 GPUs. Sampling uses a 50-step Heun ODE solver, and the CFG interval is $3.17$7 (Shin et al., 15 Mar 2026).

The method is therefore best characterized as an auxiliary training objective with a structured adapter, not as a new diffusion backbone. Its intervention is confined to the representation-alignment branch, and the paper explicitly emphasizes that this branch does not modify the main JiT denoising path and contributes no inference-time overhead. A plausible implication is that PixelREPA is intended to preserve the deployment profile of JiT while altering optimization dynamics during training (Shin et al., 15 Mar 2026).

5. Empirical performance, convergence, and diversity

The principal empirical results are reported on ImageNet $3.17$8. For JiT-B/16, the baseline achieves FID $3.17$9 and IS $275.1$0, whereas PixelREPA-B/16 achieves FID $275.1$1 and IS $275.1$2. For JiT-L/16, the baseline reaches FID $275.1$3 and IS $275.1$4, while PixelREPA-L/16 reaches FID $275.1$5 and IS $275.1$6. For JiT-H/16, the baseline reaches FID $275.1$7 and IS $275.1$8, while PixelREPA-H/16 reaches FID $275.1$9 and IS $284.6$0. The paper further notes that PixelREPA-H/16 at FID $284.6$1 outperforms JiT-G/16 at FID $284.6$2 despite using about half the parameters (Shin et al., 15 Mar 2026).

The method is also reported to produce more than $284.6$3 faster convergence than vanilla JiT. This claim is operationally important because the motivation for REPA-like objectives had originally included accelerated training, and PixelREPA is presented as recovering that benefit in pixel space without the instability observed for vanilla REPA (Shin et al., 15 Mar 2026).

The diversity analysis is a substantive part of the evaluation rather than a peripheral add-on. By comparing performance on the “Most Similar 100” and “Least Similar 100” subsets, the paper argues that PixelREPA improves average quality while avoiding the diversity collapse associated with feature hacking. This is notable because it reframes representation alignment as a problem not merely of final FID but of preserving variability when semantic embeddings underresolve the relevant pixel-level distinctions (Shin et al., 15 Mar 2026).

6. Ablations and relation to register-based pixel-space regularization

The ablations isolate the roles of masking and the Transformer adapter. Varying the mask ratio yields the following FID values: $284.6$4, $284.6$5, $284.6$6, $284.6$7, and $284.6$8, with $284.6$9 reported as the best setting. The interpretation given in the paper is that too little masking leaves shortcut paths intact, whereas too much masking removes too much supervision. The adapter-only variant improves over plain REPA but still underperforms JiT, while the full PixelREPA formulation performs best, indicating that both contextual aggregation and masking are necessary (Shin et al., 15 Mar 2026).

A separate but closely related line of work on pixel-space diffusion transformers connects PixelREPA to register tokens and dual-stream processing. “Registers Matter for Pixel-Space Diffusion Transformers” reports that pixel-space DiTs do not exhibit the high-norm patch-token outliers characteristic of ViTs, yet register tokens still improve convergence and generation quality. The paper attributes this to cleaner feature maps at high noise levels, reduced patch-token norms, and a separation in which some registers behave as norm sinks while others encode global semantic information. This paper explicitly positions PixelREPA as complementary rather than redundant: representation-alignment methods are said not to be the whole story, because pixel-space models also benefit from mechanisms that regularize and organize internal features (Starodubcev et al., 15 May 2026).

The reported joint results support that complementarity. For pDiT-B/16, adding PixelREPA improves FID, and adding registers on top improves it further: pDiT-B/16 + PixelREPA reaches $1.81$0 at 200 epochs and $1.81$1 at 600 epochs, compared with pDiT-B/16 without registers + PixelREPA at $1.81$2 and $1.81$3. For JiT-B/16, JiT + PixelREPA reaches $1.81$4 and $1.81$5, while JiT-B/16 + PixelREPA + compact dual reaches $1.81$6 and $1.81$7. The paper interprets this as evidence that REPA-like objectives align representations, whereas registers provide an orthogonal regularization and specialization mechanism that improves internal feature organization, especially in pixel space (Starodubcev et al., 15 May 2026).

This relation also clarifies a common misconception. PixelREPA does not claim that semantic alignment alone is sufficient for strong pixel-space diffusion training. The later register-token study suggests that global information routing, norm sinking, and spatial smoothing of internal features can remain beneficial even when alignment objectives are present. A plausible implication is that PixelREPA addresses one part of the pixel-space optimization problem—the semantic-target mismatch—while register-based mechanisms address another part concerning token-level specialization and spatial feature organization (Starodubcev et al., 15 May 2026).