Enhancing Image Restoration Transformer via Adaptive Translation Equivariance

TL;DR

This paper introduces TEAFormer, a novel image restoration transformer that incorporates adaptive translation equivariance through slide indexing and component stacking, improving training, efficiency, and generalization.

Contribution

It proposes adaptive sliding indexing and component stacking strategies to preserve translation equivariance in transformers, addressing computational challenges and enhancing image restoration performance.

Findings

TEAFormer outperforms existing methods in various image restoration tasks.

The adaptive sliding indexing mechanism improves efficiency and effectiveness.

TEAFormer demonstrates better training convergence and generalization.

Abstract

Translation equivariance is a fundamental inductive bias in image restoration, ensuring that translated inputs produce translated outputs. Attention mechanisms in modern restoration transformers undermine this property, adversely impacting both training convergence and generalization. To alleviate this issue, we propose two key strategies for incorporating translation equivariance: slide indexing and component stacking. Slide indexing maintains operator responses at fixed positions, with sliding window attention being a notable example, while component stacking enables the arrangement of translation-equivariant operators in parallel or sequentially, thereby building complex architectures while preserving translation equivariance. However, these strategies still create a dilemma in model design between the high computational cost of self-attention and the fixed receptive field associated with sliding window attention. To address this, we develop an adaptive sliding indexing mechanism to efficiently select key-value pairs for each query, which are then concatenated in parallel with globally aggregated key-value pairs. The designed network, called the Translation Equivariance Adaptive Transformer (TEAFormer), is assessed across a variety of image restoration tasks. The results highlight its superiority in terms of effectiveness, training convergence, and generalization.