{"id":"17fed9aa-9c88-41f2-bac7-2181b3f246b4","arxiv_id":"2505.00315","paper_version":1,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":2,"one_line_summary":"MoSA, an expert-choice style sparse attention that selects per-head top-k tokens, outperforms dense transformers on C4 language modeling under matched FLOPs with up to 27% perplexity improvement and reduces wall-clock time, memory, and KV-cache size in perplexity-matched comparisons.","lead":"MoSA is a new attention design that lets each attention head pick its own small set of tokens, so the cost of attention stops growing quadratically with sequence length. If its results hold at scale, it offers a practical route to cheaper training, faster generation, and much smaller inference memory for large language models.","discovery_kind":"extension","skeptic_critique":{"model":"deepseek-v4-flash","headline":"MoSA's practical KV-cache and inference claims are not established because exact top-k selection is non-autoregressive; Table 2 counts training-time materialized key-values, not cache usable during generation. The iso-FLOP perplexity claim is plausible; the resource claims need conditional wording.","rationale":"The reader's weakest assumption identifies the most load-bearing weakness: the non-autoregressive router prevents the direct transfer of the reported KV-cache counts to sequential inference. I agree with the CONDITIONAL verdict. The iso-FLOP perplexity comparison is otherwise well supported by full per-sparsity curves, the hybrid-model analysis in Appendix B, and the released code. The lack of error bars and the absent Native Sparse Attention baseline are secondary because the central comparison is internal and the claim is explicitly limited to 'tested variants.' The paper deserves credit for disclosing the autoregressive limitation, but it still presents Table 2 as evidence of practical inference benefits, which is not established. A prefix-selection agreement experiment or an actual generation benchmark would settle the concern. No change to the reader's verdict is needed.","tokens_in":21447,"tokens_out":6089,"duration_ms":71287,"concrete_test":"Run exact autoregressive decoding from the released Tiny and Small MoSA checkpoints: at each generation step, compute router scores from the prefix only, select top-k per head, and maintain a KV cache containing only selected prefix tokens; measure perplexity, wall-clock time, and peak KV cache against the dense baseline. Additionally, on held-out C4 sequences, compare the prefix-constrained selected set I_prefix(t) with the full-sequence selected set I_full and report the fraction of final selected tokens whose membership changes after future tokens arrive. If more than a small fraction of KV entries change or perplexity degrades, the Table 2 KV-cache savings do not transfer to autoregressive inference.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The paper's central iso-FLOP result (Table 1, Fig. 3) is a coherent empirical claim, but the resource-optimization contribution (Table 2) rests on an unverified bridge from training-time token selection to autoregressive generation. In MoSA, the top-k index set I is selected from the full sequence X (Sec. 2.2), and the paper's own Limitations section states that this 'requires adaptations to be directly applicable to the autoregressive scenario.' When decoding a prefix, the final top-k set for a head is unknown until future tokens exist; the same past token can enter or leave I once later tokens are observed. Serving such a model exactly would require either keeping KV pairs for all prefix positions until the final selection is known (eliminating the reported KV reduction) or approximating membership with an autoregressive classifier, e.g., MoD-style, whose accuracy and cost are not evaluated. The KV formula in Sec. 3.3, KV = TH_dense + kH_mosa, counts key-value pairs materialized during a full-sequence training step, and Table 2's wall-clock and memory numbers are training-step measurements, not generation latency. Therefore, the abstract's claim of drastically reduced KV-cache and the implied inference benefits are not yet demonstrated.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper introduces Mixture of Sparse Attention (MoSA), a sparse attention mechanism in which each attention head selects its own top-k tokens via expert-choice routing, reducing per-head cost from O(T^2) to O(k^2+T). The saved compute is used to add more heads, and the main models are hybrids that keep four dense heads alongside many sparse heads. Experiments on C4 language modeling at four scales (28M to 516M dense baselines) report iso-FLOP perplexity improvements over the dense baseline of up to 27%, claim that MoSA outperforms fixed sparse attention and Routing Transformer, report perplexity-matched wall-clock, memory, and KV-pair reductions, and include long-sequence and zero-shot downstream evaluations. The central empirical claim is that MoSA is the only tested sparse attention variant that beats the dense baseline under a fixed FLOP budget.","tokens_in":21651,"tokens_out":4304,"duration_ms":42555,"significance":"If the iso-FLOP result is robust, it is a valuable finding: it suggests that learned content-based sparsity with perfect load balancing can improve language modeling at fixed compute, contrary to the static-sparsity baselines tested. The paper's strengths include full per-sparsity perplexity curves across scales, explicit FLOP derivations, an ablation of hybrid versus pure MoSA, and a public code repository. However, the headline resource-optimization claims are not established: the top-k selection is non-autoregressive, so the reported KV-cache reduction does not translate directly to autoregressive inference, and the empirical protocol selects the best sparsity on test perplexity from single runs without seeds. The iso-FLOP central claim is plausible; the resource claims need substantial reframing or additional evidence.","major_comments":[{"comment":"The resource measurements are training-step numbers: 'Wall-time/step' is measured per training step and 'Memory (GB)' is training memory, not generation latency or inference-time memory. The Abstract and Conclusions nonetheless imply inference benefits ('faster in wall-clock time', 'reduce the resource usage', 'potentially enabling more economical scaling'). Please separate training-time efficiency from inference-time efficiency and state clearly which claim is supported by which measurement.","section":"Section 3.3, Table 2, Abstract, and Conclusions"}],"minor_comments":[{"comment":"The reference 'the Appendix 5' is unclear; it should point to Table 5 or a named appendix section.","section":"Section 3.2, 'Appendix 5'"},{"comment":"The x-axis label '8196' should be '8192'.","section":"Section 3.4, Figure 4"},{"comment":"The Introduction says 'In Section 3 we show the performance of different models in downstream zero-shot tasks', but this content is in Section 3.5; the cross-reference should be corrected.","section":"Section 3.5 / Introduction"},{"comment":"The sentence 'we train on the 105SB≈ 6.5B tokens from the dataset' is garbled; it should read '100k batches × batch size 64 × sequence length 1024 ≈ 6.5B tokens'.","section":"Section 3, Implementation details"},{"comment":"There is a typo: 'pretaining phase' should be 'pretraining phase'.","section":"Related Work"},{"comment":"The main claim refers to hybrid MoSA models with four dense heads; pure MoSA without dense heads underperforms the dense baseline (Figure 5). The paper does disclose this, but contribution statements such as 'MoSA improves perplexity' should consistently say 'hybrid MoSA with four dense heads' to avoid overgeneralization.","section":"Table 1 / Appendix B"}],"recommendation":"major_revision","confidential_remarks":"The central iso-FLOP contribution is within the journal's scope and appears plausible, but the paper currently overstates the practical resource benefits, and the test-set model selection weakens the quantitative headline. The authors' own Limitations section concedes the autoregressive gap, so the fix is likely a matter of framing plus additional evidence or a clear scope restriction rather than a fundamental invalidation. I recommend requiring the KV-cache and inference claims to be either demonstrated with an autoregressive approximation or explicitly restricted to training-time materialization, and requiring a validation-based model selection protocol with seeded repetitions."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Two things to know. The iso-FLOP perplexity result is the real contribution and it is plausible: MoSA beats dense transformers on perplexity at matched FLOPs across four scales, with full sparsity curves, while Fixed and Routing baselines do not. The other thing is that the paper's resource claims, especially 'drastically reduce the size of the KV-cache', are not established. The stress-test concern lands. MoSA selects its top-k over the full sequence X, so during autoregressive decoding the set of selected tokens for a prefix is not known until future tokens exist. The KV formula in Sec 3.3 counts key-value pairs materialized in a training-time forward pass; that is not a cache you can use for generation. The paper's own Limitations section says MoSA 'requires adaptations to be directly applicable to the autoregressive scenario', which directly undercuts the abstract's inference claim. This is a real flaw in the resource-optimization section, not a nitpick.\n\nThe mechanism is genuinely new: per-head top-k token selection with a learned sigmoid router, inspired by expert-choice routing. It differs from Routing Transformer by learning a direct scoring function, not online K-means, and it does not need to tie WQ to WK. The paper also deserves credit for showing the whole sparsity sweep, not hiding the failures: pure MoSA without dense heads underperforms (Appendix B), and the hybrid with four dense heads is necessary. That is honest and useful.\n\nSoft spots beyond the autoregressive gap. Table 1's headline improvement is selected as the best over a sparsity grid, and all numbers come from single runs with no seeds, so we do not know the variance. The baseline set is narrow: Native Sparse Attention, the most comparable recent trainable sparse method, is cited but not benchmarked. The downstream experiments are a mixed bag — MoSA loses to dense at the Large scale on most tasks — and the paper says why, but it tempers the 'improves quality' summary.\n\nThe central empirical claim, iso-FLOP perplexity, holds up as an experimental finding; the resource-optimization claim needs conditional wording or, better, an autoregressive adaptation with measured generation-time cache and latency. This is fixable in revision.\n\nMy recommendation: engage. Send it to peer review. Ask for multiple seeds, an NSA baseline, and either an autoregressive classifier that makes the cache claim real or a rewrite that limits the claim to training-time savings. If they can close the autoregressive gap, this is an important paper. As is, the iso-FLOP result is citable with a caveat.","headline":"The iso-FLOP perplexity result is worth taking seriously, but the KV-cache and inference claims are not supported by the experiments as reported.","tokens_in":22246,"tokens_out":3436,"would_cite":true,"duration_ms":34822,"reading_group":"yes","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"MoSA, a content-based sparse attention built on expert-choice routing, is the only sparse attention variant tested that beats dense transformers at equal compute, improving perplexity by up to 27 percent.","keywords":["sparse attention","mixture of experts","expert-choice routing","language modeling","efficient transformers","KV cache","content-based sparsity","perplexity"],"falsifier":"Train the same FLOP-matched hybrid MoSA model with an autoregressive router that may only select already-produced positions during both training and decoding, for instance by post-training a classifier to predict the non-autoregressive top-$k$ selections as the paper's discussion of Mixture-of-Depths suggests, and re-measure perplexity, wall-clock time, and decoding-time KV-cache size. The central claim stands only if a substantial share of the reported gains survives that change; if the gains evaporate, the results depend on training-time access to future tokens.","tokens_in":21195,"feed_emoji":"🎯","tokens_out":19880,"duration_ms":154414,"temperature":0.7,"pith_summary":"The paper claims that a learnable, content-based sparse attention can beat full dense attention in language modeling, provided each attention head is allowed to choose its own small set of tokens to attend to. The proposal, Mixture of Sparse Attention (MoSA), treats each head as an expert that picks its top-$k$ tokens from the sequence, cutting each head's cost from $O(T^2)$ to $O(k^2 + T)$ and freeing compute for many more, more specialized heads. In FLOP-matched experiments at four scales (28M to 516M parameters), MoSA improves perplexity — a language-modeling error measure, lower is better — by up to 27% over the dense baseline, and it is the only sparse method among those tested to do so; fixed sparse attention and the Routing Transformer both fall short. At matched perplexity, MoSA simultaneously improves wall-clock time and training memory, and cuts the number of cached key-value pairs by more than half. The practical upshot is a route to transformers that are both better and cheaper to run.","feed_headline":"Choosing tokens per head beats dense language-modeling baselines","feed_subtitle":"Perplexity drops by up to 27 percent at equal compute, and matched models cut the KV cache by more than half.","key_machinery":"The carrying mechanism is expert-choice routing transplanted into attention: each head acts as an expert that selects its own top-$k$ tokens through a learned per-token score $r = \\sigma(XW_r)$, which gives perfect load balance by construction and makes the sparsity pattern content-dependent. Queries, keys, values, and the attention matrix are computed only on the selected tokens, with the causal mask and rotary positions indexed by each token's original position in the sequence, and the outputs are scaled by their router scores before being scattered back, keeping the whole selection trainable end-to-end by the language-modeling objective. The per-head cost drops to $O(k^2 + T)$, and the paper's FLOP accounting shows that the routing overhead ($2hT + h'k$) is small relative to the projection and attention savings, so the freed budget can be spent on many more heads, up to hundreds of sparse heads in the hybrid models. A second load-bearing design choice is hybridization: keeping four dense heads alongside the sparse heads proves necessary, since pure-MoSA models without dense heads fail to beat the dense baseline in these experiments.","core_discovery":"The central claim is that dynamic, learned content-based sparsity makes sparse attention a strict improvement over dense attention rather than a compromise. MoSA implements this by giving every attention head a router that scores each token with a sigmoid and keeps the top-$k$ tokens, always including the first token; the head then computes queries, keys, values, and attention only on those $k$ tokens, re-inserts each output at its original position using the original positions for the causal mask and rotary embeddings, and scales it by the router score. Because each head now costs $O(k^2 + T)$ instead of $O(T^2)$, the saved compute is reinvested as many additional heads, and the paper finds that keeping four dense heads alongside the sparse heads performs best: in a fixed-FLOP comparison, perplexity falls by up to 27% (Tiny: 22.46 to 16.39) and improves at every model scale from 28M to 516M parameters, while perplexity-matched MoSA models run faster, use less memory, and cut the key-value cache by 51–70%. The same mechanism keeps its advantage on sequences up to 8192 tokens when combined with local attention, and a parameter-matched MoSA model (442M parameters, sparsity 8) beats a larger dense baseline (516M parameters) on perplexity, 12.16 versus 12.20.","pith_inferences":["Because the router sees the full sequence at training time, the reported KV-cache reductions describe training-time key-value materialization, not the cache maintained during greedy generation; how much the decoding-time cache shrinks depends on how faithfully an autoregressive adaptation reproduces the non-causal top-$k$ choices.","The comparison isolates the causal factor: fixed sparse attention enjoys the same projection savings as MoSA yet still loses to dense attention, so the design implies that content-based selection itself, not merely the smaller number of computed tokens, is what buys the perplexity gain.","A direct test of the mechanism is to restrict the router to causal information during training; if most of the perplexity gain persists, MoSA-style head specialization transfers cleanly to generative decoding, and if it disappears, the method's value is mostly in non-autoregressive settings.","The head-specialization principle should transfer to other modalities: in vision transformers or bidirectional encoders, where the router's non-causal selection is unproblematic, MoSA-style heads could deliver the same FLOP-matched gains without needing any autoregressive adaptation."],"forward_implications":["Within a fixed FLOP budget, dense heads can be traded for many sparse MoSA heads: perplexity keeps improving up to sparsity $\\rho \\approx 32$–$64$, turning head count and specialization into a free knob.","At matched perplexity, MoSA cuts the total number of key-value pairs by 51–70% relative to dense attention, which directly shrinks the KV cache that dominates inference memory.","MoSA composes with optimized attention kernels such as Flash Attention and with local attention for long sequences, where it beats fixed sparse attention and the Routing Transformer at sequence lengths up to 8192 while using fewer FLOPs.","The gains do not rest on FLOP accounting alone: a 442M-parameter MoSA model at sparsity 8 reaches 12.16 perplexity, below the 12.20 of the 516M-parameter dense baseline.","Ablations in the paper show sparse heads need a few dense heads to train stably, which fixes the architecture choice for anyone adopting the method."],"supporting_citations":[{"why":"Expert-choice routing: supplies the inverted selection paradigm in which each expert chooses its own tokens, giving perfect load balance without auxiliary losses.","marker":"[29]"},{"why":"Routing Transformer: the main content-based sparse attention baseline that MoSA must beat in the IsoFLOP, long-sequence, and downstream comparisons.","marker":"[25]"},{"why":"Sparse Transformer: defines fixed sparse attention, the static-pattern baseline that is a special case of MoSA with hand-set indices.","marker":"[18]"},{"why":"SwitchHead: earlier work applying mixture-of-experts inside attention; MoSA positions itself by making attention matrices sparse rather than only reducing heads.","marker":"[31]"},{"why":"StreamingLLM: documents attention sinks, the observation behind MoSA's rule of always including the first token in every head.","marker":"[48]"},{"why":"Mixture-of-Depths: cited as the proposed route to adapt non-autoregressive expert-choice routing to autoregressive models, the paper's stated main limitation.","marker":"[67]"},{"why":"FlashAttention: the optimized attention kernel that MoSA's sparse attention is compatible with, supporting the practical efficiency claim.","marker":"[39]"},{"why":"RoPE: the rotary positional embedding that MoSA re-indexes to tokens' original positions after top-$k$ selection.","marker":"[41]"}],"fun_headline_variants":["MoSA: learnable sparse attention that beats dense baselines","Expert-choice routing makes sparse attention beat dense","Sparse attention that beats dense: MoSA's token selection","MoSA cuts perplexity 27% and KV cache by half","Dynamic token selection lets sparse attention beat dense"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The load-bearing premise is that the router inspects the entire input sequence, including future tokens, when it picks each head's top-$k$ tokens; at generation time those future tokens do not exist, so the paper has not shown that the same selection, or the reported speed, memory, and KV-cache benefits, can be reproduced during one-token-at-a-time inference without an extra adaptation.","fun_headline_variants_meta":{"raw":{"variants":["MoSA: learnable sparse attention that beats dense baselines","Expert-choice routing makes sparse attention beat dense","Sparse attention that beats dense: MoSA's token selection","MoSA cuts perplexity 27% and KV cache by half","Dynamic token selection lets sparse attention beat dense"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000602,"raw_usage":{"total_tokens":2873,"prompt_tokens":1068,"completion_tokens":1805,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":684,"completion_tokens_details":{"reasoning_tokens":1726}},"tokens_in":684,"tokens_out":1805,"duration_ms":12983,"temperature":1.0,"reasoning_tokens":1726,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-16T04:46:04.340417+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Train the same FLOP-matched hybrid MoSA model with an autoregressive router that may only select already-produced positions during both training and decoding, for instance by post-training a classifier to predict the non-autoregressive top-$k$ selections as the paper's discussion of Mixture-of-Depths suggests, and re-measure perplexity, wall-clock time, and decoding-time KV-cache size. The central claim stands only if a substantial share of the reported gains survives that change; if the gains evaporate, the results depend on training-time access to future tokens.","supporting_citations":[{"cited_title":"Mixture-of-experts with expert choice routing","cited_arxiv_id":null,"evidence_quote":"Expert-choice routing: supplies the inverted selection paradigm in which each expert chooses its own tokens, giving perfect load balance without auxiliary losses."},{"cited_title":"Efficient content-based sparse attention with routing transformers","cited_arxiv_id":null,"evidence_quote":"Routing Transformer: the main content-based sparse attention baseline that MoSA must beat in the IsoFLOP, long-sequence, and downstream comparisons."},{"cited_title":"Switchhead: Accelerating transformers with mixture-of-experts attention","cited_arxiv_id":null,"evidence_quote":"SwitchHead: earlier work applying mixture-of-experts inside attention; MoSA positions itself by making attention matrices sparse rather than only reducing heads."},{"cited_title":"Efficient streaming language models with attention sinks","cited_arxiv_id":null,"evidence_quote":"StreamingLLM: documents attention sinks, the observation behind MoSA's rule of always including the first token in every head."},{"cited_title":"Fu, Stefano Ermon, Atri Rudra, and Christopher Ré","cited_arxiv_id":null,"evidence_quote":"FlashAttention: the optimized attention kernel that MoSA's sparse attention is compatible with, supporting the practical efficiency claim."}],"review_version":1}