PiD: NVIDIA Replaces the VAE Decoder with Pixel Diffusion for Fast 2K/4K Image Generation

> arXiv:2605.23902 | NVIDIA Spatial Intelligence Lab | Open source, Apache 2.0

The decoder has always been the overlooked bottleneck

Three years of diffusion model research has focused on the front of the pipeline — the U-Net or DiT: faster denoising, better text alignment, fewer malformed hands. But one component has rarely been seriously examined: the latent decoder.

It is the last gate to the pixel world. Every latent diffusion model (FLUX, SD3, Stable Diffusion) and autoregressive image generator (VAR, MAR) depends on the same thing: reconstructing a compact latent representation into pixels. The traditional approach is a VAE decoder — a convolutional network trained for reconstruction, tasked with "restoring whatever the encoder stored."

The flaws in this design have never gone away:

1. Imperfect reconstruction: encoder-decoder round trips inevitably lose fine-grained detail 2. No artifact correction: the decoder only "restores" — artifacts in the latent (blotches, color blocks) are faithfully amplified 3. Cost explosion at high resolution: at 2K/4K, VAE decoder memory and compute costs spike 4. Upsampler cascades: after decoding you still need a super-resolution pipeline, stacking latency

NVIDIA's core judgment in PiD: decoding should not be a reconstruction task — it should be a generative task.

PiD's approach: conditional pixel diffusion decoding

PiD stands for Pixel Diffusion Decoder. It does not reconstruct; it generates.

Traditional pipeline: latent → VAE decode → low-res image → super-resolution model → high-res image PiD pipeline: latent → conditional pixel diffusion (high-res output directly)

One step: decoding and upsampling are unified in a single generative module.

Architecture: PixelDiT + ControlNet-style adapter

PiD builds on PixelDiT (1.3B parameters, MMDiT backbone) — a transformer natively trained for pixel-space diffusion. Starting from a 1024×1024 pretrained checkpoint, it is extended to 2K and even 4K resolution.

Scaling to high resolution brings the classic problem of exploding patch token sequence lengths. PiD replaces RoPE with NTK-aware RoPE to improve spatial extrapolation.

A lightweight ControlNet-style adapter injects latent information into the pixel diffusion backbone:

• nearest-neighbor upsampling to align the patch grid
• Conv2d(16→512) → SiLU → Conv2d(512→512)
• 4 pre-activation residual blocks (GroupNorm 4, 512 channels)
• flatten + Linear projection into PixelDiT's hidden dimension (1536)

Latent tokens are injected every two DiT blocks. PiT pixel blocks remain untouched.

Key design 1: Noisy latent training

Training feeds not only clean latents but also noised latents:

z̃_σ = (1-σ)z + σξ, σ∼U(0, σ_max), σ_max=0.8

Two purposes:

1. Prevent the decoder from over-trusting the latent, which would suppress its ability to generate detail 2. Enable the decoder to handle partially denoised latents — the basis for early termination

Key design 2: Sigma-aware gating

How reliable a latent is depends on how much noise remains. PiD uses a learnable gate to modulate injection strength:

g_i(h_i, l_i, σ) = sigmoid(Linear_i([h_i, l_i]) - ασ)

• Content-dependent term: base injection strength determined by current hidden state and latent content
• The ασ term: noisier latent → weaker injection → the pixel prior takes over more
• Initialization: bias=2.0, α≈5, giving roughly sigmoid(2-5σ) at start
• Zero-initialization: training starts from the pretrained pixel-prior behavior and gradually learns to exploit the latent

Ablations show removing this gate significantly degrades scores, confirming its necessity.

Key design 3: DMD2 distillation to 4 steps

After training the teacher, PiD is distilled into a student via Distribution Matching Distillation (DMD2) for 4-step inference. The σ schedule is fixed at {0.999, 0.866, 0.634, 0.342}.

Classifier-free guidance is additionally distilled, so inference does not need separate conditional + unconditional forward passes.

GAN regularization uses a projected discriminator (26-block DiT, hidden 1536), with DMD loss weight 1.0, score matching 1.0, GAN loss 0.05, R1 regularization 200.0.

