Pith. sign in

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 →

arxiv 2509.10406 v4 pith:LTYDHCQP submitted 2025-09-12 cs.LG

classification cs.LG
keywords multipoleexpansionsoftmaxattentionapproximationsemanticclusteringlong-contextpretrainingK-meansdipolecorrectiontransformerefficiencycausal
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

This paper tries to show that the quadratic cost of softmax attention can be approximated during training by clustering queries and keys separately and keeping only summary statistics. The proposed method, MuSe, uses a two-stage attention scheme plus a covariance-based dipole correction to capture directional detail within clusters. In a 30M-parameter model trained on 16k-token book-length sequences, MuSe reduces wall-clock pretraining time by 12.2% while increasing final loss by only 0.36%. Because it is a drop-in replacement for the attention layer, it offers a path to longer-context pretraining without changing the architecture.

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.

Watch

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

Editorial extensions of the paper, not claims the author makes directly.

  • 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.
Share X Bluesky LinkedIn Reddit HN

Signed reviews

No signed human review yet.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

3 major / 5 minor

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)
  1. [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.
  2. [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.
  3. [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)
  1. [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'.
  2. [Section 4 title] Typo: 'desribed' should be 'described'.
  3. [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).
  4. [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.
  5. [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

0 steps flagged · score 0.0 of 10

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 4 free parameters · 4 assumptions · 3 invented entities

The method rests on standard probability theory (CGF expansions) and on the domain assumption that K-means clusters of raw query/key vectors are tight enough for a first-order expansion to hold. The partial dipole specialization is an explicit implementation compromise. No new physical entities are postulated; the dipole term and two-stage summaries are algorithmic constructs validated only in-paper.

free parameters (4)
  • Number of clusters C = 64 (32 and 128 explored)
    Controls the speed/accuracy tradeoff; chosen as 64 for square matmuls (Section 5). Error ranges from 0.2475 (C=32) to 0.1123 (C=128, table 1).
  • K-means iterations = 1 (main runs); 3 and 5 explored
    Few iterations keep clustering cost comparable to attention cost; error drops roughly 20% from 1 to 5 iterations (table 5).
  • Cluster size cap ratio = 1.5 (1.25 to 4.0 explored)
    Caps the ragged-tensor block size; lower cap reduces compute but worsens clustering quality (Section 7.1, tables 5-6).
  • Causal block size B (hierarchical decomposition) = 8192
    Only instantiation tested: one exact diagonal 8k block and one MuSe block; the O(log N) multi-level hierarchy is never swept or tested.
assumptions (4)
  • domain assumption Query and key representations are clusterable by Euclidean K-means with 1-5 iterations.
    Sections 3.1 and 5 run K-means on raw query/key vectors; Section 4's error bound requires small intra-cluster variances, and the paper concedes the ideal metric would be variance-weighted, not Euclidean.
  • 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.
    Section 4 defines M_j(q)=E exp(q.k) and expands the log MGF; this is standard probability theory. The truncation to terms linear in q_tilde is the method's defining approximation.
  • ad hoc to paper Partial dipole specialization (omitting the exponential tilt by q_i in Cov_j(v,k)) is an acceptable accuracy loss.
    Section 4 'Partially Specialized Dipoles' makes this compromise to reduce covariance cost from O(Cq N D^2) to O(N D^2), degrading the error bound from O(q_tilde^2 k_tilde^2) to O(q_bar q_tilde k_tilde^2).
  • domain assumption No key cluster can be neglected; locality in attention is captured only through exponential tilting.
    Figure 1 caption and Section 4 state that all clusters must be approximated; this is a claim about the structure of trained attention patterns and is not justified by reference to measured attention sparsity.
invented entities (3)
  • Query-specialized monopole summaries (exponentially tilted key/value centroids, K_ij and V_ij)
    purpose: Let coarse query clusters produce query-dependent key/value summaries that fine residual queries refine in the second stage.
    Central to the two-stage mechanism (algorithms 2-3); its contribution is measured only through in-paper error tables and ablations (table 3), with no external benchmark.
  • Partially specialized dipole correction (aggregated covariance Cov_j(v,k) per query cluster)
    purpose: Linear correction capturing directional variance inside key-value clusters that centroids lose.
    Ablation shows removing it raises relative squared error by 15-80% (table 3), but this evidence is internal to the paper; no independent falsifiable handle.
  • Hierarchical block decomposition for causal attention
    purpose: Splits the causal mask into exact diagonal blocks and MuSe off-diagonal blocks to reach O(NCD log N).
    Only the single-level case (block size 8k) is exercised in pretraining (Section 7.3); the multi-level version that would realize the log N scaling is untested.

how reviews work

0 comments
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

Figures reproduced from arXiv: 2509.10406 by the authors.

Figure 1
Figure 1. Idealized depiction of MuSe attention. For some particular query centroid [PITH_FULL_IMAGE:figures/full_fig_p006_1.png] view at source ↗
Figure 2
Figure 2. Hierarchical block decomposition for causal attention. Diagonal blocks (green) use Flash attention for exact computation, while below-diagonal blocks (blue) use MuSeapproxi￾mation. 6 From Acausal to Causal We extend our acausal approximation to causal attention by decomposing the lower-triangular attention matrix into blocks as shown in figure 2. The block-diagonal uses Flash attention for exact local computation, w… view at source ↗
Figure 3
Figure 3. Training negative log-likelihood loss in nats [PITH_FULL_IMAGE:figures/full_fig_p011_3.png] view at source ↗

Discussion (0). Continue with ORCID to comment.

Reference graph

Works this paper leans on

17 extracted references · 2 linked inside Pith

  1. [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

  2. [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

  3. [3]

    Reformer: The efficient transformer

    Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In ICLR . OpenReview.net, 2020

  4. [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

  5. [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

  6. [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

  7. [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

  8. [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

Show all 17 references
  1. [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

  2. [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

  3. [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

  4. [12]

    A hierarchical o(nlogn) force-calculation algorithm

    Hut Barnes. A hierarchical o(nlogn) force-calculation algorithm. Nature, 1986

  5. [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

  6. [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

  7. [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

  8. [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...

  9. [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

Pith tools

Reviewed August 4, 2026 · model on record in the stance chip above.