PixRestore: 픽셀 확산 트랜스포머를 통한 통합 이미지 복원

초록

통합 이미지 복원(UIR)은 다양한 열화가 있는 저화질(LQ) 이미지에서 고화질(HQ) 콘텐츠를 단일 모델로 복구하는 것을 목표로 한다. 최근 대부분의 방법은 강력한 용량과 생성 사전(generative prior)을 활용하기 위해 대규모 사전 학습된 텍스트-이미지(T2I) 잠재 확산 모델을 적용한다. 그러나 잠재 T2I 모델의 변분 오토인코더(VAE)는 복원에 민감한 세부 정보를 버릴 수 있으며, 개방형 합성 사전은 콘텐츠와 불일치하는 인공물(artifact)을 유발할 수 있다. 우리는 UIR을 위한 VAE 없는 픽셀 공간 확산 트랜스포머(DiT)인 PixRestore를 제시한다. 여기서 확산 백본은 T2I 사전 학습에 의존하지 않고 완전히 처음부터 학습된다. PixRestore는 패치화된 픽셀에 대해 직접 플로우 매칭(flow matching)을 수행하여 토큰 시퀀스를 관리 가능하게 유지하면서 미세한 세부 정보를 보존한다. 다양한 열화에 적응하기 위해 PixRestore는 LQ–HQ DINO 특징 유사도를 사용하여 레이어 특징의 신뢰성을 예측하도록 학습한다. 더 신뢰할 수 있는 레이어의 특징은 밀집 조건(dense conditioning)으로 융합되는 반면, 신뢰도가 낮은 레이어는 열화 제거를 장려하기 위해 더 강한 HQ 특징 감독(supervision)을 받는다. 우리는 다양한 장면과 열화로 구성된 대규모 말뭉치에서 PixRestore를 학습하고, 효율적인 추론을 위해 DINO 기반 적대적 목적 함수를 사용하여 단일 단계 생성기로 추가 미세 조정한다. 공개 벤치마크와 실제 세계 테스트 세트에 대한 실험은 약 5천만 개의 파라미터와 단일 단계 추론만으로도 PixRestore가 경쟁 UIR 모델들 중에서 전반적인 충실도, 지각 품질, 열화에 대한 견고성에서 최고의 성능을 달성하면서 훨씬 더 효율적임을 보여준다. 더 큰 PixRestore 변형은 성능을 더욱 향상시킬 수 있으며, 이는 우리 픽셀 공간 설계의 확장성을 입증한다. 코드와 선별된 벤치마크는 https://github.com/csslc/PixRestore에서 확인할 수 있다.

English

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.