[Paper Review] Pre-Trained Image Processing Transformer

Hanting Chen, Yunhe Wang+8 more|arXiv (Cornell University)|Dec 1, 2020

Advanced Image Processing TechniquesComputer Science89 references120 citations

TL;DR

IPT is a transformer-based model pre-trained on a large synthetic image processing corpus (ImageNet-derived) to handle multiple low-level vision tasks such as super-resolution, denoising, and deraining, achieving strong fine-tuned performance.

ABSTRACT

As the computing power of modern hardware is increasing strongly, pre-trained deep learning models (e.g., BERT, GPT-3) learned on large-scale datasets have shown their effectiveness over conventional methods. The big progress is mainly contributed to the representation ability of transformer and its variant architectures. In this paper, we study the low-level computer vision task (e.g., denoising, super-resolution and deraining) and develop a new pre-trained model, namely, image processing transformer (IPT). To maximally excavate the capability of transformer, we present to utilize the well-known ImageNet benchmark for generating a large amount of corrupted image pairs. The IPT model is trained on these images with multi-heads and multi-tails. In addition, the contrastive learning is introduced for well adapting to different image processing tasks. The pre-trained model can therefore efficiently employed on desired task after fine-tuning. With only one pre-trained model, IPT outperforms the current state-of-the-art methods on various low-level benchmarks. Code is available at https://github.com/huawei-noah/Pretrained-IPT and https://gitee.com/mindspore/mindspore/tree/master/model_zoo/research/cv/IPT

Motivation & Objective

• Motivate the need for cross-task pre-training in low-level vision tasks (where data can be scarce and task types vary).
• Propose a universal pre-trained transformer (IPT) with task-specific heads/tails and a shared body for image processing tasks.
• Leverage large-scale synthetic degradation of ImageNet images to create diverse training data.
• Incorporate contrastive learning to enhance patch-level representations and generalization to unseen tasks.
• Demonstrate that a single pre-trained IPT can surpass task-specific models after fine-tuning on SR, denoising, and deraining.

Proposed method

• Introduce a four-component IPT: multi-task heads for each task, a shared encoder-decoder transformer body, and multi-tail outputs for reconstruction.
• Convert input features to patches (visual words) with positional encodings; process via Transformer encoder and task-aware decoder with embeddings.
• Pre-train IPT on a large synthesized ImageNet-derived dataset with multiple degradation models (bicubic SR, Gaussian noise, rain, etc.).
• Use supervised losses for reconstruction plus a contrastive loss between patches from the same image to learn universal features (L_IPT = λ L_contrastive + L_supervised).
• Fine-tune the pre-trained IPT on specific tasks (e.g., ×2/×3/×4 SR, denoising, deraining) by freezing unused heads/tails as needed.

Experimental results

Research questions

• RQ1 Can a single pre-trained transformer model generalize across multiple low-level image processing tasks after fine-tuning?
• RQ2 Does large-scale pre-training on degraded ImageNet data improve SR, denoising, and deraining performance over task-specific models?
• RQ3 What is the impact of contrastive learning on the quality and generalization of IPT across tasks?
• RQ4 How does IPT compare to state-of-the-art CNN-based methods when pre-trained on ImageNet and fine-tuned?
• RQ5 What are the effects of multi-task training versus single-task pre-training on transfer to new tasks?

Key findings

• IPT outperforms many task-specific methods on several low-level benchmarks after fine-tuning.
• On SR, IPT achieves PSNR values of 38.37 (Set5), 34.43 (Set14), 32.48 (B100), and 33.76 (Urban100) for ×2; 34.81, 30.85, 29.38, 29.49 for ×3; and 32.64, 29.01, 27.82, 27.26 for ×4.
• On color image denoising (Gaussian noise), IPT achieves 30.75 (BSD68, σ=30) and 28.39 (Urban100, σ=50); 30.75 and 28.39 are among the best reported.
• On deraining, IPT achieves 41.62 dB PSNR on Rain100L, surpassing prior methods.
• Contrastive learning (λ > 0) improves PSNR by up to ~0.1 dB on SR when combined with supervised loss.
• Multi-task pre-training enhances generalization to unseen tasks compared to single-task pre-training.