PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2608.16793v1cs.CV

TL;DR

Unified image restoration seeks one model that can recover high-quality images across diverse degradations, but existing approaches risk oversmoothing, detail loss, and inconsistent synthesis. PixRestore uses a compact VAE-free pixel-space diffusion transformer with adaptive DINO guidance and achieves leading restoration quality with single-step inference while remaining efficient.

  • Problem

    Unified image restoration needs one model for diverse degradations, while existing regression and latent-diffusion approaches can oversmooth outputs or lose restoration-sensitive details.

  • Method

    PixRestore trains a VAE-free pixel-space DiT from scratch with patchified-pixel flow matching and adaptive hierarchical DINO guidance for conditioning and supervision.

  • Results

    PixRestore variants achieve the best overall performance across eight restoration tasks, ranking first or second on most metrics while using about 50M parameters and single-step inference.

  • Takeaways & Limitations

    PixRestore offers faithful, perceptually strong unified restoration with a compact and efficient single-step model, and larger variants further improve performance.

  • Takeaways & Limitations

    A model trained at one resolution cannot directly extend to higher resolution without retraining or architectural modification, and compact PixRestore may struggle with extremely information-scarce degradations.

Abstract

from arXiv · show

Unified image restoration (UIR) aims to recover high-quality (HQ) content from low-quality (LQ) images with different degradations using a single model. Most recent methods adapt large pretrained text-to-image (T2I) latent diffusion models for their strong capacity and generative priors. However, the variational autoencoder (VAE) in latent T2I models may discard restoration-sensitive details, while the open-ended synthesis prior can introduce content-inconsistent artifacts. We present PixRestore, a VAE-free pixel-space Diffusion Transformer (DiT) for UIR, where the diffusion backbone is trained entirely from scratch, without relying on T2I pretraining. PixRestore performs flow matching directly on patchified pixels, preserving fine-grained details while keeping the token sequence tractable. To adapt to different degradations, PixRestore learns to predict the reliability of layer features using LQ--HQ DINO feature similarity. Features from more reliable layers are fused as dense conditioning, while less reliable layers receive stronger HQ-feature supervision to encourage degradation removal. We train PixRestore on a large-scale corpus of diverse scenes and degradations, and further finetune it into a one-step generator using DINO-based adversarial objectives for efficient inference. Experiments on public benchmarks and real-world test sets show that, with only about 50M parameters and single-step inference, PixRestore achieves the best overall fidelity, perceptual quality, and robustness to degradations among competing UIR models while being far more efficient. Larger PixRestore variants can further boost performance, demonstrating the scalability of our pixel-space design. Code and the curated benchmark can be found at https://github.com/csslc/PixRestore.

1 Introduction

PixRestore addresses unified image restoration with a VAE-free pixel-space Diffusion Transformer trained from scratch, preserving pixel-aligned evidence without T2I priors. It combines adaptive hierarchical visual guidance with single-step DINO-adversarial finetuning to deliver efficient, high-quality restoration.

  • Motivation: UIR seeks one model that recovers high-quality images across diverse, co-occurring degradations instead of training separate specialist models.The degradations include noise, blur, rain, haze, and low light.
  • Motivation: Recent diffusion UIR methods often finetune pretrained T2I latent models, whose lossy latents and open-ended priors can limit restoration fidelity and introduce inconsistencies.The supplied passage identifies rich priors and perceptual realism as benefits of T2I pretraining, but notes associated costs.
  • Method: PixRestore is a VAE-free pixel-space DiT trained from scratch, applying flow matching to patchified pixels to preserve pixel-aligned evidence while keeping token sequences tractable.This design removes autoencoder overhead and targets improved content fidelity while maintaining perceptual quality.
  • Efficiency and Results: PixRestore is finetuned from a multi-step model into a single-step generator with DINO-based adversarial objectives, enabling efficient inference with little quality loss.The model is trained on a large-scale corpus of diverse scenes and degradations before single-step finetuning.
  • Efficiency and Results: With about 50M parameters and single-step inference, PixRestore achieves superior fidelity, perceptual quality, and robustness on public benchmarks and real-world test sets.The introduction presents it as the fastest among compared diffusion-based UIR methods while attaining the best quality.
  • Method: Adaptive hierarchical visual guidance provides dense conditioning and supervision to stabilize restoration across different degradations.The method guides restoration through conditioning and stabilizes training through supervision.

2 Related Work

