UCIP: A Universal Framework for Compressed Image Super-Resolution using Dynamic Prompt

Compressed Image Super-resolution aims to tackle complicated hybrid distortions, including compression artifacts and low-resolution artifacts [21, 67, 24, 26, 30, 14, 66, 7]. The first challenge for this task was held in the AIM2022 [67], where the image is first downsampled with the bicubic operation and then compressed with a JPEG codec. To solve this hard degradation, some works [26, 10, 46, 24] seek to utilize the Transformer-based architecture as their backbone. For instance, Swin2SR [10] eliminates the training instability and the requirements for large data for CSR by incorporating the Swin Transformer V2 to SwinIR [34]. HST [24] utilizes the multi-scale information flow and pre-training strategy [28] to enhance the restoration process with a hierarchical swin transformer. To further fuse the advantages of convolution and transformer, Qin et al. [46] proposes a dual-branch network, which achieves the consecutive interaction between the convolution branch and transformer branch. In contrast, to achieve the trade-off between the performance and computational cost, we aim to explore one efficient and effective framework for universal CSR problem.

As the alternative model for Transformer and Convolution Neural Networks (CNNs), MLP-like models [33, 8, 54, 64, 70, 55, 52, 51, 18, 68] have attracted great attention for their concise architectures. Typically, the noticeable success of MLP-like models stems from the well-designed token-mixing strategies [64]. The pioneering works, MLP-Mixer [54] and ResMLP [55] adopt two types of MLP layers, i.e., channel-mixing MLP and token-mixing MLP, which are responsible for the channel and spatial information interaction. To simplify the token-mixing MLP, Hou et al. [18] and Tang et al. [51] decompose the token-mixing MLP into the horizontal and vertical token-mixing MLPs. Sequentially, As-MLP [33] introduces the two-axis token shift in different channels to achieve global token mixing. There are also several works that take the hand-craft windows to enlarge the receptive field for better spatial token mixing, e.g., WaveMLP [52], and MorphMLP [70]. However, the token-mixing strategies in the above methods are restrictively fixed and lack flexibility and adaptability for different contents. To overcome this, ATM [64] is proposed to achieve the active token selection and mixing in each channel. Based on the progress of the above MLP-like models, MAXIM [56] is the first work to introduce the MLP-like model in low-level processing. However, the potential of MLP-like models is yet to be explored, as restoration model not only requires long-range token mixing but also demands efficient local feature extractions.

In the field of Natural Language Processing (NLP), prompt learning has emerged as a pivotal technique, particularly with the advent of transformer-based pre-trained models such as GPT [6, 43] and BERT [11]. Prompt learning involves providing models with specific textual cues that guide their processing of subsequent input, which helps models fast adapt to unseen tasks or applications. This approach has proven instrumental in directing models for task-specific outputs without necessitating extensive retraining or fine-tuning. Despite the success in NLP tasks, some researchers adopt prompt learning into vision tasks [20, 47, 27, 2, 22, 35, 61]. Among them, PromptIR [45] is the first to explore the low-level restoration model with prompts to facilitate multi-task learning [57, 58, 25, 48, 13]. Prompts here act as a small set of learnable parameters which interact with image features during training, providing task-specific guidance. Therefore, the prompts should be as much dynamic as possible to adapt to various degradation tasks and different pixel distributions.

As stated in Sec. 1, the universal CSR tasks entail the content/spatial- and task-adaptive modulation. An intuitive strategy is to set one prompt with the image size for each task individually or fuse them adaptively. However, it will bring severe parameter costs with the increase of the task number or image size [45]. To mitigate this, we propose the dynamic prompt strategy, and design the corresponding dynamic prompt generation module (DPM), intending to only exploit a small amount of prompt with $1\times 1\times C_{p}$ and achieve the content/spatial- and task-adaptive with the cooperation of them. To this end, we decouple the large dynamic prompt with the size of $H\times W\times C_{p}$ into two smaller matrices, i.e., the coefficients $\mathbf{w_{I}}$ with the size of $H\times W\times D$ and $D$ basic prompts with the size of $1\times 1\times C_{p}$. We can understand that for each spatial position $\{i,j\}$, there is one group of coefficients $w_{I}(i,j)$ to combine $D$ basic prompts. thereby being content/spatial-adaptive. To let the dynamic prompt perceive the task information, we generate the coefficients with the feature of input images directly, thereby being task-adaptive and suitable for any input size. Our implementation has two advantages: 1) no extra operations to adjust the spatial size of prompts, and thus the guidance information from prompts is explicit and accurate; 2) our prompts have fewer parameters and are more computationally-friendly compared to previous methods [45].

