AdaWPT: Adaptive Window Pruning for Deblurring

• The paper introduces AdaWPT, a mechanism that prunes non-blurred regions to enhance deblurring efficiency and accuracy in Vision Transformers.
• It combines window-based attention, blur confidence prediction, and Gumbel-Softmax sampling to prune approximately 66% of windows, reducing FLOPs and accelerating inference.
• Empirical results show improved PSNR and computational savings through its U-shaped encoder–decoder approach, highlighting a balance between performance and efficiency.

AdaWPT (Adaptive Window Pruning Transformer) is a mechanism designed to enhance the efficiency and fidelity of local motion deblurring in high-resolution images by dynamically selecting computation regions within a vision Transformer architecture. Local motion blur, characterized by spatially varying blur patterns due to moving objects and static backgrounds, presents challenges for contemporary deblurring networks that predominantly perform global operations. AdaWPT addresses this inefficiency by pruning non-blurred spatial windows during processing, focusing Transformer capacity on genuinely degraded regions. Developed in the context of the Local Motion Deblurring Vision Transformer (LMD-ViT), AdaWPT integrates window-based attention, blur-aware confidence prediction, Gumbel-Softmax-based discrete window selection, and loss functions guided by annotated blur masks (Li et al., 2023).

1. Structural Overview and Context

LMD-ViT employing AdaWPT follows a U-shaped encoder–decoder architecture with skip connections. AdaWPT blocks are inserted at various resolution levels, both in the encoder, bottleneck, and decoder pathways. The input image $B\in\mathbb{R}^{H\times W\times 3}$ is projected to an initial feature map $X_0$. Each encoder stage includes a down-sampling operation followed by $n_i$ AdaWPT blocks, generating successively lower-resolution ($X_i$) feature maps. After a two-block AdaWPT bottleneck, the decoder symmetrically upsamples the features, integrating encoder features via skip connections. The network outputs a residual image $R$, added to the input $B$ to yield the final sharp prediction $S' = B + R$. Only AdaWPT blocks conduct adaptive window pruning; other layers use standard convolutional or sampling operations (Li et al., 2023).

2. Window Partitioning and Feature Tokenization

An input feature map $X\in\mathbb{R}^{H\times W\times C}$ entering an AdaWPT block is partitioned into a grid of non-overlapping $P \times P$ spatial windows:

$$N = \frac{H}{P} \times \frac{W}{P}, \quad \{W_i\}_{i=1}^{N}, \quad W_i \in \mathbb{R}^{P \times P \times C}$$

Each window $X_0$0 is flattened and/or linearly projected to a token vector:

$X_0$1

This yields a set of tokens $X_0$2 representing local regions, forming the computational units for adaptivity.

3. Blurriness Confidence Prediction and Gumbel-Softmax Pruning

AdaWPT employs a learned blurriness confidence predictor. In the first block at each scale (AdaWPT-F), each $X_0$3 passes through a multi-layer perceptron stack:

$X_0$4

Softmax is applied to $X_0$5 to obtain a blur confidence $X_0$6:

$X_0$7

Subsequent blocks at the same scale (AdaWPT-P) reuse the same $X_0$8 decision map.

During training, binarization of $X_0$9 is facilitated by the Gumbel-Softmax re-parameterization. Sampling with Gumbel noise $n_i$0 and temperature $n_i$1 yields a soft, differentiable discrete decision:

$n_i$2

where $n_i$3, $n_i$4. The hard sample $n_i$5 is used in the forward pass, with gradients propagated through the soft variable $n_i$6 (straight-through estimator). At inference, $n_i$7 is set as $n_i$8 if $n_i$9 (a threshold, with $X_i$0 used in practice), and $X_i$1 otherwise (Li et al., 2023).

4. Pruning Mechanism and Transformer Application

During the forward pass, only windows with $X_i$2 are retained for subsequent Transformer computation; others are zeroed:

$X_i$3

Retained windows proceed through window-based multi-head self-attention (W-MSA) and local feed-forward (W-LeFF) operations. After processing, outputs are scattered back into the original spatial grid.

The pruning procedure is applied in both encoder and decoder stages, compounding computational savings by focusing on windows corresponding to blurred regions. A single confidence map is predicted and reused within each resolution group, reducing redundant computation (Li et al., 2023).

5. Supervision and Loss Functions

End-to-end training of LMD-ViT with AdaWPT is supervised by a composite loss:

• Reconstruction loss ($X_i$4): Per-pixel loss (e.g., $X_i$5) between deblurred prediction $X_i$6 and ground truth $X_i$7, potentially emphasizing annotated blur regions defined by mask $X_i$8.
• Pruning loss ($X_i$9): Binary cross-entropy between predicted windowwise $R$0 and downsampled ground-truth mask $R$1, enforcing $R$2 on blurred regions and $R$3 elsewhere:

$R$4

The total loss is

$R$5

with $R$6 tuned to balance deblurring and pruning accuracy ($R$7--$R$8) (Li et al., 2023).

6. Computational Complexity and Comparative Results

With $R$9 windows and local window FLOPs $B$0, the unpruned network computes $B$1 operations per block. Window pruning yields

$B$2

with relative FLOPs reduction $B$3. Empirically, AdaWPT prunes $B$4 of windows, leading to a 66\% FLOPs saving and over $B$5 inference speedup versus global Transformer baselines (Uformer, Restormer), while also increasing PSNR by $B$6 dB. Weighted PSNR improves by $B$7 dB over the best-performing CNNs with comparable or faster inference. Pruning in all 9 AdaWPT blocks provides the optimal trade-off, reducing inference time from 1.30 s to 0.56 s. Refined blur masks enhance pruning precision from 66\% to 97\%, yielding a further $B$8 dB PSNR increase (Li et al., 2023).

System	FLOPs Reduction	Speedup	PSNR Gain
vs. best CNN (ReLoBlur)	n/a	similar	+0.50 dB
vs. global Transformer	66%	>2×	+0.28 dB

7. Limitations and Prospective Enhancements

Operationally, AdaWPT is not yet real-time for very high-resolution images. While pruning precision exceeds 95%, small or subtle blur regions may be missed. A plausible implication is that precise blur-region detection governs deblurring fidelity and computational efficiency. Future extensions proposed include integrating cross-window context into the confidence predictor, employing dynamic window sizes, multi-task deblurring and mask refinement, hardware-aware real-time implementations, and expanding to temporal modules for multi-frame deblurring.

In summation, AdaWPT’s integration of region-aware token selection, blur-confidence prediction, Gumbel-Softmax binarization, and selective Transformer computation efficiently targets the spatially localized nature of real-world motion blur, elevating both restoration accuracy and resource utilization efficiency (Li et al., 2023).