Learning Representations to Guide Pixel Diffusion Transformers

AI summary

Overview

Research area: Generative computer vision — pixel-space diffusion transformers, representation learning for image synthesis.

Technical level: Advanced. The paper assumes familiarity with flow matching, classifier-free guidance, diffusion transformers, and self-supervised representation metrics.

Scope: The paper diagnoses why borrowed visual representations often fail to help pixel-space diffusion transformers and proposes a training-time mechanism, Representation Guidance (RG), that forms and exploits a model's own internal representation.

What This Paper Is About

Pixel diffusion transformers synthesize images directly from raw pixels, so representation learning and pixel synthesis happen end-to-end in one model. The paper asks what actually makes an internal representation useful for generation, rather than merely semantically rich. It finds that externally aligned semantics decay before reaching pixel prediction, and that adding specialized pixel decoders mainly shifts work between encoder and decoder, then proposes RG to make a compact internal representation both properly formed and actively used.

Key Contributions

1. Two diagnostic findings about representations in pixel diffusion transformers. Using REPA alignment inside JiT-B/16, the authors show better representation quality at an aligned layer does not translate into better generation. Comparing vanilla, PixelDiT, DiP, and DeCo decoders, they show architectural separation redistributes semantic versus spatial modeling rather than producing a representation better suited to generation.

2. Representation Guidance (RG). A low-dimensional linear bottleneck between a JiT-style encoder and decoder creates a compact implicit latent space, and a modified flow-matching target — shifted along the representation-induced velocity offset — forces the decoder to depend on that representation while gradients through it reshape the encoder.

3. Strong pixel-space results with a plain Transformer. RG-L/16 reaches FID 1.79 with 459M parameters and RG-XL/16 reaches FID 1.55 on ImageNet 256×256 with 725M parameters, reported as the lowest among pixel-space methods in the paper's comparison table.

4. Extension to higher resolution and text-to-image. RG-XL/16 achieves FID 1.74 at 512×512, and RG-T2I-L/16 reaches 0.89 on GenEval and 87.1 on DPG-Bench.

Main Findings

• Alignment gains do not survive to pixel prediction. REPA-4 improves RMSC, LDS, and linear probing accuracy at layer 4 of JiT-B/16, but these gains diminish immediately in later layers as the network approaches pixel prediction.

• REPA accelerates early training but loses in the end. REPA-4 beats the baseline early, but the baseline overtakes it after roughly 160 epochs and achieves a better final FID (4.56 vs. 4.80). Aligning deeper degrades further: REPA-6 reaches FID 5.38 and REPA-8 reaches 8.33.

• Pixel-specialized decoders reshape rather than improve representations. Linear probing accuracy drops sharply after the aligned layer and falls below a vanilla Transformer decoder. DeCo and DiP retain stronger spatial variation but weaker global semantics; vanilla and PixelDiT preserve stronger semantics but poor spatial structure.

• Decoder specialization brings no consistent generation gain. Under the same training budget with REPA at encoder layer 3, PixelDiT roughly matches vanilla (4.256 vs. 4.165) with 55% more parameters and 67% more FLOPs (203.37M / 47.38 G vs. 131.32M / 28.42 G), while DiP reaches FID 6.089 and DeCo reaches FID 9.167.

• RG scales with fewer parameters and FLOPs than plain-Transformer JiT baselines. RG-XL/16 reaches FID 1.55 with 725M parameters at 600 epochs, compared with PixelDiT-XL (1.61, 797M), DeCo-XL/16 (1.69, 682M), DiP-XL/16 (1.79, 631M), and JiT-G/16 (1.82, 2B). RG-L/16 reaches 1.79 with 459M parameters, outperforming JiT-G/16 and PixelREPA-H/16 (953M, 1.81).

• Both bottlenecks and guidance are necessary. On B/16 at 200 epochs, REPA alone gives FID 4.17 (vs. JiT-B at 4.56); adding the bottleneck gives 3.79; adding guidance gives 3.28; the full model reaches 2.64 at 600 epochs. Removing REPA or the bottleneck while keeping guidance raises FID to 7.85–8.99.

• RG is cheap. It adds 0.4M parameters and 13% FLOPs relative to the B/16 baseline, requires two gradient-free forward passes of the EMA network during training, and adds no inference cost.

• Feature structure becomes clearer and functionally split. PCA visualization shows RG-B maintains clear semantic contours and foreground–background separation in encoder layers 3–8, while high-frequency patterns dominate decoder layers 9–12, indicating functional specialization; REPA's induced structure degrades in later layers.

• Hyperparameters matter. The best RG scale is around s = 1.4, and the best guidance interval lower bound is t_min = 0.3; applying RG at all noise levels (t_min = 0) raises FID to 5.28. Bottleneck width d = 128 is chosen as the default, giving FID 2.644 at 600 epochs versus 2.891 at d = 512.

