ATD: Improved Transformer with Adaptive Token Dictionary for Image Restoration

Leheng Zhang; Wei Long; Yawei Li; Xingyu Zhou; Xiaorui Zhao; Shuhang Gu

cs.CVMar 3, 2026

ATD: Improved Transformer with Adaptive Token Dictionary for Image Restoration

Authors: Leheng Zhang, Wei Long, Yawei Li, Xingyu Zhou, Xiaorui Zhao, Shuhang Gu

Organizations: School of Computer Science and Engineering (School of Cyber Security), University of Electronic Science and Technology of China (UESTC), China · Computer Vision Lab, ETH Zürich, Zürich, Switzerland and with the Integrated Systems Laboratory, ETH Zürich, Zürich, Switzerland

Abstract

Recently, Transformers have gained significant popularity in image restoration tasks such as image super-resolution and denoising, owing to their superior performance. However, balancing performance and computational burden remains a long-standing problem for transformer-based architectures. Due to the quadratic complexity of self-attention, existing methods often restrict attention to local windows, resulting in limited receptive field and suboptimal performance. To address this issue, we propose Adaptive Token Dictionary (ATD), a novel transformer-based architecture for image restoration that enables global dependency modeling with linear complexity relative to image size. The ATD model incorporates a learnable token dictionary, which summarizes external image priors (i.e., typical image structures) during the training process. To utilize this information, we introduce a token dictionary cross-attention (TDCA) mechanism that enhances the input features via interaction with the learned dictionary. Furthermore, we exploit the category information embedded in the TDCA attention maps to group input features into multiple categories, each representing a cluster of similar features across the image and serving as an attention group. We also integrate the learned category information into the feed-forward network to further improve feature fusion. ATD and its lightweight version ATD-light, achieve state-of-the-art performance on multiple image super-resolution benchmarks. Moreover, we develop ATD-U, a multi-scale variant of ATD, to address other image restoration tasks, including image denoising and JPEG compression artifacts removal. Extensive experiments demonstrate the superiority of out proposed models, both quantitatively and qualitatively.

Figures & tables

Fig. 1 : Visual comparisons on different attention mechanisms. (a) Window-based self-attention constrain the attention to local windows due to the quadratic computational complexity, resulting in limited receptive field. (b) Our proposed token dictionary cross-attention leverages the typical image structures summarized in the token dictionary to incorporate external information to input features. (c) Our proposed adaptive category-based self-attention exploits relationships within categories, connecting distant but similar tokens across the image.

Panel of Figure 2 — Fig. 2 : Overall architectures of the proposed ATD and ATD-U networks. Each ATD block contains several consecutive transformer layers and a learnable token dictionary. The transformer layer combines three attention branches: token dictionary cross-attention ( IV-B ), adaptive category-based self-attention ( IV-C ), and window-based self-attention to enhance image feature. The attention map of TDCA branch is further utilized by the ACMSA branch for categorization and CFFN for dictionary entry selection, respectively.

Panel of Figure 3 — Fig. 3 : Illustrations of the proposed (a) Token Dictionary Cross-Attention (TDCA), (b) Adaptive Category-based Multi-head Self-Attention (AC-MSA), and (c) Category-aware Feed-Forward Network (CFFN). Descriptions of TDCA and ACMSA are provided in Sec. IV-B and IV-C , respectively. Further details of the $\operatorname{Categorize}$ operation can be found in Eq. 12 and Fig. 5 .

Fig. 4 : Distribution of the maximum attention values over image tokens with respect to the token dictionary in TDCA across different layers. Y-axis is log-scaled for better visualization.

Fig. 5 : Illustration of the proposed $\operatorname{Categorize}$ operation. The attention map is flattened and sorted based on the index of the maximum value in each row as denoted in Eq. 12 and Eq. 14 . This effectively clusters tokens sharing similar features for self-attention. Subsequently, the $\operatorname{UnCategorize}$ operation applies the inverse permutation to restore the original spatial structure.

Panel of Figure 6 — Fig. 6 : Visualization of the attention groups for window-based self-attention and the proposed adaptive category-based self-attention. Green and red regions represent the query tokens and their corresponding key tokens , respectively. In window-based self-attention (window size $n_{w}=16^{2}$ ), similarity computations are confined within local windows, which limits the ability to exploit similar regions across the image. In contrast, the proposed adaptive category-based self-attention (category size $n_{s}=256$ , equal to window size) could connect distant but similar tokens across the image. Additionally, as shown in (c) and (d), the attention categories vary across different layers, allowing the receptive field to expand progressively. Zoom in for a better view.

