💻 computer science

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

The paper proposes the Flexible Image Transformer (FIT), a unified image restoration model that enhances performance across diverse and spatially non-uniform degradations by explicitly integrating degradation awareness throughout the entire tokenization pipeline via adaptive patch deformation and a novel task-token dropout strategy.

Original authors: Zihao He, Yunfeng Wu, Xinchao Wang, Songhua Liu

Published 2026-08-10

📖 9 min read🧠 Deep dive

CC BY 4.0

Original authors: Zihao He, Yunfeng Wu, Xinchao Wang, Songhua Liu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to fix a messy room, but the mess isn't the same everywhere. In one corner, there's a pile of wet, muddy clothes; in another, a shattered window; and in a third, a layer of thick dust. If you had a robot cleaner that used the exact same sweeping pattern for the whole room, it would likely miss the mud, smash the glass, or just push the dust around. This is the daily struggle of "image restoration," a branch of computer science dedicated to fixing photos that have been ruined by rain, fog, blur, or darkness. For years, the smartest tools for this job have been like that rigid robot: they chop an image into tiny, fixed-size square tiles (like a grid of Post-it notes) to analyze them. They then try to fix each tile independently. But when a messy streak of rain crosses the boundary between two tiles, the robot gets confused. It tries to fix the top half of the rain in one tile and the bottom half in another, often leaving a visible, ugly seam where the two pieces don't quite match.

The paper you are about to read tackles this specific "seam" problem. It introduces a new way of thinking about how computers look at broken images. Instead of forcing the image into a rigid grid, the authors propose a system that lets the grid itself bend and warp to fit the mess. Think of it like a flexible rubber sheet instead of a stiff cardboard box. If a rain streak runs diagonally, the rubber sheet stretches and shifts its "tiles" to follow the streak perfectly, ensuring the computer sees the whole problem in one go. By doing this, the computer can fix the image without leaving those distracting, blocky lines behind. The researchers call their new system FIT, which stands for Flexible Image Transformer, and they show that by making the image "bend" to match the damage, they can create cleaner, sharper photos than ever before.

The Problem: The "Grid" That Won't Bend

To understand why this new method is a big deal, we have to look at how current image-fixing AI works. Most modern systems use a technique called "patch tokenization." Imagine you have a giant jigsaw puzzle, but instead of irregular shapes, every piece is a perfect square. To fix a blurry photo, the computer cuts the image into these tiny squares, studies each one, and then tries to reconstruct the picture.

The problem, as the authors point out, is that real-world damage doesn't respect these squares. Rain streaks, motion blur, and fog often cut right across the boundaries of these fixed squares. When the computer tries to fix a square that contains half a rain streak and half a clear sky, it gets confused. It mixes the "bad" pixels with the "good" ones. When it puts the squares back together, the mismatch creates a visible "grid artifact"—a faint, unnatural line or seam where the squares meet. It's like trying to patch a torn shirt with square patches of fabric that don't line up with the tear; the repair looks clumsy and obvious.

The authors argue that previous attempts to fix this were like trying to teach the robot cleaner to be smarter after it had already chopped the room into squares. They would inject information about the type of damage (e.g., "this is rain") into the computer's brain, but the computer was still stuck looking at the world through a rigid grid. The damage was already mixed up inside the squares before the computer even started thinking about how to fix it.

The Solution: A Flexible, Shape-Shifting Grid

The authors propose a radical change: don't use a fixed grid at all. Instead, let the grid bend.

They introduce a system called FIT (Flexible Image Transformer). Here is how it works, step-by-step, using a simple analogy:

The "Damage Detective" (Degradation Encoder): Before the computer even starts looking at the image pieces, it sends a quick scout to look at the whole photo. This scout doesn't just say "it's raining"; it creates a map. It draws a global picture of the damage (a "global vector") and a detailed map showing exactly where the damage is worst (a "spatial map"). It's like a firefighter looking at a burning building and immediately knowing which rooms are hottest and where the fire is spreading.
The "Bending Grid" (Deformable Tokenization): Now comes the magic. Instead of cutting the image into rigid squares, the computer uses the damage map to stretch and shift the squares. If there is a long rain streak, the computer stretches the grid lines to follow the streak. If a part of the image is very blurry, the grid shifts to group those blurry pixels together.
- The Analogy: Imagine you are cutting a cake. If the cake has a long, thin strawberry running through it, a rigid knife would cut the strawberry in half, ruining the piece. But if you have a flexible knife that can bend to follow the strawberry's curve, you can cut around it perfectly. FIT does exactly this. It bends the "cutting lines" of the image so that every piece (or "token") contains a coherent part of the damage, rather than a messy mix of clean and dirty pixels.
The "Smart Fixer" (Transformer): Once the image is cut into these flexible, damage-aware pieces, the computer processes them. Because the pieces are now organized logically (the rain is in one piece, the clear sky in another), the computer can fix them much more accurately.
The "Reverse Bend" (Unembedding): After the computer fixes the pieces, it has to put them back together. It uses the same bending map to warp the pieces back into their original positions. Because the pieces were aligned with the damage in the first place, they fit together perfectly, leaving no visible seams.

