Learning Representations to Guide Pixel Diffusion Transformers

Zihao Pan; Mengting Chen; Cheng Shi; Taihang Hu; Zhengrong Yue; Zikang Wang; Yuan Wang; Wei Li

cs.CVOct 9, 2026

Learning Representations to Guide Pixel Diffusion Transformers

Authors: Zihao Pan, Mengting Chen, Cheng Shi, Taihang Hu, Zhengrong Yue, Zikang Wang, Yuan Wang, Wei Li, +1 more

Organizations: Sun Yat-sen University · Alibaba Group · The University of Hong Kong · Shanghai Jiao Tong University

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.

Figures & tables

Figure 2: Our RG enables efficient scaling and better internal representations. (a) RG scales consistently from B to XL. It achieves lower FID with fewer parameters and FLOPs than plain-Transformer JiT baselines, while also outperforming pixel-specialized decoders at lower complexity. (b) RG learns clearer spatial structure and more localized semantic correspondence.

Figure 3: Architecture comparison. (a) JiT and (b) RG use patch-level input and standard Transformer blocks. JiT uses a single stack, while RG conditions its decoder on the bottleneck output, which is randomly dropped during training. (c)–(e) Pixel-specialized decoders take pixel-level input: patch attention with pixel-wise FFNs in PixelDiT, a pixel-wise MLP in DeCo, and a convolutional network with skip connections in DiP. Input tokens are noisy, output tokens are clean, and rounded squares denote the tokens output by the patch embedding or patch encoder.

Panel of Figure 4 — Figure 4: Representation quality does not necessarily translate into generative utility. REPA improves semantics and spatial structure at the aligned layer, but these gains quickly decay and do not improve final generation. REPA- $k$ aligns layer $k$ with DINOv2 features.

Panel of Figure 5 — Figure 5: Decoder design reshapes the encoder representation. Different pixel decoders induce different balances of global semantics and spatial structure in encoder. Vanilla and PixelDiT retain stronger global semantics, whereas DiP and DeCo preserve more structures but weaker semantics.

Architecture	Encoder depth	Decoder depth	Params (M)	FLOPs (G)	FID $\downarrow$	IS $\uparrow$
Vanilla	8	4	131.32	28.42	4.165	298.62
PixelDiT	8	4	203.37	47.38	4.256	306.05
DiP	8	$4\downarrow+4\uparrow$	93.60	33.72	6.089	295.64
DeCo	8	4	99.52	25.80	9.167	249.23

Table 1: Pixel-specialized decoders offer limited gains under the same training budget. PixelDiT achieves comparable FID with 55% more parameters and 67% more FLOPs, while DiP and DeCo perform worse. All models use REPA at encoder layer 3; for DiP,

4\downarrow+4\uparrow

denotes four down/up-sampling stages.

Figure 6: Training with Representation Guidance. The encoder maps the noisy input $\mathbf{z}_{t}$ through a bottleneck to the representation $\mathbf{r}$ . The shared decoder predicts the clean image $\hat{\mathbf{x}}$ with and without $\mathbf{r}$ , and the predictions are converted to the velocities $\mathbf{v}_{\text{with\_rep}}$ and $\mathbf{v}_{\text{without\_rep}}$ , whose difference is $\Delta\mathbf{v}_{\text{rep}}$ . The representation-free branch is supervised by the flow matching target $\mathbf{v}=\mathbf{x}-\bm{\epsilon}$ , and $\mathbf{v}_{\text{with\_rep}}$ by the guided target $\mathbf{v}_{\text{RG}}$ , which extrapolates $\mathbf{v}$ along the stop-gradient offset.

Latent-space Generative Models
Method	Epoch	Params	Pre-training	Generation@256
Method	Epoch	Params	Pre-training	gFID $\downarrow$	IS $\uparrow$	Precision $\uparrow$	Recall $\uparrow$
DiT-XL ( Peebles & Xie, 2023 )	1400	675M	SD-VAE	2.27	278.2	0.83	0.57
SiT-XL ( Ma et al., 2024 )	1400	675M	SD-VAE	2.06	270.3	0.82	0.59
REPA ( Yu et al., 2024 )	800	675M	SD-VAE, DINOv2	1.42	305.7	0.80	0.65
Lightning DiT ( Yao et al., 2025 )	800	675M	VA-VAE	1.35	295.3	0.79	0.65

Table 2: Class-conditional results on ImageNet

256\times 256

. RG-XL/16 reaches FID 1.55 at 600 epochs, narrowing the performance gap between pixel-space and latent-space diffusion models. The “Pre-training” columns list the external models. Bold = best in column.

Table 8

Figure 7: PCA visualization of layer-wise features. Visualized at $t=0.5$ . While REPA-induced structure in JiT-B degrades in later layers, RG-B keeps clear semantic contours in encoder layers 3–8, with high-frequency patterns in decoder layers 9–12, indicating functional specialization.

Figure 8: Layer-wise representation metrics in RG-B. The bottleneck (BN) improves spatial diversity and locality at the encoder output (layer 8), while representation guidance further strengthens global semantics, as measured by linear probing, and reinforces encoder–decoder specialization.

Table 11

REPA	BN	Guid.	Params (M)	FLOPs (G)	Epoch	FID $\downarrow$	IS $\uparrow$
Baseline: JiT-B/16			131.12	25.22	200	4.56	254.86
✗	✗	✓	131.32	28.42	200	8.08	230.81
✗	✓	✗	131.52	28.47	200	4.71	268.45
✗	✓	✓	131.52	28.47	200	8.99	226.20
✓	✗	✗	131.32	28.42	200	4.17	298.62
✓	✗	✓	131.32	28.42	200	7.85	220.40

