Multipole Semantic Attention: A Fast Approximation of Softmax Attention for Pretraining

2025-09-12

AI Summary

To address the O(N²) computational complexity of softmax attention in Transformers for long sequences, this paper proposes Multipole Semantic Attention (MuSe). MuSe preserves attention asymmetry by decoupling query and key spaces via semantic clustering, enhances centroid approximation accuracy using multipole expansion (monopole plus dipole correction), and supports both non-causal and causal attention through hierarchical block decomposition. As a plug-and-play module requiring no architectural modifications, MuSe reduces complexity to O(NCD) or O(NCD log N). Experiments demonstrate that on 8k-context sequences, MuSe achieves 3× speedup per layer with relative squared error under 20%; in end-to-end pretraining on 16k-context sequences, it reduces runtime by 12.2% while increasing loss by only 0.36%. These results significantly improve efficiency for long-sequence modeling without compromising model fidelity.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Natural Language Processing: Sentence-level Semantics, Textual Inference, etc.
• Computer Vision: Multi-modal Vision

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Graph Algorithms and Modeling for the Web: Algorithms and analysis for heterogeneous, signed, attributed, multi-relational, temporal, higher-order, and annotated Web-related graphs

Abstract

We present Multipole Semantic Attention (MuSe), an efficient approximation of softmax attention that combines semantic clustering with multipole expansions from computational physics. Our method addresses the quadratic computational complexity of transformers in the context length by clustering queries and keys separately in their learned representation spaces, enabling a hierarchical two-stage attention mechanism. Unlike prior clustering approaches that group only keys or use unified clustering, we maintain separate clusterings that respect attention's asymmetric treatment of these spaces. We augment centroid-based (monopole) approximations with dipole corrections that capture directional variance within clusters, preserving richer information during training. The method operates as a drop-in replacement for standard attention, requiring only hyperparameter specification without architectural modifications. Our approach achieves $mathcal{O}(NCD)$ complexity for acausal attention with $C$ clusters and $mathcal{O}(NCD log N)$ for causal attention. On isolated attention layers, we demonstrate $3 imes$ speedup over CUDNN Flash Attention at 8k context length, with relative squared errors below 20%. For causal attention, we develop a hierarchical block decomposition that combines exact local computation with efficient long-range approximation. In end-to-end pretraining of a 30M parameter model on book-length texts with 16k context, we achieve 12.2% runtime reduction with only 0.36% loss degradation, establishing the viability of multipole approximations for efficient transformer pretraining.

Problem

• Reduces quadratic complexity of transformer attention
• Approximates softmax via semantic clustering and multipole expansions
• Enables efficient long-context pretraining with minimal performance loss

Innovation

• Multipole Semantic Attention clustering queries and keys
• Dipole corrections capturing directional variance within clusters
• Hierarchical block decomposition combining exact and approximate computation