PixRestore: 画素拡散Transformerによる統合画像復元

要旨

統一的画像復元（UIR）は、単一のモデルを用いて、異なる劣化を伴う低品質（LQ）画像から高品質（HQ）コンテンツを復元することを目的とする。最近の手法の多くは、その強力な容量と生成的先験知識を活用するため、大規模な事前学習済みテキストから画像への（T2I）潜在拡散モデルを適用している。しかし、潜在T2Iモデルにおける変分オートエンコーダ（VAE）は復元に敏感な詳細を捨て去る可能性があり、一方でオープンエンドな合成の先験知識は内容と矛盾するアーティファクトを導入し得る。我々は、VAEを用いないピクセル空間の拡散Transformer（DiT）であるPixRestoreを提案する。ここでは拡散バックボーンをT2I事前学習に頼らず完全にゼロから訓練する。PixRestoreはパッチ化されたピクセル上で直接フローマッチングを行い、トークン系列を扱いやすい状態に保ちつつ、細かな詳細を保存する。異なる劣化に適応するため、PixRestoreはLQ-HQ DINO特徴類似度を用いて層特徴の信頼性を予測することを学習する。より信頼性の高い層からの特徴は密な条件付けとして融合され、一方、信頼性の低い層には劣化除去を促進するためにより強いHQ特徴の監視が与えられる。我々は多様なシーンと劣化からなる大規模コーパスでPixRestoreを訓練し、さらに効率的な推論のためにDINOに基づく敵対的目的関数を用いて一段階生成器へと微調整する。公開ベンチマークと実世界のテストセットにおける実験により、わずか約50Mパラメータと一段階推論で、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.