Neural Discrimination-Prompted Transformers for Efficient UHD Image Restoration and Enhancement

Cong Wang; Jinshan Pan; Liyan Wang; Wei Wang; Yang Yang

cs.CVMar 1, 2026

Neural Discrimination-Prompted Transformers for Efficient UHD Image Restoration and Enhancement

Authors: Cong Wang, Jinshan Pan, Liyan Wang, Wei Wang, Yang Yang

Organizations: Shenzhen Campus of Sun Yat-Sen University, China · The Hong Kong Polytechnic University, Hong Kong · Nanjing University of Science and Technology, China · Dalian University of Technology, China · University of California, San Francisco, USA

Abstract

We propose a simple yet effective UHDPromer, a neural discrimination-prompted Transformer, for Ultra-High-Definition (UHD) image restoration and enhancement. Our UHDPromer is inspired by an interesting observation that there implicitly exist neural differences between high-resolution and low-resolution features, and exploring such differences can facilitate low-resolution feature representation. To this end, we first introduce Neural Discrimination Priors (NDP) to measure the differences and then integrate NDP into the proposed Neural Discrimination-Prompted Attention (NDPA) and Neural Discrimination-Prompted Network (NDPN). The proposed NDPA re-formulates the attention by incorporating NDP to globally perceive useful discrimination information, while the NDPN explores a continuous gating mechanism guided by NDP to selectively permit the passage of beneficial content. To enhance the quality of restored images, we propose a super-resolution-guided reconstruction approach, which is guided by super-resolving low-resolution features to facilitate final UHD image restoration. Experiments show that UHDPromer achieves the best computational efficiency while still maintaining state-of-the-art performance on $3$ UHD image restoration and enhancement tasks, including low-light image enhancement, image dehazing, and image deblurring. The source codes and pre-trained models will be made available at https://github.com/supersupercong/uhdpromer.

Figures & tables

Panel of Figure 1 — Figure 1: Overview of our proposed Neural Discrimination-Prompted Transformers (UHDPromer) . Our UHDPromer contains $4$ parts: (a) High-Resolution Feature Representation (HRFR); (b) Neural Discrimination-Prompted Transformers (NDPT); (c) Feature Super-Resolution (FeaSR); and (d) SR-Guided Reconstruction (SRG-Recon). Firstly, HRFR explores hierarchical multi-scale high-resolution features, which participate in forming neural discrimination priors (NDP) with low-resolution features. Then, NDPT, which is guided by NDP, learns low-resolution features. Next, FeaSR super-resolves the output features of NDPT. Finally, SRG-Recon reconstructs final images guided by learned SR features in FeaSR.

Panel of Figure 2 — Figure 2: Illustration of the Neural Discrimination Priors (NDP) in the different layers by ( 1 ). We note that the low-resolution features exist obvious content differences compared with high-resolution ones. Specifically, low-resolution features always hand down extensive noises while some structures are blurry. Our NDP can better identify these regions, which can effectively measure whether the structures lie in shallower layers or deeper layers, serving as the discrimination function to help learn more representative features in the low-resolution domain.

Panel of Figure 3 — Figure 3: (a) Neural Discrimination-Prompted Attention ( NDPA ) and (b) Neural Discrimination-Prompted Network ( NDPN ). The NDPA first performs cross attention between NDP and query vector generated by low-resolution features to globally perceive the NDP, which would be re-performed by the attention computation between rest key and value vectors of low-resolution features to further explore useful learned information. The NDPN considers the continuous gating mechanism guided by the NDP to allow useful information to be passed to facilitate final image reconstruction.

Dataset	Training	Testing	Resolution
UHD-LL	2,000	150	3,840 $\times$ 2,160
UHD-Haze	2,290	230	3,840 $\times$ 2,160
UHD-Blur	1,964	300	3,840 $\times$ 2,160

Table 1: Datasets Statistics of UHD image restoration benchmarks .

Method	Venue	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$
Setting 1: Trained on LOL
SwinIR ( Liang et al., 2021 )	ICCVW’21	17.900	0.7379	0.5217
Restormer ( Zamir et al., 2022 )	CVPR’22	19.728	0.7703	0.4566
Uformer ( Wang et al., 2022c )	CVPR’22	18.168	0.7201	0.5593
LLFlow ( Wang et al., 2022b )	AAAI’22	19.596	0.7333	0.4606
LLformer ( Wang et al., 2023b )	AAAI’23	21.440	0.7763	0.4528