The "Dropout" Trick: Learning to Guess

There is one more clever trick the authors use called Task-Token Dropout. Usually, when you train an AI to fix a photo, you tell it exactly what's wrong: "This is rain," or "This is blur." But in the real world, you often don't know what's wrong. You just see a bad photo.

To make the AI robust enough to handle this, the authors train it to sometimes ignore the label. They randomly tell the AI, "Here is a rainy photo," and then immediately take that label away, forcing the AI to figure out it's rain just by looking at the picture. This is like teaching a student to solve a math problem by sometimes giving them the answer key and sometimes taking it away, forcing them to learn the logic instead of just memorizing the answer. This ensures that when the AI faces a photo with a mix of rain and fog (or a type of damage it has never seen before), it can still figure out the problem and fix it without getting confused.

The Results: Sharper, Cleaner, No Seams

The authors tested their new FIT system on five different types of image damage: removing noise, removing rain, removing fog (haze), fixing motion blur, and brightening dark photos. They compared it against the best existing methods, including a very strong baseline called JIT (Just Image Transformers).

The results were impressive. On a standard test of five different damage types, FIT achieved an average score of 30.72 dB (a measure of how close the restored image is to the original). This beat the previous best methods by a significant margin, improving the score by 0.5 to 1.1 dB over other top systems. In the world of image restoration, even a tiny fraction of a decibel is a huge deal; a jump of 1 dB is often considered a massive leap in quality.

Specifically, the paper highlights that FIT is particularly good at handling "spatially non-uniform" damage—where the mess is different in different parts of the photo.

On Rain: FIT removed rain streaks more thoroughly while keeping the background textures sharp.
On Fog: It cleared up foggy areas without making the clear parts look weird or washed out.
On Blur: It recovered sharp edges that other methods turned into smudges.

Perhaps the most visual proof is the "Grid Score." The authors invented a way to measure those ugly seams. The old rigid-grid methods had high grid scores (lots of seams), especially when the damage was heavy. FIT's grid score was much lower, meaning the images looked smooth and natural, with no visible blocky lines.

Why This Matters

The paper suggests that the way we slice up an image (the "tokenization") is just as important as the brain that fixes it. For a long time, researchers assumed the grid had to be fixed and rigid. This paper proves that assumption wrong. By making the grid flexible and letting it bend to the shape of the damage, the computer can see the problem more clearly and fix it more effectively.

The authors don't claim this solves every problem in image restoration, and they note that the system still needs to be tested on real-world, unpredictable scenarios. However, they demonstrate that this "bending" approach is a powerful new tool. It suggests that in the future, our photo-editing tools might not just be smarter algorithms, but also more flexible ones—tools that can reshape their own view of the world to match the mess we are trying to clean up.

In short, FIT teaches us that sometimes, to fix a broken picture, you have to stop looking at it through a rigid window and start looking through a flexible lens. And when you do, the picture looks a whole lot better.

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

1. Problem Statement

All-in-one image restoration (AiOIR) aims to train a single model capable of recovering images degraded by diverse and spatially non-uniform corruptions (e.g., noise, rain, haze, blur, low-light) without task-specific switching. While recent unified Transformer-based methods have advanced this field, they rely on fixed patch partitioning for tokenization and reconstruction.

The paper identifies a critical limitation in existing approaches (including the recent JIT baseline): fixed $P \times P$ patch boundaries often cut through degradation structures (e.g., rain streaks or blur kernels). This forces the model to mix corrupted and clean pixels within the same token. During reconstruction, mapping these mixed tokens back to a rigid grid produces visible grid artifacts at patch boundaries. As degradation severity increases, these artifacts become more pronounced, creating a bottleneck that limits restoration quality regardless of model capacity. Furthermore, existing methods typically inject degradation conditions only after tokenization (into backbone blocks), leaving the embedding and reconstruction stages insensitive to local degradation variations.

2. Methodology: Flexible Image Transformer (FIT)

The authors propose FIT, a unified restoration architecture that explicitly models degradation awareness across the entire pipeline, from patch sampling to pixel reconstruction. The core principle is to make patch boundaries and reconstruction functions dependent on local degradation severity.

2.1 Dual-Granularity Degradation Encoder (DGDE)

Instead of relying solely on a global degradation vector, FIT employs a lightweight CNN backbone to extract two complementary representations:

Global Vector ( $g$ ): Encodes image-wide degradation statistics.
Spatial Degradation Map ( $M$ ): A high-resolution map capturing location-dependent severity.
During training, auxiliary heads supervise these representations using ground-truth degradation type and strength labels, though the model operates without them at inference.

