Skip to content
Bud's Wiki

PixelDiT: Pixel Diffusion Transformers for Image Generation

TL;DR

Section titled “TL;DR”

PixelDiT is a single-stage, end-to-end pixel-space Diffusion Transformer that removes the VAE entirely while staying competitive with latent baselines. It uses a dual-level architecture: a patch-level DiT (large patch size, short sequence) carries global semantics, and a pixel-level DiT (one token per pixel) refines texture details, with two efficiency tricks — pixel-wise AdaLN modulation and a pixel-token compaction mechanism that compresses pixel tokens into patch tokens for the attention step. PixelDiT-XL hits gFID 1.61 on ImageNet 256² in 320 epochs (beating PixelFlow-XL 1.98, PixNerd-XL 1.93, JiT-G 2B at 1.82), and a T2I variant trained directly at 1024² in pixel space scores GenEval 0.74 / DPG 83.5 with 1.3B params, competitive with FLUX-dev (12B) and beating SDXL/PixArt-Σ. Releases under an NVIDIA non-commercial research license.

Key claims

Section titled “Key claims”
  • Single-stage pixel-space DiT removes the autoencoder lossy-reconstruction bottleneck and matches/beats best pixel baselines at parity model size [§1, Table 1].
  • Dual-level architecture (patch DiT for semantics + pixel DiT for detail) is the key inductive bias: vanilla DiT/16 baseline gets 9.84 gFID at 80 epochs; full PixelDiT at the same budget gets 2.36 [Table 4].
  • Pixel token compaction (compress all P² pixel tokens within a patch into one token before global attention, then decompress) cuts the attention sequence from H×W back to (H/P)(W/P), an 82,247 GFLOPs → 311 GFLOPs reduction; without it the model OOMs [Table 5].
  • Pixel-wise AdaLN — generate independent modulation parameters for each pixel within a patch from the patch-level semantic token, rather than broadcasting one set across the patch — improves gFID from 3.50 → 2.36 at 80 epochs [Table 4, Fig. 2].
  • PixelDiT-XL achieves gFID 1.61 on ImageNet 256² at 320 epochs, beating PixelFlow-XL (1.98, 320 epochs), PixNerd-XL (1.93, 320 epochs), EPG (2.04, 800 epochs), and JiT-G (1.82, 2B params, 600 epochs) [Table 1].
  • T2I variant trained directly at 1024² in pixel space (1.3B params, 26M image-text pairs, Gemma-2 text encoder, MM-DiT blocks in patch pathway only) reaches GenEval 0.74 / DPG 83.5 vs FLUX-dev’s 0.67/84.0 at 12B params [Table 3].
  • Throughput at 1024² is 0.33 samples/s on a single A100, ~8× faster than FLUX-dev (0.04) and ~2× faster than Hunyuan-DiT (0.05) despite operating directly on pixels [Table 3].
  • REPA-style alignment of mid-patch-pathway tokens with frozen DINOv2 features is part of the training recipe (loss weight λ=0.5, applied at block 8 of the patch-level pathway) [§3.4].
  • Image-editing case study (FlowEdit pipeline) shows PixelDiT preserves small scene text where Stable Diffusion 3 and FLUX visibly distort it; the distortion is traced to the VAE’s own reconstruction, not the diffusion process [§4.4, Fig. 1(b)].
  • Patch-size ablation: smaller patches help small models but the gain shrinks with scale — at PixelDiT-XL, p=8 and p=16 converge to similar gFID, so p=16 is the practical choice; for PixelDiT-B, p=4 is needed to close the gap [§4.5.3, Fig. 4(a)].

Method

Section titled “Method”

PixelDiT splits the model into two pathways. Patch level: standard DiT blocks (RMSNorm + 2D RoPE per [lgt]) operate on N = (H/P)(W/P) patch tokens of hidden size D, conditioned by global AdaLN parameters derived from the class label or text embedding plus timestep. After M blocks, the patch pathway emits semantic tokens c_pat ∈ R^{N×D}. Pixel level: the image is also embedded as one token per pixel, x_pix ∈ R^{HW×D’} where D’ is much smaller (e.g., 64), and reshaped so the H·W·D’ grid splits into N groups of P² pixel tokens each. Each PiT block runs (a) per-pixel update via pixel-wise AdaLN: project each semantic token c_pat[i] through a linear map to (6 × P² × D’) AdaLN params, partition into six (P²·D’)-vectors, and apply per-pixel-token scale/shift/gate independently within the patch; and (b) global attention via pixel token compaction: a learned linear flattening compresses each patch’s P² pixel tokens into one patch token, global self-attention runs on N tokens, and a learned expansion decompresses back to P² per patch. Residual connections preserve high-frequency information across the compress-attend-expand bottleneck. For T2I, the patch pathway uses MM-DiT blocks [SD3] that fuse text and image streams; the pixel pathway is unchanged because semantic tokens already carry text intent. Training uses Rectified Flow velocity loss plus REPA alignment.