Distillation takes only 3,000 iterations — 2 hours on 128 H100s.

Key design 4: LDM early termination

Because the decoder can handle partially denoised latents, the base latent diffusion model can stop before fully converging. The latent, still carrying residual noise σ, goes straight to PiD, which finishes the final refinement steps in pixel space.

This saves the last few denoising steps of the base model, further reducing total latency.

Training cost

| Stage | Resources | Time | Data | |-------|-----------|------|------| | Pixel prior (2K) | 128 H100 | ~1 day | 2.6M high-quality images (MultiAspect-4K-1M + PDF + internal) | | Latent-conditioned decoder | 64 H100 | ~half a day | batch 64, 30K iterations | | DMD2 distillation (4-step) | 128 H100 | ~2 hours | 3K iterations, context parallelism |

Total: roughly 2.5 days, on the order of 200+ GPU-days. Not heavy by NVIDIA standards.

The text encoder is a frozen Gemma-2-2B-it; images are auto-captioned with Qwen3-VL-8B-Instruct (long/medium/short captions at 10% sampling each).

Performance numbers

Speed

| Hardware | Configuration | Latency | |----------|---------------|---------| | RTX 5090 (consumer) | 512→2048, 13GB peak memory | <1 second | | GB200 (data center) | same | 210ms | | vs SeedVR2 cascaded diffusion | same resolution, same hardware | 5.9× faster (211.2ms vs 1237.5ms) |

Compatibility

| Latent type | Models | Upsampling factor | |-------------|--------|-------------------| | VAE | FLUX.1 / FLUX.2 / SD3 / Z-Image | 4× | | RAE semantic | DINOv2-B (DiT-DH) | 4× | | RAE semantic | SigLIP (Scale-RAE DiT 2.8B) | 8× |

This matters especially for RAE: RAE latents preserve strong semantics but leave low-level appearance underdetermined — traditional decoders simply cannot fill in the missing texture. PiD's generative capability fills exactly that gap.

Why "throw away the VAE"

Existing diffusion decoder explorations (ε-VAE, SSDD) remain low-resolution, reconstruction-oriented, and still require separate super-resolution. PiD is the first to combine the generative power of pixel-space diffusion with structural guidance from latent space, and to scale it to 2K/4K.

This brings three structural changes:

1. Simplified pipeline: decode + upsample in one step, no cascade needed 2. Artifact correction: the pixel diffusion prior can fix artifacts at the latent level 3. Consumer viability: 13GB peak memory, 2K output in under 1 second on an RTX 5090

An analogy: the VAE decoder is film developing — restoring the negative into a photo. PiD is digital painting — using the negative's composition as a hint and repainting a higher-fidelity image directly on the canvas.

Open source and deployment

• Code: github.com/nv-tlabs/PiD (Apache 2.0)
• Weights: huggingface.co/nvidia/PiD
• Released: adapters for FLUX / FLUX.2 / Z-Image / Z-Image-Turbo / SD3 / DINOv2 / SigLIP
• Coming soon: Qwen-Image

This means existing FLUX/SD3 pipelines can swap the VAE decoder for PiD without retraining the base model. Plug-and-play.

References

• Lu et al., "PiD: Fast and High-Resolution Latent Decoding with Pixel Diffusion", arXiv:2605.23902, 2026
• Project: sites.google.com/view/pid-decoder
• Code: github.com/nv-tlabs/PiD
• Weights: huggingface.co/nvidia/PiD
• Base architecture: PixelDiT (Wu et al., 2026), MMDiT (Peebles & Xie, 2023)
• Distillation: DMD2 (Yin et al., 2024)
• Related: ε-VAE, SSDD, SeedVR2, LUA, Real-ESRGAN

> PiD's implication runs deeper than the technique itself: the "latent" side of latent diffusion has been pushed to its limits (FLUX, SD3, various tokenizers), yet the final step back into pixel space has been dragged down by the VAE — a relic of the reconstruction era. With one paper, NVIDIA shows that every stage of a generative model should be generative — including decoding.