Multipole Semantic Attention: A Fast Approximation of Softmax Attention for Pretraining
Authors: Rupert Mitchell, Kristian Kersting
Organizations: Department of Computer Science, TU Darmstadt, Darmstadt, Germany · Hessian Center for AI (hessian.AI), Darmstadt, Germany · German Research Center for Artificial Intelligence (DFKI), Darmstadt, Germany · Center for Cognitive Science, TU Darmstadt, Darmstadt, Germany
Pretraining transformers on long sequences (entire code repositories, collections of related documents) is bottlenecked by quadratic attention costs. We present Multipole Semantic Attention (MuSe), which accelerates 64k-context pretraining by 36% while matching baseline loss, requiring no architectural changes. MuSe clusters queries and keys separately in representation space. This yields query-specific summaries that substantially outperform spatial blocking at matched sparsity, while also enabling drop-in compatibility with existing pretrained models; we validate on Llama 3.1-8B and 3.2-1B without retraining. We pretrain language models up to 1B parameters at 64k context on code and scientific documents, confirming that MuSe preserves quality and long-context utilization during training.