提出线性复杂度的多极注意力机制,高效处理高分辨率图像与物理模拟。
Linear Attention with Global Context: A Multipole Attention Mechanism for Vision and Physics
- 将注意力建模为网格点间的多尺度相互作用,采用距离驱动的分层计算
- 在图像分类和达西流任务中性能媲美ViT与Swin Transformer,推理速度提升数倍
- 适合需要长序列建模且对内存敏感的应用,如高分辨率视觉或科学模拟
Transformer已成为从图像分类到物理模拟等众多任务的标准架构。尽管表现优异,标准Transformer随输入长度增长呈现二次方的时间与内存复杂度,难以处理高分辨率输入。为此,已有多种变体被提出,最成功的依赖于分块、下采样或粗化技术,但常损失细粒度细节。本文另辟蹊径:受n体数值模拟前沿技术启发,将注意力视为网格点间的相互作用问题。提出多极注意力神经算子(MANO),以距离为基础进行多尺度计算,在每个注意力头中保持全局感受野,实现与网格点数量呈线性的时间与内存复杂度。在图像分类和达西流任务上的实证表明,MANO性能可媲美ViT与Swin Transformer,同时将运行时间与峰值内存使用量降低数个数量级。代码已开源,地址:https://github.com/AlexColagrande/MANO。
原文摘要 · Abstract (English)
Transformers have become the de facto standard for a wide range of tasks, from image classification to physics simulations. Despite their impressive performance, the quadratic complexity of standard Transformers in both memory and time with respect to the input length makes them impractical for processing high-resolution inputs. Therefore, several variants have been proposed, the most successful relying on patchification, downsampling, or coarsening techniques, often at the cost of losing the finest-scale details. In this work, we take a different approach. Inspired by state-of-the-art techniques in $n$-body numerical simulations, we cast attention as an interaction problem between grid points. We introduce the Multipole Attention Neural Operator (MANO), which computes attention in a distance-based multiscale fashion. MANO maintains, in each attention head, a global receptive field and achieves linear time and memory complexity with respect to the number of grid points. Empirical results on image classification and Darcy flows demonstrate that MANO rivals state-of-the-art models such as ViT and Swin Transformer, while reducing runtime and peak memory usage by orders of magnitude. We open source our code for reproducibility at https://github.com/AlexColagrande/MANO.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。