Method	Scale	Params	Set5 [ 103 ]		Set14 [ 51 ]		BSD100 [ 104 ]		Urban100 [ 105 ]		Manga109 [ 106 ]
Method	Scale	Params	PSNR	SSIM	PSNR	SSIM	PSNR	SSIM	PSNR	SSIM	PSNR	SSIM
EDSR [ 3 ]	$\times$ 2	42.6M	38.11	0.9602	33.92	0.9195	32.32	0.9013	32.93	0.9351	39.10	0.9773
RCAN [ 31 ]	$\times$ 2	15.4M	38.27	0.9614	34.12	0.9216	32.41	0.9027	33.34	0.9384	39.44	0.9786
SAN [ 75 ]	$\times$ 2	15.7M	38.31	0.9620	34.07	0.9213	32.42	0.9028	33.10	0.9370	39.32	0.9792
HAN [ 76 ]	$\times$ 2	63.6M	38.27	0.9614	34.16	0.9217	32.41	0.9027	33.35	0.9385	39.46	0.9785
SwinIR [ 4 ]	$\times$ 2	11.8M	38.42	0.9623	34.46	0.9250	32.53	0.9041	33.81	0.9433	39.92	0.9797

TABLE I : Quantitative comparison (PSNR/SSIM) with state-of-the-art methods on classical image SR task. Best and second best results are marked with bold and underline . Experimental details can be found in Sec. V-A .

Method	Scale	Params	Set5 [ 103 ]		Set14 [ 51 ]		BSD100 [ 104 ]		Urban100 [ 105 ]		Manga109 [ 106 ]
Method	Scale	Params	PSNR	SSIM	PSNR	SSIM	PSNR	SSIM	PSNR	SSIM	PSNR	SSIM
CARN [ 70 ]	$\times$ 2	1,592K	37.76	0.9590	33.52	0.9166	32.09	0.8978	31.92	0.9256	38.36	0.9765
IMDN [ 71 ]	$\times$ 2	694K	38.00	0.9605	33.63	0.9177	32.19	0.8996	32.17	0.9283	38.88	0.9774
LAPAR-A [ 108 ]	$\times$ 2	548K	38.01	0.9605	33.62	0.9183	32.19	0.8999	32.10	0.9283	38.67	0.9772
LatticeNet [ 72 ]	$\times$ 2	756K	38.15	0.9610	33.78	0.9193	32.25	0.9005	32.43	0.9302	-	-
SwinIR-light [ 4 ]	$\times$ 2	910K	38.14	0.9611	33.86	0.9206	32.31	0.9012	32.76	0.9340	39.12	0.9783

TABLE II : Quantitative comparison (PSNR/SSIM) with state-of-the-art methods on lightweight image SR task. Best and second best results are marked with bold and underline . Experimental details can be found in Sec. V-A .

Panel of Figure 9 — Fig. 7 : Qualitative comparison on $\times 4$ image super-resolution.

Model	$128\times 128$			$256\times 256$			$512\times 512$			Urban100 [ 105 ]		Manga109 [ 106 ]
Model	FLOPs	Runtime	Memory	FLOPs	Runtime	Memory	FLOPs	Runtime	Memory	PSNR	SSIM	PSNR	SSIM
ART [ 47 ]	0.39T	104ms	1565M	2.85T	1237ms	15037M	32.3T	-	OOM	27.77	0.8321	32.31	0.9283
HAT [ 5 ]	0.41T	105ms	1443M	1.65T	450ms	4057M	6.60T	1950ms	12855M	27.97	0.8368	32.48	0.9292
MambaIRv2 [ 94 ]	0.45T	175ms	1057M	1.78T	821ms	2371M	7.13T	3830ms	7727M	27.89	0.8344	32.57	0.9295
ATD (Ours)	0.51T	137ms	1025M	2.03T	612ms	2385M	8.11T	2530ms	8681M	28.24	0.8420	32.71	0.9311

TABLE III : Comparisons on computational overhead. We use single RTX4090 GPU to test the runtime and GPU memory footprint for various LQ input sizes. “OOM” denotes out-of-memory.

