How Powerful Potential of Attention on Image Restoration?

TL;DR

This paper investigates the potential of attention mechanisms in image restoration by removing the feed-forward network component, introducing Continuous Scaling Attention (CSAttn), and demonstrating competitive performance across tasks.

Contribution

The paper presents a novel attention structure, CSAttn, that operates without FFN, offering a flexible alternative for image restoration tasks and providing new insights into attention mechanisms.

Findings

CSAttn outperforms CNN and Transformer-based methods in image restoration.

Removing FFN is feasible and effective for image restoration.

Simple operations within attention significantly impact performance.

Abstract

Transformers have demonstrated their effectiveness in image restoration tasks. Existing Transformer architectures typically comprise two essential components: multi-head self-attention and feed-forward network (FFN). The former captures long-range pixel dependencies, while the latter enables the model to learn complex patterns and relationships in the data. Previous studies have demonstrated that FFNs are key-value memories \cite{geva2020transformer}, which are vital in modern Transformer architectures. In this paper, we conduct an empirical study to explore the potential of attention mechanisms without using FFN and provide novel structures to demonstrate that removing FFN is flexible for image restoration. Specifically, we propose Continuous Scaling Attention (\textbf{CSAttn}), a method that computes attention continuously in three stages without using FFN. To achieve competitive performance, we propose a series of key components within the attention. Our designs provide a closer look at the attention mechanism and reveal that some simple operations can significantly affect the model performance. We apply our \textbf{CSAttn} to several image restoration tasks and show that our model can outperform CNN-based and Transformer-based image restoration approaches.