REVIEW 3 major objections 5 minor 17 references
Multipole Semantic Attention: A Fast Approximation of Softmax Attention for Pretraining
T0 review · 3 major / 5 minor · reviewed 2026-08-04 · deepseek-v4-flash
Pith's one-line read MuSe clusters queries and keys to approximate softmax attention, cutting 16k-context pretraining time by 12.2% while degrading final loss by only 0.36%.
desk verdict A genuinely novel multipole derivation with a sound error bound, but the abstract overclaims the empirical results and the single-run loss difference is within noise. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The load-bearing identity is the cumulant expansion of attention. Treating each key-value cluster as a probability distribution, the log-normalizer and the value output are derivatives of the joint cumulant generating function of keys and values. Truncating after the first cumulants yields the monopole terms (exponentially tilted cluster centroids) and the dipole term (the key-value covariance contracted with the residual query). The two-stage scheme drops only the residual-residual term q̃·k̃, giving an error of order Tr(Cov(q̃,q̃) Cov(k̃,k̃)), which K-means is designed to minimize.
What would settle it
Run the same 30M/16k pretraining comparison at 64k context and with a 1B-parameter model, and check whether the runtime gain matches the abstract's claimed 36% and the loss gap stays near 0.36%; if the gap grows substantially, the claim that MuSe preserves quality at scale fails. Independently, instrument the backward pass to measure the relative error of gradients; a large gradient error would show that training degradation is driven by something other than forward approximation fidelity.
Extended reading notes
Core claim
The paper's central claim is that softmax attention can be faithfully approximated by a two-stage semantic clustering in the representation spaces of queries and keys. Queries are partitioned into coarse clusters; each coarse query centroid attends to fine-grained key-value clusters to form query-dependent summary centroids; then individual residual queries refine those summaries. A dipole term, computed from the within-cluster covariance of keys and values, adds the first-order directional correction. The approximation error is governed by the product of intra-cluster query and key variances, so increasing the number of clusters drives the method toward exact attention. Empirically, on isol
Load-bearing premise
The error bound assumes that a single iteration of K-means yields clusters tight enough in both query and key space for the first-order dipole truncation to be accurate; the paper's measurements show a relative squared error of about 0.19 at its chosen settings, and the backward-pass error is never measured.
Editorial extensions
If this is right
- Training on long sequences becomes nearly linear in context length instead of quadratic, so 64k and longer contexts can be trained at a fraction of the previous cost.
- Because MuSe is a drop-in attention replacement, existing pretrained transformers can be adapted to longer contexts without retraining from scratch.
- The dipole correction provides a principled way to recover directional information lost by centroid-only methods, improving the accuracy of clustering-based attention.
- Users can trade speed against accuracy by choosing the number of clusters; the paper reports a roughly fourfold increase in clusters halves the relative squared error.
- The observed 12.2% end-to-end speedup at 16k context suggests larger relative gains at longer contexts, where the exact attention cost dominates more.
Reading between the lines
- If the error bound transfers, MuSe should scale to 64k context and billion-parameter models, but the paper's own evidence for such scales is only a projection; the abstract's 36% speedup at 64k is not backed by an experiment in the text.
- The backward pass uses the same approximations but its error is never measured; training dynamics could be sensitive to gradient error even when forward error is modest.
- Because the ideal metric for clustering queries depends on key variance and vice versa, replacing Euclidean K-means with a variance-weighted clustering could reduce error at fixed cluster count.
- The separate clustering of queries and keys, justified by attention's invariance to basis changes of the keys, suggests a general principle for designing efficient attention beyond this specific implementation.
Signed reviews
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper proposes MuSe, an approximate softmax attention mechanism based on separate K-means clustering of queries and keys, using monopole (centroid) summaries plus dipole (covariance) corrections, and a hierarchical block decomposition for causal attention. The authors derive the approximation from a cumulant generating function expansion (Sections 3-4), analyze its complexity as O(NCD) for acausal and O(NCD log N) for causal attention, and report microbenchmarks on isolated attention layers plus an end-to-end pretraining experiment of a 29.4M-parameter model at 16k context (Section 7.3), where they claim a 12.2% (later 12.7%) runtime reduction with a 0.36% loss degradation. The abstract additionally claims 64k-context 36% speedups, validation on Llama 3.1-8B/3.2-1B without retraining, and pretraining up to 1B parameters at 64k context, but these results are not present in the body.
Significance. The theoretical framing is a genuine contribution: connecting softmax attention to multipole expansions via the cumulant generating function is elegant and yields a plausible error analysis, and the paper is transparent about the partial dipole specialization and its costs. The microbenchmarks show substantial speedups at long context (20-30x at 64k) and the method has the desirable property of converging to exact attention as cluster count grows. However, the central empirical claim of loss preservation rests on a single unseeded run, and several abstract claims go beyond what the body presents. If the training result were replicated over multiple seeds, the method would be a useful addition to efficient-attention literature.
major comments (3)
- [Abstract vs. Sections 5-7] The abstract claims 'accelerates 64k-context pretraining by 36% while matching baseline loss', validation on Llama 3.1-8B and 3.2-1B without retraining, and pretraining up to 1B parameters at 64k context. None of these experiments appear in the body. The only end-to-end training is a 29.4M model at 16k context (Section 7.3), and 64k appears only in isolated-layer benchmarks (Table 2). These unsupported claims overstate the scope and must be corrected or substantiated.
- [Section 7.3, Table 4] The headline '0.36% loss degradation' compares a single MuSe run (3.322 nats) with a single Pallas Flash run (3.310 nats). No seeds or variance are reported for either condition. For a 29.4M model trained on 2B tokens, run-to-run variation is typically comparable to or larger than 0.012 nats, so the difference is not statistically distinguishable from noise. Please provide multiple seeds (at least 3) or another statistical comparison to support the claim that MuSe preserves pretraining quality.
- [Section 7.2 and Section 7, Q4] The paper explicitly lists as a key question whether forward and backward approximation errors are sufficiently small, but only forward relative squared error is measured (Tables 1 and 3). Backward-pass gradient error is never quantified. The end-to-end loss is the only indirect evidence, and it is confounded by the single-run issue. Add a gradient-error measurement on recorded activations, or provide additional training runs, to support the claim that training quality is preserved.
minor comments (5)
- [Section 4, after Eq. (7)] The sentence claiming the error quantity is 'exactly the quantity minimized by K-means' is imprecise: K-means minimizes the trace of within-cluster covariance (sum of squared distances), not the product of query and key covariance traces. Rephrase as 'related to' or 'a quantity that K-means tends to reduce'.
- [Section 4 title] Typo: 'desribed' should be 'described'.
- [Algorithm 3] The covariance C[k] is computed over the U dimension unweighted by the attention probabilities. Clarify in the comment or pseudocode that this is Cov_j(v,k), not the exponentially tilted Cov_{ij}(v,k) used in Eq. (6).
- [Section 5, Clustering] The K-means++ initialization described as 'sampling points proportional to their squared norm' is equivalent to the first K-means++ sample only for zero-mean data. State whether the queries/keys are mean-centered or whether this is an approximation.
- [Figure 3] The inset showing convergence in the 1-2B token range is important but appears too small to read in the preprint. Ensure legibility in the camera-ready version.
Circularity Check
No significant circularity: MuSe's multipole approximation is derived from a CGF Taylor expansion, converges to exact attention as cluster count grows, and rests on external runtime baselines rather than self-citation.
full rationale
The paper's central derivation is self-contained. The softmax attention approximation is built from a cumulant-generating-function Taylor expansion (Eqs. 6 and 7), with the dipole term being a truncated linear correction; this is a genuine mathematical derivation, not an assumption of the desired result. The stated error bound O(Tr(Cov(q̃,q̃) Cov(k̃,k̃))) follows from the truncation and is not used as an input. The method provably reduces to exact attention when the number of clusters equals the sequence length, since singleton clusters have zero intra-cluster covariance and the two-stage algorithm reproduces the exact softmax normalization. No fitted parameter is relabeled as a prediction: clustering centroids and covariances are computed from the input data as part of the approximation, and the end-to-end pretraining result (0.36% loss degradation) is an empirical outcome, not a quantity forced by construction. The runtime comparisons use external Flash Attention implementations as baselines, so the speedup claim does not reduce to a self-citation. The main limitations—the error bound assumes K-means produces tight clusters, and the 0.012-nat loss difference rests on a single unseeded run—are robustness/reproducibility concerns, not circularity. The paper contains no load-bearing self-citations, and its novelty claims are anchored in external related work. Therefore the derivation chain is not circular.
Assumptions & free parameters
free parameters (4)
- Number of clusters C =
64 (32 and 128 explored)
- K-means iterations =
1 (main runs); 3 and 5 explored
- Cluster size cap ratio =
1.5 (1.25 to 4.0 explored)
- Causal block size B (hierarchical decomposition) =
8192
assumptions (4)
- domain assumption Query and key representations are clusterable by Euclidean K-means with 1-5 iterations.
- standard math Softmax attention can be modeled as exponential tilting of a discrete uniform distribution over keys, and its CGF Taylor expansion truncates at first order.
- ad hoc to paper Partial dipole specialization (omitting the exponential tilt by q_i in Cov_j(v,k)) is an acceptable accuracy loss.
- domain assumption No key cluster can be neglected; locality in attention is captured only through exponential tilting.
invented entities (3)
-
Query-specialized monopole summaries (exponentially tilted key/value centroids, K_ij and V_ij)
-
Partially specialized dipole correction (aggregated covariance Cov_j(v,k) per query cluster)
-
Hierarchical block decomposition for causal attention
Cite this review
Pith. "Pith review of Multipole Semantic Attention: A Fast Approximation of Softmax Attention for Pretraining." pith.science (2026). https://pith.science/paper/LTYDHCQP
@misc{pith2026250910406,
author = {Pith},
title = {Pith review of: Multipole Semantic Attention: A Fast Approximation of Softmax Attention for Pretraining},
year = {2026},
howpublished = {\url{https://pith.science/paper/LTYDHCQP}},
note = {Machine review of arXiv:2509.10406}
}
read the original abstract
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.
Figures
Reference graph
Works this paper leans on
-
[1]
Fu, Stefano Ermon, Atri Rudra, and Christopher R \' e
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher R \' e . Flashattention: Fast and memory-efficient exact attention with io-awareness. In NeurIPS, 2022
2022
-
[2]
Mahoney, Yakun Sophia Shao, Kurt Keutzer, and Amir Gholami
Coleman Hooper, Sebastian Zhao, Luca Manolache, Sehoon Kim, Michael W. Mahoney, Yakun Sophia Shao, Kurt Keutzer, and Amir Gholami. Multipole attention for efficient long context reasoning. CoRR, abs/2506.13059, 2025 a
arXiv 2025
-
[3]
Reformer: The efficient transformer
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In ICLR . OpenReview.net, 2020
2020
-
[4]
Efficient content-based sparse attention with routing transformers
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Trans. Assoc. Comput. Linguistics, 9: 0 53--68, 2021
2021
-
[5]
Mahoney, Kurt Keutzer, and Amir Gholami
Coleman Richard Charles Hooper, Sehoon Kim, Hiva Mohammadzadeh, Monishwaran Maheswaran, Sebastian Zhao, June Paik, Michael W. Mahoney, Kurt Keutzer, and Amir Gholami. Squeezed attention: Accelerating long context length LLM inference. In ACL (1) , pages 32631--32652. Association for Computational Linguistics, 2025 b
2025
-
[6]
A \( ^2 \) ats: Retrieval-based KV cache reduction via windowed rotary position embedding and query-aware vector quantization
Junhui He, Junna Xing, Nan Wang, Rui Xu, Shangyu Wu, Peng Zhou, Qiang Liu, Chun Jason Xue, and Qingan Li. A \( ^2 \) ats: Retrieval-based KV cache reduction via windowed rotary position embedding and query-aware vector quantization. In ACL (Findings) , pages 12451--12463. Association for Computational Linguistics, 2025
2025
-
[7]
Tactic: Adaptive sparse attention with clustering and distribution fitting for long-context llms
Kan Zhu, Tian Tang, Qinyu Xu, Yile Gu, Zhichen Zeng, Rohan Kadekodi, Liangyu Zhao, Ang Li, Arvind Krishnamurthy, and Baris Kasikci. Tactic: Adaptive sparse attention with clustering and distribution fitting for long-context llms. CoRR, abs/2502.12216, 2025
arXiv 2025
-
[8]
Li, Madian Khabsa, Han Fang, and Hao Ma
Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. CoRR, abs/2006.04768, 2020
arXiv 2006
Show all 17 references
-
[9]
o mformer: A nystr \
Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nystr \" o mformer: A nystr \" o m-based algorithm for approximating self-attention. In AAAI , pages 14138--14148. AAAI Press, 2021
2021
-
[10]
Fast multipole attention: A divide-and-conquer attention mechanism for long sequences
Yanming Kang, Giang Tran, and Hans De Sterck. Fast multipole attention: A divide-and-conquer attention mechanism for long sequences. arXiv preprint arXiv:2310.11960, 2023
2023
-
[11]
H-transformer-1d: Fast one-dimensional hierarchical attention for sequences
Zhenhai Zhu and Radu Soricut. H-transformer-1d: Fast one-dimensional hierarchical attention for sequences. In ACL/IJCNLP (1) , pages 3801--3815. Association for Computational Linguistics, 2021
2021
-
[12]
A hierarchical o(nlogn) force-calculation algorithm
Hut Barnes. A hierarchical o(nlogn) force-calculation algorithm. Nature, 1986
1986
-
[13]
Rapid solution of integral equations of classical potential theory
V Rokhlin. Rapid solution of integral equations of classical potential theory. Journal of Computational Physics, 60 0 (2): 0 187--207, 1985
1985
-
[14]
Perceiver: General perception with iterative attention
Andrew Jaegle, Felix Gimeno, Andy Brock, Oriol Vinyals, Andrew Zisserman, and Jo \ a o Carreira. Perceiver: General perception with iterative attention. In ICML , volume 139 of Proceedings of Machine Learning Research, pages 4651--4664. PMLR , 2021
2021
-
[15]
Luna: Linear unified nested attention
Xuezhe Ma, Xiang Kong, Sinong Wang, Chunting Zhou, Jonathan May, Hao Ma, and Luke Zettlemoyer. Luna: Linear unified nested attention. In NeurIPS, pages 2441--2453, 2021
2021
-
[16]
Colwell, and Adrian Weller
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tam \' a s Sarl \' o s, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J. Colwell, and Adrian Weller. Rethinking attention with perfo...
2021
-
[17]
Compressive transformers for long-range sequence modelling
Rae, Jack W., Potapenko, Anna, Jayakumar, Siddhant M., Lillicrap, and Timothy P. Compressive transformers for long-range sequence modelling. In ICLR, 2020
2020
Reviewed August 4, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.