Results

Section titled “Results”
  • ImageNet 256² class-conditioned, 320 epochs, PixelDiT-XL (797M): gFID 1.61, sFID 4.68, IS 292.7, P 0.78, R 0.64. At 80 epochs already hits gFID 2.36, IS 282.3 [Table 1].
  • Closest pixel-space baselines at 256²: PixelFlow-XL 1.98 (320 ep, 677M), PixNerd-XL 1.93 (320 ep, 700M), EPG 2.04 (800 ep, 583M), JiT-G 1.82 (600 ep, 2B params). Beat all of them with smaller compute budget and competitive size [Table 1].
  • Closest latent baselines: REPA 1.42 (800 ep, 675M), LightningDiT 1.35 (800 ep, 675M), DDT-XL 1.26 (400 ep, 675M), RAE-XL 1.13 (800 ep, 839M). PixelDiT narrows but doesn’t close the gap — RAE-XL is ~0.5 FID better with no autoencoder needed in the generation pathway [Table 1].
  • T2I at 512²: GenEval 0.78 / DPG 83.7 (PixelDiT-T2I 1.3B) vs PixNerd 0.73/80.9 and PixArt-Σ 0.52/79.5 [Table 3].
  • T2I at 1024²: GenEval 0.74 / DPG 83.5 (PixelDiT-T2I 1.3B) — beats FLUX-dev (0.67/84.0, 12B), SDXL (0.55/74.7, 2.6B), Playground v2.5 (0.56/75.5, 2.6B), DALLE 3 (0.67/83.5). Lower throughput per sample than the 1.3B latent models would be expected but actually higher: 0.33 samples/s vs FLUX-dev 0.04 [Table 3].
  • Ablations: dual-level design without pixel token compaction → OOM; with compaction → 3.50 gFID; with compaction + pixel-wise AdaLN → 2.36; “No Pixel-Pathway Attention” (compaction only, MLPs replace attention in pixel level) → 2.56 at 80ep, 2.22 at 160ep — worse than full model, confirming the compact global attention is doing work, not just the per-pixel modulation [Tables 4, 5].
  • Patch size: B/4 needs P=4 (gFID 22.04 at 200K iters with P=4 vs B/16’s much worse), but XL/8 and XL/16 converge to similar quality, so P=16 is the practical default [§4.5.3].

Why it’s interesting

Section titled “Why it’s interesting”

PixelDiT is the natural successor to PixelGen / JiT / PixNerd / FARMER on the “pixel-space DiT is back” arc filed under Pixel-space diffusion. Where PixelGen argued the gap was the loss function and CAFM / FD-loss extended that to learned and closed-form distributional criteria respectively, PixelDiT argues the gap is architectural: a single-pathway DiT (vanilla DiT/16) caps out at gFID 9.84 at 80 epochs no matter the loss, and the right fix is to factorize semantics and texture into two pathways with carefully designed cross-conditioning. The pixel-token-compaction trick is the load-bearing piece — without it, dense per-pixel attention OOMs at 256², and the entire dual-level design collapses. The result also reframes Asymmetric Flow Models: AsymFlow’s 1.57 FID on JiT-H/16 with rank-asymmetric velocity is the parameterization-side alternative to PixelDiT’s architecture-side solution at comparable headline numbers, suggesting the loss-side, parameterization-side, and architecture-side levers are roughly substitutable on this benchmark. PixelDiT’s 1024² native pixel T2I result is the most consequential — until now the conventional wisdom (echoed even in DiP: Taming Diffusion Models in Pixel Space) was that pixel-space diffusion couldn’t scale past 512² without cascading; the 0.33 samples/s throughput at 1024² on a single A100 puts that argument in serious doubt. Note that the released NVIDIA weights are non-commercial-only, so this is more useful as a research datapoint than as a production base.

See also

Section titled “See also”