无需训练即可替换注意力模块,实现快速图像修复
Plug-and-play linear attention with provable guarantees for training-free image restoration
- 基于Nyström方法的线性注意力,可直接替换预训练模型中的自注意力
- 在RTX 4090上提速1.8至3.6倍,CPU上提速1.8至7倍
- 保持接近原模型性能,适合实时与低资源场景部署
多头自注意力(MHSA)是现代视觉Transformer的核心组件,但其在令牌数量上的二次复杂度仍是实时和资源受限部署的主要瓶颈。我们提出PnP-Nystra,一种无需训练的基于Nyström的线性注意力模块,作为预训练图像修复Transformer中MHSA的即插即用替代方案,并具备可证明的核近似误差保证。PnP-Nystra可直接集成到窗口式架构如SwinIR、Uformer和Dehazeformer中,实现无需微调的高效推理。在图像去噪、去模糊、去雾和超分辨率任务上,PnP-Nystra在NVIDIA RTX 4090 GPU上实现1.8–3.6倍加速,在CPU推理上实现1.8–7倍加速。相比最强的无训练线性注意力基线,本方法质量下降最小,输出最接近原始模型。
原文摘要 · Abstract (English)
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ö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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。