arXiv · 2506.08520
Plug-and-play linear attention with provable guarantees for training-free image restoration
Abstract
Multi-head self-attention (MHSA) is a key building block in modern vision Transformers, yet its quadratic complexity in the number of tokens remains a major bottleneck for real-time and resource-constrained deployment. We present PnP-Nystra, a training-free Nystr\"{o}m-based linear attention module designed as a plug-and-play replacement for MHSA in {pretrained} image restoration Transformers, with provable kernel approximation error guarantees. PnP-Nystra integrates directly into window-based architectures such as SwinIR, Uformer, and Dehazeformer, yielding efficient inference without finetuning. Across denoising, deblurring, dehazing, and super-resolution on images, PnP-Nystra delivers $1.8$--$3.6\times$ speedups on an NVIDIA RTX 4090 GPU and $1.8$--$7\times$ speedups on CPU inference. Compared with the strongest training-free linear-attention baselines we evaluate, our method incurs the smallest quality drop and stays closest to the original model's outputs.
Explore related subjects
Keep this discovery
Srinivasan Kidambi, Karthik Palaniappan, Pravin Nair. 2025-06-10. Plug-and-play linear attention with provable guarantees for training-free image restoration. https://arxiv.org/abs/2506.08520
Cite the original work for its findings. Save a collection to share your selection of sources.