Methodology in Plain English

The authors first run a diagnostic study. They take a plain pixel diffusion transformer (JiT) trained with flow matching, insert REPA alignment with DINOv2 at various layers, and track both representation metrics and generation quality. They then swap in three pixel-specialized decoders (PixelDiT, DiP, DeCo) under equal training budgets to see whether separating representation from synthesis helps.

For their own method, they keep the plain Transformer but split it into an encoder and a decoder that communicate only through a representation vector. A linear layer compresses the encoder's output tokens into a much lower-dimensional latent, and another linear layer expands it back; the decoder receives this expanded representation as a token-wise conditioning signal.

Training uses the network's own predictions to sharpen the target. The decoder is run twice per step — once with the representation and once with it zeroed out (which serves as the unconditional branch) — and the difference between the two predicted velocities becomes an offset. That stop-gradient offset is added to the ordinary flow-matching target, scaled by a factor s − 1. The model is then trained to match this shifted target, so the representation is pushed to carry the semantic and spatial information that separates conditional from unconditional prediction. Because gradients flow back through the representation, the encoder is optimized for generation rather than for matching a frozen external encoder. The offset is computed with an EMA network for stability, guidance is applied only for t in (t_min, 1], and it is switched on after a short plain flow-matching warm-up. REPA alignment is still applied at an intermediate encoder layer as a semantic scaffold.

Why This Matters

Impact on research. The paper challenges the assumption that stronger pretrained features automatically help generation, and offers evidence that representation utility depends on where information is preserved and whether the generative process is forced to use it. It also argues against the trend of ever more elaborate pixel decoders, showing that a plain Transformer with good representation handling can be more parameter- and compute-efficient.

Real-world applications.

• High-resolution text-to-image systems, where the paper demonstrates 0.89 GenEval and 87.1 DPG-Bench at 512×512 and samples from 256×256 up to 1024×1024.
• Image generation pipelines that need to avoid a separate autoencoder and its reconstruction bottleneck, since RG keeps diffusion in pixel space.
• Compute-constrained deployments, because RG-L/16 reaches FID 1.79 with 459M parameters and adds no inference-time cost.
• Class-conditional image synthesis benchmarks and downstream data-generation or augmentation workflows built on ImageNet-style conditional models.

Industry relevance. The results are directly relevant to teams training diffusion transformers, since the method is a training-time change to the loss and target, adds no sampling overhead, and improves the parameter–quality and compute–quality trade-off relative to both plain JiT and pixel-specialized decoder designs.

Future Directions

• Extending RG beyond image generation to larger model scales, video generation, and other modalities such as 3D and audio, which the authors explicitly leave to future work.
• Determining more precisely why representations extracted from highly noisy inputs are unreliable, given that applying guidance at all noise levels (t_min = 0) raises FID to 5.28 while t_min = 0.3 performs best.
• Exploring whether an online teacher for the guidance offset, which the paper reports performs slightly better than the EMA teacher, is worth adopting despite the stability trade-off.
• Understanding the encoder–decoder depth trade-off further, since a 2-block decoder (10/2 split) raises FID to 4.191 while 8/4 gives 3.281, suggesting there are limits to how much synthesis work can be pushed onto the decoder.

Target Audience

Researchers and engineers working on diffusion transformers, pixel-space generative modeling, and representation learning for generation. The paper is most useful to readers already comfortable with flow matching and classifier-free guidance who want to understand the design space of representation conditioning in generative models, and to practitioners looking for parameter- and compute-efficient alternatives to specialized pixel decoders.

Authors’ abstract

Visual representations are largely developed for understanding, yet what makes them useful for generation remains unclear. We study this question in pixel diffusion transformers and uncover two key insights. First, better representation quality does not necessarily translate into greater generative utility, as externally aligned semantics can be progressively lost during pixel synthesis. Second, architectural separation with specialized pixel decoders mainly redistributes semantic and spatial modeling, without necessarily yielding representations better suited to generation. These insights suggest that a useful generative representation must be both properly formed and actively exploited. Motivated by these insights, we introduce Representation Guidance (RG), which directly forms and exploits internal representations while retaining a plain Transformer architecture. A low-dimensional bottleneck forces generation-relevant information into a compact representation, which is then explicitly used to guide decoder generation by modifying the flow-matching target according to its induced prediction offset. RG learns clearer semantic and spatial structure while scaling efficiently in both parameters and computation. On ImageNet 256x256, RG-L/16 with only 459M parameters outperforms substantially larger pixel diffusion transformers, while RG-XL/16 achieves an FID of 1.55. RG further reaches an FID of 1.74 at 512x512 and extends to text-to-image generation, achieving 0.89 on GenEval and 87.1 on DPG-Bench.