PixRestore: Unified Image Restoration via Pixel Diffusion Transformer
Lingchen Sun, Rongyuan Wu, Xiangtao Kong, Jixin Zhao, Qiaosi Yi, Yujing Sun, Shuaizheng Liu, Zhengqiang Zhang, Lei Zhang
cs.CV
2026-08-18
PolyU's PixRestore skips the VAE and T2I pretraining, runs flow matching on patchified pixels with adaptive DINO layer guidance, and beats billion-parameter restorers with 50M parameters in a single step.
Unified image restoration (UIR) wants one model to handle noise, blur, rain, haze, low light, super-resolution, and more. The recent mainstream finetunes text-to-image latent diffusion models like SDXL and FLUX: strong priors, good perceptual quality, but three hidden costs. The VAE compresses images into latent space and discards exactly the small structures and textures restoration cares about. The open-ended T2I synthesis prior invents plausible details that do not match the input. And a billion-scale backbone plus an MLLM planner plus iterative sampling makes the inference bill enormous; a FLUX-based restorer takes 5.8 seconds per image.
The Hong Kong Polytechnic University and OPPO team flips the premise: UIR starts from a low-quality image dense with visual cues, so what it needs is robustness to degradation and pixel-faithful reconstruction, not open-ended generation. The T2I baggage is unnecessary.
PixRestore is a pixel-space Diffusion Transformer with no VAE, no T2I pretraining, and a backbone trained from scratch. Three mechanisms.
Flow matching on patchified pixels. The low-quality image and the noisy state are concatenated channel-wise, cut into patches as tokens, and the DiT learns a velocity field directly on RGB pixels. A relatively large patch (8 by default) keeps the VAE-free sequence tractable. A controlled experiment establishes the choice: the same DiT-S reaches 26.62 dB PSNR in pixel space versus at most 22.80 with any of three mainstream VAEs, with more parameters and an extra 41-78 ms of VAE coding.
DINO layer-reliability routing. A frozen DINOv2 extracts multi-layer features from the low-quality input. The authors first measure a fact: reliability varies by layer and degradation. Shallow layers are sensitive to rain, blur, and snow; deep layers to noise; middle layers to low light, super-resolution, and haze. No single layer wins everywhere. A lightweight predictor therefore guesses per-layer weights from the low-quality input alone, supervised during training by low-quality-to-high-quality feature similarity. Reliable layers are fused into dense conditioning injected by cross-attention into every DiT block; unreliable layers get stronger high-quality-feature supervision instead, pushing the model to wash the degradation out of them.
Single-stepping. After multi-step training, t is fixed at 1 so the model maps pure Gaussian noise to output in one pass, and a multi-layer discriminator over frozen DINO features recovers the textures one-step generation loses.
Training used 2.83M images across eight degradations on 8 A800 GPUs: 250K iterations for the multi-step model, then 100K for single-step finetuning.
Across 15 public benchmarks and eight degradations, the two variants rank first or second on most metrics. Representative numbers:
| Method | Params | De-rainstreak PSNR | Denoise PSNR | Avg. latency |
| FAPE-IR (FLUX-based, retrained) | 21575M | 31.91 | 34.22 | 1011 ms |
| FoundIR-v2 (SDXL-based, retrained) | 16910M | 27.85 | 28.17 | 18293 ms |
| PixRestore-S | 53.7M | 32.28 | 34.87 | 44 ms |
| PixRestore-B | 210.9M | 32.85 | 34.62 | 79 ms |
On a real-world test set without ground truth, PixRestore-B takes the best average DR-Score (67.88) over retrained FAPE-IR (67.55); DR-Score is the degradation-removal metric the authors introduce.
The ablation chain is clean: plain pixel DiT averages 26.62 dB, adding single-layer DINO conditioning brings 27.36, six adaptive layers with hierarchical supervision 27.66, and single-step finetuning 28.49. Pure regression training under the same budget reaches only 27.00, so the gain comes from flow pretraining. Shrinking the patch or growing the backbone keeps helping, with a -0.96 correlation between GFLOPs and LPIPS.
The paper shakes the two-year-old assumption that restoration must ride on a T2I backbone. Restoration needs pixel fidelity, not open-ended synthesis, and the VAE's lossy compression is pure liability for it. 54M versus 21.6B parameters and 44 ms versus 5.8 s are both two-order-of-magnitude gaps, turning on-device real-time restoration from implausible to practical. The DINO reliability router is not restoration-specific: any dense prediction task with degraded inputs could borrow it. And there is a metrology bonus: MUSIQ-style no-reference metrics favor outputs that keep the rain streaks, while the proposed VLM-judged DR-Score aligns far better with human judgment.
The authors state: fixed tokenization ties the model to its training resolution, and moving to higher resolution needs retraining or architectural changes; under extreme information loss the compact backbone struggles and large-model generative priors still win; DR-Score depends on a proprietary VLM (Gemini-3.1 Pro) that may drift with updates, and scoring each image five times is not cheap.
Two more from a close read. Retrained baselines used the authors' 2.83M-image corpus, while official FAPE-IR checkpoints remain strong on several tasks, so "beats everything" needs a discount. DR-Score is a self-proposed metric whose best scores come from the proposed method; the human-alignment validation exists but independent replication does not. On super-resolution, PixRestore's 27.01 trails retrained PromptIR's 27.59; the most information-starved degradation is the weakest task, consistent with the authors' own caveat.