Fast Multipole Attention: A Scalable Multilevel Attention Mechanism for Text and Images

📄 arXiv: 2310.11960v4 📥 PDF

作者: Yanming Kang, Giang Tran, Hans De Sterck

分类: cs.CL, cs.LG

发布日期: 2023-10-18 (更新: 2025-09-18)

🔗 代码/项目: GITHUB

💡 一句话要点

提出快速多极注意力机制以解决长序列处理问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 快速多极注意力 自注意力机制 长序列处理 高分辨率输入 多模态任务 计算机视觉 自然语言处理

📋 核心要点

  1. 现有Transformer模型在处理长序列和高分辨率输入时面临二次复杂度的挑战，限制了其应用。
  2. 本文提出的快速多极注意力机制通过分治策略显著降低了自注意力的时间和内存复杂度，保持了全局交互能力。
  3. 实验表明，1D变体在语言建模任务中表现优异，2D变体在视觉分类和语义分割任务中超越了现有基线。

📝 摘要（中文）

尽管Transformer网络具有全局感受野，但其相对于序列长度的二次复杂度限制了其在长序列和高分辨率输入中的应用。本文提出了快速多极注意力（FMA），这是一种受n体物理学启发的自注意力分治机制。FMA将自注意力的时间和内存复杂度从$ ext{O}(n^2)$降低到$ ext{O}(n ext{log} n)$和$ ext{O}(n)$，同时保持全上下文交互。FMA包含一个具有$ ext{O}( ext{log} n)$级别分辨率的学习层次结构。在该层次结构中，近邻token以全分辨率交互，而远离的token通过逐渐粗糙的学习基函数进行交互。我们为语言和视觉任务分别开发了1D和2D实现。实验结果表明，FMA使得基于Transformer的模型能够扩展到更长的序列和更高的分辨率，而不损失准确性。

🔬 方法详解

问题定义：本文旨在解决Transformer模型在处理长序列和高分辨率输入时的计算复杂度问题，现有方法在时间和内存使用上存在显著限制。

核心思路：快速多极注意力（FMA）通过分治策略，将自注意力的复杂度从$ ext{O}(n^2)$降低到$ ext{O}(n ext{log} n)$，同时保持全局上下文交互。该方法利用学习的层次结构，使得近邻token以全分辨率交互，而远离token通过粗糙的基函数进行交互。

技术框架：FMA的整体架构包括多个分辨率级别的层次结构，1D实现用于语言任务，2D实现用于视觉任务。每个级别的交互通过学习的基函数进行，确保了信息的有效传递。

关键创新：FMA的主要创新在于其分治机制和层次结构设计，使得模型能够在保持准确性的同时，处理更长的序列和更高的分辨率输入。这一设计与传统的自注意力机制有本质区别。

关键设计：FMA的设计包括多个层次的分辨率设置，学习的基函数，以及针对不同任务的1D和2D实现，确保了在不同应用场景下的高效性和准确性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，1D变体在自回归和双向语言建模基准上与领先的高效注意力基线相匹配或超越，同时显著降低了内存使用。2D变体在分类和语义分割任务中表现优于强大的视觉Transformer基线，展现出线性复杂度的优势。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理和计算机视觉，尤其是在需要处理长文本和高分辨率图像的任务中。快速多极注意力机制的引入将推动更大规模的神经网络模型的发展，提升多模态任务的处理能力，具有重要的实际价值和未来影响。

📄 摘要（原文）

While Transformer networks benefit from a global receptive field, their quadratic cost relative to sequence length restricts their application to long sequences and high-resolution inputs. We introduce Fast Multipole Attention (FMA), a divide-and-conquer mechanism for self-attention inspired by the Fast Multipole Method from n-body physics. FMA reduces the time and memory complexity of self-attention from $\mathcal{O}\left(n^2\right)$ to $\mathcal{O}(n \log n)$ and $\mathcal{O}(n)$ while preserving full-context interactions. FMA contains a learned hierarchy with $\mathcal{O}(\log n)$ levels of resolution. In this hierarchy, nearby tokens interact at full resolution, while distant tokens engage through progressively coarser, learned basis functions. We have developed both 1D and 2D implementations of FMA for language and vision tasks, respectively. On autoregressive and bidirectional language modeling benchmarks, the 1D variant either matches or outperforms leading efficient attention baselines with substantially lower memory use. With linear complexity, the 2D variant demonstrates superior performance over strong vision transformer baselines in classification and semantic segmentation tasks. Our results confirm that the multilevel attention implemented by FMA allows Transformer-based models to scale to much longer sequences and higher-resolution inputs without loss in accuracy. This provides a principled, physics-inspired approach for developing scalable neural networks suitable for language, vision, and multimodal tasks. Our code will be available at https://github.com/epoch98/FMA.