Towards Vision Transformer Unrolling Fixed-Point Algorithm: a Case Study on Image Restoration

TL;DR

This paper introduces Transformer-based unrolling frameworks for fixed-point algorithms in image restoration, achieving competitive performance with fewer parameters and training time, and demonstrating the potential of Transformers in low-level vision tasks.

Contribution

It proposes FPformer, FPRformer, and FPAformer frameworks that unroll fixed-point algorithms with Transformers, incorporating parameter sharing and acceleration techniques for improved efficiency and performance.

Findings

FPAformer achieves superior performance with only 29.82% parameters of SwinIR.

The models are trained with only 26.9% of the time required for SwinIR.

Self-supervised pre-training enhances the models' effectiveness in image restoration.

Abstract

The great success of Deep Neural Networks (DNNs) has inspired the algorithmic development of DNN-based Fixed-Point (DNN-FP) for computer vision tasks. DNN-FP methods, trained by Back-Propagation Through Time or computing the inaccurate inversion of the Jacobian, suffer from inferior representation ability. Motivated by the representation power of the Transformer, we propose a framework to unroll the FP and approximate each unrolled process via Transformer blocks, called FPformer. To reduce the high consumption of memory and computation, we come up with FPRformer by sharing parameters between the successive blocks. We further design a module to adapt Anderson acceleration to FPRformer to enlarge the unrolled iterations and improve the performance, called FPAformer. In order to fully exploit the capability of the Transformer, we apply the proposed model to image restoration, using self-supervised pre-training and supervised fine-tuning. 161 tasks from 4 categories of image restoration problems are used in the pre-training phase. Hereafter, the pre-trained FPformer, FPRformer, and FPAformer are further fine-tuned for the comparison scenarios. Using self-supervised pre-training and supervised fine-tuning, the proposed FPformer, FPRformer, and FPAformer achieve competitive performance with state-of-the-art image restoration methods and better training efficiency. FPAformer employs only 29.82% parameters used in SwinIR models, and provides superior performance after fine-tuning. To train these comparison models, it takes only 26.9% time used for training SwinIR models. It provides a promising way to introduce the Transformer in low-level vision tasks.