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

Pretraining transformers on long sequences, such as entire code repositories or 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 is a training-time approximation that 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.

Modern AI language models learn by having every word in a text "pay attention" to every other word. This works well but is costly: doubling the length of the text quadruples the work, which makes training on very long inputs—entire code repositories or collections of related documents—prohibitively slow. We borrow a trick from physics. To simulate gravity among millions of stars, astronomers don't compute every pairwise pull; they group distant stars together and treat each group as a single summarized mass, computing only the nearby interactions exactly. Our method, Multipole Semantic Attention (MuSe), does the analogous thing for language: it groups words with similar meaning, replaces most word-to-word comparisons with compact group summaries, and computes exactly only the few groups that matter most for each word. Importantly, this shortcut is used only while training the model; once trained, the model runs in the normal way. In our experiments, MuSe sped up training by 36% while matching the quality of standard training. In practice, this lets anyone train on longer contexts than their current compute budget would otherwise allow.