Source-linked AI summary
Pre-Trained Image Processing Transformer
Hanting Chen, Yunhe Wang, Tianyu Guo, Chang Xu, Yiping Deng, Zhenhua Liu, Siwei Ma, Chunjing Xu, Chao Xu, Wen Gao
TL;DR
Image-processing tasks lack a broadly generalized pre-training strategy, especially because their inputs and outputs are both images and degradation models vary. IPT pre-trains a shared transformer with task-specific heads and tails on corrupted ImageNet pairs, using contrastive learning for adaptation. The resulting model outperforms state-of-the-art methods across various low-level benchmarks, while its generalization is limited by degradation types that cannot all be synthesized.
Problem
Few studies generalize pre-training across image-processing tasks, whose image-to-image mappings differ from the classification tasks used in much existing pre-training.
Method
IPT pre-trains a shared transformer body with task-specific heads and tails on synthetically corrupted ImageNet image pairs, incorporating contrastive learning for cross-task adaptation.
Results
IPT outperforms state-of-the-art methods across various low-level image-processing benchmarks using one pre-trained model after fine-tuning.
Takeaways & Limitations
One pre-trained IPT model can be fine-tuned for desired image-processing tasks after multi-task supervised and self-supervised training.
Takeaways & Limitations
IPT cannot synthesize images for all image-processing tasks because degradation models vary widely, including practical noise levels.
Abstract
from arXiv · showhide
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
1. Introduction
Image processing benefits from pre-training, but cross-task pre-training is uncommon because image-processing inputs and outputs are both images and target tasks vary. IPT addresses this with a shared transformer and task-specific heads and tails trained on synthetically corrupted ImageNet images.
- Image-processing results can substantially influence later image recognition and understanding, motivating stronger low-level vision methods.
- Few studies have generalized pretraining across image processing tasks, despite relationships among tasks suggesting transfer may be useful.
- Pre-training can address limited task-specific data, distribution shifts, and the need to prepare separate modules before the requested image-processing task is known.
- Existing pre-training strategies are difficult to apply directly because image processing maps images to images rather than class labels.
- IPT uses multiple task-specific head-tail pairs with one shared body, while ImageNet images are degraded into diverse training counterparts.
- The transformer processes flattened image patches with positional and task embeddings, and contrastive loss supports adaptation across image-processing tasks.
2. Related Works
Prior work established image-processing methods across tasks and transformer-based pre-training in language and vision. IPT builds on these directions by applying a shared transformer framework to low-level image processing.
- Image Processing: Image processing includes tasks such as super-resolution, denoising, dehazing, deraining, and deblurring, with deep-learning methods often targeting one or several tasks.
- Transformer Pre-training: Transformer variants have served as unsupervised or self-supervised pre-training frameworks for language tasks, including autoregressive prediction and masked-word prediction.
- Transformer in Vision: Vision research has explored self-attention by augmenting convolutional networks with spatial, channel, or contextual attention.
- IPT Architecture: IPT uses a shared transformer body with encoder and decoder, task-specific heads and tails, patch-based visual words, and reconstructed output patches.
3. Image Processing Transformer
IPT is a transformer-based pre-trained model for multiple image-processing tasks, using task-specific heads and tails around a shared transformer body. It pre-trains on diverse ImageNet-derived corrupted images with supervised reconstruction and contrastive objectives, then supports task-specific fine-tuning.
- IPT architecture: IPT uses multiple task-specific heads and tails with a shared encoder-decoder transformer body for different image-processing tasks.The architecture targets tasks including super-resolution, denoising, and deraining.
- IPT architecture: Input features are divided into patches, augmented with learnable position encodings, and processed by transformer layers containing multi-head self-attention and feed-forward networks.The encoder preserves the patch feature size while the decoder uses task-specific embeddings for different tasks.
- Pre-training on ImageNet: ImageNet provides over 1 million diverse natural images for synthesizing corrupted image pairs across image-processing tasks.The semantic labels are removed, and task-specific degradation models generate the corrupted inputs.
- Pre-training on ImageNet: The supervised objective reconstructs desired images from corrupted inputs while training multiple image-processing tasks simultaneously.Each task uses its corresponding head, tail, and task embedding; after pre-training, the model can be fine-tuned for a desired task.
- Limitations: The approach cannot synthesize images for all image-processing tasks because degradation models vary, including practical variation in noise levels.The paper identifies enhanced generalization as necessary for the resulting IPT.
- Pre-training on ImageNet: Contrastive learning brings patched features from the same images closer and separates patches from different images to improve adaptation to unseen tasks.The final objective combines the contrastive loss with the supervised loss using a λ-balanced formulation.
4. Experiments
Experiments evaluate IPT across super-resolution, denoising, deraining, generalization, data scaling, and contrastive-learning ablations. IPT achieves strong benchmark results, while large-scale pre-training and contrastive learning improve performance under reported conditions.
- Super-resolution: IPT achieves the best super-resolution performance at ×2, ×3, and ×4 scales across all evaluated datasets.On Urban100 at ×2, IPT reaches 33.76dB PSNR, exceeding other methods by more than ∼0.4dB.
- Super-resolution: 33.76dB PSNR on ×2-scale Urban100 exceeds other methods by more than ∼0.4dB.The paper contrasts this gain with previous SOTA improvements of less than 0.2dB.
- Denoising: IPT achieves the best color-denoising results across Gaussian noise levels on BSD68 and Urban100.The reported Urban100 improvement over the state of the art is approximately 0.3dB.
- Deraining: IPT achieves 41.62dB on Rain100L deraining, improving on state-of-the-art methods by 1.62dB.The evaluation uses the synthesized Rain100L dataset, which contains 100 rainy images.
- Generalization Ability: On unseen synthesized corruptions, IPT outperforms conventional methods in image denoising.The experiment tests corrupted images not included in the synthesized ImageNet pre-training dataset.
- Ablation Study: With large-scale data, transformer-based models outperform CNN models, whereas CNNs perform better without pre-training or below 60% of the dataset.The comparison uses 20%, 40%, 60%, 80%, and 100% of the synthesized ImageNet dataset.
- Ablation Study: 38.37dB PSNR with λ = 0.1 is about 0.1dB higher than 38.27dB with λ = 0 in ×2 super-resolution on Set4.The comparison evaluates the contrastive-learning loss balance parameter λ.
5. Conclusions and Discussions
The paper presents IPT as a pre-trained transformer for multiple low-level image-processing tasks, using synthesized ImageNet data and supervised and self-supervised training. Experiments show that one pre-trained model can outperform state-of-the-art methods after quick fine-tuning.
- IPT combines multi-heads, multi-tails, and a shared transformer body to support tasks including super-resolution and denoising.
- The model is trained on synthesized ImageNet counterparts generated by degrading original images into paired training data.
- IPT uses supervised and self-supervised approaches to capture intrinsic features for low-level image processing.
- One pre-trained IPT model outperforms state-of-the-art methods after quick fine-tuning.
A. Results on Deblurring
IPT is evaluated on GoPro for image deblurring after task-specific fine-tuning and architectural modifications. It achieves the best reported quantitative and visual results among the compared methods.
- IPT achieves the best deblurring results among all compared methods on the GoPro dataset.The model is fine-tuned and tested on GoPro, with patch size, patch dimension, and feature count modified for a higher receptive field.
- IPT produces the best visual quality among previous models in the reported deblurring visualizations.
B. Architecture of IPT
IPT uses task-specific heads and tails around a shared transformer body. The full model has 114M parameters and 33G FLOPs, with more parameters but fewer FLOPs than EDSR.
- IPT consists of task-specific heads, a shared transformer body, and task-specific tails.Each head contains one convolutional layer and two ResBlocks, while the body contains 12 encoder and 12 decoder layers.
- The denoising and deraining tails use convolutional layers, while super-resolution tails use pixelshuffle layers for different upsampling scales.
- 114M parameters and 33G FLOPs give IPT more parameters but fewer FLOPs than EDSR's 43M parameters and 99G FLOPs.
C. Impact of Multi-task Training
Multi-task training preserves performance on the six pre-trained tasks and improves transfer to other tasks. This supports learning a universal representation for image processing.
- Multi-task training covers six tasks: ×2, ×3, and ×4 super-resolution, plus denoising at noise levels 30 and 50 and deraining.
- Performance remains almost the same between multi-task and single-task training on tasks pre-trained with large-scale ImageNet data.
- Multi-task pre-training performs about 0.3dB better than single-task pre-training when transferred to other tasks.
- The transfer improvement suggests that multi-task training learns universal representations of image-processing tasks.
D. Visualization of Embeddings
The visualizations show that IPT learns position embeddings aligned with image layout and task embeddings that reflect task-specific spatial relationships. Removing or replacing these learnable embeddings reduces performance across tasks.
- Position embeddings: Position embeddings give patches with similar rows or columns similar representations, encoding useful image-location information.Learnable position embeddings outperform fixed or absent embeddings by 0.2dB to 0.3dB across different tasks.
- Task embeddings: Task embeddings vary with task scale and degradation structure, reflecting distinct spatial relationships for super-resolution, deraining, and denoising.The visualizations associate farther patch relationships with ×2 super-resolution, stronger vertical attention with deraining, and neighbour similarity with higher-noise denoising.
- Task embeddings: Removing task embeddings causes a significant accuracy drop of 0.1dB to 0.5dB across different tasks.This supports using task embeddings to adapt the shared transformer to different image-processing settings.