Model	Urban100 [ 105 ]		Manga109 [ 106 ]
Model	PSNR	SSIM	PSNR	SSIM
baseline	26.40	0.7956	30.84	0.9131
+TDCA	26.46	0.7971	30.92	0.9140
+ACMSA	26.70	0.8023	31.19	0.9165
+CFFN	26.76	0.8039	31.21	0.9169

TABLE IV : Ablation studies on proposed key components TDCA, ACMSA, and CFFN in metric of PSNR.

$\tau^{\prime}$	$M$	Urban100 [ 105 ]		Manga109 [ 106 ]
$\tau^{\prime}$	$M$	PSNR	SSIM	PSNR	SSIM
	32	26.70	0.8019	31.15	0.9162
	64	26.73	0.8029	31.15	0.9160
	128	26.70	0.8026	31.19	0.9167
$\surd$	32	26.70	0.8022	31.17	0.9163
$\surd$	64	26.72	0.8027	31.19	0.9165

TABLE V : Ablation studies on improved scaling factor

\tau^{\prime}

and choice of dictionary size

M

$n_{s}$	Urban100 [ 105 ]		Manga109 [ 106 ]
$n_{s}$	PSNR	SSIM	PSNR	SSIM
0	26.46	0.7971	30.92	0.9140
32	26.66	0.8015	31.13	0.9160
64	26.73	0.8028	31.19	0.9170
128	26.76	0.8039	31.21	0.9169
256	26.75	0.8036	31.23	0.9172

TABLE VI : Ablation studies on choice of sub-category size

n_{s}

n_{s}=0

represents removing the ACMSA branch.

Dataset	$\sigma$	BM3D [ 50 ]	DnCNN [ 8 ]	IRCNN [ 9 ]	FFDNet [ 113 ]	RNAN [ 10 ]	DRUNet [ 11 ]	SwinIR [ 4 ]	Restormer [ 12 ]	SCUNet [ 86 ]	Xformer [ 89 ]	ATD-U (ours)
CBSD68 [ 104 ]	15	33.52	33.90	33.86	33.87	-	34.30	34.42	34.40	33.40	34.43	34.47
	25	30.71	31.24	31.16	31.21	-	31.69	31.78	31.79	31.79	31.82	31.85
	50	27.38	27.95	27.86	27.96	28.27	28.51	28.56	28.60	28.61	28.63	28.66
Kodak24 [ 114 ]	15	34.28	34.60	34.69	34.63	-	35.31	35.34	35.35	35.34	35.39	35.57
	25	32.15	32.14	32.18	32.13	-	32.89	32.89	32.93	32.92	32.99	33.12
	50	28.46	28.95	28.93	28.98	29.58	29.86	29.79	29.87	29.87	29.94	30.12

TABLE VII : Quantitative comparison (PSNR/SSIM) with state-of-the-art methods on color image denoising task. Best and second best results are marked with bold and underline . Experimental details can be found in Sec. VI-A .

Dataset	$\sigma$	BM3D [ 50 ]	DnCNN [ 8 ]	IRCNN [ 9 ]	FFDNet [ 113 ]	RNAN [ 10 ]	DRUNet [ 11 ]	SwinIR [ 4 ]	Restormer [ 12 ]	SCUNet [ 86 ]	Xformer [ 89 ]	ATD-U (ours)
Set12 [ 8 ]	15	32.37	32.86	32.76	32.75	-	33.25	33.36	33.42	33.43	33.46	33.49
	25	29.97	30.44	30.37	30.43	-	30.94	31.01	31.08	31.09	31.16	31.13
	50	26.72	27.18	27.12	27.32	27.70	27.90	27.91	28.00	28.04	28.10	28.10
BSD68 [ 104 ]	15	31.08	31.73	31.63	31.63	-	31.91	31.97	31.96	31.99	31.98	32.02
	25	28.57	29.23	29.15	29.19	-	29.48	29.50	29.52	29.55	29.55	29.56
	50	25.60	26.23	26.19	26.29	26.48	26.59	26.58	26.62	26.67	26.65	26.65

TABLE VIII : Quantitative comparison (PSNR/SSIM) with state-of-the-art methods on grayscale image denoising task. Best and second best results are marked with bold and underline . Experimental details can be found in Sec. VI-A .

Fig. 8 : Qualitative comparison on color Gaussian image denoising with noise level $\sigma=50$ .

Fig. 9 : Qualitative comparison on grayscale Gaussian image denoising with noise level $\sigma=50$ .