Table 8: Component ablation on B/16.

Figure 9: Ablations on RG and the alignment target. (a) A moderate RG scale (

s=1.4

) yields the lowest FID at 200 epochs. (b) Excluding highly noisy inputs from guidance improves generation, with

t_{\min}=0.3

performing best. (c) DINOv2 with registers and DINOv3 yield comparable FIDs, both lower than standard DINOv2.

Appendix figures & tables16 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 14

Architecture
Configuration	RG-B/16	RG-L/16	RG-XL/16
encoder depth	8	16	20
decoder depth	4	8	10
hidden dim $D$	768	1024	1152
heads	12	16	16
patch size	16

Appendix

Table 9: Configurations and implementation details of RG on ImageNet.

Method	Weak prediction lacks	Guidance applied at	Extra inference cost
CFG ( Ho & Salimans, 2022 )	condition $\mathbf{c}$	sampling	second forward pass
Autoguidance ( Karras et al., 2024 )	capacity or training	sampling	second (smaller) model
IG ( Zhou et al., 2026 )	depth (intermediate head)	sampling or training	auxiliary head
RG (ours)	learned representation $\mathbf{r}$	training	none

Appendix

Table 10: Self-guidance methods. Top: how each method obtains its weak prediction. Bottom: B/16 models at 200 epochs on ImageNet

256\times 256

, evaluated with the same protocol; CFG is the optimal CFG scale.

Figure 11: Layer-wise representation of RG-B, JiT-B, and JiT-B with IG. The shaded band marks block 4, where IG attaches its auxiliary output head. IG collapses spatial structure at this block and leaves global semantics unchanged.

Figure 12: FID versus CFG scale at epoch 200. (a) Different RG scales

s

. (b) Different lower bounds

t_{\min}

of the RG interval. Larger RG scales favor smaller CFG scales.

Figure 13: Guidance teacher and warm-up. RG-B at 200 epochs. Left: FID versus CFG with

\Delta\mathbf{v}_{\text{rep}}

from the EMA (default) or online network, both with stop-gradient. Right: results at the optimal CFG for each teacher (top) and each epoch at which RG is enabled (bottom).

Figure 14: FID during training. RG-B matches the 600-epoch FID of JiT-B (3.66) after only 150 epochs (3.64), a

4\times

training speed-up. REPA accelerates JiT-B early but is overtaken by the baseline, whereas RG-B keeps improving to 2.64 at 600 epochs.

Figure 15: Layer-wise representation of RG-B and JiT-B. RG-B layers 1–8 form the encoder and 9–12 the decoder. RG-B stays semantic and spatially diverse through the last encoder layer.

Panel of Figure 28 — Figure 28: Text-to-image $512\times 512$ RG-T2I-L/16 samples. The prompt is shown below each image.

Panel of Figure 29 — Figure 29: Text-to-image $1024\times 1024$ samples of RG-T2I-L/16 fine-tuned from the $512\times 512$ model. The prompt is shown below each image.

Explore similar work

PixelDiT2: Representation-Grounded Pixel Diffusion Transformers
Sep 21, 2026Yongsheng Yu, Wei Xiong, Yichen Sheng +2Diffusion Model Guidance Diffusion Transformer
NVIDIA · University of Rochester
Pixel-Space Diffusion Transformers
Jul 20, 2026Renye Yan, Jikang Cheng, You Wu +8Diffusion Transformer Pixel-Space Diffusion Models
Peking University, China · Nanjing University, China · Stanford University, USA +3
Persistence Forcing: Exploiting Feature Specialization in Pixel-Space Diffusion
Sep 28, 2026Chong Wang, Zixuan Fu, Shiqi Huang +3Diffusion Transformer Diffusion Guidance
Nanyang Technological University · KTH Royal Institute of Technology · Hebei University of Technology

Sep 21, 2026J/K move · Enter open · S save

PixelDiT2: Representation-Grounded Pixel Diffusion Transformers

Yongsheng Yu, Wei Xiong, Yichen Sheng, Shiqiu Liu+1

NVIDIA · University of Rochester

Diffusion Model Guidance Diffusion Transformer Pixel-Space Diffusion Models Diffusion-Based Image Generation

Open paper AskPDF arXiv

Recent advances in pixel-space diffusion models have narrowed the image quality gap with latent-space diffusion, but still converge more slowly and lag behind in final image quality. We argue that a key reason is the lack of an explicit representation prior: unlike latent diffusion, which usually denoises in a compact and structured latent space, pixel diffusion needs to learn denoising-friendly representations and pixel generation simultaneously from raw RGB space. To address this problem, we propose PixelDiT2, an end-to-end pixel-space diffusion model designed to decouple representation learning from pixel generation without introducing an autoencoder or latent reconstruction bottleneck. We propose representation grounding that uses a frozen pretrained vision foundation model to provide explicit per-patch representation guidance throughout denoising, allowing the pixel diffusion transformer to focus more on pixel generation. On ImageNet-256x256, PixelDiT2 achieves an FID of 1.46 after 600 epochs; at 512x512 resolution, PixelDiT2 achieves an FID of 1.48 after 680 epochs. Project page: https://pixeldit.github.io/pixeldit2/

Abstract

Explore similar work

PixelDiT2: Representation-Grounded Pixel Diffusion Transformers

Pixel-Space Diffusion Transformers

Persistence Forcing: Exploiting Feature Specialization in Pixel-Space Diffusion