Table 2: UHD low-light image enhancement . Our UHDPromer achieves the state-of-the-art performance under different training settings. The best and second best are marked in bold and underlined , respectively.

Panel of Figure 4 — Figure 4: UHD low-light enhancement on UHD-LL ( Li et al., 2023 ) under Setting 1 . UHDPromer is able to produce clearer results with vivid colors.

Panel of Figure 5 — Figure 5: UHD low-light enhancement on UHD-LL ( Li et al., 2023 ) under Setting 2 . UHDPromer is capable of generating clearer results with more natural colors.

Method	Venue	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$
Setting 1: Trained on SOTS-ITS
GridNet ( Liu et al., 2019 )	ICCV’19	14.783	0.8466	0.2518
MSBDN ( Dong et al., 2020 )	CVPR’20	15.043	0.8570	0.2709
UHD ( Zheng et al., 2021 )	ICCV’21	11.708	0.6569	0.5948
Restormer ( Zamir et al., 2022 )	CVPR’22	13.875	0.6405	0.4477
Uformer ( Wang et al., 2022c )	CVPR’22	15.264	0.6724	0.4288

Table 3: UHD image dehazing . Our UHDPromer outperforms SOTA approaches.

Panel of Figure 6 — Figure 6: UHD image dehazing on UHD-Haze ( Wang et al., 2024b ) under Setting 1 . UHDPromer is able to generate much cleaner results.

Panel of Figure 7 — Figure 7: UHD dehazing on UHD-Haze ( Wang et al., 2024b ) under Setting 2 . UHDPromer can generate much clearer dehazing results.

Method	Venue	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$
Setting 1: Trained on GoPro
DMPHN ( Zhang et al., 2019 )	CVPR’19	26.490	0.7985	0.2696
MIMO-Unet++ ( Cho et al., 2021 )	ICCV’21	24.290	0.7354	0.3723
MPRNet ( Zamir et al., 2021 )	CVPR’21	24.571	0.7426	0.3597
Restormer ( Zamir et al., 2022 )	CVPR’22	24.872	0.7484	0.3550
Uformer ( Wang et al., 2022c )	CVPR’22	24.382	0.7209	0.3901

Table 4: Image deblurring . UHDformer outperforms SOTAs.

Panel of Figure 8 — Figure 8: UHD image deblurring on UHD-Blur ( Wang et al., 2024b ) under Setting 1 . UHDformer is able to generate results with finer structures.

Panel of Figure 9 — Figure 9: UHD image deblurring on UHD-Blur ( Wang et al., 2024b ) under Setting 2 . UHDformer is capable of producing results with much clearer characters.

ID	Experiment	SR Branch		Main Branch		Parameters $\downarrow$
ID	Experiment	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	Parameters $\downarrow$
(a)	w/o NDP in NDPA&NDPN	26.487	0.8974	26.811	0.9282	0.5322M
(b)	w/o NDP in NDPA	25.962	0.8952	26.183	0.9252	0.7348M
(c)	w/o NDP in NDPN	26.745	0.8984	27.026	0.9283	0.5404M
(d)	NDP $\rightarrow$ Direct Feature	25.999	0.9021	26.138	0.9272	0.7430M
(e)	Using NDP before NDPTB	25.790	0.8974	26.161	0.9263	0.5322M

Table 5: Ablation study on neural discrimination-prompted Transformers . Inserting NDP in both attention and forward-network produces a positive effect.

Panel of Figure 10 — Figure 10: Visual effect on neural discrimination priors. Using NDP is capable of helping produce visually pleasing results with more natural colors.

ID	Experiment	SR Branch		Main Branch		Parameters $\downarrow$
ID	Experiment	PSNR $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	Parameters $\downarrow$
(a)	Cascaded Reconstruction	-	-	26.375	0.9043	0.7414M
(b)	w/o Second Term in ( 5 )	-	-	27.125	0.9051	0.7425M
(c)	w/ SR Branch ( Ours )	26.859	0.8985	27.159	0.9285	0.7430M

Table 6: Effect on super-resolution-guided reconstruction . The introduced SR-guided reconstruction significantly improves restoration quality in terms of SSIM.

Panel of Figure 11 — Figure 11: Visual effect on SR-guided reconstruction. The proposed SR-guided reconstruction is able to improve the visual results with more vivid colors.

Figure 12: Effect on number of channels . Our model, with only $2$ channels and 0.0168M parameters, achieves a performance comparable to UHDFour ( Li et al., 2023 ) , while being approximately 1,000 times smaller in model size (UHDFour has $17.5$ M parameters, as seen in Tab. 7 ).