2.2 Degradation-Guided Deformable Tokenization (DGDT)

This is the core innovation, replacing fixed grid partitioning with adaptive sampling:

Deformable Patch Embedding: The input image is processed alongside the upsampled spatial map $M$ and global vector $g$ . An offset network, modulated by $g$ via FiLM (Feature-wise Linear Modulation), predicts sampling offsets ( $\Delta$ ). These offsets allow patch boundaries to align with local corruption patterns (e.g., shifting boundaries to avoid cutting through rain streaks). Pixels are sampled at these offset locations via bilinear interpolation to form tokens.
Deformable Patch Unembedding: After processing tokens through the Transformer backbone, the tokens are projected back to pixel space. An inverse offset network predicts offsets ( $\Delta'$ ) to warp the reconstructed tokens back to the original grid, ensuring smooth reconstruction without grid discontinuities.

2.3 Transformer Backbone and Task-Token Dropout

The token sequence is processed by an AdaLN Transformer where the global vector $g$ modulates attention and feed-forward layers.

Task-Token Dropout (TTD): To improve robustness against spatially mixed or unseen degradation combinations, the model employs a stochastic strategy during training. It randomly switches between using explicit "hard" task tokens (if labels are available) and "soft" tokens derived from the inferred global vector $g$ . This encourages the model to rely on content-driven degradation inference rather than over-relying on oracle labels.

2.4 Training Objectives

The model is trained with a weighted loss function including:

Reconstruction Loss ( $L_{recon}$ ): Pixel-wise fidelity ( $L_1$ ).
Auxiliary Losses: Type classification and strength regression (training only).
Seam Loss ( $L_{seam}$ ): Explicitly penalizes discontinuities across patch boundaries by matching ground-truth cross-boundary gradients and discouraging seam-like jumps.
Offset Regularization ( $L_{off}$ ): Encourages small, spatially smooth offsets to prevent degenerate sampling.

3. Key Contributions

FIT Architecture: A unified restoration framework that conditions deformable embedding and reconstruction on global-local degradation representations, making the pipeline end-to-end degradation-aware.
Task-Token Dropout: A regularization strategy that bridges the gap between labeled and label-free inference, improving robustness to mixed degradations.
State-of-the-Art Performance: Demonstration of superior performance across five restoration tasks (denoising, deraining, dehazing, deblurring, low-light enhancement) on standard benchmarks (BSD68, Rain100L, SOTS, GoPro, LOLv1).
Interpretability: The learned offsets provide a direct mechanism for visualizing how the model adapts spatially to degradation structures.

4. Experimental Results

The paper evaluates FIT on five standard benchmarks under both three-degradation and five-degradation settings.

Quantitative Performance:
- Five-Degradation Setting: FIT achieves an average PSNR of 30.72 dB, outperforming recent unified methods like AdaIR (+0.52 dB) and the JIT baseline (+1.44 dB).
- Three-Degradation Setting: FIT achieves an average PSNR of 32.83 dB, surpassing the previous best (VLU-Net) by +0.13 dB and JIT by +1.11 dB.
- Significant gains are observed in tasks with spatially non-uniform degradations (e.g., +2.52 dB over JIT on dehazing).
Qualitative Analysis:
- Grid Artifact Mitigation: Visual comparisons show that FIT eliminates the visible grid artifacts (seams) present in fixed-patch methods like JIT, particularly in high-degradation scenarios.
- Offset Visualization: The learned offsets demonstrate interpretable behavior: they align with rain streak directions, correlate with scene depth in hazy images, and concentrate around edges in blurry images.
Ablation Studies:
- Removing deformable tokenization causes a significant drop in PSNR and an increase in grid artifacts.
- Task-Token Dropout is shown to be crucial for maintaining performance when task labels are incorrect or unavailable, reducing performance drops in mixed-degradation scenarios.
- Scaling analysis reveals that while fixed-patch models (JIT) plateau due to grid artifacts, FIT continues to benefit from increased model capacity.

5. Significance and Claims

The paper claims that tokenization geometry is an active and critical component of restoration quality. By adapting the image-token interface itself to the degradation structure, FIT addresses a fundamental bottleneck in unified restoration: the mismatch between rigid grid partitioning and spatially varying corruptions.

The authors argue that previous methods, which condition degradation only during token transformation, are suboptimal for non-uniform degradations. FIT demonstrates that making the embedding and reconstruction stages "deformation-aware" allows for artifact-free restoration and better handling of mixed degradation types. The work suggests that future unified restoration architectures should consider the flexibility of the tokenization process itself, rather than treating it as a fixed preprocessing step.

The paper concludes that FIT achieves state-of-the-art results by effectively mitigating grid artifacts and enabling spatially adaptive processing, offering a new direction for all-in-one image restoration.