Dataset	$q$	ARCNN [ 121 ]	DnCNN [ 8 ]	RNAN [ 10 ]	FBCNN [ 122 ]	DRUNet [ 11 ]	SwinIR [ 4 ]	ART [ 47 ]	MambaIR [ 92 ]	ATD-U (ours)
Classic5 [ 119 ]	10	29.03/0.7929	29.40/0.8026	29.96/0.8178	30.12/0.8223	30.16/0.8234	30.27/0.8249	30.27 / 0.8258	30.27 /0.8256	30.34 / 0.8264
	20	31.15/0.8517	31.63/0.8610	32.11/0.8693	32.31/0.8724	32.39/0.8734	32.52 / 0.8748	-/-	-/-	32.58 / 0.8755
	30	32.51/0.8806	32.91/0.8861	33.38/0.8924	33.54/0.8943	33.59/0.8949	33.73/0.8961	33.74 /0.8964	33.74 / 0.8965	33.80 / 0.8967
	40	33.32/0.8953	33.77/0.9003	34.27/0.9061	34.35/0.9070	34.41/0.9075	34.52/0.9082	34.55 / 0.9086	34.53/0.9084	34.59 / 0.9088
LIVE1 [ 120 ]	10	28.96/0.8076	29.19/0.8123	29.63/0.8239	29.75/0.8268	29.79/0.8278	29.86/0.8287	29.89 /0.8300	29.88/ 0.8301	29.91 / 0.8304
	20	31.29/0.8733	31.59/0.8802	32.03/0.8877	32.13/0.8893	32.17/0.8899	32.25 / 0.8909	-/-	-/-	32.29 / 0.8917

TABLE IX : Quantitative comparison (PSNR/SSIM) with state-of-the-art methods on JPEG compression artifacts removal task. Best and second best results are marked with bold and underline . Experimental details can be found in Sec. VI-A .

Fig. 10 : Qualitative comparison on JPEG CAR with quality factor $q=10$ .

Explore similar work

Fast and Accurate Image Restoration and Generation with Rank Enhanced Linear Attention
May 22, 2025Yuang AiEfficient ViTs Linear Attention
University of Chinese Academy of Sciences, Beijing, China
In-Token Learning for High-Fidelity Image Restoration via Diffusion Transformers
Sep 27, 2026Xingfu Yi, Xiaoxue YuImage-to-Image Translation Diffusion Transformer
Independent Researcher, Hangzhou, China · Zhejiang University, Hangzhou, China
SAT: Selective Aggregation Transformer for Image Super-Resolution
Apr 9, 2026Dinh Phu Tran, Thao Do, Saad Wazir +3Image Super-Resolution Efficient ViTs
School of Computing, KAIST, Republic of Korea

May 22, 2025J/K move · Enter open · S save

Fast and Accurate Image Restoration and Generation with Rank Enhanced Linear Attention

Yuang Ai

University of Chinese Academy of Sciences, Beijing, China

Efficient ViTs Linear Attention Image Restoration

Open paper AskPDF arXiv

Transformer-based models have made remarkable progress in image restoration (IR) tasks. However, the quadratic complexity of self-attention in Transformer hinders its applicability to high-resolution images. Existing methods mitigate this issue with sparse or window-based attention, yet inherently limit global context modeling. Linear attention, a variant of softmax attention, demonstrates promise in global context modeling while maintaining linear complexity, offering a potential solution to the above challenge. Despite its efficiency benefits, vanilla linear attention suffers from a significant performance drop in IR, largely due to the low-rank nature of its attention map. To counter this, we propose Rank Enhanced Linear Attention (RELA), a simple yet effective method that enriches feature representations by integrating a lightweight depthwise convolution. Building upon RELA, we propose an efficient and effective Vision Transformer, named LAformer. LAformer eliminates hardware-inefficient operations such as softmax and window shifting, enabling efficient processing of high-resolution images. Extensive experiments across 7 IR tasks and 21 benchmarks demonstrate that LAformer outperforms SOTA methods and offers significant computational advantages. Furthermore, we extend LAformer to diffusion-based and flow-based visual generation, showcasing its strong potential as a competitive alternative to DiT and SiT. Code and models are available at https://github.com/shallowdream204/LAformer.

Abstract

Explore similar work

Fast and Accurate Image Restoration and Generation with Rank Enhanced Linear Attention

In-Token Learning for High-Fidelity Image Restoration via Diffusion Transformers

SAT: Selective Aggregation Transformer for Image Super-Resolution