Panel of Figure 13 — Figure 13: Effect on shuffle down factor . It is noteworthy that while the model using a shuffle down factor of 4 achieves a larger receptive field compared to the baseline model with a factor of 8, there is no significant improvement in SSIM.

Method	LMAR	UHD	UHDFour	FFTformer	DehazeFormer	UHDformer	UHDPromer
Method	( Yu et al., 2024 )	( Zheng et al., 2021 )	( Li et al., 2023 )	( Kong et al., 2023 )	( Song et al., 2023 )	( Wang et al., 2024b )	( Ours )
Parameters (M)	1.314	34.5	17.5	16.6	2.5	0.3393	0.7430
Reduction	43.5%	97.8%	95.8%	95.5%	70.3%	119.0%	-

Table 7: Parameters comparisons . Our UHDPromer significantly reduces the number of parameters compared to the advanced methods, but is slightly above UHDformer.

Method	LMAR	UHD	UHDFour	FFTformer	DehazeFormer	UHDformer	UHDPromer
Method	( Yu et al., 2024 )	( Zheng et al., 2021 )	( Li et al., 2023 )	( Kong et al., 2023 )	( Song et al., 2023 )	( Wang et al., 2024b )	( Ours )
FLOPs (G)	208.34	113.46	75.63	2107.94	375.40	51.63	32.56
Reduction	84.4%	71.3%	56.9%	98.5%	91.3%	36.9%	-

Table 8: FLOPs comparisons on

1024

\times

1024

pixels. Our UHDPromer significantly reduces the computational complexity in terms of FLOPs, indicating the efficiency of our approach.

Method	Restormer	LMAR	DehazeFormer	FFTformer	UHDformer	UHDPromer
Method	( Zamir et al., 2022 )	( Yu et al., 2024 )	( Song et al., 2023 )	( Kong et al., 2023 )	( Wang et al., 2024b )	( Ours )
RT (s)	1.86	3.40	0.45	1.27	0.16	0.12

Table 9: Running Time (RT) (

1024\times 1024

pixels) on the same GPU. Our UHDPromer runs faster than Transformer-based methods as well as the CNN-based approach ( Yu et al., 2024 ) .

Panel of Figure 14 — Figure 14: Illustration of Neural Discrimination Priors (NDP) at the first and last layer. Our NDP can better measure the salient structure at both the first and last layer, which can serve as the discrimination function to help the learning of Transformers to facilitate restoration.

Method	Venue	PSNR $\uparrow$	SSIM $\uparrow$	Parameters $\downarrow$
Retinex-Net ( Wei et al., 2018 )	BMVC’18	16.77	0.54	0.45M
Zero-DCE ( Guo et al., 2020 )	CVPR’20	16.79	0.67	79.416K
AGLLNet ( Lv et al., 2021 )	IJCV’21	17.52	0.77	3.438K
Zhao et al. ( Zhao et al., 2021 )	ICCV’21	21.67	0.87	11.560M
RUAS ( Liu et al., 2021a )	CVPR’21	16.44	0.70	3K
SCI ( Ma et al., 2022 )	CVPR’22	14.78	0.62	0.258K

Table 10: Low-light image enhancement on LOL ( Wei et al., 2018 ) . Although our UHDPromer is designed for UHD images, where most operations are performed in the low-resolution space, our UHDPromer still achieves the best performance on general image enhancement benchmarks ( Wei et al., 2018 ) . The term ‘UHDFour-2/8’ refers to a variant of UHDFour with a shuffle-down factor of 2/8. Conversely, ‘UHDFour-8’ denotes the UHDFour version utilizing a shuffle-down factor of 8, specifically tailored for enhancing UHD low-light images.

Method	Venue	PSNR $\uparrow$	SSIM $\uparrow$	Parameters $\downarrow$
GridNet ( Liu et al., 2019 )	ICCV’19	33.8244	0.9924	0.96M
MSBDN ( Dong et al., 2020 )	CVPR’20	34.4219	0.9899	31.4M
UHD ( Zheng et al., 2021 )	CVPR’21	18.7488	0.8486	34.5M
PFDN ( Dong and Pan, 2020 )	ECCV’20	31.2308	0.9853	12M
D4 ( Yang et al., 2022 )	CVPR’22	20.5449	0.9236	22.9M
DeHamer ( Guo et al., 2022 )	CVPR’22	38.2672	0.9936	132.4M

Table 11: Comparisons on SOTS-Indoor ( Li et al., 2019a ) . Although our UHDPromer is designed for UHD images, where most operations are performed in low-resolution space, our method is comparable with general image dehazing methods in terms of SSIM but lower in terms of PSNR.

