Plug-and-play linear attention with provable guarantees for training-free image restoration

TL;DR

This paper introduces PnP-Nystra, a training-free, Nyström-based linear attention module that replaces multi-head self-attention in vision Transformers, providing efficient, provably accurate image restoration with significant speedups.

Contribution

The paper proposes PnP-Nystra, a novel plug-and-play linear attention method with provable guarantees, compatible with pretrained models for image restoration tasks without additional training.

Findings

Achieves 1.8–3.6× speedup on GPU and 1.8–7× on CPU.

Maintains high image quality with minimal drop compared to original models.

Compatible with window-based architectures like SwinIR, Uformer, Dehazeformer.

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.