Multipole Attention for Efficient Long Context Reasoning

Large Reasoning Models (LRMs) have shown promising accuracy improvements on complex problem-solving tasks. While these models have attained high accuracy by leveraging additional computation at test time, they need to generate long chain-of-thought reasoning in order to think before answering, which requires generating thousands of tokens. While sparse attention methods can help reduce the KV cache pressure induced by this long autoregressive reasoning, these methods can introduce errors which disrupt the reasoning process. Additionally, prior methods often pre-process the input to make it easier to identify the important prompt tokens when computing attention during generation, and this pre-processing is challenging to perform online for newly generated reasoning tokens. Our work addresses these challenges by introducing Multipole Attention, which accelerates autoregressive reasoning by only computing exact attention for the most important tokens, while maintaining approximate representations for the remaining tokens. Our method first performs clustering to group together semantically similar key vectors, and then uses the cluster centroids both to identify important key vectors and to approximate the remaining key vectors in order to retain high accuracy. We design a fast cluster update process to quickly re-cluster the input and previously generated tokens, thereby allowing for accelerating attention to the previous output tokens. We evaluate our method using emerging LRMs such as Qwen-8B, demonstrating that our approach can maintain accuracy on complex reasoning tasks even with aggressive attention sparsity settings. We also provide kernel implementations to demonstrate the practical efficiency gains from our method, achieving up to 4.5$\times$ speedup for attention in long-context reasoning applications. Our code is available at https://github.com/SqueezeAILab/MultipoleAttention.

💡 Research Summary

The paper introduces “Multipole Attention,” a novel attention mechanism designed to make long‑context reasoning in Large Reasoning Models (LRMs) both memory‑efficient and accurate. LRMs such as Qwen‑8B and DeepSeek‑R1‑Distil‑Qwen2.5‑14B achieve high performance on complex tasks by generating thousands of chain‑of‑thought tokens, but this leads to a massive KV‑cache that strains GPU memory during inference. Existing sparse‑attention techniques reduce KV‑cache load by selecting a subset of tokens, yet they often cause substantial accuracy loss because important tokens may be omitted, and they rely on preprocessing that cannot be applied online to newly generated tokens.

Multipole Attention addresses these issues through three core ideas:

1. Semantic Clustering of Keys – All key vectors in the KV‑cache are clustered using k‑means based on semantic similarity (not positional proximity). Each cluster is represented by a key centroid (Kc) and a value centroid (Vc).

2. Centroid‑Based Importance Estimation – For a given query q, the model computes a soft‑max‑style score S_i = exp(q·Kc_i) / Σ_j N_j·exp(q·Kc_j), where N_j is the number of keys in cluster j. Clusters are sorted by S_i, and the top clusters that fit within a predefined token budget are deemed “important.” Exact attention is computed for all keys inside these important clusters.

3. Multipole Approximation for the Rest – For less important clusters, the attention contribution is approximated as N_i·exp(q·Kc_i)·Vc_i (the denominator of the soft‑max is omitted because it is shared across all clusters). This retains the aggregate influence of the entire cluster while avoiding per‑token computation.

Because rotary positional embeddings (RoPE) rotate keys differently depending on position, the authors adopt a “Windowed RoPE” strategy that assumes a fixed relative offset when computing centroids, improving clusterability.

A hierarchical extension uses progressively coarser centroids for clusters farther from the query, further reducing the number of centroid‑query dot products.

Fast Online Clustering – During autoregressive generation, new tokens continuously extend the KV‑cache. Re‑clustering the whole sequence each step would be prohibitive. The authors propose a block‑wise clustering scheme: the sequence is divided into blocks of W tokens; only the final block is reclustered when new tokens arrive. A sliding window with an overlap α ensures the final block never contains fewer than α tokens, guaranteeing enough data for a stable k‑means run. For the newly added tokens, a quick initial assignment is performed by sampling random centroids and using a batched version of the sequential k‑means update (MacQueen, 1967), followed by a few refinement iterations. This yields an O(new‑tokens) update cost.

System Implementation – Custom Triton kernels implement three stages: (1) centroid‑query comparison to select important clusters, (2) exact attention on the selected keys, and (3) centroid‑based approximate attention for the remaining keys. The implementation keeps the full KV‑cache in memory but loads only a small subset of keys for exact computation, dramatically reducing memory bandwidth.

Experimental Results – On benchmarks requiring chain‑of‑thought reasoning (MATH, Codeforces, ARC‑E), Multipole Attention maintains near‑full accuracy while operating with less than 10 % of the original KV‑cache budget. Compared to prior sparse‑attention methods, the accuracy drop is ≤0.5 % versus 5‑7 % for those baselines. Performance-wise, the attention layer achieves up to 4.5× speedup, and end‑to‑end decoding sees about 2.8× acceleration. Ablation studies show that removing the hierarchical refinement reduces accuracy by ~15 %, confirming its importance.

Limitations and Future Work – The approach relies on k‑means clustering, which may struggle with very high‑dimensional key spaces or highly non‑Gaussian distributions. Dynamic block sizing and alternative nearest‑neighbor structures (e.g., IVF, HNSW) could further improve scalability.

Conclusion – Multipole Attention offers a practical solution for long‑context inference: it preserves the full contextual information of the KV‑cache through centroid approximations while drastically cutting memory traffic and compute. The fast online clustering and hierarchical design make it suitable for real‑time generation, and the demonstrated speedups suggest it could become a standard component in future LRM deployments that require extensive reasoning over thousands of tokens.