The overall architecture of $\operatorname{DPM}$ is shown in Fig. 2, where the learnable basic prompts $\textbf{P}_{\text{I}}\in\mathbb{R}^{D\times 1\times 1\times C_{P}}$ are set. Here, the $D$ and $C_{P}$ are the number of base prompts and the channel dimension of prompts. To generate dynamic prompt coefficients from input features $\textbf{F}_{\text{X}}\in\mathbb{R}^{H\times W\times C}$, an MLP layer is applied to extract the degradation prior and transform the channel dimension from $C$ to the number of basic prompts $D$. Then, the $\operatorname{softmax}$ operation is exploited to generate the composable coefficients $\textbf{w}_{\text{I}}\in\mathbb{R}^{D\times H\times W\times 1}$ for basic prompts. Based on the inversion of the above dynamic prompt decomposition, we can obtain the dynamic prompt as:

As shown in Fig. 1, we build our UCIP following the popular pipeline of compressed super-resolution backbones, which is composed of shallow feature extraction, deep feature restoration, and HR reconstruction modules. Given a low-resolution input image $\textbf{X}_{\text{LR}}\in\mathbb{R}^{H\times W\times 3}$, UCIP first extracts the shallow features $\textbf{F}_{\text{X}}\in\mathbb{R}^{H\times W\times C}$ using a patch-embedding layer, where $H$, $W$ are the spatial dimensions of features. Then, we pass $\textbf{F}_{\text{X}}$ through several PTMB to recursively remove the compression artifacts and generate the restored features $\textbf{F}_{\text{X}_{r}}$. Finally, following [63, 34], we use a series of convolution layers and nearest interpolation operations to obtain the final high-resolution output $\textbf{X}_{\text{HR}}$, which can be represented as:

To facilitate current and future research in CSR, we propose the first benchmark dataset for universal CSR, dubbed UCSR dataset, which not only considers traditional compression methods but also learning-based compression methods. We consider 6 types of compression codecs, including 3 most representative traditional codecs JPEG [59], HM [49], VTM [5], and 3 open-sourced learning-based codecs $\text{Cheng}_{\text{PSNR}}$ [9], $\text{Cheng}_{\text{SSIM}}$ [9] (abbreviated as $\text{C}_{\text{PSNR}}$ and $\text{C}_{\text{SSIM}}$ in the following paper), HIFIC [42]. Thesse three learning-based codecs are PSNR-oriented and SSIM-oriented variants from [9] and perceptual-oriented GAN-based codecs from [42], respectively. To cover the prominent compression types in real scenarios, we consider four different compression qualities for each codec, except for HIFIC, since only the weights for three bitrate points are released.

To generate the training dataset, we choose the popular DF2K [1, 53], which contains 3450 high-quality images. Each image is downsampled by a scale factor of 4 using MATLAB bicubic algorithm. Then, we compress the downsampled images with six different compression algorithms to yield the training dataset of all competitive methods and our UCIP. The quality factors we used for different codecs are respectively as: (i) [10, 20, 30, 40] for JPEG, where the smaller value means poorer image quality. (ii) [32, 37, 42, 47] for HM, VTM, where value denotes the quantization parameter (QP), and larger value means poor quality. (iii) [1, 2, 3, 4] for $\text{C}_{\text{PSNR}}$, $\text{C}_{\text{SSIM}}$, where the smaller value indicates poorer quality. We adopt the implementation in the popular open-sourced compression tools compressai [3]. (iv) [‘low’, ‘med’, ‘high’] for HIFIC, where ‘low’ indicates the poorest image quality. We use the PyTorch implementation [16] to compress images. All the methods are trained from scratch on our proposed benchmarks. We adopt the same process to generate the evaluation datasets based on five commonly used benchmarks: Set5 [4], Set14 [69], BSD100 [40], Urban100 [19] and Manga109 [41].