Related work spans regression-based unified restoration, which adapts to multiple degradations through representations, prompts, or expert routing, and generative approaches that synthesize HQ images from LQ inputs. PixRestore instead uses conditional pixel-space flow modeling with degradation-aware hierarchical DINO guidance.

  • Regression-based Unified Restoration: UIR methods use learned degradation representations, prompts, or expert routing to adapt one model across multiple degradations.Conventional restoration methods typically target specific degradations such as noise, blur, rain, or haze.
  • Generative Unified Restoration: Generative UIR methods synthesize HQ images conditioned on LQ inputs, either learning restoration-native conditional generation from scratch or using restoration-specific diffusion and degradation-aware conditioning.Examples include DiffUIR and DA-CLIP, which target fidelity across different degradations.
  • Generative Unified Restoration: PixRestore combines pixel-space flow modeling with degradation-aware hierarchical DINO guidance for unified restoration.Its guidance is derived densely and per image from hierarchical DINO features.

3 Method

PixRestore formulates unified image restoration as conditional flow matching directly on RGB pixels with a VAE-free Diffusion Transformer, preserving spatial details while keeping token sequences tractable. It uses adaptive DINO-based conditioning and hierarchical supervision, then distills the model into a one-step generator with feature-level adversarial training.

  • Pixel-space diffusion: PixRestore performs conditional flow matching directly on RGB pixels, avoiding VAE compression and its potential loss of small structures and fine textures.The model patchifies full-resolution pixels to retain spatial structure while controlling token-sequence length.
  • Adaptive visual guidance: A frozen DINOv2 encoder supplies multi-layer LQ features, and an adaptive router predicts per-layer reliability weights for fused cross-attention conditioning.The router estimates weights from the LQ image and projected features because HQ images are unavailable at inference.
  • Pixel space vs. latent space: 26.62 dB PSNR, 0.1593 LPIPS, and 54.32 MUSIQ are achieved by pixel DiT across eight restoration tasks, outperforming latent DiTs while using 23.41M parameters.Removing the VAE also reduces cost and latency by 41–78 ms under matched training and testing settings.
  • Pixel-space diffusion: The noisy state follows x_t = (1 − t)y_hq + tε, while the DiT predicts the clean image and is trained through velocity matching with a denominator clip at 0.05.The objective minimizes the squared difference between predicted and target velocities.
  • Adaptive visual guidance: Hierarchical supervision assigns stronger restoration emphasis to unreliable layers while reliable layers provide conditioning, using complementary reliability and restoration weights.The complete objective combines flow matching with router prediction and feature-fidelity losses, with both auxiliary weights set to 0.5.
  • One-step generation: The multi-step model is finetuned into a one-step generator at t = 1, using a DINO-based multi-layer adversarial objective to preserve fine textures.Inference restores an image from one Gaussian-noise sample conditioned on the LQ input.

4 Experiments

PixRestore is evaluated across eight restoration tasks using large-scale training data, paired public benchmarks, and real-world no-reference tests. It achieves strong overall restoration quality and degradation removal while remaining highly efficient, with ablations and scaling studies supporting its design.

  • Datasets and Evaluation: The training corpus contains about 2.83M images spanning eight restoration tasks, sampled with equal probability during training.The tasks include deblur, dehaze, denoise, de-rainstreak, de-raindrop, desnow, low-light enhancement, and super-resolution.
  • Datasets and Evaluation: Evaluation combines paired public benchmarks for fidelity and perceptual quality with real-world test sets containing 100 LQ images per degradation without ground truth.Public metrics include PSNR, SSIM, LPIPS, and DISTS, while DR-Score is used to assess degradation removal in addition to no-reference metrics.
  • Public-Benchmark Results: Across eight public tasks, PixRestore variants rank first or second on most reference-based metrics and DR-Score, achieving the best overall performance among compared methods.Retraining competing models on PixRestore’s diverse dataset generally improves their results, indicating the value of the training corpus.
  • Public-Benchmark Results: FoundIR-v2 obtains high DR-Scores but lower fidelity and perceptual quality, while FAPE-IR is strong on selected degradations, showing that stronger generative priors alone do not ensure robust UIR.The paper attributes FoundIR-v2’s fidelity limitation to input-detail loss from latent VAE compression.
  • Efficiency and Scalability: With 53.7M parameters, 658G FLOPs, and 44 ms per image, PixRestore runs about 7–23× faster than PromptIR, DiffUIR, FoundIR, and FAPE-IR.Its computation is about 1/181 of FoundIR-v2 and 1/1209 of Flux-IR, while PixRestore-B further improves performance while remaining compact.
  • Ablations and Scalability: Averaging six DINO layers raises PSNR to 27.72 and MUSIQ to 55.01, while increasing image-token computation improves LPIPS, with GFLOPs and LPIPS correlation −0.96.The S model’s LPIPS improves from about 0.129 with p = 16 to about 0.101 with p = 4.