Method	Venue	PSNR $\uparrow$	SSIM $\uparrow$	Parameters $\downarrow$
SRN ( Tao et al., 2018 )	CVPR’18	31.751	0.9144	6.8M
DMPHN ( Zhang et al., 2019 )	CVPR’19	31.734	0.9130	21.7M
MIMO-Unet++ ( Cho et al., 2021 )	ICCV’21	33.277	0.9348	16.1M
MPRNet ( Zamir et al., 2021 )	CVPR’21	34.195	0.9454	20.1M
Restormer ( Zamir et al., 2022 )	CVPR’22	34.483	0.9488	26.1M
Uformer ( Wang et al., 2022c )	CVPR’22	34.668	0.9502	20.6M

Table 12: Comparisons on GoPro ( Nah et al., 2017 ) . Our UHDPromer is less effective in handling general image deblurring benchmarks, such as GoPro ( Nah et al., 2017 ) . Note that the PSNR and SSIM may be different from existing papers, as we use the commonly-used IQA PyTorch Toolbox (https://github.com/chaofengc/IQA-PyTorch) to recompute the results for all benchmarks in this paper.

Explore similar work

Ultra-High-Definition Restoration Transformers with Correlation Matching Transformation
Aug 18, 2026Cong Wang, Liyan Wang, Jinshan Pan +4Image Restoration All-in-One Image Restoration
School of Computing and Communications, Lancaster University
From Zero to Detail: A Progressive Spectral Decoupling Paradigm for UHD Image Restoration with New Benchmark
Apr 17, 2026Chen Zhao, Yunzhe Xu, Zhizhou Chen +5Frequency-Domain Image Processing Image Restoration
State Key Laboratory of Novel Software Technology, Nanjing University, Nanjing 210023, China · School of Computer Science and Technology, Dalian University of Technology, Dalian 116023, China · Department of Computer Science and Engineering, Michigan State University, East Lansing, MI 48824, U.S.A.
PixRestore: Unified Image Restoration via Pixel Diffusion Transformer
Aug 17, 2026Lingchen Sun, Rongyuan Wu, Xiangtao Kong +6All-in-One Image Restoration Pixel-Space Diffusion Models
The Hong Kong Polytechnic University · OPPO Research Institute

Aug 18, 2026J/K move · Enter open · S save

Ultra-High-Definition Restoration Transformers with Correlation Matching Transformation

Cong Wang, Liyan Wang, Jinshan Pan, Wei Wang+3

School of Computing and Communications, Lancaster University

Image Restoration All-in-One Image Restoration Image Deblurring Image Super-Resolution

Open paper AskPDF arXiv

We propose UHDformer++, a general Transformer-based framework to solve numerous Ultra-High-Definition (UHD) image restoration tasks. UHDformer++ operates across

4

coordinated learning spaces: 1) a high-resolution space (HR) for multi-level feature extraction, 2) a low-resolution space (LR) for learning compact, representative features, 3) a super-resolution space (SR) for upsampling low-resolution features from SR, and 4) a low-high fusion and reconstruction space (LHFR) for final image restoration. Specifically, HR extracts multi-scale high-resolution features and fuses them with low-resolution cues to produce residual images, while LR distills complementary representations from HR to improve restoration quality. To supply LHFR with richer features, SR super-resolves LR outputs before fusion. We further introduce two modules to bridge the high- and low-resolution spaces. The Feature-Refined Correlation Matching Transformation (FR-CMT) module selects the top

C/r~(C~\text{denotes the number of channels;~}r\geq1~\text{controls the squeezing level})

, from the fusion between max- and mean-pooled high-resolution features to replace less informative channels in the low-resolution Transformer. The Adaptive Channel Modulator (ACM) adaptively recalibrates multi-scale high-resolution features, ensuring that only task-relevant information propagates to LR. Extensive experiments demonstrate that UHDformer++ reduces model parameters by at least 86% compared with recent state-of-the-art methods while achieving substantial performance gains across

5

UHD restoration tasks, including low-light image enhancement, dehazing, deblurring, deraining, and desnowing. Code will be released at https://github.com/supersupercong/uhdformerplus.

Abstract

Explore similar work

Ultra-High-Definition Restoration Transformers with Correlation Matching Transformation

From Zero to Detail: A Progressive Spectral Decoupling Paradigm for UHD Image Restoration with New Benchmark

PixRestore: Unified Image Restoration via Pixel Diffusion Transformer