TSFormer: A Robust Framework for Efficient UHD Image Restoration

TL;DR

TSFormer is an all-in-one framework that integrates Trusted learning with S Parsification to boost both generalization capability and computational efficiency in UHD image restoration, and achieves state-of-the-art restoration quality while enhancing generalization and reducing computational demands.

Abstract

Ultra-high-definition (UHD) image restoration is vital for applications demanding exceptional visual fidelity, yet existing methods often face a trade-off between restoration quality and efficiency, limiting their practical deployment. In this paper, we propose TSFormer, an all-in-one framework that integrates \textbf{T}rusted learning with \textbf{S}parsification to boost both generalization capability and computational efficiency in UHD image restoration. The key is that only a small amount of token movement is allowed within the model. To efficiently filter tokens, we use Min-$p$ with random matrix theory to quantify the uncertainty of tokens, thereby improving the robustness of the model. Our model can run a 4K image in real time (40fps) with 3.38 M parameters. Extensive experiments demonstrate that TSFormer achieves state-of-the-art restoration quality while enhancing generalization and reducing computational demands. In addition, our token filtering method can be applied to other image restoration models to effectively accelerate inference and maintain performance.

Figures (11)

• Figure 1: Illustration of the Min-$p$ sampling and stability filtering process.(a) The original input image used as a reference. (b) Result after Min-$p$ sampling, where high-probability regions are retained and highlighted in a distinct color, indicating sparsified yet significant areas. (c) Result after stability filtering, where only the stable, high-confidence features are preserved, with stable regions marked in red and unstable regions suppressed.
• Figure 2: The overall architecture of the proposed TSFormer for UHD image restoration, which main consists of Trusted Sparse Blocks (TSB) and a Feed-forward Network (FFN). Our core is Min-p Sapre Attention (MSA), which incorporates the ability of dynamic token filtering and sparse representation, effectively reducing the burden of running a UHD image. In addition, MSA is also a plug-and-play algorithm that can be used in any image restoration network based on the Transformer architecture.
• Figure 3: Comparison of Low-light Enhancement Methods. The top row is from the UHD-LL dataset, and the bottom is from the UHD-LOL dataset.
• Figure 4: Image deblurring on UHD-Blur. TSFormer is able to generate deblurring results with sharper structures.
• Figure 5: Image dehazing on UHD-Haze. TSFormer is capable of producing clearer results.