5 Conclusion

PixRestore is a compact, VAE-free pixel-space diffusion transformer that performs faithful, perceptually strong unified image restoration with single-step inference and scales to larger variants. Its limitations include resolution inflexibility, weaker performance under extremely information-scarce degradations, and reliance on a proprietary VLM for DR-Score.

  • Conclusion: PixRestore is a VAE-free pixel-space diffusion transformer for unified image restoration.It performs flow matching directly on patchified pixels and uses adaptive hierarchical DINO guidance.
  • Conclusion: PixRestore achieves faithful restoration with strong perceptual quality across diverse degradations while remaining compact and efficient through single-step inference.Larger variants further improve performance, demonstrating the scalability of the design.
  • Limitations: A model trained at one resolution cannot be directly extended to higher resolution without retraining or architectural modification.This limitation follows from the DiT architecture setting.
  • Limitations: Compared with billion-scale T2I models, the compact PixRestore backbone may struggle under extremely information-scarce degradations, while DR-Score relies on a proprietary VLM.The proprietary VLM may introduce differences across model updates.

Appendix

The appendix provides supplementary details on PixRestore’s data, diffusion-space design, visual foundation prior, DR-Score, and expanded benchmark comparisons.

  • Data: It documents the training and real-world test data collection procedures.These details correspond to Sec. 4.1 of the main paper.
  • Diffusion space: It compares pixel-space and latent-space diffusion models for unified image restoration.The comparison supports the discussion in Sec. 3.1 of the main paper.
  • Visual foundation prior: It explains the visual foundation prior used for unified image restoration.This material corresponds to Sec. 3.2 of the main paper.
  • Evaluation metric: It provides further details of DR-Score.The appendix expands on the metric introduced in Sec. 4.1 of the main paper.
  • Benchmark comparisons: It includes additional public benchmark comparisons with per-dataset numerical and visual results.These comparisons extend the experiments reported in Sec. 4.2 of the main paper.

A Training Data and Real-World Test Data

PixRestore is trained on a diverse 2.83M-image corpus spanning eight restoration tasks with balanced sampling. Its real-world evaluation uses six-degradation benchmarks of real photographs without ground-truth references, excluding denoising and super-resolution.

  • Training Data: About 2.83M training images cover eight restoration tasks, with samples drawn from each task at equal probability.The tasks are deblur, dehaze, denoise, de-rainstreak, de-raindrop, desnowing, low-light enhancement, and super-resolution.
  • Training Data: Depth-dependent haze, rainstreak, and snow degradations use image-depth datasets with synthetic pipelines, while denoising adds Gaussian noise at σ=15, 25, and 50.The synthesis follows MioIR [55] and RealRestorer [56] for depth-dependent degradations and DnCNN [1] for noise.
  • Real-World Test Data: Real-world evaluation covers six degradation types with 100 real photographs per type and no ground-truth references.These benchmarks assess restoration-method generalization to real images.
  • Real-World Test Data: Real-world denoising and super-resolution are excluded because isolated degradations are difficult to define, while real low-quality images commonly contain mixed degradations.The benchmark therefore focuses on degradation types represented in real photographs without ground-truth references.

B Pixel-space vs. Latent-space

Across eight degradation types, pixel-space diffusion shows a clear advantage for unified image restoration over latent-space alternatives. Under matched backbones and compression, Pixel DiT leads most degradation tasks, while Qwen-VAE is stronger on SR and deblurring perceptual metrics.

  • Overall comparison: Pixel-space diffusion shows a clear advantage over latent-space diffusion for unified image restoration when averaged across all eight degradation types.Per-degradation results are reported in Table S.3.
  • Experimental setup: The comparison uses the same LightningDiT-S backbone and training protocol, with matched spatial compression ratios across RGB-patch and VAE-latent models.Pixel DiT uses RGB patches of size 8, whereas latent models use FLUX-VAE, Qwen-VAE, or SD2-VAE with latent patch size 1.
  • Per-degradation results: Qwen-VAE achieves better LPIPS and MUSIQ scores than Pixel DiT on SR and deblurring, including LPIPS 0.1992 vs. 0.2067 on SR.The corresponding deblurring LPIPS values are 0.1545 vs. 0.1850, while MUSIQ values are 62.25 vs. 57.80 on SR and 46.69 vs. 41.76 on deblurring.
  • Per-degradation results: Pixel DiT performs best on all three metrics for dehazing, de-raindrop removal, de-rainstreak removal, and low-light enhancement.Its dehazing results include 24.78 PSNR, 0.1020 LPIPS, and 58.17 MUSIQ; low-light enhancement reaches 20.80 on the reported metric.

