Bend the Basics: Degradation-Aware Deformable Tokenization for All-in-One Image Restoration

2026-08-07

AI Summary

This work addresses the limitation of existing unified Transformers, which struggle to handle spatially non-uniform image degradation due to fixed patch partitioning strategies. To overcome this, the authors propose the Flexible Image Transformer (FIT), the first framework to integrate a degradation-aware mechanism throughout the entire image restoration pipeline. FIT employs a lightweight degradation encoder to produce both a global degradation vector and a spatial degradation map, which jointly guide deformable patch partitioning, adaptive embedding and de-embedding, and pixel-level reconstruction. Additionally, a task token dropout strategy is introduced to enhance model generalization. Extensive experiments demonstrate that FIT achieves state-of-the-art performance across five benchmarks, attaining average PSNR values of 30.72 dB and 32.83 dB under five-degradation and three-degradation settings, respectively—outperforming current methods by 0.5–1.1 dB.

Abstract

All-in-one image restoration seeks a single model that can recover images degraded by diverse and spatially non-uniform corruptions. However, many unified Transformers rely on fixed patch partitioning: task/degradation condition is injected only into the backbone blocks after tokenization, leaving the embedding and reconstruction stages insensitive to local degradation variations. In contrast to previous approaches, we present Flexible Image Transformer (FIT) that explicitly models degradation awareness across the entire pipeline, from patch sampling to pixel reconstruction. Specifically, FIT employs a lightweight Degradation Encoder to predict a global degradation vector $\mathbf{g}$ and a spatial degradation map $\mathbf{M}$ from local degradation severity, which jointly condition the patch embedding and unembedding through adaptive deformation. Moreover, to improve robustness across degradation types, we introduce a task-token dropout strategy that regularizes task conditioning during training. On five standard benchmarks (BSD68, Rain100L, SOTS, GoPro, and LOLv1), FIT achieves state-of-the-art performance with 30.72 dB average PSNR on the five-degradation setting and 32.83 dB on the three-degradation setting, outperforming recent unified restoration methods by +0.5$\sim$1.1 dB. Moreover, the learned offsets provide a direct handle for visualizing degradation-aware spatial adaptation.

Problem

• image restoration
• degradation-aware
• non-uniform corruption
• patch tokenization
• all-in-one model

Innovation

• Degradation-Aware
• Deformable Tokenization
• All-in-One Image Restoration
• Task-Token Dropout
• Adaptive Embedding