We evaluate UCIP with six state-of-the-art models on our CSR benchmarks which composes of five commonly adopted datasets: Set5 [4], Set14 [69], BSD100 [40], Urban100 [19] and Manga109 [41]. The compared models include the fully-convolutional network RRDB [63], the transformer-based image restoration model SwinIR [34] and its upgraded version Swin2SR [10], the MLP-like model MAXIM [56] and two multi-task models AIRNet [25] and PromptIR [45]. We add the HR reconstruction module to last three models, enabling them to perform super-resolution tasks. All compared methods are trained from scratch with our proposed UCSR dataset for fair comparisons.

As demonstrated in Table. 1 and Table. 2, our UCIP outperforms all other methods on almost all codecs and compression qualities. Particularly, UCIP achieves PSNR gain of up to 0.45dB against PromptIR [45] with only one-third the number of parameters. Another intriguing observation is that the gains provided by UCIP become more significant as the compression ratio decreases. We attribute this to the preservation of more high-frequency information at milder compression levels. The abundance of high-frequency details further enhances the capability of PTMM to conduct global-wise informative tokens extraction, thus leads to a better performance.

As illustrated in Fig. 3, UCIP leverages the implicit guidance of the dynamic prompt to recover more textural details while avoiding the generation of artifacts. Specifically, as observed in the first row, our model recovers the clearest texture of the monarch. Besides, in the second and final rows, images reconstructed by our method exhibit clearer edges and fewer distorted lines. For the third row, our method successfully removes compression artifacts, while other methods suffer from blocked and blurry outputs. We attribute these performances to the generation of the dynamic prompt and the fusion of global tokens with local features.

It is noteworthy that, though we do not specifically tailor prompts for various compression qualities within certain codec, experimental evidence suggests that our dynamic prompt not only possesses task-specific adaptability but is also capable of handling different distortion degrades. As shown in Fig. 4, our method maintains robust image restoration capabilities across three levels of compression qualities (e.g., always recovers straight lines on the right side of image) .

Prompt learning can be utilized in two popular ways: (i) One is to utilize prompt learning for multi-task learning, e.g., PromptIR [45], ProRes [39], PIP [32] in low-level vision, which needs to train whole model from scratch; (ii) another is prompt tuning, which requires a strong baseline model and aims to optimize only a small part of parameters for downstream tasks. Notably, in the CSR field, there are no pre-trained baseline models on multiple types of compression artifacts existed, which prevents us to study prompt tuning in the beginning. And thus, we build the first Universal CSR framework and corresponding dataset with the first way, which follows existing prompt learning works in low-level vision [45, 39, 32, 38]. However, training a model from scratch is time consuming. To further explore the potential of our proposed UCIP in prompt tuning, we choose two unseen codecs, including one traditional codec WebP [15] and one learning-based codec ELIC [17], to fine-tune UCIP. In Tab. 3, we explore four ways of fine-tuning: i) directly evaluated pre-trained model without fine-tuning. ii) pre-training without prompt, then adding prompt and only training prompt parameters on new tasks. iii) pre-training without prompt, then adding prompt and only training prompt parameters on new tasks. iv) pre-training with prompt, then fine-tuning only prompt parameters on new tasks. All the experiments are conducted under the same settings with the same training iterations. As shown in Tab. 3, compared between ii) and iii), tuning only prompt achieves comparable performance on ELIC codec against tuning full model. Compared between iii) and iv), tuning only prompt based on UCIP achieves comparable and even better performances against tuning full model after adding prompt. The experimental results indicate that our proposed UCIP can serve as a strong baseline model in CSR field, which will also benefit the prompt tuning for new codecs in future works.