C Visual Foundation Prior

PixRestore uses a frozen visual foundation encoder to provide dense visual cues for restoration. Its encoder comparison controls architecture, layer position, model, data, optimization, sampling, and guidance settings so differences reflect pretrained vision models for UIR.

  • Visual foundation prior: A frozen visual foundation encoder provides dense visual cues for PixRestore restoration.The encoder is kept frozen during use as the visual prior.
  • Encoder comparison: The study compares CLIP, DINOv2, MAE, and SigLIP as frozen visual foundation candidates.All candidates use the ViT-B architecture and share the same parameter count.
  • Evaluation metric: DR-Score evaluates whether the target degradation has been removed, assigning higher scores to better restorations.The VLM judges the restored result using the LQ input and task description.
  • Controlled evaluation: The comparison extracts dense tokens from the 11th layer while fixing the pixel DiT, training data, optimization schedule, sampling steps, and guidance scale.Inference uses 10 sampling steps and a guidance scale of 1.0.

D DR-Score Details

DR-Score is a VLM-based auxiliary metric designed to assess degradation removal when ground-truth images are unavailable, using task-specific restoration criteria. Its method-level rankings are stable and align strongly with human preferences, despite image-level score fluctuations and occasional subtle-case failures.

  • Metric Design: DR-Score uses a vision-language model to evaluate degradation removal in UIR tasks without replacing standard full-reference metrics.It addresses the lack of ground-truth images in many real-world benchmarks and supplements existing no-reference metrics.
  • Metric Design: The metric judges task completion using degradation-specific criteria, including artifact removal, background restoration, brightness, denoising, texture preservation, sharpness, and artifact suppression.Criteria cover de-rainstreak, de-raindrop, desnow, low-light enhancement, denoising, and super-resolution.
  • Stability: Each image is evaluated five times and averaged because VLM outputs can be stochastic.Higher scores indicate better task completion, while lower scores indicate remaining degradation or restoration artifacts.
  • Stability: Method-level run standard deviation is about 0.10–0.43, while image-level standard deviation is about 5.3–6.4, indicating reproducible averages despite per-image variation.The reported method-level standard deviation remains below 0.5, supporting stable relative rankings and mean DR-Score reporting.
  • Human Alignment: DR-Score achieves the highest agreement with human preference at 90.7% in pairwise comparisons of restored results.Human judgments consider degradation removal, scene-content preservation, and suppression of obvious artifacts.
  • Limitations: DR-Score may fail to distinguish restored results with subtle appearance differences when the images are visually similar.This limitation is illustrated among representative failure cases in the human-alignment analysis.

E More Public Benchmark Comparisons

Across diverse public restoration benchmarks, PixRestore delivers strong, balanced performance, remaining competitive on no-reference metrics while maintaining fidelity-oriented results. Larger PixRestore variants generally improve restoration fidelity, and retrained baselines still lag behind, indicating that both training data and model design matter.

  • Benchmark-wide performance: PixRestore achieves strong and balanced performance across diverse restoration benchmarks rather than specializing in a single test dataset.The comparisons cover deblurring, dehazing, denoising, rain streak and raindrop removal, low-light enhancement, desnowing, and super-resolution.
  • Benchmark-wide performance: PixRestore remains competitive on no-reference metrics such as MUSIQ and AFINE-NR while emphasizing fidelity-oriented full-reference performance.Some prior methods achieve higher no-reference scores through sharper or more contrastive outputs but can fall behind on fidelity-oriented metrics.
  • Model scaling: Scaling from PixRestore-S through PixRestore-XL generally improves average performance, especially for deblur, dehaze, desnow, and de-raindrop restoration.Some individual datasets and no-reference metrics show non-monotonic behavior, but larger models repeatedly provide stronger restoration fidelity in these tasks.
  • Retrained baselines: Retrained existing methods improve on almost all datasets but remain behind PixRestore on the main full-reference metrics, showing that better training data alone is insufficient.The appendix attributes the remaining gap to model design as well as training data.
  • Visual comparisons: Visual comparisons show PixRestore producing cleaner, more balanced outputs with stronger degradation removal and detail preservation across synthetic restoration tasks.Examples include more thorough snow removal while preserving fine fence structures and sharper pole boundaries in deblurring.
Loading 2608.16793v1…