Linear Attention with Global Context: A Multipole Attention Mechanism for Vision and Physics
作者: Alex Colagrande, Paul Caillon, Eva Feillet, Alexandre Allauzen
分类: cs.CV, cs.AI, cs.LG
发布日期: 2025-07-03
备注: Accepted at ECLR Workshop at ICCV 2025
🔗 代码/项目: GITHUB
💡 一句话要点
提出基于多极子展开的线性注意力机制MANO,用于视觉和物理模拟任务。
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱六:视频提取与匹配 (Video Extraction)
关键词: 线性注意力机制 多极子展开 Transformer 图像分类 物理模拟
📋 核心要点
- 传统Transformer的二次复杂度限制了其在高分辨率输入上的应用,现有方法如分块或降采样会损失细节。
- MANO将注意力机制建模为网格点间的交互,借鉴n体模拟的多极子展开思想,实现线性复杂度。
- 实验表明,MANO在图像分类和Darcy流任务上性能媲美SOTA模型,且显著降低了运行时间和内存占用。
📝 摘要(中文)
Transformer模型在图像分类到物理模拟等多种任务中已成为事实标准。尽管性能出色,但标准Transformer在内存和时间上的复杂度都随输入长度呈二次方增长,使其难以处理高分辨率输入。因此,许多变体被提出,其中最成功的方法依赖于分块、降采样或粗化技术,但通常以牺牲最精细的细节为代价。本文采取了一种不同的方法,受到n体数值模拟中最新技术的启发,将注意力机制视为网格点之间的交互问题。我们引入了多极子注意力神经算子(MANO),它以基于距离的多尺度方式计算注意力。MANO在每个注意力头中保持全局感受野,并实现了关于网格点数量的线性时间和内存复杂度。在图像分类和Darcy流上的实验结果表明,MANO可以与ViT和Swin Transformer等最先进的模型相媲美,同时将运行时间和峰值内存使用量降低几个数量级。我们在https://github.com/AlexColagrande/MANO开源了我们的代码,以保证可重复性。
🔬 方法详解
问题定义:Transformer模型在处理高分辨率图像或大规模物理模拟时,其计算复杂度和内存需求呈平方级增长,成为瓶颈。现有的解决方案,如ViT和Swin Transformer,通常采用图像分块或降采样等方法来降低计算量,但这些方法不可避免地会损失图像的细节信息,影响最终的性能。
核心思路:MANO的核心思路是将注意力机制视为网格点之间的相互作用,并借鉴n体模拟中的多极子展开方法来近似计算这些相互作用。通过将远距离的网格点进行聚类,并用一个“多极子”来代表这些点的整体影响,从而避免了对每个点对之间进行单独计算,大大降低了计算复杂度。
技术框架:MANO的整体框架包括以下几个主要步骤:1. 输入嵌入:将输入数据(例如图像像素或物理场网格点)嵌入到高维空间中。2. 多极子展开:根据网格点之间的距离,将远距离的点聚类成多极子,并计算每个多极子的中心和权重。3. 局部交互:计算每个网格点与其邻近网格点之间的直接交互。4. 全局交互:计算每个网格点与所有多极子之间的交互。5. 输出融合:将局部交互和全局交互的结果融合,得到最终的注意力输出。
关键创新:MANO最重要的创新点在于将多极子展开方法引入到注意力机制中,从而实现了线性复杂度的注意力计算。与传统的注意力机制需要计算所有点对之间的相似度不同,MANO只需要计算每个点与其邻近点以及所有多极子之间的相似度,大大减少了计算量。
关键设计:MANO的关键设计包括:1. 多极子聚类策略:如何有效地将远距离的网格点聚类成多极子,并保证聚类的精度。2. 多极子中心和权重的计算方法:如何计算每个多极子的中心和权重,以准确地代表其所包含的网格点的整体影响。3. 局部交互和全局交互的融合方法:如何有效地融合局部交互和全局交互的结果,以获得最佳的性能。
🖼️ 关键图片
📊 实验亮点
MANO在图像分类和Darcy流模拟任务上取得了显著的成果。在图像分类任务中,MANO的性能与ViT和Swin Transformer等SOTA模型相当,同时将运行时间和峰值内存使用量降低了几个数量级。在Darcy流模拟任务中,MANO也表现出了优异的性能,能够准确地模拟复杂的流体行为。
🎯 应用场景
MANO具有广泛的应用前景,包括高分辨率图像处理、视频分析、三维重建、物理模拟(如流体动力学、电磁场模拟)等领域。其线性复杂度使其能够处理大规模的数据,为解决传统Transformer难以处理的问题提供了新的思路。未来,MANO有望在更多需要处理长序列数据的任务中发挥重要作用。
📄 摘要(原文)
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.