{"id":"d54c7eae-067d-40e4-835e-cb10be30b69f","arxiv_id":"2506.13541","paper_version":1,"verdict":"CONDITIONAL","confidence":"HIGH","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":3,"one_line_summary":"Under equal KV cache budgets, routing each token to weight-shared grouped-attention experts with group sizes 1, 2, and 4 yields higher ROUGE-L and lower perplexity than static GQA and CLA baselines.","lead":"A new mixture-of-experts attention layer for large language models routes each token to one of several KV-cache grouping schemes (full, half, quarter heads), keeping the same average memory but shifting granularity to where it helps. This paper tests it on instruction tuning and continued pretraining, reporting quality gains over static grouped-query attention under equal KV budgets.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"Router receives no language-modeling gradient: the claimed 'learned importance' partition is trained only by a self-referential consistency loss, so the same-budget gains may reflect task-specific hidden-state clustering rather than a transferable token-importance mechanism.","rationale":"The reader's weakest assumption and my concern coincide: the router is trained only by the auxiliary consistency loss because the hard routing masks block gradients from L_model, so the decode-time argmax must inherit a meaningful, capacity-respecting token assignment from a self-referential objective. I agree this is the most load-bearing point because the entire contribution—dynamic token-wise KV optimization guided by learned importance—rests on the router encoding something beyond self-consistency. My critique adds two concrete sharpenings: (1) the hidden-state features x do carry LM-trained signal, so the router is not necessarily arbitrary, but the paper provides no evidence that the consistency-trained partition is importance-aligned; and (2) the 'same KV budget' claim requires the argmax frequencies to match the nominal expert ratios at decode, which is not guaranteed by the auxiliary loss and is not directly reported. The proposed test would settle both points: if the empirical expert frequencies stay at 3:1:6 and expert 1 tokens are measurably more important, the concern does not land; if either fails, the central comparison and the mechanism description require revision. The reader's CONDITIONAL verdict remains appropriate: the experiments are suggestive, but the router-training assumption is not yet supported. No code or formal verification is provided, and L_sce is undefined, which further supports conditionality rather than rejection or acceptance.","tokens_in":15472,"tokens_out":12730,"duration_ms":145878,"concrete_test":"Reproduce or obtain the trained Llama3.1-8B mixSGA model and run greedy decoding on 1,000 held-out Dolly prompts plus WikiText-2, recording the argmax expert for every token and the actual KV bytes stored. Then (i) compare the empirical expert frequencies and average KV size against the nominal 3:1:6 ratios and 50% budget, and (ii) compute an independent per-token importance score, e.g., the increase in sequence perplexity when that token's KV is replaced by the group-averaged KV, as in Figure 1, and test whether tokens routed to expert 1 have significantly higher importance than tokens routed to expert 3 using a permutation test on the rank-biserial correlation.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The load-bearing mechanism is the router S(x)=sigmoid(xφ+β) (Eq. 1). The hard masks m_e in Eq. 2 are piecewise-constant functions of S, so the language-modeling loss L_model in Eq. 5 cannot backpropagate into φ or β. The only gradient for the router comes from the auxiliary consistency loss in Eq. 4, L_sce(S(x), argmax T(x)), and that target is itself a function of S via the greedy top-⌈ρ_e L⌉ construction. The router is therefore optimized purely for self-consistency between the full-sequence greedy assignment and the per-token argmax; no term in the objective rewards assigning important tokens to high-KV experts. The claim that mixSGA performs 'token-wise expert-choice routing guided by learned importance scores' thus depends on the unverified assumption that hidden states x, which do receive LM gradients through the backbone, happen to separate tokens by importance in a way that the consistency-trained linear layer can latch onto. The same assumption underlies the decode-time budget: nothing in Eq. 4 explicitly enforces the 3:1:6 expert proportions at inference, and if argmax frequencies drift on new data, the actual KV cache size is no longer equal to the GQA baseline, invalidating the 'same KV budget' comparison. Table 6 shows that removing L_aux produces near-random expert ratios, but it does not show that the learned ratios are stable, task-independent, or importance-aligned. Additionally, L_sce is never defined in the paper, so the training signal itself is not fully specified.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper proposes mixSGA, an attention mechanism that assigns each token to one of several weight-shared KV-group experts with different group sizes (e.g., group sizes 1, 2, 4) via a learned router, while retaining all tokens instead of evicting low-importance ones. A progressive top-k routing is used during prefill/training (Eq. 2), and decode-time routing uses the argmax of the same score function (Eq. 3), with an auxiliary loss (Eq. 4) intended to align the two. Experiments on instruction-following (Dolly-15k fine-tuning) and continued pretraining (MiniPile on TinyLlama-1.1B) report consistent improvements in ROUGE-L and perplexity over GQA and CLA under matched KV budgets, plus compatibility with H2O eviction.","tokens_in":15848,"tokens_out":7769,"duration_ms":77429,"significance":"If the reported results are reproducible, mixSGA offers a simple and parameter-efficient way to allocate KV resources non-uniformly across tokens while preserving full context, with only a modest decoding-throughput overhead (~3-4%). The paper covers multiple model families (OPT, Llama3, Gemma2, TinyLlama), reports memory and FLOP overheads, and includes ablations over expert ratios, number of experts, and the auxiliary loss. However, the central claim that routing is guided by 'learned importance scores' is not supported by the training objective as written: the router receives no gradient from the language-modeling loss and its auxiliary target is derived from the router's own scores. The same-budget comparisons also depend on the decode-time argmax maintaining the predefined expert ratios, which is not verified on held-out data. Several dynamic KV-optimization baselines cited in the motivation are never compared, and no uncertainty estimates are provided.","major_comments":[{"comment":"The auxiliary loss L_sce is never defined, so the training signal for the router is not fully specified and the method is not reproducible as written. Moreover, the target argmax T(x) in Eq. (4) is itself the routing produced from S(x) by the progressive procedure in Eq. (2), making the objective self-referential: it only encourages the per-token argmax to agree with the sequential top-(rho_e L) assignment. No term in Eq. (5) rewards routing based on token importance, so the abstract's and Section 3.2's claim of 'learned importance scores' is not justified by the given training objective. Please define L_sce, and either add an external importance-supervision signal or empirically demonstrate, for example by correlating expert assignments with the per-token sensitivity measure used in Figure 1, that the learned routing is indeed importance-aligned.","section":"Section 3.3, Eq. (4)"},{"comment":"Because the masks in Eq. (2) are hard indicators and Eq. (3) uses argmax, no gradient from L_model in Eq. (5) reaches the router parameters phi and beta; the router is trained only by the auxiliary consistency loss. This is an assumption about training dynamics, and the paper never verifies that the trained router maintains the intended 3:1:6 expert proportions at decoding time on the evaluation sets. Table 6 shows that without the auxiliary loss the ratios become nearly uniform, but the paper does not report the measured ratios for the full model on held-out data. If the decode-time ratios drift, the 'same KV budget' comparisons in Tables 1-3 are not guaranteed to hold. Please report the actual expert-assignment distribution for the trained router on the test sets, or explicitly enforce the target ratios at decoding (e.g., via capacity-limited sampling) and describe how the KV budget is then computed.","section":"Section 3.1-3.2; Tables 1-3"},{"comment":"The paper motivates mixSGA against DynamicKV [40] and PyramidKV [5] and in Section 1 claims 'superior efficiency and performance over static and dynamic baselines', yet none of the experiments compare against these or other token-level KV-compression methods. The only dynamic baseline is H2O, which is applied as an eviction post-processing on top of both GQA and mixSGA. Please either add comparisons to at least one of the cited dynamic methods (DynamicKV or PyramidKV) under matched budgets, or revise the contributions and abstract to claim superiority over static-grouping baselines only.","section":"Section 4.1-4.2; Section 2.1"},{"comment":"All reported results are from single runs, with no standard deviations, confidence intervals, or multiple seeds. Given that the fine-tuning uses a relatively small Dolly-15k subset (14,000 training samples) and the continued-pretraining run is one epoch, differences on the order of a few tenths of a percent (e.g., Table 2: HellaSwag 37.00 vs 36.70; ARC-E 54.84 vs 54.92) are within plausible random variation. Please report means and standard deviations over at least three independent seeds for the main comparisons, and indicate how many seeds were used for each reported number.","section":"Tables 1-3"},{"comment":"The main text states that the grid search for the optimal expert density ratio was conducted on OPT-355M, while Appendix C.2 reports the grid search on Gemma2-2B. This inconsistency must be corrected. The paper should also explain how the same 3:1:6 ratio is justified across different architectures and whether the optimal ratio was re-tuned per model family.","section":"Section 4.1 vs Appendix C.2"}],"minor_comments":[{"comment":"In the Llama3.2-3B row, the GQA Vicuna DeepSeek-v3 feedback score is listed as 35.4, which is an order of magnitude larger than all neighboring values (typically 2-5); this appears to be a typo for 3.54.","section":"Table 1"},{"comment":"The group-size notation is inconsistent: the text says 'group size of 2^e' and 'H/2^e groups', but the example with H=4 and E=3 gives G1 as four groups of size 1, G2 as two groups of size 2, and G3 as one group of size 4. This corresponds to group sizes 2^(e-1), not 2^e. Please fix the notation or the indexing.","section":"Section 3.4"},{"comment":"The sentence 'while efficiently handling less critical tokens with the third [expert]' describes the third expert as having 'group size 3', but the method defines the third expert with group size 4; correct the typo.","section":"Appendix C.2"},{"comment":"The indicator notation 1[...] is used without defining its output; please state explicitly that it returns a zero-one tensor of the same shape as its argument, and clarify the behavior when the ceiling operation makes the sum of capacities exceed L.","section":"Eq. (2)"},{"comment":"Reference [33] is cited as a general KV-cache reference, but the cited work ('KV-cache: A scalable high-performance web-object cache for manycore') is about a web-object cache, not the transformer key-value cache; this citation is misplaced.","section":"References"}],"recommendation":"major_revision","confidential_remarks":"The self-referential nature of the auxiliary loss is the most serious conceptual gap: the router is trained only to agree with itself, so the 'learned importance' claim is not established. A strong revision would add an oracle-routing ablation (e.g., route by the Figure 1 sensitivity measure) or at least report the measured decode-time ratios and show they match the budget. The undefined L_sce and missing dynamic baselines are additional reproducibility and positioning issues that should be addressed before publication. The work is potentially interesting, but the current evidence and presentation do not yet justify the central mechanism claim."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Quick take: this is a real extension, not a restatement. The mixSGA idea—route each token to one of three weight-shared attention experts with different KV group sizes, keep all tokens, and add an auxiliary loss to align prefill and decode assignments—is new relative to GQA, CLA, and the ECR literature. The group-size mechanism is simple and parameter-efficient: same projections, averaged KV heads, scatter/gather, 3–4% overhead. The experiments cover several model families, and the gains over static GQA are large and fairly consistent. On Llama3.1-8B average ROUGE-L goes from 19.97 to 27.08 at the same nominal KV budget; TinyLlama Wikitext-2 PPL drops from 22.66 to 20.46; and the H2O combination table shows benefit under pressure. That is genuine evidence.\n\nSoft spots, in proportion. The biggest is the router training. Equation 4 is an auxiliary consistency loss matching S to the greedy assignment, and the greedy assignment is computed from S itself. No language-modeling gradient reaches the router. So the learned importance score is not grounded in an external notion of token utility; the paper shows the auxiliary loss matters (Table 6: removing it collapses routing to near random), but not that the assignments are importance-aligned or stable across tasks. The decode-time argmax could drift from the 3:1:6 ratios on new data, which would quietly break the same-KV-budget comparison. Also, L_sce in Eq. 4 is never defined. These are fixable in revision: define L_sce, add a small head that receives LM gradient, or at least report measured decode-time ratios and per-seed variance.\n\nMissing baselines are more than cosmetic. DynamicKV and PyramidKV are cited in Related Work, and the abstract claims superiority over dynamic baselines, but neither appears in the tables. If the point is dynamic token-wise KV allocation, omitting the two most direct dynamic KV competitors makes the headline claim provisional. There is also an obvious typo in Table 1 (Llama3.2-3B Vicuna DSv3 = 35.4), and no error bars anywhere; five seeds would cost little here.\n\nMy own read: the central empirical claim probably holds against static GQA, because the effects are large and mechanical—giving some tokens full KV and others quarter KV is rarely worse than uniform half-KV—but the paper's stronger claim about learned importance is not yet supported. It deserves a serious referee; the fixes are substantial but the mechanism is worth engaging. If I were advising, I would send it out and ask for DynamicKV/PyramidKV comparisons, a defined L_sce, and variance reporting.","headline":"A useful KV-grouping variant with large clean gains over static GQA, but the router's training signal is self-referential and key dynamic baselines are missing.","tokens_in":16332,"tokens_out":2020,"would_cite":true,"duration_ms":20934,"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":"Under a fixed KV-cache budget, routing each token to a weight-shared attention expert chosen by learned importance improves language-model quality over static grouped-query attention.","keywords":["KV cache compression","mixture of experts","token-wise routing","grouped query attention","causal language modeling","dynamic inference","attention efficiency"],"falsifier":"Take a trained mixSGA model and, on held-out data, compare the router's assignments against an oracle importance measure such as the per-token perplexity change when a token's KV heads are averaged away, as in the paper's own Figure 1. If the router's chosen expert for a token does not track that oracle ordering, then the reported gains must come from something other than importance-aware allocation, and the claim that dynamic token importance drives the improvement fails.","tokens_in":15293,"feed_emoji":"⚡","tokens_out":10575,"duration_ms":94165,"temperature":0.7,"pith_summary":"This paper claims that the KV cache of a causal language model can be compressed as aggressively as grouped-query attention, but with less quality loss, by making the grouping decision per token instead of per layer. mixSGA routes each token to one of several attention \"experts\" whose key and value projections are shared but whose KV groups differ in size, so important tokens keep many KV heads while unimportant tokens get few, and no token is dropped. The router is trained so that the same assignment is made during prefill and during decode, and the claim is that under the same average cache size the method improves instruction-following quality (Llama3.1-8B average ROUGE-L rises from 19.97 to 27.08) and lowers perplexity on continued pretraining (22.66 to 20.46 on Wikitext-2 with TinyLlama-1.1B). If right, the practical upshot is cheaper, longer-context inference that does not pay for memory savings by evicting information.","feed_headline":"Dynamic per-token KV grouping beats static GQA at equal cache size","feed_subtitle":"Same KV memory, lower perplexity: Llama3.1-8B ROUGE-L rises from 19.97 to 27.08 after instruction tuning.","key_machinery":"The load-bearing mechanism is token-wise expert-choice routing with heterogeneous expert capacities, welded to weight-shared grouped-attention projections. A trainable linear layer maps each token embedding to $E$ sigmoid scores $S(x)=\\sigma(x\\phi+\\beta)$; during prefill and training, expert $e$ receives the top $\\lceil \\rho_e L \\rceil$ tokens not already claimed by earlier experts, producing a hard, exclusive one-hot mask $m_e(x)$; during decoding, each token simply picks $\\arg\\max_e S(x)$. To make those two assignments coincide, the auxiliary loss $\\mathcal{L}_{\\mathrm{aux}}$, a cross-entropy between $S(x)$ and the prefill assignment, is added to the language-modeling loss at every layer. The KV side uses one shared set of key and value projections for all experts and obtains expert $e$'s representation by averaging heads into groups of size $2^e$, so switching experts changes only memory and arithmetic, not parameters. That combination is what lets the paper claim a per-token allocation of KV size proportional to importance without adding parameters or discarding tokens.","core_discovery":"On the paper's own terms, the discovery is that token importance in a sequence is not just a ranking, but a continuous spectrum that shifts with context, and the right response is not to evict low-importance tokens but to give them a coarser KV representation. mixSGA realizes this with a set of heterogeneous grouped-attention experts: expert e averages the shared key and value projections into groups of size 2^e, so its KV cache per token is H/2^e heads, and a learned router assigns each token to exactly one expert. Prefill assigns tokens greedily by capacity ratios, decode uses an argmax of the same scoring function, and an auxiliary cross-entropy loss is the only training signal for the router, because the hard one-hot masks block gradients from the language-modeling loss. With a 3:1:6 capacity split among group sizes 1, 2, and 4 (50% of the original cache), the paper reports consistent improvements over GQA across OPT, Llama3, and Gemma2 models on instruction-following tasks and over GQA and CLA in continued pretraining, and shows the method combines with H2O-style eviction.","pith_inferences":["If the learned importance scores are genuinely transferable, they could serve as a cheap token-importance signal for other memory-saving devices the paper does not test, such as prefix caching, prompt compression, or deciding where to place tokens in hierarchical caches; the paper only demonstrates routing learned jointly with fine-tuning or continued pretraining.","Because the routing weights are trained only by the auxiliary loss, the approach is easy to probe causally: freeze the model, supervise the router directly with an oracle importance ranking, and compare quality against the auxiliary-loss-trained router, a clean way to measure how much of the gain is routing accuracy versus the richer token-importance prior itself.","The paper fixes expert group sizes to powers of two (1, 2, 4); nothing in the weight-sharing design forces that, so one could interpolate to fractional or per-head granularities or condition capacity ratios on sequence length, which might extend the budget-linear scaling shown in Figure 3 to more extreme compression."],"forward_implications":["Under the same average KV cache size as a GQA baseline, mixSGA raises average ROUGE-L across five instruction-following datasets, with the largest reported jump on Llama3.1-8B from 19.97 (GQA) to 27.08 (mixSGA 3:1:6).","In continued pretraining on TinyLlama-1.1B, mixSGA lowers Wikitext-2 perplexity from 22.66 (GQA) to 20.46 and matches or beats GQA and CLA on HellaSwag, PIQA, Winogrande, and ARC-Challenge.","The method composes with KV eviction: with H2O keep ratios from 80% down to 20%, mixSGA combined with H2O beats GQA with H2O on average accuracy and perplexity at every tested ratio.","The routing consistency loss is load-bearing: removing the auxiliary loss drops average ROUGE-L on Gemma2-2B from 21.20 to 7.35, and replacing the learned router with a random router costs about 3.5 points on the same average.","The optimal expert mix under a 50% KV budget is consistently the 3:1:6 ratio of group sizes 1, 2, and 4, which allocates 30% of tokens to fine, 10% to medium, and 60% to coarse representation."],"supporting_citations":[{"why":"Defines grouped query attention, the static KV-grouping baseline all experiments are gated against; the weight-shared experts are built by extending GQA's head grouping.","marker":"[1]"},{"why":"Cross-layer attention is the second static baseline in the continued-pretraining comparison, against which mixSGA reports lower perplexity across KV budgets.","marker":"[4]"},{"why":"H2O heavy-hitter eviction is the baseline for the combined experiment, establishing that mixSGA layers on top of token eviction and still outperforms GQA+H2O at every keep ratio.","marker":"[39]"},{"why":"Expert-choice routing supplies the capacity-based assignment paradigm that mixSGA adapts to autoregressive decoding, and its prefill-decode mismatch motivates the consistency loss.","marker":"[41]"},{"why":"The Dolly-15k instruction-following dataset is used for fine-tuning and for the ROUGE-L benchmarks in Table 1.","marker":"[10]"},{"why":"MiniPile is the 1.6-billion-token corpus on which TinyLlama is continued-pretrained for the perplexity and downstream-accuracy comparisons.","marker":"[19]"},{"why":"DynamicKV exemplifies token-level KV allocation that prioritizes high-value tokens; the paper contrasts mixSGA with its rigid allocation and motivates retaining all tokens.","marker":"[40]"}],"fun_headline_variants":["mixSGA: dynamic KV groups outperform static GQA at same cache","Per-token KV experts: keep all tokens, cut cache size","Adaptive KV grouping: lower perplexity, no token eviction","Token importance drives KV allocation in mixSGA, beats GQA","Same KV budget, better scores: mixSGA dynamic grouping"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The routing assignments that decide which tokens get fine or coarse KV representations are learned only from the auxiliary consistency loss, since the hard one-hot masks block any gradient from the language-modeling loss; the whole method therefore rests on that single signal learning a meaningful, transferable ranking of token importance.","fun_headline_variants_meta":{"raw":{"variants":["mixSGA: dynamic KV groups outperform static GQA at same cache","Per-token KV experts: keep all tokens, cut cache size","Adaptive KV grouping: lower perplexity, no token eviction","Token importance drives KV allocation in mixSGA, beats GQA","Same KV budget, better scores: mixSGA dynamic grouping"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000249,"raw_usage":{"total_tokens":1599,"prompt_tokens":1045,"completion_tokens":554,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":661,"completion_tokens_details":{"reasoning_tokens":463}},"tokens_in":661,"tokens_out":554,"duration_ms":5663,"temperature":1.0,"reasoning_tokens":463,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-15T19:58:57.418624+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Take a trained mixSGA model and, on held-out data, compare the router's assignments against an oracle importance measure such as the per-token perplexity change when a token's KV heads are averaged away, as in the paper's own Figure 1. If the router's chosen expert for a token does not track that oracle ordering, then the reported gains must come from something other than importance-aware allocation, and the claim that dynamic token importance drives the improvement fails.","supporting_citations":[{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"Defines grouped query attention, the static KV-grouping baseline all experiments are gated against; the weight-shared experts are built by extending GQA's head grouping."},{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"Cross-layer attention is the second static baseline in the continued-pretraining comparison, against which mixSGA reports lower perplexity across KV budgets."},{"cited_title":null,"cited_arxiv_id":null,"evidence_quote":"H2O heavy-hitter eviction is the baseline for the combined experiment, establishing that mixSGA layers on top of token eviction and still outperforms GQA+H2O at every keep ratio."},{"cited_title":"Dai, Zhifeng Chen, Quoc V Le, and James Laudon","cited_arxiv_id":null,"evidence_quote":"Expert-choice routing supplies the capacity-based assignment paradigm that mixSGA adapts to autoregressive decoding, and its prefill-decode mismatch motivates the consistency loss."}],"review_version":2}